test_model.py

import torch
import torch.nn as nn
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from models.fma_model import FMAEncoderModel
from models.baseline_model import StandardTransformerModel
import os

class ModelTester:
    def __init__(self, device='cuda'):
        self.device = torch.device(device if torch.cuda.is_available() else 'cpu')
        print(f"Using device: {self.device}")
        
        # Model configs (must match training)
        self.vocab_size = None
        self.d_model = 128
        self.num_layers = 2
        self.num_modes = 32
        self.max_len = 128
        self.num_classes = 2
        self.nhead = 4
        
        self.tokenizer = None
        self.teacher = None
        self.student = None
        self.teacher_hf = None  # For HuggingFace pretrained
        
    def load_distilled_models(self):
        """Load models trained with train_distill.py (BERT-Tiny teacher + FMA student)"""
        print("\n=== Loading Distilled Models ===")
        
        # Load tokenizer
        teacher_name = "prajjwal1/bert-tiny"
        print(f"Loading tokenizer: {teacher_name}")
        self.tokenizer = AutoTokenizer.from_pretrained(teacher_name)
        self.vocab_size = self.tokenizer.vocab_size
        
        # Load HuggingFace teacher
        print(f"Loading Teacher (HF): {teacher_name}")
        self.teacher_hf = AutoModelForSequenceClassification.from_pretrained(
            teacher_name, num_labels=2
        ).to(self.device)
        self.teacher_hf.eval()
        
        # Load FMA student
        print("Loading Student (FMA)")
        self.student = FMAEncoderModel(
            vocab_size=self.vocab_size,
            d_model=self.d_model,
            num_layers=self.num_layers,
            num_modes=self.num_modes,
            max_len=self.max_len,
            num_classes=self.num_classes
        ).to(self.device)
        
        # Load checkpoint if exists
        ckpt_path = "checkpoints/student_fma_distilled.pt"
        if os.path.exists(ckpt_path):
            self.student.load_state_dict(torch.load(ckpt_path, map_location=self.device))
            print(f"✓ Loaded checkpoint: {ckpt_path}")
        else:
            print(f"⚠ Checkpoint not found: {ckpt_path}")
            print("  Run: conda run -n fsnn python -m training.train_distill")
        
        self.student.eval()
        
    def load_tiny_models(self):
        """Load models trained with train_tiny.py (synthetic data)"""
        print("\n=== Loading Tiny Models ===")
        
        self.vocab_size = 1000
        max_len = 64
        
        # Load teacher
        print("Loading Teacher (Standard Transformer)")
        self.teacher = StandardTransformerModel(
            vocab_size=self.vocab_size,
            d_model=self.d_model,
            num_layers=self.num_layers,
            nhead=self.nhead,
            max_len=max_len,
            num_classes=self.num_classes
        ).to(self.device)
        
        ckpt_path = "checkpoints/teacher.pt"
        if os.path.exists(ckpt_path):
            self.teacher.load_state_dict(torch.load(ckpt_path, map_location=self.device))
            print(f"✓ Loaded checkpoint: {ckpt_path}")
        else:
            print(f"⚠ Checkpoint not found: {ckpt_path}")
            
        self.teacher.eval()
        
        # Load student
        print("Loading Student (FMA)")
        self.student = FMAEncoderModel(
            vocab_size=self.vocab_size,
            d_model=self.d_model,
            num_layers=self.num_layers,
            num_modes=self.num_modes,
            max_len=max_len,
            num_classes=self.num_classes
        ).to(self.device)
        
        ckpt_path = "checkpoints/student_kd.pt"
        if os.path.exists(ckpt_path):
            self.student.load_state_dict(torch.load(ckpt_path, map_location=self.device))
            print(f"✓ Loaded checkpoint: {ckpt_path}")
        else:
            print(f"⚠ Checkpoint not found: {ckpt_path}")
            
        self.student.eval()
    
    def predict(self, text):
        """Run inference on text input"""
        if self.tokenizer is None:
            print("Error: No tokenizer loaded. Use load_distilled_models() first.")
            return
        
        # Tokenize
        inputs = self.tokenizer(
            text,
            padding="max_length",
            truncation=True,
            max_length=self.max_len,
            return_tensors="pt"
        )
        input_ids = inputs['input_ids'].to(self.device)
        attention_mask = inputs['attention_mask'].to(self.device)
        
        with torch.no_grad():
            # Teacher prediction
            if self.teacher_hf is not None:
                teacher_outputs = self.teacher_hf(input_ids=input_ids, attention_mask=attention_mask)
                teacher_logits = teacher_outputs.logits
                teacher_probs = torch.softmax(teacher_logits, dim=-1)
                teacher_pred = torch.argmax(teacher_logits, dim=-1).item()
            else:
                teacher_pred = None
                teacher_probs = None
            
            # Student prediction
            student_logits = self.student(input_ids)
            student_probs = torch.softmax(student_logits, dim=-1)
            student_pred = torch.argmax(student_logits, dim=-1).item()
        
        # Format results
        labels = ["Negative", "Positive"]
        
        print(f"\n{'='*60}")
        print(f"Text: {text}")
        print(f"{'='*60}")
        
        if teacher_pred is not None:
            print(f"\n📊 Teacher (BERT-Tiny):")
            print(f"   Prediction: {labels[teacher_pred]}")
            print(f"   Confidence: {teacher_probs[0][teacher_pred].item():.2%}")
            print(f"   Probs: Neg={teacher_probs[0][0].item():.2%}, Pos={teacher_probs[0][1].item():.2%}")
        
        print(f"\n⚡ Student (FMA):")
        print(f"   Prediction: {labels[student_pred]}")
        print(f"   Confidence: {student_probs[0][student_pred].item():.2%}")
        print(f"   Probs: Neg={student_probs[0][0].item():.2%}, Pos={student_probs[0][1].item():.2%}")
        
        if teacher_pred is not None:
            match = "✓" if teacher_pred == student_pred else "✗"
            print(f"\n{match} Agreement: {'YES' if teacher_pred == student_pred else 'NO'}")
        
        return {
            'teacher_pred': teacher_pred,
            'teacher_probs': teacher_probs[0].cpu().numpy() if teacher_probs is not None else None,
            'student_pred': student_pred,
            'student_probs': student_probs[0].cpu().numpy()
        }

def main():
    tester = ModelTester()
    
    # Check which checkpoints exist
    print("\n=== Checking Available Checkpoints ===")
    checkpoints = [
        "checkpoints/teacher.pt",
        "checkpoints/student_kd.pt",
        "checkpoints/student_fma_distilled.pt"
    ]
    
    available = []
    for ckpt in checkpoints:
        if os.path.exists(ckpt):
            size = os.path.getsize(ckpt) / (1024**2)  # MB
            print(f"✓ {ckpt} ({size:.2f} MB)")
            available.append(ckpt)
        else:
            print(f"✗ {ckpt}")
    
    if not available:
        print("\n⚠ No checkpoints found. Please train a model first:")
        print("   conda run -n fsnn python -m training.train_distill")
        return
    
    # Load appropriate models
    if "checkpoints/student_fma_distilled.pt" in available:
        tester.load_distilled_models()
        
        # Test prompts
        test_texts = [
            "This movie was absolutely amazing! I loved every minute of it.",
            "Terrible film. Complete waste of time and money.",
            "The acting was great but the plot was confusing.",
            "Best movie of the year! Highly recommend!",
            "Boring and predictable. Don't bother watching."
        ]
        
        print("\n\n" + "="*60)
        print("TESTING DISTILLED MODELS ON REAL PROMPTS")
        print("="*60)
        
        for text in test_texts:
            tester.predict(text)
            
    else:
        print("\n⚠ Distilled model not found. Run full training:")
        print("   conda run -n fsnn python -m training.train_distill")

if __name__ == "__main__":
    main()