File size: 4,143 Bytes
210f9af | 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 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 | """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
|