| """The fit corpus: 16 short passages (8 roleplay, 8 novel prose), split 4 + 4 of each into a fit split and a |
| held-out split. |
| |
| Each item becomes one sequence: a chat-template user turn asking the model to continue a passage (128 tokens |
| of it), then 128 reference tokens as the model's turn. The fit averages the fine-tune's edit over every |
| position of the fit split; the score compares next-token distributions at the reference positions. |
| The item texts and reference token ids are stored in data/fit_corpus.json, so the prompts are identical |
| wherever this runs. |
| """ |
| from __future__ import annotations |
|
|
| import json |
| from pathlib import Path |
|
|
| import numpy as np |
|
|
| DATA = Path(__file__).parent / "data" / "fit_corpus.json" |
| SPLITS = ("fit", "held_out") |
|
|
|
|
| class Corpus: |
| """Token arrays and position masks for both splits. |
| |
| ids[s] (8, T) token ids, right-padded |
| valid[s] per sequence, the number of real tokens T_valid (the sequences are cut to the shortest) |
| starts[s] per sequence, where the reference turn starts |
| fit_mask[s] (8, T) 1 at the positions a vector is fitted over: 1 <= t < T_valid (BOS and padding excluded) |
| rows[s] (sequence, first, last): the positions whose next-token prediction is scored |
| spans[s] per sequence, its [start, end) in the concatenated scored rows |
| """ |
|
|
| def __init__(self, tokenizer, path: Path = DATA, turns: dict[str, list[int]] | None = None): |
| data = json.loads(Path(path).read_text()) |
| self.items = data["items"] |
| n_ref = data["reference_tokens"] |
| pad = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0 |
| self.ids, self.valid, self.starts, self.fit_mask, self.rows, self.spans = {}, {}, {}, {}, {}, {} |
| for s in SPLITS: |
| seqs, starts, valid = [], [], [] |
| for reg in ("roleplay", "narrative"): |
| group = [] |
| for it in (i for i in self.items if i["split"] == s and i["register"] == reg): |
| user = it["framing"].replace("{{TEXT}}", it["context"]) |
| prompt = tokenizer.apply_chat_template([{"role": "user", "content": user}], |
| add_generation_prompt=True, return_dict=True)["input_ids"] |
| ref = (turns[it["id"]] if turns else it["reference_ids"])[:n_ref] |
| group.append(list(prompt) + list(ref)) |
| starts.append(len(prompt)) |
| |
| T = min(len(x) for x in group) |
| seqs += [x[:T] for x in group] |
| valid += [T] * len(group) |
| self.valid[s] = valid |
| self.starts[s] = starts |
| width = max(self.valid[s]) |
| self.ids[s] = np.array([x + [pad] * (width - len(x)) for x in seqs]) |
| self.fit_mask[s] = np.array([[1.0 if 1 <= t < v else 0.0 for t in range(width)] for v in self.valid[s]], |
| np.float32) |
| self.rows[s] = [(i, starts[i] - 1, self.valid[s][i] - 1) for i in range(len(seqs))] |
| spans, o = [], 0 |
| for _, a, b in self.rows[s]: |
| spans.append((o, o + b - a)) |
| o += b - a |
| self.spans[s] = spans |
|
|
