File size: 3,861 Bytes
0cf8154 | 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 104 105 106 107 108 | """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|>"
|