From 7bd2389c1fe302b59f7ddc7dcacdb5f11e7765e2 Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Sun, 19 May 2024 14:23:51 +0000 Subject: [PATCH] add mmlu as well and refine both a bit --- dev/hellaswag.py | 20 +++--- dev/mmlu.py | 156 +++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 168 insertions(+), 8 deletions(-) create mode 100644 dev/mmlu.py diff --git a/dev/hellaswag.py b/dev/hellaswag.py index a0fea2a..cdfe364 100644 --- a/dev/hellaswag.py +++ b/dev/hellaswag.py @@ -14,6 +14,9 @@ endings: a list of 4 endings. The correct index is given by label (0,1,2, or 3) split: train, val, or test. split_type: indomain if the activity label is seen during training, else zeroshot source_id: Which video or WikiHow article this example came from + +gpt2 (124M) = 28.2% +gpt2-xl (1558M) = 39.22% """ import os @@ -28,7 +31,7 @@ from torch.nn import functional as F from transformers import GPT2LMHeadModel -DATA_CACHE_DIR = "data" +DATA_CACHE_DIR = os.path.join("data", "hellaswag") hellaswags = { "train": "https://raw.githubusercontent.com/rowanz/hellaswag/master/data/hellaswag_train.jsonl", @@ -95,23 +98,24 @@ def render_example(example): return tokens, mask, label def iterate_examples(split): + # there are 10,042 examples in total in val + download(split) with open(os.path.join(DATA_CACHE_DIR, f"hellaswag_{split}.jsonl"), "r") as f: for line in f: example = json.loads(line) - rendered = render_example(example) - yield example, rendered + yield example @torch.no_grad() -def evaluate(model_type="gpt2-xl", device="cuda"): +def evaluate(model_type, device): model = GPT2LMHeadModel.from_pretrained(model_type) model.to(device) num_correct = 0 num_total = 0 - data_it = iterate_examples("val") - for example, (tokens, mask, label) in data_it: + for example in iterate_examples("val"): + tokens, mask, label = render_example(example) tokens = tokens.to(device) mask = mask.to(device) @@ -154,7 +158,7 @@ def evaluate(model_type="gpt2-xl", device="cuda"): if __name__ == "__main__": import argparse parser = argparse.ArgumentParser() - parser.add_argument("-m", type=str, default="gpt2-xl", help="the model type to use") - parser.add_argument("-d", type=str, default="cuda", help="the device to use") + parser.add_argument("-m", "--model_type", type=str, default="gpt2", help="the model type to use") + parser.add_argument("-d", "--device", type=str, default="cuda", help="the device to use") args = parser.parse_args() evaluate(args.model_type, args.device) diff --git a/dev/mmlu.py b/dev/mmlu.py new file mode 100644 index 0000000..f15b785 --- /dev/null +++ b/dev/mmlu.py @@ -0,0 +1,156 @@ +""" +Downloads and evaluates MMLU in Python. +This then acts as the reference file for llm.c +https://github.com/hendrycks/test + +gpt2 (124M) ~= 25% (chance) +gpt2-xl (1558M) = 27.00% ... +""" + +import os +import requests +import tiktoken +import pandas as pd +from tqdm import tqdm + +import torch +import torch.nn as nn +from torch.nn import functional as F + +from transformers import GPT2LMHeadModel + +DATA_CACHE_DIR = os.path.join("data", "mmlu") + +enc = tiktoken.get_encoding("gpt2") +data_url = "https://people.eecs.berkeley.edu/~hendrycks/data.tar" + +def download_file(url: str, fname: str, chunk_size=1024): + """Helper function to download a file from a given url""" + resp = requests.get(url, stream=True) + total = int(resp.headers.get("content-length", 0)) + with open(fname, "wb") as file, tqdm( + desc=fname, + total=total, + unit="iB", + unit_scale=True, + unit_divisor=1024, + ) as bar: + for data in resp.iter_content(chunk_size=chunk_size): + size = file.write(data) + bar.update(size) + +def download(): + """Downloads MMLU to DATA_CACHE_DIR""" + os.makedirs(DATA_CACHE_DIR, exist_ok=True) + data_filename = os.path.join(DATA_CACHE_DIR, f"data.tar") + if not os.path.exists(data_filename): + print(f"Downloading {data_url} to {data_filename}...") + download_file(data_url, data_filename) + os.system(f"tar -xf {data_filename} -C {DATA_CACHE_DIR}") # untar + # creates a directory "data" inside it, with e.g. data/test/*csv + else: + print(f"{data_filename} already exists, skipping download...") + +def iterate_examples(): + # there are 14,042 examples in total in the test set + + download() + test_dir = os.path.join(DATA_CACHE_DIR, "data", "test") + csv_files = [f for f in os.listdir(test_dir) if f.endswith(".csv")] + for csv_file in csv_files: + csv_path = os.path.join(test_dir, csv_file) + print(csv_path) + df = pd.read_csv(csv_path, header=None) + n = df.shape[0] + for idx in range(n): + example = { + "question": df.iloc[idx, 0], + "endings": [df.iloc[idx, 1], df.iloc[idx, 2], df.iloc[idx, 3], df.iloc[idx, 4]], + "label": df.iloc[idx, 5], + } + yield example + +def render_example(example): + """ + Given the example as a dictionary, render it as three torch tensors: + - tokens (the tokens of context + completion, of size 4xN, as there are always 4 candidates) + - mask (is 1 in the region of the candidate completion, where we evaluate likelihoods) + - label (the index of the correct completion, which we hope has the highest likelihood) + """ + ctx = f"Question: {example['question']}\n\nAnswer:" + ctx_tokens = enc.encode(ctx) + + tok_rows = [] + mask_rows = [] + for end in example["endings"]: + end_tokens = enc.encode(" " + str(end)) # note: prepending " " because GPT-2 tokenizer + tok_rows.append(ctx_tokens + end_tokens) + mask_rows.append([0]*len(ctx_tokens) + [1]*len(end_tokens)) + + # have to be careful during the collation because the number of tokens in each row can differ + max_len = max(len(row) for row in tok_rows) + tokens = torch.zeros((4, max_len), dtype=torch.long) + mask = torch.zeros((4, max_len), dtype=torch.long) + for i, (tok_row, mask_row) in enumerate(zip(tok_rows, mask_rows)): + tokens[i, :len(tok_row)] = torch.tensor(tok_row) + mask[i, :len(mask_row)] = torch.tensor(mask_row) + + label = "ABCD".index(example["label"]) + return tokens, mask, label + +@torch.no_grad() +def evaluate(model_type, device): + + model = GPT2LMHeadModel.from_pretrained(model_type) + model.to(device) + + num_correct = 0 + num_total = 0 + for example in iterate_examples(): + tokens, mask, label = render_example(example) + tokens = tokens.to(device) + mask = mask.to(device) + + # get the logits + logits = model(tokens).logits + # evaluate the autoregressive loss at all positions + shift_logits = (logits[..., :-1, :]).contiguous() + shift_tokens = (tokens[..., 1:]).contiguous() + flat_shift_logits = shift_logits.view(-1, shift_logits.size(-1)) + flat_shift_tokens = shift_tokens.view(-1) + shift_losses = F.cross_entropy(flat_shift_logits, flat_shift_tokens, reduction='none') + shift_losses = shift_losses.view(tokens.size(0), -1) + # now get the average loss just for the completion region (where mask == 1), in each row + shift_mask = (mask[..., :-1]).contiguous() + masked_shift_losses = shift_losses * shift_mask + # sum and divide by the number of 1s in the mask + sum_loss = masked_shift_losses.sum(dim=1) + avg_loss = sum_loss / shift_mask.sum(dim=1) + # now we have a loss for each of the 4 completions + # the one with the lowest loss should be the most likely + # to think through more carefully: sum or average? sum is more right probabilistically + use_loss = sum_loss + # ok predict what the model thinks is the most likely completion + pred = use_loss.argmin().item() + + # accumulate stats + num_total += 1 + num_correct += int(pred == label) + print(f"accuracy: {num_correct/num_total:.4f} ({num_correct}/{num_total})") + + # debug prints + if num_total < 10: + print("---") + print(f"Context:\n {example['question']}") + print(f"Endings:") + for i, end in enumerate(example["endings"]): + print(f"{i} (loss: {use_loss[i].item():.4f}) {end}") + print(f"predicted: {pred}, actual: {label}") + +if __name__ == "__main__": + import argparse + parser = argparse.ArgumentParser() + parser.add_argument("-m", "--model_type", type=str, default="gpt2", help="the model type to use") + parser.add_argument("-d", "--device", type=str, default="cuda", help="the device to use") + args = parser.parse_args() + evaluate(args.model_type, args.device)