Update app.py

app.py

@@ -3,47 +3,34 @@ import sys
 import pickle
 import numpy as np
 import gradio as gr
-from typing import List
 import fitz  # PyMuPDF
 from docx import Document
-from transformers import AutoModel, AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import AutoModel, AutoTokenizer
 import faiss
 
 # =============================================
-# FIX FOR HUGGINGFACE HUB IMPORT ISSUE
+# EMBEDDING MODEL SETUP (NO sentence-transformers dependency)
 # =============================================
-try:
-    from huggingface_hub import cached_download
-except ImportError:
-    from huggingface_hub.utils import cached_download
-    import huggingface_hub
-    sys.modules['huggingface_hub'].cached_download = cached_download
-
-# Now we can safely import sentence-transformers
-from sentence_transformers import SentenceTransformer
+model_name = "sentence-transformers/all-MiniLM-L6-v2"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+embedding_model = AutoModel.from_pretrained(model_name)
+
+def get_embeddings(texts):
+    if isinstance(texts, str):
+        texts = [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()
 
 # =============================================
-# INITIALIZE MODELS
+# DOCUMENT STORAGE SETUP
 # =============================================
-# Initialize embedding model (using direct transformers as fallback)
-try:
-    embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
-except Exception as e:
-    print(f"Failed to load SentenceTransformer, falling back to direct transformers: {e}")
-    model_name = "sentence-transformers/all-MiniLM-L6-v2"
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    embedding_model = AutoModel.from_pretrained(model_name)
-    
-    def get_embeddings(texts):
-        inputs = tokenizer(texts, padding=True, truncation=True, return_tensors="pt")
-        outputs = embedding_model(**inputs)
-        return outputs.last_hidden_state.mean(dim=1).detach().numpy()
-
-# Initialize FAISS index
 index_path = "faiss_index.pkl"
 document_texts_path = "document_texts.pkl"
 document_texts = []
 
+embedding_dim = 384  # Dimension for all-MiniLM-L6-v2
 if os.path.exists(index_path) and os.path.exists(document_texts_path):
     try:
         with open(index_path, "rb") as f:
@@ -51,10 +38,10 @@ if os.path.exists(index_path) and os.path.exists(document_texts_path):
         with open(document_texts_path, "rb") as f:
             document_texts = pickle.load(f)
     except Exception as e:
-        print(f"Error loading FAISS index: {e}")
-        index = faiss.IndexFlatIP(384)  # 384 is dim for all-MiniLM-L6-v2
+        print(f"Error loading index: {e}")
+        index = faiss.IndexFlatIP(embedding_dim)
 else:
-    index = faiss.IndexFlatIP(384)
+    index = faiss.IndexFlatIP(embedding_dim)
 
 # =============================================
 # DOCUMENT PROCESSING FUNCTIONS
@@ -93,14 +80,12 @@ def upload_files(files):
             else:
                 continue
 
-            sentences = [s for s in text.split("\n") if s.strip()]
-            
-            if hasattr(embedding_model, 'encode'):
-                embeddings = embedding_model.encode(sentences, normalize_embeddings=True)
-            else:
-                embeddings = get_embeddings(sentences)
-
-            index.add(np.array(embeddings))
+            sentences = [s.strip() for s in text.split("\n") if s.strip()]
+            if not sentences:
+                continue
+                
+            embeddings = get_embeddings(sentences)
+            index.add(embeddings)
             document_texts.extend(sentences)
 
         # Save updated index
@@ -115,12 +100,9 @@ def upload_files(files):
 
 def query_text(query):
     try:
-        if hasattr(embedding_model, 'encode'):
-            query_embedding = embedding_model.encode([query], normalize_embeddings=True)
-        else:
-            query_embedding = get_embeddings([query])
-
-        D, I = index.search(np.array(query_embedding), k=3)
+        query_embedding = get_embeddings(query)
+        D, I = index.search(query_embedding, k=3)
+        
         results = []
         for idx in I[0]:
             if 0 <= idx < len(document_texts):
@@ -149,4 +131,5 @@ with gr.Blocks() as demo:
     upload_btn.click(upload_files, inputs=file_input, outputs=upload_output)
     search_btn.click(query_text, inputs=query_input, outputs=results_output)
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()