File size: 3,306 Bytes
e51b495 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 | """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))
# each register's sequences are cut to that register's shortest
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
|