run.py fix

model/model.py

@@ -33,7 +33,7 @@ class FactChecker(nn.Module):
 class TruthfulnessSayer(nn.Module):
     def __init__(self, model_name, num_labels=3, weights_path=None):
         super().__init__()
-        self.fact_checker = FactChecker(model_name, num_labels)
+        self.fact_checker = FactChecker(model_name, num_labels).float()
         if weights_path:
             self.fact_checker.load_state_dict(torch.load(weights_path, map_location="cpu"))
         self.fact_checker.eval()

model/run.py

@@ -19,7 +19,6 @@ from transformers import AutoTokenizer
 from sentence_transformers import SentenceTransformer
 
 args, device = parse_gpu_args(extra_args=[
-    (["claim"], {"type": str, "help": "Claim to check"}),
     (["--weights"], {"type": str, "default": None, "help": "Path to model weights"}),
     (["--model"], {"type": str, "default": MODEL_NAME}),
     (["--pq"], {"type": str, "default": PQ_NAME}),
@@ -29,6 +28,7 @@ weights = args.weights or os.path.join(
     "checkpoints", f"{args.model.rsplit('/', 1)[-1]}__{args.pq}", "best.pt"
 )
 
+print("Loading models...")
 embedder = SentenceTransformer(EMBEDDER, device=device, model_kwargs={"dtype": "auto"})
 tokenizer = AutoTokenizer.from_pretrained(args.model)
 model = TruthfulnessSayer(args.model, NUM_LABELS, weights).to(device)
@@ -36,36 +36,45 @@ model = TruthfulnessSayer(args.model, NUM_LABELS, weights).to(device)
 index = faiss.read_index(faiss_path(args.pq))
 faiss.extract_index_ivf(index).nprobe = NPROBE
 conn = sqlite3.connect(DB_PATH)
+print("Ready.\n")
 
-query_emb = embedder.encode(
-    [args.claim], prompt_name=QUERY_PROMPT, normalize_embeddings=True
-).astype(np.float32)
-
-_, ids_found = index.search(query_emb, TOP_K)
-chunk_ids = [i for i in ids_found[0].tolist() if i >= 0]
-
-evidence_texts = []
-if chunk_ids:
-    ph = ",".join("?" * len(chunk_ids))
-    rows = conn.execute(f"SELECT id, chunk FROM chunks WHERE id IN ({ph})", chunk_ids).fetchall()
-    id_to_text = {r[0]: r[1] for r in rows}
-    evidence_texts = [id_to_text.get(i, "") for i in chunk_ids]
-while len(evidence_texts) < TOP_K:
-    evidence_texts.append("")
-conn.close()
+while True:
+    try:
+        claim = input("Claim: ").strip()
+    except (EOFError, KeyboardInterrupt):
+        break
+    if not claim:
+        continue
 
-import torch
-pairs = tokenizer(
-    [args.claim] * TOP_K, evidence_texts,
-    truncation=True, max_length=MAX_LENGTH, padding="longest", return_tensors="pt",
-)
-input_ids = pairs["input_ids"].unsqueeze(0).to(device)
-attention_mask = pairs["attention_mask"].unsqueeze(0).to(device)
+    query_emb = embedder.encode(
+        [claim], prompt_name=QUERY_PROMPT, normalize_embeddings=True
+    ).astype(np.float32)
+
+    _, ids_found = index.search(query_emb, TOP_K)
+    chunk_ids = [i for i in ids_found[0].tolist() if i >= 0]
+
+    evidence_texts = []
+    if chunk_ids:
+        ph = ",".join("?" * len(chunk_ids))
+        rows = conn.execute(f"SELECT id, chunk FROM chunks WHERE id IN ({ph})", chunk_ids).fetchall()
+        id_to_text = {r[0]: r[1] for r in rows}
+        evidence_texts = [id_to_text.get(i, "") for i in chunk_ids]
+    while len(evidence_texts) < TOP_K:
+        evidence_texts.append("")
 
-score = model(input_ids, attention_mask).item()
+    pairs = tokenizer(
+        [claim] * TOP_K, evidence_texts,
+        truncation=True, max_length=MAX_LENGTH, padding="longest", return_tensors="pt",
+    )
+    input_ids = pairs["input_ids"].unsqueeze(0).to(device)
+    attention_mask = pairs["attention_mask"].unsqueeze(0).to(device)
 
-print(f"\nClaim: {args.claim}")
-print(f"Truthfulness: {score:.4f}")
-for i, ev in enumerate(evidence_texts):
-    if ev:
-        print(f"  Evidence {i+1}: {ev[:120]}...")
+    score = model(input_ids, attention_mask).item()
+
+    print(f"Truthfulness: {score:.4f}")
+    for i, ev in enumerate(evidence_texts):
+        if ev:
+            print(f"  Evidence {i+1}: {ev[:120]}...")
+    print()
+
+conn.close()