Something went wrong. Try again.
End-to-end Tinker model training and serving with Eve.
Something went wrong. Try again.
4.9 kB · 124 lines
Python
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125from __future__ import annotations
import argparseimport hashlibimport jsonfrom dataclasses import asdictfrom pathlib import Pathfrom typing import Any
from training.recipes import all_recipes, get_recipefrom training.recipes.base import ROOT, SPLITS, Case, Recipe, Split
DEFAULT_RECIPE = "capability-placement"
def validate_cases(recipe: Recipe, cases: dict[Split, tuple[Case, ...]]) -> None: if set(cases) != set(SPLITS): raise ValueError(f"{recipe.id} does not contain every split") seen_ids: set[str] = set() seen_slugs: set[str] = set() for split in SPLITS: rows = cases[split] expected_count = len(recipe.kinds) * recipe.split_counts[split] if len(rows) != expected_count: raise ValueError(f"{recipe.id}/{split} has {len(rows)} rows, expected {expected_count}") kinds = {row.kind for row in rows} if kinds != set(recipe.kinds): raise ValueError(f"{recipe.id}/{split} does not cover every kind") for kind in recipe.kinds: count = sum(row.kind == kind for row in rows) if count != recipe.split_counts[split]: raise ValueError(f"{recipe.id}/{split}/{kind} has {count} rows") for row in rows: if row.split != split: raise ValueError(f"split drift in {row.id}") if row.id in seen_ids or row.slug in seen_slugs: raise ValueError(f"duplicate id or slug in {recipe.id}: {row.id} / {row.slug}") seen_ids.add(row.id) seen_slugs.add(row.slug) if not row.prompt.strip() or not row.target: raise ValueError(f"empty prompt or target in {row.id}") if not recipe.validate_target(row, row.target): raise ValueError(f"target schema drift in {recipe.id}/{row.id}")
def _jsonl_bytes(rows: tuple[Case, ...]) -> bytes: return ("".join(json.dumps(asdict(row), separators=(",", ":")) + "\n" for row in rows)).encode()
def rendered_files(recipe: Recipe) -> dict[Path, bytes]: cases = recipe.generate_cases() validate_cases(recipe, cases) files = {recipe.data_dir / f"{split}.jsonl": _jsonl_bytes(cases[split]) for split in SPLITS} test_fixture = [ {"id": row.id, "prompt": row.prompt, "expected": row.target} for row in cases["test"] ] files[recipe.eve_eval_path] = ( json.dumps(test_fixture, indent=2, ensure_ascii=False) + "\n" ).encode() return files
def dataset_sha256(recipe: Recipe) -> str: digest = hashlib.sha256() for path, content in sorted(rendered_files(recipe).items(), key=lambda item: str(item[0])): if path == recipe.eve_eval_path: continue digest.update(path.name.encode()) digest.update(b"\0") digest.update(content) return digest.hexdigest()
def write_files(recipes: tuple[Recipe, ...], check: bool) -> dict[str, str]: mismatches: list[str] = [] hashes: dict[str, str] = {} for recipe in recipes: for path, content in rendered_files(recipe).items(): if check: if not path.exists() or path.read_bytes() != content: mismatches.append(str(path.relative_to(ROOT))) continue path.parent.mkdir(parents=True, exist_ok=True) path.write_bytes(content) hashes[recipe.id] = dataset_sha256(recipe) if mismatches: raise SystemExit("generated data is stale: " + ", ".join(mismatches)) return hashes
def load_split(recipe: Recipe, split: Split) -> tuple[Case, ...]: rows: list[Case] = [] path = recipe.data_dir / f"{split}.jsonl" for line in path.read_text().splitlines(): payload: dict[str, Any] = json.loads(line) rows.append(Case(**payload)) return tuple(rows)
def parser() -> argparse.ArgumentParser: result = argparse.ArgumentParser(description="Generate deterministic Eve recipe corpora") selection = result.add_mutually_exclusive_group() selection.add_argument( "--recipe", choices=tuple(recipe.id for recipe in all_recipes()), help=f"one recipe to generate; defaults to {DEFAULT_RECIPE}", ) selection.add_argument("--all", action="store_true", help="generate every recipe") action = result.add_mutually_exclusive_group(required=True) action.add_argument("--write", action="store_true", help="write generated files") action.add_argument("--check", action="store_true", help="fail when generated files are stale") return result
def main() -> None: args = parser().parse_args() recipes = all_recipes() if args.all else (get_recipe(args.recipe or DEFAULT_RECIPE),) hashes = write_files(recipes, check=args.check) for recipe_id, digest in hashes.items(): print(f"{recipe_id} dataset_sha256={digest}")
if __name__ == "__main__": main()