Update main.py

main.py

@@ -2,60 +2,72 @@ import os
 import faiss
 import torch
 import numpy as np
+import pandas as pd
 from fastapi import FastAPI
 from pydantic import BaseModel
-from transformers import AutoModelForTokenClassification, AutoTokenizer
+from transformers import AutoModel, AutoTokenizer
 
+# Hugging Face Cache Directory
 os.environ["HF_HOME"] = "/app/huggingface"
 
 app = FastAPI()
 
-# Load Model for NER
-model_name = "priyanandanwar/fine-tuned-gatortron"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForTokenClassification.from_pretrained(model_name)
+# --- Load Clinical Trials CSV (for metadata lookup) ---
+csv_path = "clinical_trials.csv"  # Ensure this file is uploaded
+df_trials = pd.read_csv(csv_path)
 
-# Dummy FAISS Retrieval System
+# --- Load FAISS Index ---
 dimension = 768
-index = faiss.IndexFlatL2(dimension)
-db_vectors = np.random.rand(10, dimension).astype('float32')
-index.add(db_vectors)
+faiss_index_path = "clinical_trials.index"
 
-# Request Model
+if os.path.exists(faiss_index_path):
+    index = faiss.read_index(faiss_index_path)
+    print("✅ FAISS Index Loaded!")
+else:
+    index = faiss.IndexFlatL2(dimension)
+    print("⚠️ FAISS Index Not Found! Using Empty Index.")
+
+# --- 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)
+
+# --- Request Model ---
 class QueryRequest(BaseModel):
     text: str
-    temperature: float = 0.7
-    max_tokens: int = 256
-    top_p: float = 0.9
-    top_k: int = 50
-
-@app.post("/ner")
-async def predict_ner(request: QueryRequest):
-    """Perform Named Entity Recognition (NER)"""
-    tokens = tokenizer(request.text, return_tensors="pt", truncation=True, padding=True, max_length=request.max_tokens)
-    outputs = model(**tokens)
-    predictions = torch.argmax(outputs.logits, dim=-1)[0].tolist()
-    tokenized_text = tokenizer.tokenize(request.text)
-    
-    return {
-        "tokens": tokenized_text, 
-        "labels": predictions,
-        "temperature": request.temperature,
-        "top_p": request.top_p,
-        "top_k": request.top_k
-    }
+    top_k: int = 5
+
+# --- 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)
+    return outputs.last_hidden_state[:, 0, :].numpy()  # CLS Token Embedding
 
+# --- Retrieve Clinical Trial Info ---
+def get_trial_info(nct_id):
+    trial_info = df_trials[df_trials["NCT_ID"] == nct_id].to_dict(orient="records")
+    return trial_info[0] if trial_info else None
+
+# --- Retrieval Endpoint ---
 @app.post("/retrieve")
 async def retrieve_trial(request: QueryRequest):
     """Retrieve Clinical Trial based on text"""
-    query_vector = np.random.rand(1, dimension).astype('float32')  # Dummy Query Encoding
-    _, indices = index.search(query_vector, request.top_k)  # Retrieve Top K Matches
-    
-    return {
-        "matched_trial_ids": indices.tolist(),
-        "top_k": request.top_k
-    }
+    query_vector = generate_embedding(request.text)
+    distances, indices = index.search(query_vector, request.top_k)
+
+    matched_trials = []
+    for idx, dist in zip(indices[0], distances[0]):
+        nct_id = df_trials.iloc[idx]["NCT_ID"]  # Get NCT_ID using FAISS index mapping
+        trial_data = get_trial_info(nct_id)  # Fetch complete trial details
+        
+        if trial_data:
+            trial_data["similarity"] = round(100 / (1 + dist), 2)  # Convert similarity
+            matched_trials.append(trial_data)
+
+    return {"matched_trials": matched_trials}
 
+# --- Root Endpoint ---
 @app.get("/")
 async def root():
-    return {"message": "TrialGPT API is Running with Parameterized Inputs!"}
+    return {"message": "TrialGPT API is Running with FAISS-based Retrieval!"}