Update main.py

main.py

@@ -1,6 +1,4 @@
 import os
-os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf_cache"  # cache before importing model
-
 import torch
 import pandas as pd
 from fastapi import FastAPI
@@ -8,136 +6,79 @@ from pydantic import BaseModel
 from sentence_transformers import SentenceTransformer, util
 from langdetect import detect
 from huggingface_hub import hf_hub_download
+import threading
+
+os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf_cache"
 
-# --- Configuration ---
 MODEL_PATH = './muril_combined_multilingual_model'
 CSV_PATH = './muril_multilingual_dataset.csv'
-EMBEDDINGS_PATH = './answer_embeddings.pt'
-HF_DATASET_REPO = "Sp2503/muril-dataset"  # your HF dataset repo
-HF_FILE_NAME = "answer_embeddings.pt"
-
-
+HF_REPO = "Sp2503/muril-dataset"
+HF_FILE = "answer_embeddings.pt"
 
+app = FastAPI(title="MuRIL Multilingual QA API")
 
+model = None
+df = None
+answer_embeddings = None
 
-def load_or_download_embeddings():
-    CACHE_DIR = "/tmp"
-    EMBEDDING_FILENAME = "answer_embeddings.pt"
-    LOCAL_PATH = os.path.join(CACHE_DIR, EMBEDDING_FILENAME)
 
+def load_embeddings():
     print("📥 Downloading embeddings from Hugging Face...")
+    hf_path = hf_hub_download(repo_id=HF_REPO, filename=HF_FILE, repo_type="dataset", cache_dir="/tmp")
+    print(f"✅ Embeddings available at {hf_path}")
+    return torch.load(hf_path, map_location="cpu")
 
-    try:
-        # Download (stays in cache_dir)
-        hf_path = hf_hub_download(
-            repo_id="Sp2503/muril-dataset",
-            filename=EMBEDDING_FILENAME,
-            repo_type="dataset",
-            token=os.getenv("HF_TOKEN"),
-            cache_dir=CACHE_DIR
-        )
-
-        print(f"✅ Embeddings available at {hf_path}")
-
-        # Load directly from hf_path — no rename, no copy
-        if not os.path.exists(hf_path):
-            raise FileNotFoundError(f"{hf_path} not found after download!")
-
-        embeddings = torch.load(hf_path, map_location="cpu")
-        print("✅ Embeddings loaded successfully.")
-        return embeddings
-
-    except Exception as e:
-        print(f"❌ Failed to load embeddings: {e}")
-        print("⚙️ Computing new embeddings from scratch...")
-
-        # === Compute your embeddings here ===
-        # Example:
-        # from sentence_transformers import SentenceTransformer
-        # model = SentenceTransformer("muril_combined_multilingual_model")
-        # embeddings = model.encode(sentences)
-        #
-        # torch.save(embeddings, LOCAL_PATH)
-        # =====================================
-
-        raise RuntimeError("Embeddings not available and could not be regenerated.") from e
-
-
-# === Call this during app startup ===
-answer_embeddings = load_or_download_embeddings()
 
-# --- Load Model + Data ---
 def load_resources():
+    global model, df, answer_embeddings
     try:
-        # Load model
+        print("⚙️ Loading model and dataset in background...")
         model = SentenceTransformer(MODEL_PATH)
-        # Load dataset
         df = pd.read_csv(CSV_PATH).dropna(subset=['question', 'answer'])
-
-        # Use already loaded embeddings from HF
-        if answer_embeddings is None:
-            print("⚙️ Computing new embeddings from scratch...")
-            answers = df['answer'].tolist()
-            embeddings = model.encode(answers, convert_to_tensor=True)
-            torch.save(embeddings, EMBEDDINGS_PATH)
-            print("✅ Computed and saved embeddings")
-        else:
-            embeddings = answer_embeddings
-            print("✅ Using embeddings loaded from Hugging Face")
-
-        return model, df, embeddings
-
+        answer_embeddings = load_embeddings()
+        print("✅ Model and embeddings ready.")
     except Exception as e:
         print(f"❌ Error loading resources: {e}")
-        return None, None, None
 
 
-model, df, answer_embeddings = load_resources()
+# --- Fast startup ---
+@app.on_event("startup")
+def schedule_background_load():
+    thread = threading.Thread(target=load_resources)
+    thread.start()
 
-# --- FastAPI Setup ---
-app = FastAPI(title="MuRIL Multilingual QA API")
 
 class QueryRequest(BaseModel):
     question: str
-    lang: str = None  # optional: en, hi, mr, etc.
+    lang: str = None
 
 class QAResponse(BaseModel):
     answer: str
 
-# --- API Endpoint ---
+
+@app.get("/")
+def root():
+    ready = model is not None and df is not None and answer_embeddings is not None
+    return {"status": "✅ Running", "model_loaded": ready}
+
+
 @app.post("/get-answer", response_model=QAResponse)
 def get_answer_endpoint(request: QueryRequest):
-    if not model:
-        return {"answer": "❌ Model not loaded properly."}
+    if model is None or df is None or answer_embeddings is None:
+        return {"answer": "⏳ Model still loading, please try again shortly."}
 
     question_text = request.question.strip()
-    lang_filter = request.lang
+    lang_filter = request.lang or detect(question_text)
 
-    # Detect language if not given
-    if not lang_filter:
-        try:
-            lang_filter = detect(question_text)
-        except Exception:
-            lang_filter = None
-
-    # Filter dataframe by language if column exists
     filtered_df = df
     filtered_embeddings = answer_embeddings
-    if lang_filter and 'lang' in df.columns:
+    if 'lang' in df.columns and lang_filter:
         mask = df['lang'] == lang_filter
         filtered_df = df[mask].reset_index(drop=True)
-        if len(filtered_df) == 0:
-            return {"answer": f"⚠️ No data found for language '{lang_filter}'."}
         filtered_embeddings = answer_embeddings[mask.values]
 
-    # Encode question + find best match
     question_emb = model.encode(question_text, convert_to_tensor=True)
     cosine_scores = util.pytorch_cos_sim(question_emb, filtered_embeddings)
     best_idx = torch.argmax(cosine_scores).item()
     answer = filtered_df.iloc[best_idx]['answer']
-
     return {"answer": answer}
-
-@app.get("/")
-def root():
-    return {"status": "✅ MuRIL Multilingual QA API running successfully!"}