| """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()) |
|
|
