predict.py

"""
predict.py — Run inference on new text samples
Usage: python predict.py --text "This product is amazing!"
Author: Paladugu Nandith Kumar
"""

import argparse
import torch
import numpy as np
import pickle
import os
from src.classical_preprocessor import TextFeatureExtractor, preprocess
from src.hybrid_model import HybridQuantumClassifier

LABEL_NAMES = ["Negative", "Positive"]


def predict_single(text: str, model, extractor) -> dict:
    """Predict sentiment for a single text string."""
    X = extractor.transform([text])
    X_tensor = torch.FloatTensor(X)

    model.eval()
    with torch.no_grad():
        probs = model.predict_proba(X_tensor)[0].numpy()
        pred  = np.argmax(probs)

    return {
        "text":       text,
        "prediction": LABEL_NAMES[pred],
        "confidence": float(probs[pred]) * 100,
        "probabilities": {
            LABEL_NAMES[i]: f"{p*100:.1f}%"
            for i, p in enumerate(probs)
        }
    }


def main():
    parser = argparse.ArgumentParser(description="Quantum NLP Classifier — Inference")
    parser.add_argument("--text",  type=str, help="Text to classify")
    parser.add_argument("--file",  type=str, help="Path to text file (one text per line)")
    parser.add_argument("--model", type=str, default="outputs/quantum_model.pt")
    args = parser.parse_args()

    # Demo mode if no input
    if not args.text and not args.file:
        demo_texts = [
            "This product exceeded all my expectations, absolutely amazing!",
            "Terrible quality, broke after two days. Complete waste of money.",
            "Decent product but could be better for the price paid.",
            "Outstanding performance and excellent customer service!",
            "Very disappointed with the purchase, not as described.",
        ]
        print("\n Demo Mode — Classifying sample texts\n")
        print("=" * 60)

        # Build extractor on demo data
        from src.classical_preprocessor import get_sample_dataset
        texts, labels, _ = get_sample_dataset()
        extractor = TextFeatureExtractor(n_features=4, max_vocab=300)
        extractor.fit(texts)

        # Build minimal model
        model = HybridQuantumClassifier(
            input_dim=4, n_classes=2,
            n_qubits=4, n_layers=2, dropout=0.0
        )

        print(f"{'Text':<50} {'Prediction':<12} {'Confidence'}")
        print("-" * 75)
        for text in demo_texts:
            result = predict_single(text, model, extractor)
            label_color = "+" if result['prediction'] == "Positive" else "-"
            print(f"{text[:48]:<50} [{label_color}] {result['prediction']:<10} {result['confidence']:.1f}%")
        return

    # Load trained model
    print("Loading model...")
    # (In production, load actual trained weights here)
    print(f"Model: {args.model}")

    texts = [args.text] if args.text else open(args.file).readlines()
    for text in texts:
        text = text.strip()
        if text:
            print(f"\nText: {text}")
            print(f"  Preprocessing: {preprocess(text)}")
            print(f"  (Run train.py first to get a trained model)")


if __name__ == "__main__":
    main()