#!/usr/bin/env python3 """ModelFarm runtime — chat with any persona package produced by persona_lab. A package is a directory (exported by `persona_lab.py export`) containing: runtime_config.json method + base model + steering parameters persona_vector.pt (engram arm) or neuron_recipe.json (neurons arm) adapter/ (lora arm) or nothing for pure-prompt packages Usage: python farm_runtime.py --chat python farm_runtime.py "Say something in character." python farm_runtime.py --alpha 0.5 "Who are you?" # dose override Requires: torch, transformers (+ peft only for lora packages). """ from __future__ import annotations import argparse import json import re import sys from pathlib import Path import torch from transformers import AutoModelForCausalLM, AutoTokenizer DECODER_LAYER_RE = re.compile(r"^(?:.*\.)?layers\.(\d+)$") def locate_layers(model): found = {} for name, module in model.named_modules(): match = DECODER_LAYER_RE.fullmatch(name) if not match: continue try: module.get_submodule("mlp.down_proj") except AttributeError: continue idx = int(match.group(1)) if idx not in found or len(name) < found[idx]._farm_path: # type: ignore[attr-defined] module._farm_path = name # type: ignore[attr-defined] found[idx] = module if not found: raise RuntimeError("No decoder layers with mlp.down_proj found.") return found def install_engram(model, layers, config, package: Path, alpha: float): blob = torch.load(package / config["vector_file"], weights_only=False) layer = blob["chosen"]["layer"] vec = blob["vectors"][str(layer)].float() unit = vec / vec.norm() median = blob["medians"][str(layer)] vector = (unit * (alpha * median)).to(next(model.parameters()).device, torch.bfloat16) def hook(_mod, _args, output): if isinstance(output, tuple): return (output[0] + vector,) + tuple(output[1:]) return output + vector handle = layers[layer].register_forward_hook(hook) print(f"[modelfarm] engram: layer {layer}, |v·alpha·median| = " f"{vector.norm():.2f} (alpha={alpha})") return handle def install_neurons(model, layers, config, package: Path, dose: float): recipe = json.loads((package / config["recipe_file"]).read_text()) handles = [] by_layer = {} for entry in recipe["recipe"]: by_layer.setdefault(entry["layer"], []).append({**entry, "delta": entry["delta"] * dose}) for layer, entries in by_layer.items(): down_proj = layers[layer].get_submodule("mlp.down_proj") def make_hook(items): def hook(_mod, args): hidden = args[0].clone() for item in items: hidden[..., item["neuron"]] = hidden[..., item["neuron"]] + item["delta"] return (hidden,) return hook handles.append(down_proj.register_forward_pre_hook(make_hook(entries))) print(f"[modelfarm] neurons: {len(recipe['recipe'])} neurons, dose x{dose}") return handles def install_lora(model, package: Path): from peft import PeftModel model = PeftModel.from_pretrained(model, str(package / "adapter")) print("[modelfarm] lora adapter loaded") return model @torch.no_grad() def generate(model, tokenizer, messages, device, max_new_tokens=160): try: inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt", return_dict=True, enable_thinking=False).to(device) except TypeError: inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt", return_dict=True).to(device) out = model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=False, pad_token_id=tokenizer.pad_token_id or tokenizer.eos_token_id) return tokenizer.decode(out[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True).strip() def main() -> None: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("package", type=Path) parser.add_argument("prompt", nargs="?", default=None) parser.add_argument("--chat", action="store_true") parser.add_argument("--alpha", type=float, default=None, help="engram dose override") parser.add_argument("--dose", type=float, default=None, help="neuron dose override") parser.add_argument("--max-new-tokens", type=int, default=160) parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") args = parser.parse_args() config = json.loads((args.package / "runtime_config.json").read_text()) tokenizer = AutoTokenizer.from_pretrained(config["base_model"]) dtype = torch.bfloat16 if args.device != "cpu" else torch.float32 model = AutoModelForCausalLM.from_pretrained(config["base_model"], dtype=dtype).to(args.device) model.eval() layers = locate_layers(model) handles = [] method = config["method"] if method == "engram": handles.append(install_engram(model, layers, config, args.package, args.alpha if args.alpha is not None else config.get("alpha", 1.0))) elif method == "neurons": handles.extend(install_neurons(model, layers, config, args.package, args.dose if args.dose is not None else config.get("dose", 1.0))) elif method == "lora": model = install_lora(model, args.package) elif method in ("sysprompt", "prompt"): pass # plain system-prompt package, nothing to install else: raise SystemExit(f"unknown method {method}") if method in ("sysprompt", "prompt"): print("[modelfarm] plain system-prompt package") system = config.get("system_prompt") or "You are a helpful assistant." if args.chat: print(f"chatting with {args.package.name} ({method}); Ctrl+C to exit\n") try: while True: user = input("you: ").strip() if user: print(f"{args.package.name}: {generate(model, tokenizer, [{'role': 'system', 'content': system}, {'role': 'user', 'content': user}], args.device, args.max_new_tokens)}\n") except (KeyboardInterrupt, EOFError): print("\nbye") else: if not args.prompt: raise SystemExit("give a prompt or --chat") messages = [{"role": "system", "content": system}, {"role": "user", "content": args.prompt}] print(generate(model, tokenizer, messages, args.device, args.max_new_tokens)) for handle in handles: handle.remove() if __name__ == "__main__": main()