Update app.py

app.py

@@ -16,6 +16,7 @@ from langchain.chains import RetrievalQA
 from langchain_google_genai import ChatGoogleGenerativeAI
 from PyPDF2 import PdfReader
 from gtts import gTTS
+from sentence_transformers import SentenceTransformer
 
 temp_file_map = {}  
 
@@ -43,7 +44,13 @@ vector_store = None
 file_id_map = {}
 temp_dir = "./temp_downloads"
 os.makedirs(temp_dir, exist_ok=True)
-embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+
+# ✅ Define Sentence-Transformers for both models
+def get_embedding_model(file_size_category):
+    if file_size_category in ["small", "medium"]:
+        return SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
+    else:
+        return SentenceTransformer('sentence-transformers/all-MiniLM-L3-v2')
 
 # ✅ Get list of files from Google Drive
 def get_files_from_drive():
@@ -102,10 +109,13 @@ def process_documents(selected_files):
     # ✅ Dynamically adjust chunk size for efficiency  
     if total_words < 1000:
         chunk_size, chunk_overlap = 500, 50   # Small  
+        file_size_category = "small"
     elif total_words < 5000:
         chunk_size, chunk_overlap = 1000, 100  # Medium  
+        file_size_category = "medium"
     else:
         chunk_size, chunk_overlap = 2000, 200  # Large  
+        file_size_category = "large"
 
     logging.info(f"📄 Document Size: {total_words} words | Chunk Size: {chunk_size}, Overlap: {chunk_overlap}")
 
@@ -113,6 +123,10 @@ def process_documents(selected_files):
     text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
     split_docs = text_splitter.split_documents(docs)
 
+    # ✅ Choose embedding model based on file size category
+    embedding_model = get_embedding_model(file_size_category)
+    embeddings = HuggingFaceEmbeddings(model_name=embedding_model)
+
     # ✅ Store efficiently in vector database  
     vector_store = Chroma.from_documents(split_docs, embeddings)
 
@@ -120,10 +134,6 @@ def process_documents(selected_files):
 
 
 # ✅ Query document
-
-
-
-
 # ✅ Ensure temp_file_map exists
 temp_file_map = {}
 
@@ -215,4 +225,3 @@ with gr.Blocks() as demo:
     submit_button.click(query_document, inputs=user_input, outputs=[response_output, audio_output])
 
 demo.launch()
-