Update app.py

app.py

@@ -4,16 +4,29 @@ import faiss
 import torch
 import numpy as np
 import pandas as pd
+import zipfile
 from fastapi import FastAPI
 from pydantic import BaseModel
 from transformers import AutoModel, AutoTokenizer
 from sklearn.feature_extraction.text import TfidfVectorizer
 from sklearn.metrics.pairwise import cosine_similarity
 
-# 🌍 Set Hugging Face Cache Directory
+# 🌍 Set Hugging Face Cache Directory (if needed)
 os.environ["HF_HOME"] = "/app/huggingface"
 os.environ["HF_HUB_DOWNLOAD_TIMEOUT"] = "60"
 
+# ✅ Unzip model if not already unzipped
+model_path = "my_model"
+if not os.path.exists(model_path):
+    with zipfile.ZipFile("my_model.zip", "r") as zip_ref:
+        zip_ref.extractall(model_path)
+    print("✅ Model unzipped!")
+
+# 🤗 Load tokenizer and model from local directory
+retrieval_tokenizer = AutoTokenizer.from_pretrained(model_path)
+retrieval_model = AutoModel.from_pretrained(model_path)
+
+# ✅ Start FastAPI
 app = FastAPI()
 
 # 📄 Load Clinical Trials CSV
@@ -45,11 +58,6 @@ else:
     index = faiss.IndexFlatL2(dimension)
     print("⚠ FAISS Index Not Found. Using Empty Index.")
 
-# 🤗 Load Public Model from Hugging Face
-retrieval_model_name = "distilbert-base-uncased"
-retrieval_tokenizer = AutoTokenizer.from_pretrained(retrieval_model_name)
-retrieval_model = AutoModel.from_pretrained(retrieval_model_name)
-
 # 📦 Request Models
 class QueryRequest(BaseModel):
     text: str
@@ -137,4 +145,4 @@ async def get_trial_details(nct_id: str):
 
 @app.get("/")
 async def root():
-    return {"message": "🌟 TrialGPT API is Running with Public Model & Timeline Extraction!"}
+    return {"message": "🌟 TrialGPT API is Running with Local Model & Timeline Extraction!"}