py/test_model.py

import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import numpy as np
import pandas as pd
from sklearn.metrics import classification_report, confusion_matrix
import matplotlib.pyplot as plt
import seaborn as sns
import argparse

class SentimentTester:
    def __init__(self, model_path="./vietnamese_sentiment_finetuned"):
        self.model_path = model_path
        self.tokenizer = None
        self.model = None
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        
    def load_model(self):
        """Load the fine-tuned model and tokenizer"""
        print(f"Loading model from: {self.model_path}")
        print(f"Using device: {self.device}")
        
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_path)
        self.model = AutoModelForSequenceClassification.from_pretrained(self.model_path)
        self.model.to(self.device)
        self.model.eval()
        
        print("Model loaded successfully!")
        print(f"Number of labels: {self.model.config.num_labels}")
        
    def predict_sentiment(self, text, return_probabilities=False):
        """Predict sentiment for a single text"""
        # Tokenize the text
        inputs = self.tokenizer(
            text,
            return_tensors="pt",
            truncation=True,
            padding=True,
            max_length=512
        )
        
        # Move to device
        inputs = {k: v.to(self.device) for k, v in inputs.items()}
        
        # Get predictions
        with torch.no_grad():
            outputs = self.model(**inputs)
            logits = outputs.logits
            probabilities = torch.softmax(logits, dim=-1)
            predicted_class = torch.argmax(probabilities, dim=-1).item()
        
        if return_probabilities:
            return predicted_class, probabilities.cpu().numpy()[0]
        else:
            return predicted_class
    
    def predict_batch(self, texts):
        """Predict sentiment for a batch of texts"""
        predictions = []
        probabilities = []
        
        for text in texts:
            pred, probs = self.predict_sentiment(text, return_probabilities=True)
            predictions.append(pred)
            probabilities.append(probs)
        
        return np.array(predictions), np.array(probabilities)
    
    def test_custom_texts(self):
        """Test the model with custom Vietnamese texts"""
        test_texts = [
            "Giảng viên dạy rất hay và tâm huyết.",
            "Môn học này quá khó và nhàm chán.",
            "Lớp học ổn định, không có gì đặc biệt.",
            "Tôi rất thích cách giảng dạy của thầy cô.",
            "Chương trình học cần cải thiện nhiều.",
            "Thời gian biểu hợp lý, dễ theo kịp.",
            "Bài tập quá nhiều và khó.",
            "Môi trường học tập tốt, bạn bè thân thiện."
        ]
        
        print("\n" + "="*60)
        print("TESTING WITH CUSTOM VIETNAMESE TEXTS")
        print("="*60)
        
        label_names = ["Negative", "Neutral", "Positive"]  # Assuming 3 classes
        
        for i, text in enumerate(test_texts, 1):
            pred, probs = self.predict_sentiment(text, return_probabilities=True)
            confidence = np.max(probs)
            
            print(f"\n{i}. Text: {text}")
            print(f"   Predicted: {label_names[pred]} (Class {pred})")
            print(f"   Confidence: {confidence:.4f}")
            print(f"   Probabilities: {probs}")
    
    def interactive_test(self):
        """Interactive testing mode"""
        print("\n" + "="*60)
        print("INTERACTIVE SENTIMENT ANALYSIS")
        print("="*60)
        print("Enter Vietnamese text to analyze sentiment (type 'quit' to exit):")
        
        label_names = ["Negative", "Neutral", "Positive"]  # Assuming 3 classes
        
        while True:
            text = input("\nEnter text: ").strip()
            
            if text.lower() in ['quit', 'exit', 'q']:
                break
            
            if not text:
                continue
            
            try:
                pred, probs = self.predict_sentiment(text, return_probabilities=True)
                confidence = np.max(probs)
                
                print(f"Predicted: {label_names[pred]} (Class {pred})")
                print(f"Confidence: {confidence:.4f}")
                print(f"Probabilities: {probs}")
                
            except Exception as e:
                print(f"Error: {e}")
    
    def evaluate_from_file(self, file_path, text_column, label_column=None):
        """Evaluate model on a dataset from file"""
        print(f"\nEvaluating on dataset from: {file_path}")
        
        try:
            # Load dataset
            if file_path.endswith('.csv'):
                df = pd.read_csv(file_path)
            elif file_path.endswith('.json'):
                df = pd.read_json(file_path)
            else:
                print("Unsupported file format. Please use CSV or JSON.")
                return
            
            print(f"Loaded {len(df)} samples")
            
            # Get texts and labels
            texts = df[text_column].tolist()
            
            if label_column and label_column in df.columns:
                true_labels = df[label_column].tolist()
                has_labels = True
            else:
                true_labels = None
                has_labels = False
            
            # Make predictions
            print("Making predictions...")
            predictions, probabilities = self.predict_batch(texts)
            
            # Display results
            if has_labels:
                print("\nClassification Report:")
                print(classification_report(true_labels, predictions))
                
                # Confusion matrix
                cm = confusion_matrix(true_labels, predictions)
                plt.figure(figsize=(8, 6))
                sns.heatmap(cm, annot=True, fmt='d', cmap='Blues')
                plt.title('Confusion Matrix')
                plt.xlabel('Predicted')
                plt.ylabel('Actual')
                plt.savefig('test_confusion_matrix.png', dpi=300, bbox_inches='tight')
                plt.show()
                
                # Calculate accuracy
                accuracy = np.mean(np.array(predictions) == np.array(true_labels))
                print(f"Overall Accuracy: {accuracy:.4f}")
            
            # Show some examples
            print("\nSample predictions:")
            label_names = ["Negative", "Neutral", "Positive"]
            for i in range(min(5, len(texts))):
                pred_label = label_names[predictions[i]]
                confidence = np.max(probabilities[i])
                true_label = f" (True: {label_names[true_labels[i]]})" if has_labels else ""
                print(f"{i+1}. {texts[i][:50]}...")
                print(f"   Predicted: {pred_label} (Confidence: {confidence:.3f}){true_label}")
            
        except Exception as e:
            print(f"Error evaluating file: {e}")
    
    def compare_with_original(self):
        """Compare fine-tuned model with original model"""
        print("\n" + "="*60)
        print("COMPARING WITH ORIGINAL MODEL")
        print("="*60)
        
        test_texts = [
            "Giảng viên dạy rất hay và tâm huyết.",
            "Môn học này quá khó và nhàm chán.",
            "Lớp học ổn định, không có gì đặc biệt."
        ]
        
        original_model = "5CD-AI/Vietnamese-Sentiment-visobert"
        
        try:
            # Load original model
            print("Loading original model...")
            original_tokenizer = AutoTokenizer.from_pretrained(original_model)
            original_model_instance = AutoModelForSequenceClassification.from_pretrained(original_model)
            original_model_instance.to(self.device)
            original_model_instance.eval()
            
            print("\nComparison Results:")
            print("-" * 50)
            
            label_names = ["Negative", "Neutral", "Positive"]
            
            for i, text in enumerate(test_texts, 1):
                # Fine-tuned model prediction
                ft_pred, ft_probs = self.predict_sentiment(text, return_probabilities=True)
                
                # Original model prediction
                inputs = original_tokenizer(
                    text,
                    return_tensors="pt",
                    truncation=True,
                    padding=True,
                    max_length=512
                )
                inputs = {k: v.to(self.device) for k, v in inputs.items()}
                
                with torch.no_grad():
                    outputs = original_model_instance(**inputs)
                    orig_logits = outputs.logits
                    orig_probs = torch.softmax(orig_logits, dim=-1)
                    orig_pred = torch.argmax(orig_probs, dim=-1).item()
                    orig_probs = orig_probs.cpu().numpy()[0]
                
                print(f"\n{i}. Text: {text}")
                print(f"   Fine-tuned: {label_names[ft_pred]} (Conf: {np.max(ft_probs):.3f})")
                print(f"   Original:    {label_names[orig_pred]} (Conf: {np.max(orig_probs):.3f})")
                
                if ft_pred != orig_pred:
                    print(f"   *** DIFFERENT PREDICTION ***")
            
        except Exception as e:
            print(f"Error in comparison: {e}")

def main():
    parser = argparse.ArgumentParser(description='Test fine-tuned Vietnamese sentiment analysis model')
    parser.add_argument('--model_path', type=str, default='./vietnamese_sentiment_finetuned',
                       help='Path to the fine-tuned model')
    parser.add_argument('--mode', type=str, choices=['custom', 'interactive', 'file', 'compare'],
                       default='custom', help='Testing mode')
    parser.add_argument('--file_path', type=str, help='Path to test file (for file mode)')
    parser.add_argument('--text_column', type=str, default='text', help='Text column name (for file mode)')
    parser.add_argument('--label_column', type=str, help='Label column name (for file mode)')
    
    args = parser.parse_args()
    
    # Initialize tester
    tester = SentimentTester(args.model_path)
    
    # Load model
    tester.load_model()
    
    # Run tests based on mode
    if args.mode == 'custom':
        tester.test_custom_texts()
    elif args.mode == 'interactive':
        tester.interactive_test()
    elif args.mode == 'file':
        if not args.file_path:
            print("Error: --file_path required for file mode")
            return
        tester.evaluate_from_file(args.file_path, args.text_column, args.label_column)
    elif args.mode == 'compare':
        tester.compare_with_original()

if __name__ == "__main__":
    main()