"""Minimal FactGuard-General-4B-v2 inference: label probabilities for one claim against evidence lines."""
import json, torch
from huggingface_hub import hf_hub_download
from transformers import AutoTokenizer, AutoModelForCausalLM

REPO = "ThakiCloud/FactGuard-General-4B-v2"
contract = json.load(open(hf_hub_download(REPO, "label_contract.json")))
HEAD = contract["head_general"]
LABELS = ["SUPPORTED", "CONTRADICTED", "NOT_VERIFIABLE", "STALE"]

tok = AutoTokenizer.from_pretrained(REPO)
model = AutoModelForCausalLM.from_pretrained(REPO, dtype=torch.bfloat16, device_map="auto").eval()
label_ids = [tok.encode(" " + t, add_special_tokens=False)[0] for t in "ABCD"]

packet = ("대상: 회사 공지\n근거:\n"
          "F1: [공지] 서울 사무소는 2026년 3월 2일부터 판교로 이전한다.\n"
          "F2: [공지] 이전 기간 중 고객센터 전화는 정상 운영한다.\n"
          "주장: 이전 기간에는 고객센터 전화가 중단된다.")
ids = tok.encode(HEAD + packet + "\n\n판정:", add_special_tokens=False)[-1024:]
with torch.no_grad():
    logits = model(torch.tensor([ids], device=model.device)).logits[0, -1, label_ids].float()
print(dict(zip(LABELS, [round(p, 4) for p in torch.softmax(logits, -1).tolist()])))