Update app.py

app.py

@@ -9,19 +9,24 @@ import faiss
 import torch
 
 # ===============================
-# EMBEDDING MODEL
+# EMBEDDING MODEL (E5)
 # ===============================
-model_name = "sentence-transformers/all-MiniLM-L6-v2"
+model_name = "intfloat/e5-small-v2"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 embedding_model = AutoModel.from_pretrained(model_name)
 
-def get_embeddings(texts):
+def get_embeddings(texts, is_query=False):
     if isinstance(texts, str):
         texts = [texts]
+    prefix = "query: " if is_query else "passage: "
+    texts = [prefix + t for t in texts]
+
     inputs = tokenizer(texts, padding=True, truncation=True, return_tensors="pt", max_length=512)
     with torch.no_grad():
-        outputs = embedding_model(**inputs)
-    return outputs.last_hidden_state[:, 0].cpu().numpy()
+        model_output = embedding_model(**inputs)
+
+    embeddings = model_output.last_hidden_state[:, 0]  # CLS token embeddings
+    return embeddings.cpu().numpy()
 
 # ===============================
 # TEXT CHUNKING
@@ -41,7 +46,7 @@ def chunk_text(text, chunk_size=800, overlap=100):
 index_path = "faiss_index.pkl"
 document_texts_path = "document_texts.pkl"
 document_texts = []
-embedding_dim = 384
+embedding_dim = 768  # E5-small-v2 embedding dimension
 
 if os.path.exists(index_path) and os.path.exists(document_texts_path):
     try:
@@ -90,7 +95,7 @@ def upload_document(file):
         return "Unsupported file type."
 
     chunks = chunk_text(text)
-    chunk_embeddings = get_embeddings(chunks)
+    chunk_embeddings = get_embeddings(chunks, is_query=False)
     index.add(np.array(chunk_embeddings).astype('float32'))
     document_texts.extend(chunks)
 
@@ -110,7 +115,7 @@ def generate_answer_from_file(query, top_k=10):
     if not document_texts:
         return "No documents indexed yet."
 
-    query_vector = get_embeddings(query).astype("float32")
+    query_vector = get_embeddings(query, is_query=True).astype("float32")
     scores, indices = index.search(query_vector, k=top_k)
     retrieved_chunks = [document_texts[i] for i in indices[0]]
     context = "\n\n".join(retrieved_chunks)
@@ -155,5 +160,3 @@ search_interface = gr.Interface(
 
 app = gr.TabbedInterface([upload_interface, search_interface], ["Upload", "Ask"])
 app.launch()
-
-