"""Minimal FactGuard-9B inference: returns label probabilities for one claim against evidence lines."""
import json, torch
from transformers import AutoTokenizer, AutoModelForCausalLM

REPO = "ThakiCloud/FactGuard-9B"
HEAD = ("공식 근거 사실(F1~F9)과 AI 답변의 주장 하나를 대조해, 보여 준 근거만으로 주장을 판정한다(조건부 함의).\n"
        "A=지지(근거가 주장 전체를 함의: 대상 제품·플랜·수치·조건·기간이 모두 근거에 있다) B=모순(근거가 주장과 양립할 수 없다고 명시) "
        "C=확인불가(근거가 말하지 않거나 플랜·열 귀속·조건을 확정할 수 없다. 빈 칸·흐린 칸·헤더와 안 맞는 행은 근거가 아니다) D=과거 근거로만 지지\n"
        "결정 근거: a=없음 b..j=F1..F9\n\n")
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 = ("대상 브랜드: ExampleSaaS\n근거:\n"
          "F1: [Pricing | ExampleSaaS] Pro — $20 / month, billed monthly\n"
          "F2: [Pricing | ExampleSaaS] Pro — Up to 10 projects\n"
          "주장: ExampleSaaS Pro 플랜은 월 결제 시 월 20달러이다.")
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()
probs = torch.softmax(logits, -1).tolist()
print(dict(zip(LABELS, [round(p, 4) for p in probs])))