Update main.py

main.py

@@ -4,59 +4,58 @@ import torch
 import numpy as np
 from fastapi import FastAPI
 from pydantic import BaseModel
-from transformers import AutoModelForTokenClassification, AutoTokenizer, AutoModel
+from transformers import AutoModelForTokenClassification, AutoTokenizer
 
-# Hugging Face Cache Directory (For HF Spaces)
 os.environ["HF_HOME"] = "/app/huggingface"
 
 app = FastAPI()
 
-# --- Load NER Model ---
+# Load Model for NER
 model_name = "priyanandanwar/fine-tuned-gatortron"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-ner_model = AutoModelForTokenClassification.from_pretrained(model_name)
+model = AutoModelForTokenClassification.from_pretrained(model_name)
 
-# --- Load FAISS Index ---
+# Dummy FAISS Retrieval System
 dimension = 768
-faiss_index_path = "clinical_trials.index"
+index = faiss.IndexFlatL2(dimension)
+db_vectors = np.random.rand(10, dimension).astype('float32')
+index.add(db_vectors)
 
-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 for Embeddings ---
-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 ---
+# Request Model
 class QueryRequest(BaseModel):
     text: str
-    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
+    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
+    }
 
-# --- 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)
-    
-    # Convert retrieved indices to NCT IDs directly
-    results = [{"NCT_ID": str(int(idx)), "similarity": float(round(100 / (1 + dist), 2))} for idx, dist in zip(indices[0], distances[0])]
+    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_trials": results}
+    return {
+        "matched_trial_ids": indices.tolist(),
+        "top_k": request.top_k
+    }
 
-# --- Root Endpoint ---
 @app.get("/")
 async def root():
-    return {"message": "TrialGPT API is Running with FAISS-based Retrieval!"}
+    return {"message": "TrialGPT API is Running with Parameterized Inputs!"}