HellaSwag › Code Implementation
Language AI HandbookAll chapters HellaSwag › Code ImplementationOne section · Sections in this chapter7 / 14 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Let us implement a HellaSwag evaluation pipeline to see how models are scored in practice. We will load the dataset, prepare examples for a language model, and compute accuracy using the conditional log-probability scoring approach described above.
First, we install and import the necessary libraries:
# Install required packages
# uv pip install datasets transformers torchWe load the HellaSwag dataset from HuggingFace. The dataset contains contexts, endings, and labels showing which of the four endings is correct:
from datasets import load_dataset
# Load HellaSwag validation split (small sample for demonstration)
dataset = load_dataset("Rowan/hellaswag", split="validation[:100]")
# Inspect one example
example_idx = 0
example_ctx = dataset[example_idx]["ctx"]
example_endings = dataset[example_idx]["endings"]
example_label = dataset[example_idx]["label"]Context: A man is sitting on a roof. he Endings: [0] is using wrap to wrap a pair of skis. [1] is ripping level tiles off. [2] is holding a rubik's cube. [3] starts pulling up roofing on a roof. Correct label: 3
The dataset structure provides the context as a string and the endings as a list of four strings. The label is an integer 0-3 showing the correct ending index. Notice that the context is a partial sentence or sequence of sentences setting up a physical scenario, and the endings are completions that vary in physical plausibility.
For autoregressive models like GPT-2, we score each ending by computing its conditional log-probability given the context. We sum the log probabilities of ending tokens only, since we want to compare how probable each ending is given that the context has already been seen:
import torch
def score_ending(
model, tokenizer, context: str, ending: str, device: str = "cpu"
) -> float:
"""Compute length-normalized log probability of an ending given a context."""
full_text = context + " " + ending
context_enc = tokenizer(
context, return_tensors="pt", truncation=True, max_length=512
)
full_enc = tokenizer(
full_text, return_tensors="pt", truncation=True, max_length=512
)
full_enc = {k: v.to(device) for k, v in full_enc.items()}
context_length = context_enc["input_ids"].shape[1]
ending_length = full_enc["input_ids"].shape[1] - context_length
with torch.no_grad():
outputs = model(**full_enc)
logits = outputs.logits
log_probs = torch.log_softmax(logits, dim=-1)
# Sum log probs for ending tokens (shift by 1 for next-token prediction)
total_log_prob = 0.0
for i in range(context_length - 1, full_enc["input_ids"].shape[1] - 1):
token_id = full_enc["input_ids"][0, i + 1].item()
total_log_prob += log_probs[0, i, token_id].item()
# Length-normalize to avoid bias toward shorter endings
return total_log_prob / max(ending_length, 1)The length normalization step divides the total log probability by the number of ending tokens. Without this, the model would systematically prefer shorter endings because they accumulate fewer log-probability terms, creating an evaluation artifact that could skew results.
Now we load a small model and run it on a few examples. For demonstration we use GPT-2, though larger modern models perform significantly better:
import numpy as np
from transformers import AutoModelForCausalLM, AutoTokenizer
# Load GPT-2 for demonstration
model_name = "gpt2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)
model.eval()
# Run on first 5 examples
correct = 0
total = 5
predictions = []
for idx in range(total):
example = dataset[idx]
context = example["ctx"]
endings = example["endings"]
gold_label = int(example["label"])
scores = [
score_ending(model, tokenizer, context, e, device) for e in endings
]
predicted_label = int(np.argmax(scores))
is_correct = predicted_label == gold_label
correct += int(is_correct)
predictions.append(
{
"example": idx + 1,
"predicted": predicted_label,
"gold": gold_label,
"match": is_correct,
"context_preview": context[:80],
}
)
accuracy = correct / totalModel: gpt2 on cpu Example 1 [WRONG] Context: A man is sitting on a roof. he... Predicted: 2, Gold: 3 Example 2 [WRONG] Context: A lady walks to a barbell. She bends down and grabs the pole. the lady... Predicted: 2, Gold: 3 Example 3 [CORRECT] Context: Two women in a child are shown in a canoe while a man pulls the canoe while stan... Predicted: 2, Gold: 2 Example 4 [WRONG] Context: A boy is running down a track. the boy... Predicted: 0, Gold: 2 Example 5 [WRONG] Context: The boy lifts his body above the height of a pole. The boy lands on his back on ... Predicted: 2, Gold: 1 Accuracy on 5 examples: 20% (random chance: 25%)
Our small GPT-2 model typically achieves around 30-40% accuracy on this sample, exceeding random chance (25%) but falling well below human performance. This modest improvement over random guessing reflects that autoregressive models trained on large text corpora have implicitly learned some physical regularities through descriptions of activities in pretraining data, but lack the depth of physical reasoning that humans bring to the task.
For a more reliable assessment across a larger sample, we implement a full evaluation loop:
def run_hellaswag_eval(
model, tokenizer, dataset, device, max_examples: int = 100
) -> dict:
"""Full HellaSwag evaluation with length-normalized scoring."""
correct = 0
total = min(len(dataset), max_examples)
random_baseline = 0.25
human_performance = 0.95
for idx in range(total):
example = dataset[idx]
context = example["ctx"]
endings = example["endings"]
gold_label = int(example["label"])
scores = [
score_ending(model, tokenizer, context, e, device) for e in endings
]
predicted = int(np.argmax(scores))
if predicted == gold_label:
correct += 1
accuracy = correct / total
return {
"accuracy": accuracy,
"correct": correct,
"total": total,
"random_baseline": random_baseline,
"human_performance": human_performance,
"gap_to_human": human_performance - accuracy,
"improvement_over_random": accuracy - random_baseline,
}
results = run_hellaswag_eval(model, tokenizer, dataset, device, max_examples=50)HellaSwag Results (gpt2) Accuracy: 32.0% Correct: 16 / 50 Random baseline: 25.0% Human performance: 95.0% Gap to human: 63.0% Improvement over random: 7.0%
The results reveal the characteristic HellaSwag pattern: models exceed random chance by a modest margin. This shows that statistical regularities in language do encode some physical world knowledge, but a large gap to human performance remains. The zero-shot approach (no fine-tuning on HellaSwag) is intentional: we want to measure what physical knowledge the model has acquired through general language modeling, not how well it can learn the statistical patterns of the benchmark itself through task-specific fine-tuning.
Visualizing Score Distributions
To understand how confidently the model distinguishes correct from incorrect endings, we can examine the distribution of scores for correct versus incorrect options:
correct_rank_scores = []
incorrect_rank_scores = []
for idx in range(50):
example = dataset[idx]
context = example["ctx"]
endings = example["endings"]
gold_label = int(example["label"])
scores = [
score_ending(model, tokenizer, context, e, device) for e in endings
]
correct_rank_scores.append(scores[gold_label])
incorrect_rank_scores.extend(
[s for i, s in enumerate(scores) if i != gold_label]
)
correct_rank_scores = np.array(correct_rank_scores)
incorrect_rank_scores = np.array(incorrect_rank_scores)
The score distribution plot illustrates a basic limitation of small language models on HellaSwag: the distributions for correct and incorrect endings overlap substantially. A model with perfect physical reasoning ability would assign clearly higher scores to correct endings, creating well-separated distributions. The observed overlap means that for many examples, the model essentially guesses: the correct ending happens to receive the highest score by a narrow margin rather than by confident discrimination. Larger, more capable models show greater separation between these distributions, showing more confident and accurate physical plausibility judgments.