wehub-resource-sync
223 lines
8.5 KiB
Python
223 lines
8.5 KiB
Python
"""
|
|
Base class for all Tasks.
|
|
A Task is basically a dataset of conversations, together with some
|
|
metadata and often also evaluation criteria.
|
|
Example tasks: MMLU, ARC-Easy, ARC-Challenge, GSM8K, HumanEval, SmolTalk.
|
|
"""
|
|
|
|
import os
|
|
import json
|
|
import random
|
|
import urllib.request
|
|
|
|
import numpy as np
|
|
import pyarrow as pa
|
|
import pyarrow.parquet as pq
|
|
from filelock import FileLock
|
|
|
|
from nanochat.common import get_base_dir
|
|
|
|
|
|
class HubDataset:
|
|
"""
|
|
Minimal stand-in for a HuggingFace datasets Dataset: wraps a pyarrow
|
|
Table and offers lazy row access and a seeded shuffle.
|
|
"""
|
|
|
|
def __init__(self, table, permutation=None):
|
|
self.table = table
|
|
self.permutation = permutation
|
|
|
|
def __len__(self):
|
|
return self.table.num_rows
|
|
|
|
def shuffle(self, seed):
|
|
# matches datasets.Dataset.shuffle(seed=seed) exactly, row order comes out identical
|
|
permutation = np.random.default_rng(seed).permutation(len(self))
|
|
return HubDataset(self.table, permutation)
|
|
|
|
def __getitem__(self, index):
|
|
physical_index = index if self.permutation is None else int(self.permutation[index])
|
|
row = {column: self.table[column][physical_index].as_py() for column in self.table.column_names}
|
|
return row
|
|
|
|
|
|
def load_hub_dataset(repo_id, subset="default", split="train"):
|
|
"""
|
|
Minimal stand-in for HuggingFace datasets.load_dataset(repo_id, subset, split=split).
|
|
Every dataset on the hub has an auto-generated parquet export. We list the parquet
|
|
shards via the hub API, download them (once) into the local cache directory, and
|
|
read them with pyarrow. Under torchrun, only one rank downloads, the others wait.
|
|
"""
|
|
base_dir = get_base_dir()
|
|
slug = repo_id.replace("/", "--")
|
|
shards_dir = os.path.join(base_dir, "task_data", slug, subset, split)
|
|
# the manifest is written last, so its existence means the download completed
|
|
manifest_path = os.path.join(shards_dir, "manifest.json")
|
|
if not os.path.exists(manifest_path):
|
|
os.makedirs(shards_dir, exist_ok=True)
|
|
with FileLock(manifest_path + ".lock"):
|
|
# only a single rank acquires the lock and downloads, the others block
|
|
# here and then skip the download because they recheck the manifest
|
|
if not os.path.exists(manifest_path):
|
|
listing_url = f"https://huggingface.co/api/datasets/{repo_id}/parquet/{subset}/{split}"
|
|
with urllib.request.urlopen(listing_url) as response:
|
|
shard_urls = json.loads(response.read())
|
|
filenames = []
|
|
for shard_index, shard_url in enumerate(shard_urls):
|
|
filename = f"{shard_index:05d}.parquet"
|
|
print(f"Downloading {shard_url} ...")
|
|
with urllib.request.urlopen(shard_url) as response:
|
|
content = response.read()
|
|
with open(os.path.join(shards_dir, filename), "wb") as f:
|
|
f.write(content)
|
|
filenames.append(filename)
|
|
with open(manifest_path, "w") as f:
|
|
json.dump(filenames, f)
|
|
with open(manifest_path, "r") as f:
|
|
filenames = json.load(f)
|
|
shard_paths = [os.path.join(shards_dir, filename) for filename in filenames]
|
|
tables = [pq.read_table(path) for path in shard_paths]
|
|
table = pa.concat_tables(tables)
|
|
return HubDataset(table)
|
|
|
|
|
|
class Task:
|
|
"""
|
|
Base class of a Task. Allows for lightweight slicing of the underlying dataset.
|
|
"""
|
|
|
|
def __init__(self, start=0, stop=None, step=1):
|
|
# allows a lightweight logical view over a dataset
|
|
assert start >= 0, f"Start must be non-negative, got {start}"
|
|
assert stop is None or stop >= start, f"Stop should be greater than or equal to start, got {stop} and {start}"
|
|
assert step >= 1, f"Step must be strictly positive, got {step}"
|
|
self.start = start
|
|
self.stop = stop # could be None here
|
|
self.step = step
|
|
|
|
@property
|
|
def eval_type(self):
|
|
# one of 'generative' | 'categorical'
|
|
raise NotImplementedError
|
|
|
|
def num_examples(self):
|
|
raise NotImplementedError
|
|
|
|
def get_example(self, index):
|
|
raise NotImplementedError
|
|
|
|
def __len__(self):
|
|
start = self.start
|
|
stop = self.num_examples() if self.stop is None else self.stop
|
|
step = self.step
|
|
span = stop - start
|
|
num = (span + step - 1) // step # ceil_div(span, step)
|
|
assert num >= 0, f"Negative number of examples???: {num}" # prevent footguns
|
|
return num
|
|
|
|
def __getitem__(self, index: int):
|
|
assert isinstance(index, int), f"Index must be an integer, got {type(index)}"
|
|
physical_index = self.start + index * self.step
|
|
conversation = self.get_example(physical_index)
|
|
return conversation
|
|
|
|
def evaluate(self, problem, completion):
|
|
raise NotImplementedError
|
|
|
|
|
|
class TaskMixture(Task):
|
|
"""
|
|
For SFT Training it becomes useful to train on a mixture of datasets.
|
|
Fun trick: if you wish to oversample any task, just pass it in multiple times in the list.
|
|
"""
|
|
|
|
def __init__(self, tasks, **kwargs):
|
|
super().__init__(**kwargs)
|
|
# tasks is a list of Task objects
|
|
self.tasks = tasks
|
|
self.lengths = [len(task) for task in self.tasks]
|
|
self.num_conversations = sum(self.lengths)
|
|
# Build list of all (task_idx, local_idx) pairs
|
|
self.index_map = []
|
|
for task_idx, task_length in enumerate(self.lengths):
|
|
for local_idx in range(task_length):
|
|
self.index_map.append((task_idx, local_idx))
|
|
# Deterministically shuffle to mix tasks throughout training
|
|
rng = random.Random(42)
|
|
rng.shuffle(self.index_map)
|
|
# Note: this is not the most elegant or best solution, but it's ok for now
|
|
|
|
def num_examples(self):
|
|
return self.num_conversations
|
|
|
|
def get_example(self, index):
|
|
"""
|
|
Access conversations according to a deterministic shuffle of all examples.
|
|
This ensures tasks are mixed throughout training, regardless of dataset size.
|
|
"""
|
|
assert 0 <= index < self.num_conversations, f"Index {index} out of range for mixture with {self.num_conversations} conversations"
|
|
task_idx, local_idx = self.index_map[index]
|
|
return self.tasks[task_idx][local_idx]
|
|
|
|
|
|
class TaskSequence(Task):
|
|
"""
|
|
For SFT Training sometimes we want to sequentially train on a list of tasks.
|
|
This is useful for cases that require a training curriculum.
|
|
"""
|
|
|
|
def __init__(self, tasks, **kwargs):
|
|
super().__init__(**kwargs)
|
|
self.tasks = tasks
|
|
self.lengths = [len(task) for task in self.tasks]
|
|
self.num_conversations = sum(self.lengths)
|
|
|
|
def num_examples(self):
|
|
return self.num_conversations
|
|
|
|
def get_example(self, index):
|
|
assert 0 <= index < self.num_conversations, f"Index {index} out of range for sequence with {self.num_conversations} conversations"
|
|
for task_idx, task_length in enumerate(self.lengths):
|
|
if index < task_length:
|
|
return self.tasks[task_idx][index]
|
|
index -= task_length
|
|
|
|
|
|
def render_mc(question, letters, choices):
|
|
"""
|
|
The common multiple choice rendering format we will use.
|
|
|
|
Note two important design decisions:
|
|
1)
|
|
Bigger models don't care as much, but smaller models prefer to have
|
|
the letter *after* the choice, which results in better binding.
|
|
2)
|
|
There is no whitespace between the delimiter (=) and the letter.
|
|
This is actually critical because the tokenizer has different token ids
|
|
for " A" vs. "A". The assistant responses will be just the letter itself,
|
|
i.e. "A", so it is important that here in the prompt it is the exact same
|
|
token, i.e. "A" with no whitespace before it. Again, bigger models don't care
|
|
about this too much, but smaller models do care about some of these details.
|
|
"""
|
|
query = f"Multiple Choice question: {question}\n"
|
|
query += "".join([f"- {choice}={letter}\n" for letter, choice in zip(letters, choices)])
|
|
query += "\nRespond only with the letter of the correct answer."
|
|
return query
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# very lightweight test of slicing
|
|
from tasks.mmlu import MMLU
|
|
|
|
ds = MMLU(subset="all", split="auxiliary_train")
|
|
print("Length of MMLU: ", len(ds))
|
|
ex = ds[5]
|
|
print("5th example: ", ex)
|
|
|
|
ds = MMLU(subset="all", split="auxiliary_train", start=5, stop=10)
|
|
print("Length of sliced MMLU[5:10]: ", len(ds))
|
|
print("0th example of sliced MMLU: ", ds[0])
|
|
|
|
print("They match: ", ex == ds[0])
|