"""Predict EuroVoc tags from the HF layout (encoder model.safetensors + classifier.pt).
Shipped into the HF repo as predict.py. Run from the model directory:

    uv run python predict.py "some EU document text"     # or no arg -> demo doc

Thresholds are the tuned ship operating points, not 0.5 (see MODEL_CARD.md).
"""
import json
import pickle
import sys
from pathlib import Path

import torch
from transformers import AutoModel, AutoTokenizer

HERE = Path(__file__).parent
CONCEPT_THRESHOLD = 0.9          # concept level; coarser levels use MT 0.80, DO 0.70


def load():
    tok = AutoTokenizer.from_pretrained(HERE)
    enc = AutoModel.from_pretrained(HERE).eval()
    head = torch.load(HERE / "classifier.pt", weights_only=True)
    mlb = pickle.load(open(HERE / "mlb.pickle", "rb"))
    names = json.loads((HERE / "names.json").read_text())
    return tok, enc, head, mlb, names


def predict(text, thresh=CONCEPT_THRESHOLD, model=None):
    if not any(c.isalnum() for c in text):   # no letters/digits in any script: only special tokens would
        return []                            # be pooled, which yields hundreds of spurious tags. Skip it.
    tok, enc, head, mlb, names = model or load()   # reloads per call; pass model= in a loop
    b = tok(text, truncation=True, max_length=2048, return_tensors="pt")
    with torch.no_grad():
        h = enc(**b).last_hidden_state
        m = b["attention_mask"].unsqueeze(-1).float()
        p = ((h * m).sum(1) / m.sum(1).clamp(min=1e-9) @ head["weight"].T + head["bias"]).sigmoid()[0]
    tags = [(names.get(str(mlb.classes_[i]), str(mlb.classes_[i])), float(p[i]))
            for i in (p > thresh).nonzero().flatten().tolist()]
    return sorted(tags, key=lambda t: -t[1])


if __name__ == "__main__":
    demo = "The European Parliament adopted a regulation on renewable energy and climate policy."
    text = sys.argv[1] if len(sys.argv) > 1 else demo
    hits = predict(text)
    print(f"{len(hits)} tags over {CONCEPT_THRESHOLD}:")
    for name, score in hits:
        print(f"  {score:.2f}  {name}")