| |
| """A thegrid `odin` operator backed by gridlm instead of a general LLM. |
| |
| Same queue protocol as thegrid's tools/operator/grid_operator.py: the colony writes |
| <queue>/request.json (with "genome", a list of instruction names, and "request_id"); this |
| writes <queue>/proposal.json ({"request_id", "authored_at", "genome"}) atomically. The colony |
| validates every proposal with its own parser, so a proposal it can't use costs one birth (it |
| falls back to blind variation) and never the colony. |
| |
| How a proposal is made: the generator samples K children for the parent; the screen scores |
| each; the best-scoring child that differs from the parent is proposed. Runs on CPU in |
| milliseconds, so there is no GPU queue to wait for. |
| |
| Honest scope: trained on ONE colony (kannaka-grid's colony-one, an older thegrid ISA and a |
| starvation economy). Instruction names that colony never used are unknown to the model, and |
| anything it proposes that a destination colony can't run is rejected there by design. |
| |
| usage: |
| python gridlm_operator.py --model-dir . --queue /path/to/colony/operator [--queue ...] |
| python gridlm_operator.py --model-dir . --once-genome "harvest harvest alloc copy ifnotdone jmpb fork scan move" |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import os |
| import sys |
| import time |
| from pathlib import Path |
|
|
| import torch |
|
|
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) |
| import train as g |
|
|
|
|
| def load(model_dir: Path): |
| vocab = json.loads((model_dir / "vocab.json").read_text()) |
| screen = g.Screen(len(vocab)) |
| gen = g.Generator(len(vocab)) |
| for model, name in ((screen, "screen"), (gen, "generator")): |
| st = model_dir / f"{name}.safetensors" |
| if st.exists(): |
| from safetensors.torch import load_file |
| model.load_state_dict(load_file(str(st))) |
| else: |
| model.load_state_dict(torch.load(model_dir / f"{name}.pt", map_location="cpu", weights_only=True)) |
| screen.eval() |
| gen.eval() |
| return vocab, {t: i for i, t in enumerate(vocab)}, screen, gen |
|
|
|
|
| def propose(names, vocab, idx, screen, gen, k=8, temperature=0.9, seed=None): |
| parent = [idx.get(n, g.UNK) for n in names][: g.MAXG] |
| kids = g.sample_children(gen, [parent] * k, vocab, k, temperature=temperature, |
| rng_seed=seed if seed is not None else int(time.time() * 1000) % (2 ** 31)) |
| cands = [c for c in kids if 1 <= len(c) <= g.MAXG and c != list(names)] |
| if not cands: |
| return None, None |
| rows = [{"parent": list(names)[: g.MAXG], "child": c, "label": 0, "child_id": ""} for c in cands] |
| X, _ = g.screen_tensor(rows, idx) |
| p = g.screen_scores(screen, X) |
| best = int(p.argmax()) |
| return cands[best], float(p[best]) |
|
|
|
|
| def serve(queue: Path, model, stats) -> bool: |
| req_p, prop_p = queue / "request.json", queue / "proposal.json" |
| if prop_p.exists() or not req_p.exists(): |
| return False |
| try: |
| req = json.loads(req_p.read_text()) |
| except (OSError, ValueError): |
| return False |
| child, score = propose(req.get("genome", []), *model) |
| if child is None: |
| stats["rejected"] += 1 |
| return True |
| tmp = prop_p.with_suffix(".tmp") |
| tmp.write_text(json.dumps({"request_id": req.get("request_id"), "authored_at": time.time(), |
| "genome": child, "gridlm_screen_p": score})) |
| tmp.replace(prop_p) |
| stats["served"] += 1 |
| print(f"{time.strftime('%H:%M:%S')} {queue}: {len(req.get('genome', []))} -> {len(child)} " |
| f"(screen p={score:.3f}; served {stats['served']} rejected {stats['rejected']})", flush=True) |
| return True |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--model-dir", default=".") |
| ap.add_argument("--queue", action="append", default=[]) |
| ap.add_argument("--once-genome", help="print one proposal for this genome (space-separated names) and exit") |
| ap.add_argument("--poll", type=float, default=1.0) |
| a = ap.parse_args() |
| torch.set_num_threads(2) |
| model = load(Path(a.model_dir)) |
| if a.once_genome: |
| child, score = propose(a.once_genome.split(), *model) |
| print(json.dumps({"parent": a.once_genome.split(), "child": child, "screen_p": score})) |
| return |
| stats = {"served": 0, "rejected": 0} |
| queues = [Path(q) for q in a.queue] |
| print("gridlm operator; queues:", ", ".join(map(str, queues)), flush=True) |
| while True: |
| worked = False |
| for q in queues: |
| if q.exists(): |
| worked = serve(q, model, stats) or worked |
| if not worked: |
| time.sleep(a.poll) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
