ecce_vector/examples.py

"""Generate side-by-side example replies with llama-server: stock, and stock plus control vectors.

One llama-server is started per set of vectors (llama.cpp loads control vectors at launch), each prompt is
sent once with a fixed seed, and the reply is saved with llama-server's own timings for prompt processing
(pp) and token generation (tg) and the request's wall-clock time. The first request in a fresh process pays
for Metal/CUDA kernel compilation, so one short warmup request is sent and discarded first.

    python -m ecce_vector.examples --server llama-server --model gemma-4-31B-it-Q4_K_M.gguf \\
        --prompts examples/prompts.json --out examples/runs/musica \\
        --control-vector-scaled vectors/musica.gguf:1.0

A fine-tune released as a LoRA adapter runs on the same base with --lora ADAPTER.gguf (converted with llama.cpp's
convert_lora_to_gguf.py); a full-weight fine-tune runs from its own GGUF as --model, with no vector.

Prompts file: a JSON list of {"id", "system", "turns": [{"role": "user"|"model", "text"}]}, rendered in Gemma 4's
turn format with reasoning off, byte-identical to the model's chat template for enable_thinking=false.
Sampling is Google's recommended setting for Gemma 4 (generation_config.json: temperature 1.0, top-k 64,
top-p 0.95), with llama.cpp's other samplers switched off.
"""
from __future__ import annotations

import argparse
import json
import os
import subprocess
import time
import urllib.request
from pathlib import Path

TURN_OPEN, TURN_CLOSE = "<|turn>", "<turn|>"
THOUGHT_CLOSED = "<|channel>thought\n<channel|>"          # reasoning off: an empty, closed thought channel
SAMPLER = {"temperature": 1.0, "top_k": 64, "top_p": 0.95, "min_p": 0.0, "typical_p": 1.0,
           "repeat_penalty": 1.0, "presence_penalty": 0.0, "frequency_penalty": 0.0}


def render(prompt: dict) -> str:
    parts = [f"{TURN_OPEN}system\n{prompt['system']}{TURN_CLOSE}\n"]
    parts += [f"{TURN_OPEN}{t['role']}\n{t['text']}{TURN_CLOSE}\n" for t in prompt["turns"]]
    parts.append(f"{TURN_OPEN}model\n{THOUGHT_CLOSED}")
    return "".join(parts)


def post(url: str, payload: dict, timeout: float = 900) -> dict:
    req = urllib.request.Request(url, json.dumps(payload).encode(), {"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return json.loads(r.read())


def wait_healthy(url: str, proc: subprocess.Popen, timeout: float = 600) -> None:
    t0 = time.time()
    while time.time() - t0 < timeout:
        if proc.poll() is not None:
            raise RuntimeError(f"llama-server exited with {proc.returncode}")
        try:
            with urllib.request.urlopen(url + "/health", timeout=5) as r:
                if json.loads(r.read()).get("status") == "ok":
                    return
        except Exception:
            pass
        time.sleep(1)
    raise TimeoutError("llama-server did not become healthy")


def main(argv=None):
    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    ap.add_argument("--server", default="llama-server", help="llama-server binary")
    ap.add_argument("--model", required=True, help="the base model GGUF (stock gemma-4-31B-it), or a fine-tune's own GGUF")
    ap.add_argument("--prompts", required=True, type=Path)
    ap.add_argument("--out", required=True, type=Path, help="output directory for this run")
    ap.add_argument("--control-vector-scaled", action="append", default=[], metavar="FILE:SCALE")
    ap.add_argument("--lora", action="append", default=[], metavar="FILE", help="a LoRA adapter GGUF, at scale 1")
    ap.add_argument("--seed", default="24601", help="a seed, or a comma list: the first seed's replies are ID.txt, "
                    "each further seed's ID.seedN.txt (run.json: records, then extra_seeds)")
    ap.add_argument("--max-tokens", type=int, default=320)
    ap.add_argument("--ctx", type=int, default=8192)
    ap.add_argument("--port", type=int, default=8741)
    ap.add_argument("--ik", action="store_true", help="--server is ik_llama.cpp's llama-server (its flag syntax)")
    a = ap.parse_args(argv)

    prompts = json.loads(a.prompts.read_text())
    seeds = [int(s) for s in str(a.seed).split(",")]
    a.out.mkdir(parents=True, exist_ok=True)
    cmd = [a.server, "-m", a.model, "-c", str(a.ctx), "-ngl", "99", "--port", str(a.port), "--parallel", "1"]
    cmd += ["--webui", "none"] if a.ik else ["--no-webui"]
    if a.ik:                        # ik_llama.cpp: the flag repeats, FILE and SCALE as two arguments
        for cv in a.control_vector_scaled:
            f, s = cv.rsplit(":", 1)
            cmd += ["--control-vector-scaled", f, s]
    elif a.control_vector_scaled:   # llama.cpp: one flag, a comma-separated FILE:SCALE list
        cmd += ["--control-vector-scaled", ",".join(a.control_vector_scaled)]
    for f in a.lora:                # the same flag in both engines
        cmd += ["--lora", f]
    url = f"http://127.0.0.1:{a.port}"
    log = open(a.out / "server.log", "w")
    t_launch = time.time()
    proc = subprocess.Popen(cmd, stdout=log, stderr=subprocess.STDOUT)
    records, extra = [], {}
    try:
        wait_healthy(url, proc)
        load_s = time.time() - t_launch
        post(url + "/completion", {"prompt": render(prompts[0]), "n_predict": 8, "seed": seeds[0], **SAMPLER,
                                   "cache_prompt": False})           # warmup, discarded
        for seed, p in ((s, p) for s in seeds for p in prompts):
            t0 = time.time()
            r = post(url + "/completion", {"prompt": render(p), "n_predict": a.max_tokens, "seed": seed,
                                           "cache_prompt": False, "stop": [TURN_CLOSE, TURN_OPEN], **SAMPLER})
            wall = time.time() - t0
            tm = r.get("timings", {})
            text = r["content"].strip()
            (a.out / (f"{p['id']}.txt" if seed == seeds[0] else f"{p['id']}.seed{seed}.txt")).write_text(text + "\n")
            (records if seed == seeds[0] else extra.setdefault(str(seed), [])).append({
                            "id": p["id"], "seed": seed, "text": text, "wall_s": wall,
                            "pp_tokens": tm.get("prompt_n"), "pp_ms": tm.get("prompt_ms"),
                            "pp_tok_s": tm.get("prompt_per_second"),
                            "tg_tokens": tm.get("predicted_n"), "tg_ms": tm.get("predicted_ms"),
                            "tg_tok_s": tm.get("predicted_per_second"),
                            "stop": r.get("stop_type") or ("limit" if r.get("truncated") else None)})
            print(f"{p['id']:24s} {seed:6d} pp {tm.get('prompt_n')} tok @ {tm.get('prompt_per_second', 0):.1f} tok/s   "
                  f"tg {tm.get('predicted_n')} tok @ {tm.get('predicted_per_second', 0):.2f} tok/s   wall {wall:.1f}s",
                  flush=True)
    finally:
        proc.terminate()
        try:
            proc.wait(30)
        except subprocess.TimeoutExpired:
            proc.kill()
        log.close()
    # record paths without the local machine's layout: binaries and the model by file name, vectors relative
    def portable(x: str) -> str:
        if ":" in x and Path(x.rpartition(":")[0]).suffix == ".gguf":      # FILE:SCALE
            head, sep, scale = x.rpartition(":")
        else:
            head, sep, scale = x, "", ""
        p = Path(head)
        if p.is_absolute():
            p = Path(os.path.relpath(p)) if str(p).startswith(os.getcwd()) else Path(p.name)
        return str(p) + sep + scale
    (a.out / "run.json").write_text(json.dumps({
        "model": Path(a.model).name, "control_vectors": [portable(c) for c in a.control_vector_scaled],
        "lora": [Path(f).name for f in a.lora],
        "seed": seeds[0], "max_tokens": a.max_tokens, "sampler": SAMPLER,
        "server_cmd": [Path(cmd[0]).name] + [",".join(portable(p) for p in c.split(",")) if "/" in c else c
                                             for c in cmd[1:]],
        "server_load_s": load_s, "records": records, **({"extra_seeds": extra} if extra else {})}, indent=1))


if __name__ == "__main__":
    main()