"""Predict topics from raw text — no training stack required.

Loads the two artifacts this repo ships and the public base model:
  weights/seed0/lora/       LoRA adapter (~10.1M params)  -> corrects the encoder
  weights/seed0/head/head.pt  56 prototypes + scale        -> turns vectors into topics
  Qwen/Qwen3-Embedding-0.6B  NOT stored here; transformers downloads it
                             from huggingface.co on first load (that's the point:
                             we ship only the parts we created).

Usage:
  python predict.py --weights weights/seed0 "Cold brew, oat milk, double"
  echo "PSA: picture day moved to Friday" | python predict.py --weights weights/seed0
"""
import argparse
import json
import sys

import torch
import torch.nn.functional as F
from peft import PeftModel
from transformers import AutoModel, AutoTokenizer

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"


def load(weights_dir):
    hs = torch.load(f"{weights_dir}/head/head.pt", map_location="cpu")
    tok = AutoTokenizer.from_pretrained(hs["base_model"])
    tok.padding_side = "right"
    base = AutoModel.from_pretrained(hs["base_model"], torch_dtype=torch.bfloat16).to(DEVICE)
    base.config.use_cache = False
    model = PeftModel.from_pretrained(base, f"{weights_dir}/lora").to(DEVICE).eval()
    protos = F.normalize(hs["protos"].to(DEVICE).float(), dim=-1)
    return tok, model, protos, float(hs["scale"]), hs["topics"]


@torch.no_grad()
def encode(model, tok, texts, max_length=512):
    enc = tok(texts, padding=True, truncation=True, max_length=max_length,
              return_tensors="pt").to(DEVICE)
    out = model(**enc).last_hidden_state
    last = enc["attention_mask"].sum(1) - 1
    return out[torch.arange(out.size(0), device=DEVICE), last].float()


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--weights", default="weights/seed0")
    ap.add_argument("--topk", type=int, default=3)
    ap.add_argument("text", nargs="?", default=None)
    a = ap.parse_args()

    texts = [a.text] if a.text else [l.strip() for l in sys.stdin if l.strip()]
    tok, model, protos, scale, topics = load(a.weights)
    v = F.normalize(encode(model, tok, texts), dim=-1)
    probs = F.softmax(scale * (v @ protos.t()), dim=-1)
    for t, p in zip(texts, probs):
        top = torch.topk(p, a.topk)
        print(json.dumps({
            "text": (t[:70] + "…") if len(t) > 70 else t,
            "topic": topics[top.indices[0]],
            "confidence": round(float(top.values[0]), 4),
            "top3": [{"topic": topics[i], "p": round(float(x), 4)}
                     for i, x in zip(top.indices, top.values)],
        }, ensure_ascii=False))


if __name__ == "__main__":
    main()