train.py

"""
train.py — Main training script for Hybrid Quantum-Classical NLP Classifier
Run: python train.py
Author: Paladugu Nandith Kumar
"""

import os
import sys
import warnings
import torch
import numpy as np
import json
warnings.filterwarnings('ignore')

from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, f1_score, confusion_matrix, classification_report

from src.classical_preprocessor import TextFeatureExtractor, get_sample_dataset, load_custom_dataset
from src.hybrid_model import HybridQuantumClassifier, ClassicalBaseline
from src.trainer import Trainer
from src.visualizer import (
    plot_training_history, plot_confusion_matrix,
    plot_quantum_vs_classical, plot_circuit_diagram,
    plot_feature_importance
)

# ── Config ─────────────────────────────────────────────────────
N_QUBITS      = 8
N_LAYERS      = 3
CIRCUIT_TYPE  = "basic"        # "basic" or "strongly_entangling"
EPOCHS        = 50
BATCH_SIZE    = 8
LR_QUANTUM    = 0.02
LR_CLASSICAL  = 0.01
PATIENCE      = 15
DROPOUT       = 0.2
RANDOM_STATE  = 42
DATASET       = "sample"       # "sample" or path to CSV
OUTPUT_DIR    = "outputs"
MODEL_PATH    = "outputs/quantum_model.pt"

os.makedirs(OUTPUT_DIR, exist_ok=True)


def main():
    print("=" * 60)
    print("  Hybrid Quantum-Classical NLP Classifier")
    print("  Author: Paladugu Nandith Kumar")
    print("=" * 60)

    # ── Step 1: Load Data ──────────────────────────────────────
    print("\n[1/6] Loading dataset...")
    if DATASET == "sample":
        texts, labels, label_names = get_sample_dataset()
    else:
        texts, labels = load_custom_dataset(DATASET)
        label_names = [str(i) for i in sorted(set(labels))]
    n_classes = len(set(labels))
    print(f"Samples: {len(texts)} | Classes: {label_names}")

    # ── Step 2: Feature Extraction ────────────────────────────
    print("\n[2/6] Extracting TF-IDF + PCA features...")
    extractor = TextFeatureExtractor(n_features=N_QUBITS, max_vocab=1000)
    X = extractor.fit_transform(texts).astype(np.float32)
    y = np.array(labels)
    print(f"Feature shape: {X.shape} | Range: [{X.min():.3f}, {X.max():.3f}]")

    # ── Step 3: Train/Val/Test Split ──────────────────────────
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=RANDOM_STATE, stratify=y
    )
    X_train, X_val, y_train, y_val = train_test_split(
        X_train, y_train, test_size=0.2, random_state=RANDOM_STATE, stratify=y_train
    )
    print(f"\nSplit — Train: {len(X_train)} | Val: {len(X_val)} | Test: {len(X_test)}")

    # ── Step 4: Build Hybrid Quantum Model ────────────────────
    print("\n[3/6] Building Hybrid Quantum-Classical model...")
    q_model = HybridQuantumClassifier(
        input_dim=N_QUBITS, n_classes=n_classes,
        n_qubits=N_QUBITS, n_layers=N_LAYERS,
        circuit_type=CIRCUIT_TYPE, dropout=DROPOUT
    )
    params = q_model.count_parameters()
    print(f"Parameters — Classical: {params['classical']} | "
          f"Quantum: {params['quantum']} | Total: {params['total']}")

    # ── Step 5: Train Quantum Model ───────────────────────────
    print("\n[4/6] Training Hybrid Quantum-Classical model...")
    q_trainer = Trainer(q_model, lr=LR_QUANTUM)
    q_history = q_trainer.fit(
        X_train, y_train, X_val, y_val,
        epochs=EPOCHS, batch_size=BATCH_SIZE,
        patience=PATIENCE, verbose=True
    )
    q_trainer.save(MODEL_PATH)

    # ── Step 6: Build and Train Classical Baseline ────────────
    print("\n[5/6] Training Classical baseline for comparison...")
    c_model   = ClassicalBaseline(input_dim=N_QUBITS, n_classes=n_classes, dropout=DROPOUT)
    c_trainer = Trainer(c_model, lr=LR_CLASSICAL)
    c_history = c_trainer.fit(
        X_train, y_train, X_val, y_val,
        epochs=EPOCHS, batch_size=BATCH_SIZE,
        patience=PATIENCE, verbose=True
    )

    # ── Step 7: Evaluate ──────────────────────────────────────
    print("\n[6/6] Evaluating on test set...")
    q_results = q_trainer.evaluate(X_test, y_test, label_names)
    c_results = c_trainer.evaluate(X_test, y_test, label_names)

    print("\n" + "=" * 60)
    print("  RESULTS SUMMARY")
    print("=" * 60)
    print(f"\n{'Metric':<15} {'Quantum':>12} {'Classical':>12}")
    print("-" * 40)
    for metric in ['accuracy', 'f1', 'precision', 'recall']:
        print(f"{metric:<15} {q_results[metric]:>12.4f} {c_results[metric]:>12.4f}")

    print(f"\nFull Classification Report (Quantum Model):")
    print(q_results['report'])

    # ── Save Results ──────────────────────────────────────────
    summary = {
        "config": {
            "n_qubits": N_QUBITS, "n_layers": N_LAYERS,
            "circuit_type": CIRCUIT_TYPE, "dataset_size": len(texts)
        },
        "quantum": {k: float(q_results[k]) for k in ['accuracy','f1','precision','recall']},
        "classical": {k: float(c_results[k]) for k in ['accuracy','f1','precision','recall']},
    }
    with open(f"{OUTPUT_DIR}/results_summary.json", "w") as f:
        json.dump(summary, f, indent=2)

    # ── Generate Plots ────────────────────────────────────────
    print("\nGenerating visualization plots...")
    plot_training_history(q_history,   f"{OUTPUT_DIR}/training_history.png")
    plot_confusion_matrix(q_results['confusion_matrix'], label_names, f"{OUTPUT_DIR}/confusion_matrix.png")
    plot_quantum_vs_classical(q_history, c_history,      f"{OUTPUT_DIR}/quantum_vs_classical.png")
    plot_circuit_diagram(N_QUBITS, N_LAYERS,             f"{OUTPUT_DIR}/circuit_diagram.png")
    plot_feature_importance(
        extractor.feature_names,
        extractor.pca.explained_variance_ratio_,
        f"{OUTPUT_DIR}/feature_importance.png"
    )

    print(f"\nAll outputs saved to: {OUTPUT_DIR}/")
    print("Training complete!")


if __name__ == "__main__":
    main()