demo_inference.py

"""
Interactive Demo: DistilBERT Sentiment Analysis
==============================================

Demonstrates the best-performing model (LR=1e-5) on sample reviews.
Perfect for presentations and live demonstrations.

Student: Martynas Prascevicius
Student ID: 001263199
Course: COMP1818 Artificial Intelligence Applications
"""

import torch
import torch.nn as nn
from transformers import DistilBertTokenizer, DistilBertModel
from pathlib import Path
import json

# Sample reviews for demonstration
SAMPLE_REVIEWS = [
    {
        "text": "This movie was absolutely fantastic! The acting was superb and the plot kept me engaged from start to finish. Highly recommended!",
        "expected": "POSITIVE",
        "source": "Custom (clearly positive)"
    },
    {
        "text": "Terrible waste of time. Poor acting, nonsensical plot, and awful special effects. I want my money back!",
        "expected": "NEGATIVE",
        "source": "Custom (clearly negative)"
    },
    {
        "text": "The cinematography was breathtaking and the soundtrack was amazing, but the story felt rushed and the ending was disappointing.",
        "expected": "MIXED",
        "source": "Custom (mixed sentiment)"
    },
    {
        "text": "An instant classic! This film will be remembered for generations. The director's vision is truly masterful.",
        "expected": "POSITIVE",
        "source": "Custom (high praise)"
    },
    {
        "text": "I've seen better films at a middle school talent show. The only good thing about this movie was the popcorn.",
        "expected": "NEGATIVE",
        "source": "Custom (harsh criticism)"
    }
]


class DistilBERTClassifier(nn.Module):
    """DistilBERT model for binary sentiment classification."""

    def __init__(self, model_name='distilbert-base-uncased', num_classes=2):
        super(DistilBERTClassifier, self).__init__()
        self.distilbert = DistilBertModel.from_pretrained(model_name)
        self.dropout = nn.Dropout(0.1)
        self.classifier = nn.Linear(768, num_classes)

    def forward(self, input_ids, attention_mask):
        outputs = self.distilbert(
            input_ids=input_ids,
            attention_mask=attention_mask
        )
        # Use [CLS] token representation
        pooled_output = outputs.last_hidden_state[:, 0, :]
        pooled_output = self.dropout(pooled_output)
        logits = self.classifier(pooled_output)
        return logits


def load_best_model(model_dir: Path, device: str = 'mps'):
    """
    Load the best-performing model (lr_1e5) for inference.

    Args:
        model_dir: Directory containing saved models
        device: Device to load model on ('mps', 'cuda', or 'cpu')

    Returns:
        Tuple of (model, tokenizer, device)
    """
    print("=" * 70)
    print("Loading Best Model Configuration")
    print("=" * 70)

    # Check device availability
    if device == 'mps' and not torch.backends.mps.is_available():
        print("⚠️  MPS not available, falling back to CPU")
        device = 'cpu'
    elif device == 'cuda' and not torch.cuda.is_available():
        print("⚠️  CUDA not available, falling back to CPU")
        device = 'cpu'

    device = torch.device(device)

    # Load tokenizer
    tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')

    # Initialize model
    model = DistilBERTClassifier(
        model_name='distilbert-base-uncased',
        num_classes=2
    )

    # Load best model weights (lr_1e5)
    model_path = model_dir / 'best_model_lr_1e5.pt'

    if not model_path.exists():
        print(f"\n⚠️  Pre-trained model not found at: {model_path}")
        print("Using freshly initialized DistilBERT (will give random predictions)")
        print("Note: For actual demo, train the lr_1e5 experiment first!\n")
    else:
        checkpoint = torch.load(model_path, map_location=device)
        model.load_state_dict(checkpoint['model_state_dict'])
        print(f"\n✓ Loaded best model from: {model_path}")
        print(f"  Epoch: {checkpoint.get('epoch', 'N/A')}")
        print(f"  Test Accuracy: {checkpoint.get('test_accuracy', 'N/A'):.2%}")

    model.to(device)
    model.eval()

    print(f"  Device: {device}")
    print(f"  Parameters: {sum(p.numel() for p in model.parameters()):,}")
    print()

    return model, tokenizer, device


def predict_sentiment(text: str, model, tokenizer, device, max_length=256):
    """
    Predict sentiment for a single text.

    Args:
        text: Input review text
        model: Trained DistilBERT model
        tokenizer: DistilBERT tokenizer
        device: Device to run inference on
        max_length: Maximum sequence length

    Returns:
        Tuple of (predicted_label, confidence)
    """
    # Tokenize
    encoding = tokenizer(
        text,
        add_special_tokens=True,
        max_length=max_length,
        padding='max_length',
        truncation=True,
        return_tensors='pt'
    )

    input_ids = encoding['input_ids'].to(device)
    attention_mask = encoding['attention_mask'].to(device)

    # Predict
    with torch.no_grad():
        logits = model(input_ids, attention_mask)
        probs = torch.softmax(logits, dim=1)
        confidence, predicted = torch.max(probs, dim=1)

    label = "POSITIVE" if predicted.item() == 1 else "NEGATIVE"
    confidence_pct = confidence.item() * 100

    return label, confidence_pct, probs[0].cpu().numpy()


def run_demo():
    """Run interactive sentiment analysis demo."""
    project_root = Path(__file__).parent
    model_dir = project_root / 'models'

    print("\n")
    print("╔" + "=" * 68 + "╗")
    print("║" + " " * 15 + "DistilBERT Sentiment Analysis Demo" + " " * 19 + "║")
    print("║" + " " * 20 + "COMP1818 - CW2 Presentation" + " " * 21 + "║")
    print("╚" + "=" * 68 + "╝")
    print()

    # Load model
    model, tokenizer, device = load_best_model(model_dir, device='mps')

    print("=" * 70)
    print("Running Inference on Sample Reviews")
    print("=" * 70)
    print()

    # Analyze sample reviews
    results = []
    for i, sample in enumerate(SAMPLE_REVIEWS, 1):
        print(f"Review {i}/{len(SAMPLE_REVIEWS)}")
        print("-" * 70)
        print(f"Text: \"{sample['text'][:80]}...\"" if len(sample['text']) > 80 else f"Text: \"{sample['text']}\"")
        print(f"Expected: {sample['expected']}")

        label, confidence, probs = predict_sentiment(
            sample['text'], model, tokenizer, device
        )

        print(f"Predicted: {label} ({confidence:.1f}% confidence)")
        print(f"Probabilities: Negative={probs[0]:.1%}, Positive={probs[1]:.1%}")

        # Determine correctness
        correct = "✓" if (label == sample['expected'] or sample['expected'] == "MIXED") else "✗"
        print(f"Status: {correct}")
        print()

        results.append({
            'text': sample['text'],
            'expected': sample['expected'],
            'predicted': label,
            'confidence': confidence,
            'correct': correct == "✓"
        })

    # Summary
    print("=" * 70)
    print("Demo Summary")
    print("=" * 70)

    total = len(results)
    correct = sum(1 for r in results if r['correct'])
    avg_confidence = sum(r['confidence'] for r in results) / total

    print(f"Total Samples: {total}")
    print(f"Correct Predictions: {correct}/{total} ({correct/total:.1%})")
    print(f"Average Confidence: {avg_confidence:.1f}%")
    print()

    print("Key Findings from Full Experiments:")
    print("-" * 70)
    print("  ✓ Best Configuration: LR=1e-5, Batch=16, Epochs=3")
    print("  ✓ Test Accuracy: 91.04% (25,000 IMDB reviews)")
    print("  ✓ Improvement: +0.27% vs. baseline (2e-5 learning rate)")
    print("  ✓ Training Time: 160 minutes (Mac M4)")
    print()

    print("=" * 70)
    print("Demo Complete!")
    print("=" * 70)
    print()


def interactive_mode():
    """Interactive mode for custom review input."""
    project_root = Path(__file__).parent
    model_dir = project_root / 'models'

    print("\n")
    print("╔" + "=" * 68 + "╗")
    print("║" + " " * 10 + "Interactive DistilBERT Sentiment Analyzer" + " " * 17 + "║")
    print("╚" + "=" * 68 + "╝")
    print()

    # Load model
    model, tokenizer, device = load_best_model(model_dir, device='mps')

    print("Enter movie reviews to analyze (or 'quit' to exit)")
    print("-" * 70)
    print()

    while True:
        try:
            user_input = input("Review: ").strip()

            if user_input.lower() in ['quit', 'exit', 'q']:
                print("\nThank you for using the demo!")
                break

            if not user_input:
                print("Please enter a review.\n")
                continue

            label, confidence, probs = predict_sentiment(
                user_input, model, tokenizer, device
            )

            print(f"\nPrediction: {label} ({confidence:.1f}% confidence)")
            print(f"Probabilities: Negative={probs[0]:.1%}, Positive={probs[1]:.1%}")
            print()

        except KeyboardInterrupt:
            print("\n\nDemo interrupted. Goodbye!")
            break
        except Exception as e:
            print(f"\nError: {e}")
            print("Please try again.\n")


if __name__ == '__main__':
    import sys

    if len(sys.argv) > 1 and sys.argv[1] == '--interactive':
        interactive_mode()
    else:
        run_demo()
        print("\nTip: Run 'python demo_inference.py --interactive' for custom reviews!")