Spaces:
Paused
Paused
File size: 3,833 Bytes
3abf967 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 | """
model_b.py - Model B (Claim Verification) Inference Wrapper for TruthLens AI.
Task:
Input: claim + candidate evidence passage
Output: SUPPORTS / REFUTES / NOT_ENOUGH_INFO with probability distribution.
Hardware: CUDA accelerated with model caching.
"""
import os
import sys
import json
from typing import Dict, Any
import torch
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModelForSequenceClassification
DEFAULT_MODEL_B_DIR = os.path.join(
os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
"models",
"verification",
"scifact_deberta"
)
ID2LABEL = {
0: "SUPPORTS",
1: "REFUTES",
2: "NOT_ENOUGH_INFO",
}
FALLBACK_MODEL_B = "roberta-base"
class ModelB:
_instance = None
def __init__(self, model_dir: str = DEFAULT_MODEL_B_DIR):
has_weights = os.path.exists(model_dir) and (
os.path.exists(os.path.join(model_dir, "model.safetensors")) or
os.path.exists(os.path.join(model_dir, "pytorch_model.bin"))
)
if not has_weights:
print(f"Model B local weights not found in '{model_dir}'. Falling back to HF Hub: '{FALLBACK_MODEL_B}'")
model_dir = FALLBACK_MODEL_B
if torch.cuda.is_available():
self.device = torch.device("cuda")
self.device_name = torch.cuda.get_device_name(0)
self.cuda_available = True
else:
self.device = torch.device("cpu")
self.device_name = "CPU"
self.cuda_available = False
self.tokenizer = AutoTokenizer.from_pretrained(model_dir)
self.model = AutoModelForSequenceClassification.from_pretrained(model_dir)
self.model.to(self.device)
self.model.eval()
@classmethod
def get_instance(cls, model_dir: str = DEFAULT_MODEL_B_DIR):
if cls._instance is None:
cls._instance = cls(model_dir=model_dir)
return cls._instance
def predict_verification(self, claim: str, evidence: str) -> Dict[str, Any]:
claim_clean = str(claim).strip()
evidence_clean = str(evidence).strip() if evidence else "None"
input_text = f"[CLAIM] {claim_clean} [SEP] {evidence_clean}"
inputs = self.tokenizer(
input_text,
max_length=256,
padding=True,
truncation=True,
return_tensors="pt"
)
inputs = {k: v.to(self.device) for k, v in inputs.items()}
with torch.no_grad():
with torch.amp.autocast("cuda", dtype=torch.float16):
outputs = self.model(**inputs)
probs = F.softmax(outputs.logits, dim=-1).squeeze(0).cpu().numpy()
prob_supports = float(round(float(probs[0]), 4))
prob_refutes = float(round(float(probs[1]), 4))
prob_nei = float(round(float(probs[2]), 4))
prob_dict = {
"SUPPORTS": prob_supports,
"REFUTES": prob_refutes,
"NOT_ENOUGH_INFO": prob_nei,
}
# Winning class
pred_idx = int(probs.argmax())
pred_label = ID2LABEL.get(pred_idx, "NOT_ENOUGH_INFO")
confidence = float(round(float(probs[pred_idx]), 4))
return {
"label": pred_label,
"confidence": confidence,
"probabilities": prob_dict,
}
def predict_verification(
claim: str,
evidence: str,
model_dir: str = DEFAULT_MODEL_B_DIR
) -> Dict[str, Any]:
wrapper = ModelB.get_instance(model_dir=model_dir)
return wrapper.predict_verification(claim, evidence)
if __name__ == "__main__":
print("Testing Model B wrapper...")
c = "Smoking causes lung cancer."
e = "Cigarette smoking is proven to cause lung cancer and respiratory illness."
res = predict_verification(c, e)
print(json.dumps(res, indent=2))
|