File size: 3,771 Bytes
3d7c238 | 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 | """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
|