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")