Update main.py

main.py

@@ -7,18 +7,19 @@ from fastapi import FastAPI
 from pydantic import BaseModel
 from transformers import AutoModel, AutoTokenizer
 
+# Hugging Face Cache Directory
 os.environ["HF_HOME"] = "/app/huggingface"
 os.environ["HF_HUB_DOWNLOAD_TIMEOUT"] = "60"
 
 app = FastAPI()
 
-# --- Load CSV ---
+# --- Load Clinical Trials CSV ---
 csv_path = "ctg-studies.csv"
 if os.path.exists(csv_path):
     df_trials = pd.read_csv(csv_path)
-    print("✅ CSV Loaded!")
+    print("✅ CSV File Loaded Successfully!")
 else:
-    raise FileNotFoundError("❌ CSV File Not Found!")
+    raise FileNotFoundError(f"❌ CSV File Not Found: {csv_path}. Upload it first.")
 
 # --- Load FAISS Index ---
 dimension = 768
@@ -29,9 +30,9 @@ if os.path.exists(faiss_index_path):
     print("✅ FAISS Index Loaded!")
 else:
     index = faiss.IndexFlatL2(dimension)
-    print("⚠️ FAISS Index Empty!")
+    print("⚠ FAISS Index Not Found! Using Empty Index.")
 
-# --- Load Model ---
+# --- Load Retrieval Model ---
 retrieval_model_name = "priyanandanwar/fine-tuned-gatortron"
 retrieval_tokenizer = AutoTokenizer.from_pretrained(retrieval_model_name)
 retrieval_model = AutoModel.from_pretrained(retrieval_model_name)
@@ -41,46 +42,55 @@ class QueryRequest(BaseModel):
     text: str
     top_k: int = 5
 
-# --- Generate Embedding ---
+# --- Generate Embedding for Query ---
 def generate_embedding(text):
     inputs = retrieval_tokenizer(text, return_tensors="pt", truncation=True, padding="max_length", max_length=512)
     with torch.no_grad():
         outputs = retrieval_model(**inputs)
-    emb = outputs.last_hidden_state[:, 0, :].numpy()
+    return outputs.last_hidden_state[:, 0, :].numpy()  # CLS Token Embedding
 
-    # ✅ Normalize Embeddings
-    emb = emb / np.linalg.norm(emb)
-    return emb
+# --- Retrieve Clinical Trial Info ---
+def get_trial_info(nct_id):
+    trial_info = df_trials[df_trials["NCT Number"] == nct_id].to_dict(orient="records")
+    return trial_info[0] if trial_info else None
 
-# --- Compute Similarity ---
-def compute_similarity(distance):
-    return round(np.exp(-distance) * 100, 2)  # ✅ Softmax similarity fix
-
-# --- Retrieve Trials ---
+# --- Retrieval Endpoint ---
 @app.post("/retrieve")
 async def retrieve_trial(request: QueryRequest):
+    """Retrieve Clinical Trial based on text"""
     query_vector = generate_embedding(request.text)
     distances, indices = index.search(query_vector, request.top_k)
 
-    # ✅ Fix: Remove NaN/Inf from distances
-    distances = np.nan_to_num(distances, nan=1e9, posinf=1e9, neginf=1e9)
-
     matched_trials = []
     for idx, dist in zip(indices[0], distances[0]):
         if idx < len(df_trials):
             nct_id = df_trials.iloc[idx]["NCT Number"]
-            trial_data = df_trials[df_trials["NCT Number"] == nct_id].to_dict(orient="records")[0]
-
-            # ✅ Ensure similarity is calculated safely
-            similarity = round(np.exp(-dist) * 100, 2)
-            similarity = max(0.0, similarity)  # No negative similarity
-
-            trial_data["similarity"] = f"{similarity}%"
-            matched_trials.append(trial_data)
-
-    return {"matched_trials": matched_trials}
-
-
+            trial_data = get_trial_info(nct_id)
+
+            if trial_data:
+                # *Ensure distance is valid*
+                if np.isnan(dist) or np.isinf(dist):  
+                    dist = 1e9  # Set large number instead of NaN/Inf
+
+                # ✅ *Proper Similarity Calculation*
+                similarity = round((1 / (1 + dist)) * 100, 2)  
+                trial_data["similarity"] = f"{similarity}%"  # Convert to percentage format
+                
+                matched_trials.append(trial_data)
+
+    # *Convert NaN values to None before returning JSON*
+    def clean_json(data):
+        if isinstance(data, dict):
+            return {k: clean_json(v) for k, v in data.items()}
+        elif isinstance(data, list):
+            return [clean_json(v) for v in data]
+        elif isinstance(data, float) and (np.isnan(data) or np.isinf(data)):
+            return None
+        return data
+
+    return {"matched_trials": clean_json(matched_trials)}
+
+# --- Root Endpoint ---
 @app.get("/")
 async def root():
-    return {"message": "TrialGPT API Running!"}
+    return {"message": "TrialGPT API is Running with FAISS-based Retrieval!"}