verify_model.py

#!/usr/bin/env python3
"""
Model verification script for DeBERTa v3 Small Explicit Classifier v2.0
"""

import json
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from pathlib import Path

def verify_model_integrity():
    """Verify all model files and configurations"""
    print("🔍 Verifying DeBERTa v3 Small Explicit Classifier v2.0")
    print("=" * 60)
    
    model_path = Path(".")
    
    # Check required files
    required_files = [
        "model.safetensors",
        "config.json", 
        "tokenizer.json",
        "spm.model",
        "label_mapping.json",
        "README.md"
    ]
    
    print("📁 Checking required files...")
    missing_files = []
    for file_name in required_files:
        if (model_path / file_name).exists():
            print(f"  ✅ {file_name}")
        else:
            print(f"  ❌ {file_name} - MISSING")
            missing_files.append(file_name)
    
    if missing_files:
        print(f"\n⚠️  Missing files: {missing_files}")
        return False
    
    # Load and verify model
    print("\n🤖 Loading model...")
    try:
        model = AutoModelForSequenceClassification.from_pretrained(".")
        tokenizer = AutoTokenizer.from_pretrained(".")
        print("  ✅ Model loaded successfully")
    except Exception as e:
        print(f"  ❌ Model loading failed: {e}")
        return False
    
    # Verify configuration
    print("\n⚙️  Verifying configuration...")
    config = model.config
    
    expected_labels = {
        0: "EXPLICIT-DISCLAIMER",
        1: "EXPLICIT-OFFENSIVE", 
        2: "EXPLICIT-SEXUAL",
        3: "EXPLICIT-VIOLENT",
        4: "NON-EXPLICIT",
        5: "SEXUAL-REFERENCE",
        6: "SUGGESTIVE"
    }
    
    # Check label mappings
    config_labels = {int(k): v for k, v in config.id2label.items()}
    if config_labels == expected_labels:
        print("  ✅ Label mappings correct")
    else:
        print("  ❌ Label mappings incorrect")
        print(f"    Expected: {expected_labels}")
        print(f"    Got: {config_labels}")
        return False
    
    # Verify model parameters
    total_params = sum(p.numel() for p in model.parameters())
    expected_params = 141_900_000  # Approximately 141.9M
    
    if abs(total_params - expected_params) < 1_000_000:  # Within 1M tolerance
        print(f"  ✅ Parameter count: {total_params:,} (~{total_params/1_000_000:.1f}M)")
    else:
        print(f"  ⚠️  Unexpected parameter count: {total_params:,}")
    
    # Test inference
    print("\n🧪 Testing inference...")
    try:
        test_text = "This is a test sentence for classification."
        inputs = tokenizer(test_text, return_tensors="pt", truncation=True, max_length=512)
        
        with torch.no_grad():
            outputs = model(**inputs)
            logits = outputs.logits
            probabilities = torch.softmax(logits, dim=-1)
        
        # Check output shape
        if probabilities.shape == (1, 7):  # Batch size 1, 7 classes
            print("  ✅ Inference successful")
            
            # Show predictions
            predicted_class = torch.argmax(probabilities, dim=-1).item()
            confidence = probabilities[0][predicted_class].item()
            predicted_label = config.id2label[predicted_class]
            
            print(f"    Test prediction: {predicted_label} ({confidence:.3f})")
        else:
            print(f"  ❌ Unexpected output shape: {probabilities.shape}")
            return False
            
    except Exception as e:
        print(f"  ❌ Inference failed: {e}")
        return False
    
    # Check evaluation files
    print("\n📊 Checking evaluation files...")
    eval_files = [
        "improved_classification_report.txt",
        "recommended_thresholds.json",
        "confusion_matrix.png",
        "pr_curves.png", 
        "roc_curves.png",
        "calibration.png"
    ]
    
    for file_name in eval_files:
        if (model_path / file_name).exists():
            print(f"  ✅ {file_name}")
        else:
            print(f"  ⚪ {file_name} - Optional")
    
    # Verify thresholds file
    try:
        with open("recommended_thresholds.json", "r") as f:
            thresholds = json.load(f)
        
        if len(thresholds) == 7:  # 7 classes
            print("  ✅ Thresholds file valid")
        else:
            print(f"  ⚠️  Unexpected threshold count: {len(thresholds)}")
    except Exception as e:
        print(f"  ⚠️  Could not verify thresholds: {e}")
    
    print("\n🎉 Model verification complete!")
    print("✅ All core components verified and working correctly")
    print("\n📦 Ready for deployment!")
    
    return True

def show_model_info():
    """Display model information summary"""
    print("\n📋 Model Information Summary")
    print("-" * 40)
    
    try:
        model = AutoModelForSequenceClassification.from_pretrained(".")
        config = model.config
        
        print(f"Model Type: {config.model_type}")
        print(f"Architecture: {config.architectures[0]}")
        print(f"Parameters: {sum(p.numel() for p in model.parameters()):,}")
        print(f"Layers: {config.num_hidden_layers}")
        print(f"Hidden Size: {config.hidden_size}")
        print(f"Attention Heads: {config.num_attention_heads}")
        print(f"Max Length: {config.max_position_embeddings}")
        print(f"Vocabulary Size: {config.vocab_size:,}")
        print(f"Classes: {len(config.id2label)}")
        
        print(f"\nClass Labels:")
        for id_str, label in config.id2label.items():
            print(f"  {id_str}: {label}")
            
    except Exception as e:
        print(f"Error loading model info: {e}")

if __name__ == "__main__":
    success = verify_model_integrity()
    
    if success:
        show_model_info()
    else:
        print("\n❌ Verification failed - please check the issues above")
        exit(1)