#!/usr/bin/env python3
"""Provider-side TOPLOC prover for the HOOKAGENT testnet desk (public how-to version).
pip install torch transformers toploc==0.1.6   (CPU is fine)
stdin: {"spec": {"model","revision","prompt","max_new_tokens"}}  ->  stdout: {"output","output_ids","proofs","topk","dbs"}
Rules the worker re-checks: the prompt is ONE user turn in the model's own chat template (add_generation_prompt=True);
greedy decoding (do_sample=False); fp32 weights; activations = LAST hidden layer, prefill rows + one row per decode step,
cast to bfloat16; TOPLOC build_proofs_base64 with topk=128, decode_batching_size=8, skip_prefill=False."""
import sys, json, os
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from toploc import build_proofs_base64
TOPK, DBS = 128, 8
torch.manual_seed(0); torch.set_num_threads(4)
def prove(spec, swap_model=None):
    model = swap_model or spec["model"]; rev = "main" if swap_model else spec["revision"]
    tok = AutoTokenizer.from_pretrained(model, revision=rev)
    m = AutoModelForCausalLM.from_pretrained(model, revision=rev, dtype=torch.float32); m.eval()
    ids = tok.apply_chat_template([{"role": "user", "content": spec["prompt"]}], add_generation_prompt=True, return_tensors="pt", return_dict=True)["input_ids"]
    with torch.no_grad():
        g = m.generate(ids, max_new_tokens=int(spec["max_new_tokens"]), do_sample=False, output_hidden_states=True, return_dict_in_generate=True, pad_token_id=tok.eos_token_id)
    out = g.sequences[0, ids.shape[1]:].tolist()
    acts = [g.hidden_states[0][-1][0].to(torch.bfloat16)] + [g.hidden_states[i][-1][0, -1].to(torch.bfloat16) for i in range(1, len(g.hidden_states))]
    proofs = build_proofs_base64(acts, decode_batching_size=DBS, topk=TOPK, skip_prefill=False)
    return {"output": tok.decode(out, skip_special_tokens=True), "output_ids": out, "proofs": proofs, "topk": TOPK, "dbs": DBS}
if __name__ == "__main__":
    d = json.load(sys.stdin); print(json.dumps(prove(d["spec"], d.get("swap_model"))))