Update main.py

main.py

@@ -3,20 +3,48 @@ import os
 os.environ["HF_HOME"] = "/app/huggingface"
 
 from fastapi import FastAPI
+from pydantic import BaseModel
 from transformers import AutoModelForTokenClassification, AutoTokenizer
 import torch
+import faiss
+import numpy as np
 
 app = FastAPI()
 
-# Load Model
+# Load Model for NER
 model_name = "priyanandanwar/fine-tuned-gatortron"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForTokenClassification.from_pretrained(model_name)
 
-@app.post("/predict")
-async def predict(text: str):
+# Dummy Retrieval System (Replace with Real FAISS Index)
+dimension = 768
+index = faiss.IndexFlatL2(dimension)
+db_vectors = np.random.rand(10, dimension).astype('float32')
+index.add(db_vectors)
+
+# Define Request Model
+class QueryRequest(BaseModel):
+    text: str
+
+@app.post("/ner")
+async def predict_ner(request: QueryRequest):
+    """Perform Named Entity Recognition (NER)"""
+    text = request.text
     tokens = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
     outputs = model(**tokens)
     predictions = torch.argmax(outputs.logits, dim=-1)[0].tolist()
+    tokenized_text = tokenizer.tokenize(text)
     
-    return {"tokens": tokenizer.tokenize(text), "labels": predictions}
+    return {"tokens": tokenized_text, "labels": predictions}
+
+@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, 1)  # Retrieve Top 1 Match
+    return {"matched_trial_id": int(indices[0][0])}
+
+@app.get("/")
+async def root():
+    return {"message": "TrialGPT API is Running!"}
+