Update api.py

api.py

@@ -1,10 +1,14 @@
-# api.py yes
-
 import os
 from fastapi import FastAPI
 from pydantic import BaseModel
-from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForSequenceClassification
+from transformers import (
+    pipeline,
+    AutoTokenizer,
+    AutoModelForSeq2SeqLM,
+    AutoModelForSequenceClassification,
+)
 
+# Ensure Hugging Face cache path
 os.environ["HF_HOME"] = "/tmp/hf_home"
 os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf_home"
 os.makedirs("/tmp/hf_home", exist_ok=True)
@@ -14,23 +18,32 @@ app = FastAPI(title="Agent Truth API")
 # Hugging Face token (optional if models are private)
 HF_TOKEN = os.environ.get("HF_TOKEN")
 
-# Load NLI model (text-classification)
+# ---------------------------
+# Load NLI model (sequence classification)
+# ---------------------------
 nli_model_id = os.environ.get("NLI_MODEL", "swajall/nli-model")
-nli_model = AutoModelForSequenceClassification.from_pretrained(nli_model_id, token=HF_TOKEN)
+nli_model = AutoModelForSequenceClassification.from_pretrained(
+    nli_model_id, token=HF_TOKEN
+)
 nli_tokenizer = AutoTokenizer.from_pretrained(nli_model_id, token=HF_TOKEN)
+
 nli_pipe = pipeline(
     "text-classification",
-    model=nli_model_id,
-    tokenizer=nli_model_id,
+    model=nli_model,
+    tokenizer=nli_tokenizer,
     device=-1,
-    use_auth_token=HF_TOKEN
 )
 
-# Load Seq2Seq model (T5)
+# ---------------------------
+# Load Seq2Seq model (T5 family)
+# ---------------------------
 seq2_model_id = os.environ.get("SEQ2_MODEL", "swajall/seq2seq-model")
-tokenizer = AutoTokenizer.from_pretrained(seq2_model_id, use_auth_token=HF_TOKEN)
-seq2_model = AutoModelForSeq2SeqLM.from_pretrained(seq2_model_id, use_auth_token=HF_TOKEN)
+tokenizer = AutoTokenizer.from_pretrained(seq2_model_id, token=HF_TOKEN)
+seq2_model = AutoModelForSeq2SeqLM.from_pretrained(seq2_model_id, token=HF_TOKEN)
 
+# ---------------------------
+# Request Schemas
+# ---------------------------
 class NLIRequest(BaseModel):
     premise: str
     hypothesis: str
@@ -38,13 +51,17 @@ class NLIRequest(BaseModel):
 class Seq2SeqRequest(BaseModel):
     transcript: str
 
+# ---------------------------
+# Routes
+# ---------------------------
 @app.get("/")
 def root():
-    return {"msg": "Agent Truth API is running"}
+    return {"msg": "Agent Truth API is running 🚀"}
 
 @app.post("/nli")
 def nli(req: NLIRequest):
-    res = nli_pipe((req.premise, req.hypothesis))
+    # Correct input format for text + hypothesis
+    res = nli_pipe({"text": req.premise, "text_pair": req.hypothesis})
     return {"result": res}
 
 @app.post("/seq2seq")