test_model.py

"""
Testing script to evaluate the model
"""

import torch
from model_setup import load_base_model, configure_lora
from preprocess_data import create_dataloaders,load_cuad_data,prepare_train_val_split
from train import TrainingConfig
import json


def test_model():
    print("Loading the trained model")

    model,tokenizer=load_base_model('distilbert-base-uncased')

    model=configure_lora(model,r=32,lora_alpha=64,lora_dropout=0.1)

    checkpoint = torch.load('checkpoints/best_model.pt', map_location=torch.device('cpu'))
    model.load_state_dict(checkpoint["model_state_dict"])
    model.eval()

    print("Model loaded successfully")

    print(f"Best validation loss: {checkpoint['val_loss']:.4f}")
    print(f"Best epoch: {checkpoint['epoch']}")
    return model,tokenizer

def evaluate_model(model,test_dataloader,device):
    """
    Evaluating the model on the test data
    """
    model.eval()

    total_loss=0
    correct_predictions=0
    total_predictions=0

    with torch.no_grad():
        for i,batch in enumerate(test_dataloader):
            print(f"Processing batch {i+1} of {len(test_dataloader)}")
            
            input_ids=batch['input_ids'].to(device)
            attention_mask=batch['attention_mask'].to(device)
            start_positions=batch['start_positions'].to(device)
            end_positions=batch['end_positions'].to(device)

            outputs=model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                start_positions=start_positions,
                end_positions=end_positions
            )

            total_loss+=outputs.loss.item()

            start_pred = torch.argmax(outputs.start_logits,dim=1)
            end_pred = torch.argmax(outputs.end_logits,dim=1)

            correct= (start_pred == start_positions) & (end_pred == end_positions)
            correct_predictions+= correct.sum().item()
            total_predictions+= len(start_pred)

    avg_loss=total_loss/len(test_dataloader)
    accuracy=correct_predictions/total_predictions
    return avg_loss, accuracy
def create_test_dataloader(test_data, tokenizer, batch_size=8, max_len=512):
    """
    Creating test dataloader
    """
    from torch.utils.data import DataLoader, Dataset
    
    class TestDataset(Dataset):
        def __init__(self, data, tokenizer, max_len):
            self.data = data
            self.tokenizer = tokenizer
            self.max_len = max_len
        
        def __len__(self):
            return len(self.data)
        
        def __getitem__(self, idx):
            item = self.data[idx]
            question = item['question']
            context = item['context']
             
            if item['answers'] and len(item['answers']) > 0:
                answer = item['answers'][0]['text']
            else:
                answer = ""
            
            encoding = self.tokenizer(
                question,
                context,
                truncation=True,
                padding='max_length',
                max_length=self.max_len,
                return_tensors='pt'
            )
            
            answer_start = context.find(answer)
            if answer_start == -1:
                start_pos = 0
                end_pos = 0
            else:
                start_pos = answer_start
                end_pos = answer_start + len(answer)
            
            return {
                'input_ids': encoding['input_ids'].squeeze(0),
                'attention_mask': encoding['attention_mask'].squeeze(0),
                'start_positions': torch.tensor(start_pos),
                'end_positions': torch.tensor(end_pos)
            }
    
    test_dataset = TestDataset(test_data, tokenizer, max_len)
    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
    return test_loader
            
def main():
    """Main fn to run the test"""

    print("Starting model evaluation")

    model,tokenizer= test_model()

    device= torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model =model.to(device)

    print("Loading test data")

    train_data,test_data=load_cuad_data(
        'data/train_separate_questions.json',
        'data/test.json'
    )
    from preprocess_data import parse_cuad_examples
    test_data=parse_cuad_examples(test_data)

    print("Creating test dataloader...")
    test_loader = create_test_dataloader(test_data, tokenizer, batch_size=8, max_len=512)

    print(f"Test batches: {len(test_loader)}")

    print("Evaluating model")
    test_loss,test_accuracy=evaluate_model(model,test_loader,device)
    print(f"Test Loss: {test_loss:.4f}")
    print(f"Test Accuracy: {test_accuracy:.4f}")

    print("Testing completed successfully")

if __name__ == "__main__":
    main()