FactGuard-General-4B-v2 / example_inference.py
thaki-AI's picture
FactGuard-General-4B-v2 research preview (C1 merged BF16)
e25fbe9 verified
Raw History Blame Contribute Delete
1.26 kB
"""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()])))