File size: 4,154 Bytes
35700de | 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 109 110 111 112 113 114 115 116 | """glint-router-1m: route prompts, edit the policy, teach it new categories.
python infer.py "write a python function that reverses a linked list"
python infer.py --file prompts.txt
python infer.py --policy my_policy.json "summarize this thread"
python infer.py --add-category refunds examples.txt
python infer.py --demo
everything runs on cpu. one forward pass per prompt, around 0.2 ms.
"""
from __future__ import annotations
import sys
import time
from pathlib import Path
import torch
from modeling import encode_batch, load_router
from policy import Policy, Prototypes, decide
HERE = Path(__file__).parent
PROTOTYPES = HERE / "prototypes.json"
def route(prompts: list[str], policy_path: Path | None = None,
prototypes_path: Path | None = None, device: str = "cpu"):
model, tokenizer = load_router(HERE, device)
policy = Policy.load(policy_path or HERE / "policy.json")
prototypes = Prototypes.load(prototypes_path)
tokens = encode_batch(tokenizer, prompts, model.config.max_len).to(device)
return [decide(model, tokens[i:i + 1], prompts[i], policy, prototypes)
for i in range(len(prompts))]
def add_category(name: str, examples: list[str], device: str = "cpu") -> None:
"""teach it a label it was never trained on. one forward pass, no optimizer."""
model, tokenizer = load_router(HERE, device)
tokens = encode_batch(tokenizer, examples, model.config.max_len).to(device)
with torch.no_grad():
projections = model(tokens)["proj"]
prototypes = Prototypes.load(PROTOTYPES)
prototypes.add(name, projections)
prototypes.save(PROTOTYPES)
print(f"added category {name!r} from {len(examples)} examples -> {PROTOTYPES}")
def show(decisions, prompts) -> None:
for decision, prompt in zip(decisions, prompts, strict=True):
flags = [k for k, v in decision.signals.items() if v]
head = prompt if len(prompt) <= 58 else prompt[:55] + "..."
print(f"{decision.tier:<9} {decision.domain:<20} c={decision.complexity} "
f"d={decision.difficulty:.2f} {head}")
print(f" {decision.reason}"
+ (f" | signals={flags}" if flags else "")
+ (f" | category={decision.category}" if decision.category else ""))
DEMO = [
"What's the capital of France?",
"Write a short poem about a Python snake in the garden.",
"Write a Python function that reverses a linked list in place.",
"A train leaves at 3pm going 60mph, another at 4pm going 80mph, when do they meet?",
"Summarize this email thread in two sentences.",
"What are the latest news headlines today?",
"Explain the difference between a mutex and a semaphore, with an example.",
]
def main() -> int:
args = sys.argv[1:]
if not args or args[0] in ("-h", "--help"):
print(__doc__)
return 0
if args[0] == "--demo":
started = time.perf_counter()
decisions = route(DEMO)
elapsed = time.perf_counter() - started
show(decisions, DEMO)
print(f"\n{len(DEMO)} prompts in {elapsed * 1000:.0f} ms "
f"({elapsed * 1000 / len(DEMO):.2f} ms each, cpu, includes model load)")
return 0
if args[0] == "--add-category":
if len(args) < 3:
print("usage: python infer.py --add-category <name> <examples.txt>")
return 1
examples = [line.strip() for line in Path(args[2]).read_text().splitlines()
if line.strip()]
add_category(args[1], examples)
return 0
policy_path = None
if args[0] == "--policy":
policy_path, args = Path(args[1]), args[2:]
if args and args[0] == "--file":
prompts = [line.strip() for line in Path(args[1]).read_text().splitlines()
if line.strip()]
else:
prompts = [" ".join(args)]
decisions = route(prompts, policy_path,
PROTOTYPES if PROTOTYPES.is_file() else None)
show(decisions, prompts)
if len(prompts) == 1:
print()
print(decisions[0].as_line())
return 0
if __name__ == "__main__":
raise SystemExit(main())
|