Update app.py

app.py

@@ -4,36 +4,18 @@ from pydantic import BaseModel
 import faiss
 import pickle
 import numpy as np
-from transformers import AutoTokenizer, AutoModel
+from sentence_transformers import SentenceTransformer
 import torch
 
-# Set cache to custom dir to avoid /.cache issues
+# Optional: set dummy cache paths if needed by other internal libs
 os.environ["HF_HOME"] = "/app/hf_cache"
 os.environ["TRANSFORMERS_CACHE"] = "/app/hf_cache"
 
-# Get Hugging Face token from environment (must be set in HF Space secrets)
-hf_token = os.getenv("HF_TOKEN")
-
-# Initialize FastAPI
+# Initialize FastAPI app
 app = FastAPI()
 
-import shutil
-
-model_cache_path = "/app/hf_cache/models--viswadarshan06--paraphrase-multilingual-MiniLM-L12-v2-local"
-if os.path.exists(model_cache_path):
-    shutil.rmtree(model_cache_path)
-
-# ✅ Load your private model
-tokenizer = AutoTokenizer.from_pretrained(
-    "viswadarshan06/paraphrase-multilingual-MiniLM-L12-v2-local",
-    cache_dir="/app/hf_cache",
-    token=hf_token
-)
-model = AutoModel.from_pretrained(
-    "viswadarshan06/paraphrase-multilingual-MiniLM-L12-v2-local",
-    cache_dir="/app/hf_cache",
-    token=hf_token
-)
+# ✅ Load your locally uploaded SentenceTransformer model
+model = SentenceTransformer("/app/model")
 
 # Load Thirukkural data
 with open("thirukkural_data.pkl", "rb") as f:
@@ -43,35 +25,23 @@ with open("thirukkural_data.pkl", "rb") as f:
 english_index = faiss.read_index("thirukkural_english_index.faiss")
 tamil_index = faiss.read_index("thirukkural_tamil_index.faiss")
 
-# Request schema
+# Define request schema
 class QueryRequest(BaseModel):
     query: str
     lang: str  # "en" or "ta"
     top_k: int = 3
 
-# Mean pooling (same as SentenceTransformer style)
-def mean_pooling(model_output, attention_mask):
-    token_embeddings = model_output[0]
-    input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
-    return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
-
-# POST endpoint
+# Search endpoint
 @app.post("/search/")
 def search_kural(req: QueryRequest):
-    encoded_input = tokenizer(req.query, padding=True, truncation=True, return_tensors="pt")
-
-    with torch.no_grad():
-        model_output = model(**encoded_input)
-        query_embedding = mean_pooling(model_output, encoded_input["attention_mask"])
-
-    query_embedding = query_embedding.detach().cpu().numpy()
+    query_embedding = model.encode([req.query])
     index = tamil_index if req.lang == "ta" else english_index
-    D, I = index.search(query_embedding.astype("float32"), req.top_k)
+    D, I = index.search(np.array(query_embedding).astype("float32"), req.top_k)
 
     results = [kural_data[i] for i in I[0]]
     return {"results": results}
 
-# Health check
+# Health check endpoint
 @app.get("/")
 def root():
-    return {"message": "Thirukkural FastAPI RAG is running with private model."}
+    return {"message": "Thirukkural FastAPI RAG is running with local model."}