farm_runtime.py

#!/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 <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:  # 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()