benchmark.py

import torch
import torch.nn as nn
import json
import tiktoken
from datasets import load_dataset
from model import GPTModel
from safetensors.torch import load_file
import tqdm

def load_custom_model():
    with open("config.json") as f:
        cfg = json.load(f)
    model = GPTModel(cfg)
    model.load_state_dict(load_file("model.safetensors"), strict=False)
    model.cuda()
    model.eval()
    return model

def evaluate_wikitext_perplexity(model, tokenizer):
    max_length = model.pos_emb.weight.shape[0]
    print(f"Loading WikiText-2 test set (Context Window: {max_length})...")
    dataset = load_dataset("Salesforce/wikitext", "wikitext-2-raw-v1", split="test")
    encodings = tokenizer.encode("\n\n".join(dataset["text"]))
    
    seq_len = len(encodings)
    nlls = []
    
    print(f"Evaluating Perplexity on {seq_len} tokens...")
    loss_fn = nn.CrossEntropyLoss()
    
    # Process in chunks of max_length
    with torch.no_grad():
        with torch.amp.autocast("cuda"):
            for begin_loc in tqdm.tqdm(range(0, seq_len, max_length)):
                end_loc = min(begin_loc + max_length, seq_len)
                
                # We need at least 2 tokens to compute loss
                if end_loc - begin_loc < 2:
                    break
                    
                input_ids = torch.tensor(encodings[begin_loc:end_loc]).unsqueeze(0).cuda()
                target_ids = input_ids.clone()
                
                logits = model(input_ids)
                
                # Shift logits and labels for next token prediction
                shift_logits = logits[..., :-1, :].contiguous()
                shift_labels = target_ids[..., 1:].contiguous()
                
                # Compute loss
                loss = loss_fn(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
                
                # Multiply by the number of tokens in this batch to get sum of losses
                nlls.append(loss * (end_loc - begin_loc - 1))
                
    # Calculate overall perplexity
    total_tokens = sum([(min(b + max_length, seq_len) - b - 1) for b in range(0, seq_len, max_length) if min(b + max_length, seq_len) - b >= 2])
    ppl = torch.exp(torch.stack(nlls).sum() / total_tokens)
    return ppl.item()

def evaluate_sciq_accuracy(model, tokenizer):
    print("Loading SciQ test set for Multiple Choice Accuracy...")
    dataset = load_dataset("sciq", split="test")
    
    correct_count = 0
    total = len(dataset)
    loss_fn = nn.CrossEntropyLoss(reduction='none')
    
    print(f"Evaluating Zero-Shot Accuracy on {total} questions...")
    
    with torch.no_grad():
        with torch.amp.autocast("cuda"):
            for item in tqdm.tqdm(dataset):
                q = item['question']
                choices = [item['correct_answer'], item['distractor1'], item['distractor2'], item['distractor3']]
                
                best_loss = float('inf')
                best_idx = -1
                
                for i, choice in enumerate(choices):
                    prompt_text = (
                        "Below is an instruction that describes a task. "
                        "Write a response that appropriately completes the request.\n\n"
                        f"### Instruction:\nAnswer the following question:\n{q}\n\n### Response:\n"
                    )
                    full_text = f"{prompt_text}{choice}"
                    
                    prompt_tokens = tokenizer.encode(prompt_text)
                    full_tokens = tokenizer.encode(full_text)
                    
                    # The tokens we actually want to measure the loss on (the answer)
                    answer_len = len(full_tokens) - len(prompt_tokens)
                    
                    if answer_len == 0:
                        continue
                        
                    input_ids = torch.tensor(full_tokens).unsqueeze(0).cuda()
                    
                    logits = model(input_ids)
                    
                    # Shift for next token prediction
                    shift_logits = logits[..., :-1, :].contiguous()
                    shift_labels = input_ids[..., 1:].contiguous()
                    
                    # Compute token-wise loss
                    token_losses = loss_fn(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
                    
                    # We only care about the loss of the ANSWER tokens!
                    # The answer tokens are at the very end of the sequence.
                    answer_losses = token_losses[-answer_len:]
                    avg_answer_loss = answer_losses.mean().item()
                    
                    if avg_answer_loss < best_loss:
                        best_loss = avg_answer_loss
                        best_idx = i
                        
                if best_idx == 0: # 0 is the correct_answer
                    correct_count += 1
                    
    accuracy = (correct_count / total) * 100
    return accuracy

def evaluate_mmlu_accuracy(model, tokenizer):
    print("Loading MMLU test set...")
    # Load all subjects, but sample 1000 questions to keep benchmark time reasonable
    dataset = load_dataset("cais/mmlu", "all", split="test")
    dataset = dataset.shuffle(seed=42).select(range(1000))
    
    correct_count = 0
    total = len(dataset)
    loss_fn = nn.CrossEntropyLoss(reduction='none')
    
    print(f"Evaluating Zero-Shot Accuracy on {total} MMLU questions...")
    
    with torch.no_grad():
        with torch.amp.autocast("cuda"):
            for item in tqdm.tqdm(dataset):
                q = item['question']
                choices = item['choices']
                correct_idx = item['answer']
                
                best_loss = float('inf')
                best_idx = -1
                
                # Check if this question will exceed our context window before evaluating
                skip_question = False
                
                for i, choice in enumerate(choices):
                    prompt_text = (
                        "Below is an instruction that describes a task. "
                        "Write a response that appropriately completes the request.\n\n"
                        f"### Instruction:\nAnswer the following multiple choice question:\n{q}\n\n### Response:\n"
                    )
                    full_text = f"{prompt_text}{choice}"
                    
                    prompt_tokens = tokenizer.encode(prompt_text)
                    full_tokens = tokenizer.encode(full_text)
                    
                    # Prevent CUDA crash if the question is longer than the model's brain (512 tokens)
                    if len(full_tokens) > model.pos_emb.weight.shape[0]:
                        skip_question = True
                        break
                        
                    answer_len = len(full_tokens) - len(prompt_tokens)
                    
                    if answer_len == 0:
                        continue
                        
                    input_ids = torch.tensor(full_tokens).unsqueeze(0).cuda()
                    logits = model(input_ids)
                    
                    shift_logits = logits[..., :-1, :].contiguous()
                    shift_labels = input_ids[..., 1:].contiguous()
                    
                    token_losses = loss_fn(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
                    answer_losses = token_losses[-answer_len:]
                    avg_answer_loss = answer_losses.mean().item()
                    
                    if avg_answer_loss < best_loss:
                        best_loss = avg_answer_loss
                        best_idx = i
                        
                if skip_question:
                    # Subtract from total since we couldn't evaluate it
                    total -= 1
                    continue
                        
                if best_idx == correct_idx:
                    correct_count += 1
                    
    accuracy = (correct_count / total) * 100 if total > 0 else 0
    return accuracy

if __name__ == "__main__":
    print("Loading 124M Custom Model...")
    model = load_custom_model()
    tokenizer = tiktoken.get_encoding("gpt2")
    
    ppl = evaluate_wikitext_perplexity(model, tokenizer)
    sciq_acc = evaluate_sciq_accuracy(model, tokenizer)
    mmlu_acc = evaluate_mmlu_accuracy(model, tokenizer)
    
    print(f"\n" + "="*50)
    print(f"WikiText-2 Perplexity: {ppl:.2f}")
    print(f"SciQ (Zero-Shot) Accuracy: {sciq_acc:.2f}%")
    print(f"MMLU (Zero-Shot) Accuracy: {mmlu_acc:.2f}%")
    print("="*50 + "\n")