Correct the model card (pooling is at the LABELS suffix's last token); add code/classifier from project commit 4ced7554f
3d7c238 verified | """Export loading, seeded subsets, steps of similar-length posts and token-budget micro-batches (pure Python). | |
| One optimiser step takes ``posts_per_step`` posts. Steps are drawn from length-sorted pools so their posts have | |
| similar lengths, then shuffled; each step is split into micro-batches whose padded size (rows x longest row) | |
| stays within the token budget. The loss is normalised per step (``step_counts``), so how a step is split into | |
| micro-batches does not change its gradient. No packing: the Gated DeltaNet state does not reset between | |
| packed posts. | |
| """ | |
| from __future__ import annotations | |
| import hashlib | |
| import json | |
| from pathlib import Path | |
| import random | |
| def load_jsonl(path) -> list[dict]: | |
| with open(path, encoding='utf-8') as fh: | |
| return [json.loads(line) for line in fh if line.strip()] | |
| def write_jsonl(path, rows) -> None: | |
| path = Path(path) | |
| path.parent.mkdir(parents=True, exist_ok=True) | |
| tmp = path.with_suffix(path.suffix + '.tmp') | |
| with open(tmp, 'w', encoding='utf-8') as fh: | |
| for row in rows: | |
| fh.write(json.dumps(row, ensure_ascii=False, allow_nan=False) + '\n') | |
| tmp.replace(path) | |
| def load_split(data_dir, split: str) -> list[dict]: | |
| return load_jsonl(Path(data_dir) / ('%s.jsonl' % split)) | |
| def seeded_order(ids, seed) -> list: | |
| return sorted(ids, key=lambda i: hashlib.sha256(json.dumps([str(seed), i]).encode()).hexdigest()) | |
| def subset(rows: list[dict], n: int, seed) -> list[dict]: | |
| """The first n rows in seeded hash order of their ids (independent of file order); all rows if n >= len.""" | |
| if not n or n >= len(rows): | |
| return list(rows) | |
| keep = set(seeded_order([r['id'] for r in rows], seed)[:n]) | |
| return [r for r in rows if r['id'] in keep] | |
| def steps(lengths: list[int], posts_per_step: int, seed: int, pool: int = 64) -> list[list[int]]: | |
| """Shuffled steps of similar-length posts: shuffle, sort within pools of posts_per_step*pool, chunk, shuffle.""" | |
| rng = random.Random(seed) | |
| idx = list(range(len(lengths))) | |
| rng.shuffle(idx) | |
| out = [] | |
| for s in range(0, len(idx), posts_per_step * pool): | |
| chunk = sorted(idx[s:s + posts_per_step * pool], key=lambda i: (lengths[i], i)) | |
| out += [chunk[i:i + posts_per_step] for i in range(0, len(chunk), posts_per_step)] | |
| rng.shuffle(out) | |
| return out | |
| def micro_batches(step: list[int], lengths: list[int], token_budget: int) -> list[list[int]]: | |
| """Split one step so each micro-batch's padded size (rows x longest) is at most token_budget.""" | |
| order = sorted(step, key=lambda i: (-lengths[i], i)) | |
| out, current = [], [] | |
| for i in order: | |
| width = max([lengths[j] for j in current] + [lengths[i]]) | |
| if current and width * (len(current) + 1) > token_budget: | |
| out.append(current) | |
| current = [] | |
| current.append(i) | |
| if current: | |
| out.append(current) | |
| return out | |
| def step_counts(labels: list[list[int]]) -> list[int]: | |
| """Non-null targets per field over one step's label rows (-1 marks a null); at least 1 per field.""" | |
| if not labels: | |
| return [] | |
| return [max(1, sum(1 for row in labels if row[f] >= 0)) for f in range(len(labels[0]))] | |
| def length_order(lengths: list[int]) -> list[int]: | |
| return sorted(range(len(lengths)), key=lambda i: (lengths[i], i)) | |
| def inference_batches(lengths: list[int], token_budget: int) -> list[list[int]]: | |
| """Length-sorted batches for prediction, each within the padded token budget.""" | |
| out, current = [], [] | |
| for i in length_order(lengths): | |
| if current and lengths[i] * (len(current) + 1) > token_budget: | |
| out.append(current) | |
| current = [] | |
| current.append(i) | |
| if current: | |
| out.append(current) | |
| return out | |