Update app.py

app.py

@@ -21,14 +21,18 @@ tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 print("[INFO] Loading model...")
 model = AutoModelForSequenceClassification.from_pretrained(
     MODEL_ID,
-    num_labels=2
+    num_labels=2,
+    torch_dtype=torch.float32
 )
+
 model.to(device)
 model.eval()
 
+torch.set_num_threads(1)
+
 print("[INFO] Warming up model...")
 with torch.no_grad():
-    dummy = tokenizer("Warmup text.", return_tensors="pt").to(device)
+    dummy = tokenizer("Warmup text.", return_tensors="pt")
     _ = model(**dummy)
 
 print("[INFO] Model ready.")
@@ -36,8 +40,7 @@ print("[INFO] Model ready.")
 # ================== HELPERS ==================
 
 def split_sentences(text: str):
-    sentences = re.split(r'(?<=[.!?])\s+', text)
-    return [s.strip() for s in sentences if len(s.strip()) > 10]
+    return re.split(r'(?<=[.!?])\s+', text)
 
 
 def score_sentences(sentences):
@@ -47,34 +50,33 @@ def score_sentences(sentences):
         truncation=True,
         max_length=MAX_LENGTH,
         return_tensors="pt"
-    ).to(device)
+    )
 
     with torch.no_grad():
         logits = model(**inputs).logits
         probs = F.softmax(logits, dim=-1)
 
-    return probs[:, 1].tolist()  # claim probability
+    return probs[:, 1].tolist()
 
 
 def extract_claims(text: str, threshold: float):
     if not text or len(text.strip()) < 20:
         return []
 
-    sentences = split_sentences(text)
+    sentences = [s.strip() for s in split_sentences(text) if len(s.strip()) > 10]
     if not sentences:
         return []
 
     scores = score_sentences(sentences)
 
-    results = []
-    for sent, score in zip(sentences, scores):
-        results.append({
-            "sentence": sent,
-            "claim_probability": round(score, 4),
-            "is_claim": score >= threshold
-        })
-
-    return results
+    return [
+        {
+            "sentence": s,
+            "claim_probability": round(sc, 4),
+            "is_claim": sc >= threshold
+        }
+        for s, sc in zip(sentences, scores)
+    ]
 
 # ================== FASTAPI ==================
 
@@ -100,19 +102,5 @@ def extract_api(payload: AnalyzeRequest):
 def home():
     return {
         "message": "Claim Extraction API running",
-        "model": MODEL_ID,
-        "endpoints": {
-            "POST /extract": {
-                "body": {
-                    "text": "news article text",
-                    "threshold": "optional float (default 0.6)"
-                }
-            }
-        }
-    }
-
-# ================== ENTRYPOINT ==================
-
-if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+        "model": MODEL_ID
+    }