| """Bundle loading, seeded subsets, length-bucketed steps and token-budget micro-batches (pure Python). |
| |
| One optimiser step takes ``posts_per_step`` posts (classifier-plan §4.3: 32). 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 (16k). The loss is normalised per step (``step_counts``), so how a |
| step is split into micro-batches does not change its gradient. No packing. |
| """ |
| from __future__ import annotations |
|
|
| import hashlib |
| import json |
| import os |
| import random |
| from pathlib import Path |
| from typing import Iterable |
|
|
| SPLITS = ("train", "val", "test") |
|
|
|
|
| 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: Iterable[dict], mode: int = 0o600) -> None: |
| """Atomic JSONL write (temp file + rename), file mode 0600 by default (rows may hold post text).""" |
| path = Path(path) |
| path.parent.mkdir(parents=True, exist_ok=True) |
| tmp = path.with_name(path.name + ".tmp") |
| fd = os.open(tmp, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, mode) |
| with os.fdopen(fd, "w", encoding="utf-8") as fh: |
| for row in rows: |
| fh.write(json.dumps(row, ensure_ascii=False, allow_nan=False, sort_keys=True) + "\n") |
| os.chmod(tmp, mode) |
| tmp.replace(path) |
|
|
|
|
| def load_split(data_dir, split: str) -> list[dict]: |
| return load_jsonl(Path(data_dir) / f"{split}.jsonl") |
|
|
|
|
| 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: list[list[int]] = [] |
| current: list[int] = [] |
| 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]: |
| """Unmasked targets per field over one step's label rows (-1 marks a masked target); 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 inference_batches(lengths: list[int], token_budget: int, max_rows: int = 0) -> list[list[int]]: |
| """Length-sorted batches for prediction, each within the padded token budget (and max_rows if given).""" |
| out: list[list[int]] = [] |
| current: list[int] = [] |
| for i in sorted(range(len(lengths)), key=lambda i: (lengths[i], i)): |
| full = max_rows and len(current) >= max_rows |
| if current and (full or lengths[i] * (len(current) + 1) > token_budget): |
| out.append(current) |
| current = [] |
| current.append(i) |
| if current: |
| out.append(current) |
| return out |
|
|
