| """Canonical text format shared by SFT data, training and inference. |
| |
| One turn is `<|user|>QUESTION<|spec|>SPEC<|result|>RESULT<|answer|>ANSWER<|end|>`; a document is |
| `<|bos|>` + turns + `<|eos|>`. Only SPEC and `ANSWER<|end|>` are trained; RESULT is always |
| written by the engine. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import math |
| from dataclasses import dataclass |
|
|
| from zenith.engines.base import Result |
| from zenith.engines.registry import schemas |
| from zenith.tokenizer.bpe import BPETokenizer |
|
|
| IGNORE = -100 |
|
|
|
|
| def fmt_num(x: float, sig: int = 4) -> str: |
| """`sig` significant digits, no trailing zeros; exponent form (1.988e30, 1.23e-4) outside [1e-3, 1e6).""" |
| x = float(x) |
| if x == 0: |
| return "0" |
| rounded = float(f"{x:.{sig - 1}e}") |
| if abs(rounded) >= 1e6 or abs(rounded) < 1e-3: |
| mant, exp = f"{x:.{sig - 1}e}".split("e") |
| mant = mant.rstrip("0").rstrip(".") if "." in mant else mant |
| return f"{mant}e{int(exp)}" |
| digits = max(sig - 1 - math.floor(math.log10(abs(rounded))), 0) |
| s = f"{rounded:.{digits}f}" |
| return s.rstrip("0").rstrip(".") if "." in s else s |
|
|
|
|
| def _value(v) -> str: |
| if isinstance(v, str): |
| return '"' + v.replace("\\", "\\\\").replace('"', '\\"') + '"' |
| return fmt_num(v) |
|
|
|
|
| def _obj(pairs) -> str: |
| return "{" + ",".join(f'"{k}":{v}' for k, v in pairs) + "}" |
|
|
|
|
| def spec_json(spec: dict) -> str: |
| """Compact JSON with keys in order sim, op, then the op's schema order.""" |
| pairs = [("sim", _value(spec["sim"]))] |
| if "op" in spec: |
| pairs.append(("op", _value(spec["op"]))) |
| order = [p.name for p in schemas()[spec["sim"]][spec["op"]]] |
| pairs += [(k, _value(spec[k])) for k in order if k in spec] |
| return _obj(pairs) |
|
|
|
|
| def result_json(result: Result) -> str: |
| d = result.prompt_dict() |
| values = _obj((k, _value(v)) for k, v in d["values"].items()) |
| notes = "[" + ",".join(_value(n) for n in d["notes"]) + "]" |
| return _obj([("values", values), ("notes", notes)]) |
|
|
|
|
| @dataclass |
| class Turn: |
| question: str |
| spec: dict |
| result: Result |
| answer: str |
|
|
|
|
| def _segments(tok: BPETokenizer, turn: Turn) -> list[tuple[list[int], bool]]: |
| """(ids, trainable) pieces of one turn.""" |
| sp = tok.special_ids |
| return [ |
| ([sp["<|user|>"], *tok.encode(turn.question), sp["<|spec|>"]], False), |
| (tok.encode(spec_json(turn.spec)), True), |
| ([sp["<|result|>"], *tok.encode(result_json(turn.result)), sp["<|answer|>"]], False), |
| ([*tok.encode(turn.answer), sp["<|end|>"]], True), |
| ] |
|
|
|
|
| def encode_turns(tok: BPETokenizer, turns: list[Turn], train_last_only: bool = True) -> tuple[list[int], list[int]]: |
| """Token ids and next-token targets; targets are IGNORE except where the next token is trained.""" |
| ids = [tok.special_ids["<|bos|>"]] |
| trained = [False] |
| for i, turn in enumerate(turns): |
| train_turn = not train_last_only or i == len(turns) - 1 |
| for seg, trainable in _segments(tok, turn): |
| ids += seg |
| trained += [trainable and train_turn] * len(seg) |
| ids.append(tok.special_ids["<|eos|>"]) |
| trained.append(False) |
| targets = [ids[i + 1] if trained[i + 1] else IGNORE for i in range(len(ids) - 1)] + [IGNORE] |
| return ids, targets |
|
|
|
|
| def prompt_ids(tok: BPETokenizer, history: list[Turn], question: str) -> list[int]: |
| """Prefix for inference: history turns, then the new question, ending at `<|spec|>`.""" |
| ids = [tok.special_ids["<|bos|>"]] |
| for turn in history: |
| for seg, _ in _segments(tok, turn): |
| ids += seg |
| return ids + [tok.special_ids["<|user|>"], *tok.encode(question), tok.special_ids["<|spec|>"]] |
|
|
|
|
| def continuation_text(spec_text: str, result: Result) -> str: |
| """Text that follows the prompt before answer generation.""" |
| return f"{spec_text}<|result|>{result_json(result)}<|answer|>" |
|
|
