EryriLabs

Card: compare with the leaderboard's specialist rows; Jev and Decider 1 are zero-shot generalists there. LF line endings.

b4e5b8b verified 26 days ago

Raw History Blame Contribute Delete

8.35 kB

"""On-disk feature cache: one safetensors file per example plus a JSONL manifest."""
from __future__ import annotations

import json
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Optional

import torch
from safetensors.torch import load_file, save_file

from .schema import TYPE_INDEX


@dataclass
class Decision:
    qid: str
    type: str
    slot: int
    spans: list[list[int]]
    alias: list[int]
    prior: list[float]
    gold: Optional[list[float]]


@dataclass
class Record:
    id: str
    source: str
    path: str
    n_prompt: int
    n_canvas: int
    decisions: list[Decision] = field(default_factory=list)

    def to_json(self) -> str:
        return json.dumps(asdict(self))

    @classmethod
    def from_json(cls, line: str) -> "Record":
        d = json.loads(line)
        d["decisions"] = [Decision(**x) for x in d["decisions"]]
        return cls(**d)


class CacheWriter:
    """Writes one safetensors file per example and a JSONL manifest.
    A writer always starts a fresh manifest, so re-running into the same folder does not duplicate records.
    """

    def __init__(self, out_dir):
        self.dir = Path(out_dir)
        self.dir.mkdir(parents=True, exist_ok=True)
        self._f = (self.dir / "manifest.jsonl").open("w", encoding="utf-8")

    def write(self, record: Record, feats: torch.Tensor) -> None:
        save_file({"feats": feats.detach().to("cpu", torch.float16).contiguous()}, str(self.dir / record.path))
        self._f.write(record.to_json() + "\n")
        self._f.flush()

    def close(self) -> None:
        self._f.close()


class FeatureCache:
    def __init__(self, cache_dir):
        self.dir = Path(cache_dir)
        with (self.dir / "manifest.jsonl").open("r", encoding="utf-8") as f:
            self.records = [Record.from_json(line) for line in f if line.strip()]

    def __len__(self) -> int:
        return len(self.records)

    def __getitem__(self, i: int) -> tuple[Record, torch.Tensor]:
        rec = self.records[i]
        return rec, load_file(str(self.dir / rec.path))["feats"]


@dataclass
class Batch:
    """A padded feature batch plus two views of its decisions.
    `decisions` is the per-decision list (metadata, spans, prior, gold) that evaluation, serving and the tests
    read. The flat tensors are what the head and the losses compute on: options are numbered decision by
    decision in that same order, so one gather, one index_add and one padded log-softmax cover the whole batch
    instead of a Python loop issuing hundreds of tiny ops per decision (session 3's CPU-bound trainer).
    """
    feats: torch.Tensor
    seg: torch.Tensor
    mask: torch.Tensor
    decisions: list[dict]
    dec_b: torch.Tensor  # (n_dec,) batch row of each decision
    dec_slot: torch.Tensor  # (n_dec,) answer-slot row of each decision
    dec_type: torch.Tensor  # (n_dec,) question type id
    dec_has_gold: torch.Tensor  # (n_dec,) bool
    opt_dec: torch.Tensor  # (n_opt,) decision index of each option
    opt_pos: torch.Tensor  # (n_opt,) position of each option within its decision
    opt_prior: torch.Tensor  # (n_opt,) backbone log-prior per option
    opt_gold: torch.Tensor  # (n_opt,) gold probability per option, 0 where the decision has no gold
    span_rows: torch.Tensor  # (n_span_rows,) row index into the padded (B * L) feature rows, option by option
    span_opt: torch.Tensor  # (n_span_rows,) option index of each span row
    span_len: torch.Tensor  # (n_opt,) span length per option, float
    counts: list[int]  # options per decision; plain Python so torch.split never syncs on it
    n_gold: int  # decisions with gold; plain Python so the losses never sync on it

    _FLAT = ("dec_b", "dec_slot", "dec_type", "dec_has_gold", "opt_dec", "opt_pos", "opt_prior", "opt_gold",
             "span_rows", "span_opt", "span_len")

    def to(self, device) -> "Batch":
        # spans stay on the CPU: consumers read them back as Python ints, so a device copy would only add a
        # device-to-host sync per decision.
        return Batch(self.feats.to(device), self.seg.to(device), self.mask.to(device),
                     [{**d, "prior": d["prior"].to(device),
                       "gold": None if d["gold"] is None else d["gold"].to(device)} for d in self.decisions],
                     *(getattr(self, name).to(device) for name in self._FLAT), self.counts, self.n_gold)


def _check_record(rec: Record, n_rows: int) -> None:
    """Raise a ValueError naming the record and question when a record does not fit its feature rows.
    Rows [0, n_prompt) hold the prompt and [n_prompt, n_prompt + n_canvas) the canvas. Without this check a bad
    record would be read silently from padding or from the wrong segment.
    """
    if n_rows != rec.n_prompt + rec.n_canvas:
        raise ValueError(f"record {rec.id!r}: {n_rows} feature rows, but n_prompt {rec.n_prompt} + n_canvas "
                         f"{rec.n_canvas} = {rec.n_prompt + rec.n_canvas}")
    for d in rec.decisions:
        where = f"record {rec.id!r}, question {d.qid!r}"
        if not rec.n_prompt <= d.slot < n_rows:
            raise ValueError(f"{where}: slot {d.slot} is outside the canvas rows [{rec.n_prompt}, {n_rows})")
        for s, e in d.spans:
            if not 0 <= s < e <= rec.n_prompt:
                raise ValueError(f"{where}: option span [{s}, {e}) is empty or outside the prompt rows "
                                 f"[0, {rec.n_prompt})")
        if len(d.prior) != len(d.spans):
            raise ValueError(f"{where}: {len(d.spans)} option spans but {len(d.prior)} prior values")
        if d.gold is not None and len(d.gold) != len(d.spans):
            raise ValueError(f"{where}: {len(d.spans)} option spans but {len(d.gold)} gold values")


def collate(items: list[tuple[Record, torch.Tensor]]) -> Batch:
    lmax = max(f.shape[0] for _, f in items)
    dim = items[0][1].shape[1]
    feats = torch.zeros(len(items), lmax, dim, dtype=torch.float16)
    seg = torch.zeros(len(items), lmax, dtype=torch.long)
    mask = torch.zeros(len(items), lmax, dtype=torch.bool)
    decisions = []
    dec_b, dec_slot, dec_type, dec_has_gold, counts = [], [], [], [], []
    opt_dec, opt_pos, opt_prior, opt_gold, span_rows, span_opt, span_len = [], [], [], [], [], [], []
    for b, (rec, f) in enumerate(items):
        n = f.shape[0]
        _check_record(rec, n)
        feats[b, :n] = f.to(torch.float16)
        mask[b, :n] = True
        seg[b, rec.n_prompt:n] = 1
        for d in rec.decisions:
            i = len(decisions)
            decisions.append({
                "b": b, "slot": d.slot, "type_id": TYPE_INDEX[d.type], "qid": d.qid, "type": d.type,
                "spans": torch.tensor(d.spans, dtype=torch.long),
                "prior": torch.tensor(d.prior, dtype=torch.float32),
                "gold": None if d.gold is None else torch.tensor(d.gold, dtype=torch.float32),
            })
            dec_b.append(b)
            dec_slot.append(d.slot)
            dec_type.append(TYPE_INDEX[d.type])
            dec_has_gold.append(d.gold is not None)
            counts.append(len(d.spans))
            for j, (s, e) in enumerate(d.spans):
                o = len(opt_dec)
                opt_dec.append(i)
                opt_pos.append(j)
                opt_prior.append(d.prior[j])
                opt_gold.append(0.0 if d.gold is None else d.gold[j])
                span_rows.extend(range(b * lmax + s, b * lmax + e))
                span_opt.extend([o] * (e - s))
                span_len.append(float(e - s))
    long, f32 = torch.long, torch.float32
    return Batch(feats=feats, seg=seg, mask=mask, decisions=decisions,
                 dec_b=torch.tensor(dec_b, dtype=long), dec_slot=torch.tensor(dec_slot, dtype=long),
                 dec_type=torch.tensor(dec_type, dtype=long), dec_has_gold=torch.tensor(dec_has_gold, dtype=torch.bool),
                 opt_dec=torch.tensor(opt_dec, dtype=long), opt_pos=torch.tensor(opt_pos, dtype=long),
                 opt_prior=torch.tensor(opt_prior, dtype=f32), opt_gold=torch.tensor(opt_gold, dtype=f32),
                 span_rows=torch.tensor(span_rows, dtype=long), span_opt=torch.tensor(span_opt, dtype=long),
                 span_len=torch.tensor(span_len, dtype=f32), counts=counts, n_gold=sum(dec_has_gold))

	"""On-disk feature cache: one safetensors file per example plus a JSONL manifest."""
	from __future__ import annotations

	import json
	from dataclasses import asdict, dataclass, field
	from pathlib import Path
	from typing import Optional

	import torch
	from safetensors.torch import load_file, save_file

	from .schema import TYPE_INDEX


	@dataclass
	class Decision:
	qid: str
	type: str
	slot: int
	spans: list[list[int]]
	alias: list[int]
	prior: list[float]
	gold: Optional[list[float]]


	@dataclass
	class Record:
	id: str
	source: str
	path: str
	n_prompt: int
	n_canvas: int
	decisions: list[Decision] = field(default_factory=list)

	def to_json(self) -> str:
	return json.dumps(asdict(self))

	@classmethod
	def from_json(cls, line: str) -> "Record":
	d = json.loads(line)
	d["decisions"] = [Decision(**x) for x in d["decisions"]]
	return cls(**d)


	class CacheWriter:
	"""Writes one safetensors file per example and a JSONL manifest.

	A writer always starts a fresh manifest, so re-running into the same folder does not duplicate records.
	"""

	def __init__(self, out_dir):
	self.dir = Path(out_dir)
	self.dir.mkdir(parents=True, exist_ok=True)
	self._f = (self.dir / "manifest.jsonl").open("w", encoding="utf-8")

	def write(self, record: Record, feats: torch.Tensor) -> None:
	save_file({"feats": feats.detach().to("cpu", torch.float16).contiguous()}, str(self.dir / record.path))
	self._f.write(record.to_json() + "\n")
	self._f.flush()

	def close(self) -> None:
	self._f.close()


	class FeatureCache:
	def __init__(self, cache_dir):
	self.dir = Path(cache_dir)
	with (self.dir / "manifest.jsonl").open("r", encoding="utf-8") as f:
	self.records = [Record.from_json(line) for line in f if line.strip()]

	def __len__(self) -> int:
	return len(self.records)

	def __getitem__(self, i: int) -> tuple[Record, torch.Tensor]:
	rec = self.records[i]
	return rec, load_file(str(self.dir / rec.path))["feats"]


	@dataclass
	class Batch:
	"""A padded feature batch plus two views of its decisions.

	`decisions` is the per-decision list (metadata, spans, prior, gold) that evaluation, serving and the tests
	read. The flat tensors are what the head and the losses compute on: options are numbered decision by
	decision in that same order, so one gather, one index_add and one padded log-softmax cover the whole batch
	instead of a Python loop issuing hundreds of tiny ops per decision (session 3's CPU-bound trainer).
	"""
	feats: torch.Tensor
	seg: torch.Tensor
	mask: torch.Tensor
	decisions: list[dict]
	dec_b: torch.Tensor # (n_dec,) batch row of each decision
	dec_slot: torch.Tensor # (n_dec,) answer-slot row of each decision
	dec_type: torch.Tensor # (n_dec,) question type id
	dec_has_gold: torch.Tensor # (n_dec,) bool
	opt_dec: torch.Tensor # (n_opt,) decision index of each option
	opt_pos: torch.Tensor # (n_opt,) position of each option within its decision
	opt_prior: torch.Tensor # (n_opt,) backbone log-prior per option
	opt_gold: torch.Tensor # (n_opt,) gold probability per option, 0 where the decision has no gold
	span_rows: torch.Tensor # (n_span_rows,) row index into the padded (B * L) feature rows, option by option
	span_opt: torch.Tensor # (n_span_rows,) option index of each span row
	span_len: torch.Tensor # (n_opt,) span length per option, float
	counts: list[int] # options per decision; plain Python so torch.split never syncs on it
	n_gold: int # decisions with gold; plain Python so the losses never sync on it

	_FLAT = ("dec_b", "dec_slot", "dec_type", "dec_has_gold", "opt_dec", "opt_pos", "opt_prior", "opt_gold",
	"span_rows", "span_opt", "span_len")

	def to(self, device) -> "Batch":
	# spans stay on the CPU: consumers read them back as Python ints, so a device copy would only add a
	# device-to-host sync per decision.
	return Batch(self.feats.to(device), self.seg.to(device), self.mask.to(device),
	[{**d, "prior": d["prior"].to(device),
	"gold": None if d["gold"] is None else d["gold"].to(device)} for d in self.decisions],
	*(getattr(self, name).to(device) for name in self._FLAT), self.counts, self.n_gold)


	def _check_record(rec: Record, n_rows: int) -> None:
	"""Raise a ValueError naming the record and question when a record does not fit its feature rows.

	Rows [0, n_prompt) hold the prompt and [n_prompt, n_prompt + n_canvas) the canvas. Without this check a bad
	record would be read silently from padding or from the wrong segment.
	"""
	if n_rows != rec.n_prompt + rec.n_canvas:
	raise ValueError(f"record {rec.id!r}: {n_rows} feature rows, but n_prompt {rec.n_prompt} + n_canvas "
	f"{rec.n_canvas} = {rec.n_prompt + rec.n_canvas}")
	for d in rec.decisions:
	where = f"record {rec.id!r}, question {d.qid!r}"
	if not rec.n_prompt <= d.slot < n_rows:
	raise ValueError(f"{where}: slot {d.slot} is outside the canvas rows [{rec.n_prompt}, {n_rows})")
	for s, e in d.spans:
	if not 0 <= s < e <= rec.n_prompt:
	raise ValueError(f"{where}: option span [{s}, {e}) is empty or outside the prompt rows "
	f"[0, {rec.n_prompt})")
	if len(d.prior) != len(d.spans):
	raise ValueError(f"{where}: {len(d.spans)} option spans but {len(d.prior)} prior values")
	if d.gold is not None and len(d.gold) != len(d.spans):
	raise ValueError(f"{where}: {len(d.spans)} option spans but {len(d.gold)} gold values")


	def collate(items: list[tuple[Record, torch.Tensor]]) -> Batch:
	lmax = max(f.shape[0] for _, f in items)
	dim = items[0][1].shape[1]
	feats = torch.zeros(len(items), lmax, dim, dtype=torch.float16)
	seg = torch.zeros(len(items), lmax, dtype=torch.long)
	mask = torch.zeros(len(items), lmax, dtype=torch.bool)
	decisions = []
	dec_b, dec_slot, dec_type, dec_has_gold, counts = [], [], [], [], []
	opt_dec, opt_pos, opt_prior, opt_gold, span_rows, span_opt, span_len = [], [], [], [], [], [], []
	for b, (rec, f) in enumerate(items):
	n = f.shape[0]
	_check_record(rec, n)
	feats[b, :n] = f.to(torch.float16)
	mask[b, :n] = True
	seg[b, rec.n_prompt:n] = 1
	for d in rec.decisions:
	i = len(decisions)
	decisions.append({
	"b": b, "slot": d.slot, "type_id": TYPE_INDEX[d.type], "qid": d.qid, "type": d.type,
	"spans": torch.tensor(d.spans, dtype=torch.long),
	"prior": torch.tensor(d.prior, dtype=torch.float32),
	"gold": None if d.gold is None else torch.tensor(d.gold, dtype=torch.float32),
	})
	dec_b.append(b)
	dec_slot.append(d.slot)
	dec_type.append(TYPE_INDEX[d.type])
	dec_has_gold.append(d.gold is not None)
	counts.append(len(d.spans))
	for j, (s, e) in enumerate(d.spans):
	o = len(opt_dec)
	opt_dec.append(i)
	opt_pos.append(j)
	opt_prior.append(d.prior[j])
	opt_gold.append(0.0 if d.gold is None else d.gold[j])
	span_rows.extend(range(b * lmax + s, b * lmax + e))
	span_opt.extend([o] * (e - s))
	span_len.append(float(e - s))
	long, f32 = torch.long, torch.float32
	return Batch(feats=feats, seg=seg, mask=mask, decisions=decisions,
	dec_b=torch.tensor(dec_b, dtype=long), dec_slot=torch.tensor(dec_slot, dtype=long),
	dec_type=torch.tensor(dec_type, dtype=long), dec_has_gold=torch.tensor(dec_has_gold, dtype=torch.bool),
	opt_dec=torch.tensor(opt_dec, dtype=long), opt_pos=torch.tensor(opt_pos, dtype=long),
	opt_prior=torch.tensor(opt_prior, dtype=f32), opt_gold=torch.tensor(opt_gold, dtype=f32),
	span_rows=torch.tensor(span_rows, dtype=long), span_opt=torch.tensor(span_opt, dtype=long),
	span_len=torch.tensor(span_len, dtype=f32), counts=counts, n_gold=sum(dec_has_gold))