| |
| """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 <package_dir> --chat |
| python farm_runtime.py <package_dir> "Say something in character." |
| python farm_runtime.py <package_dir> --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: |
| module._farm_path = name |
| 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 |
| 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() |
|
|
