Update app.py

app.py

@@ -9,8 +9,11 @@ from docx import Document as DocxDocument
 import fitz
 import os
 import torch
-os.environ['CUDA_LAUNCH_BLOCKING']="1"
-os.environ['PYTORCH_USE_CUDA_DSA'] = "1"
+from HybridRetriever import HybridRetriever
+from ChatEngine import ChatEngine
+from llama_index.retrievers.bm25 import BM25Retriever 
+from llama_index.core.retrievers import VectorIndexRetriever
+
 lm_list = {
     "google/gemma-2-9b-it": "google/gemma-2-9b-it",
     "mistralai/Mistral-7B-Instruct-v0.3": "mistralai/Mistral-7B-Instruct-v0.3"
@@ -60,13 +63,17 @@ def handle_file_upload(file, llm_name):
     index = VectorStoreIndex.from_documents(
         documents, transformations=[text_splitter], embed_model=Settings.embed_model
     )
-
-    return index.as_query_engine()
+    
+    bm25_retriever = BM25Retriever(nodes=documents, similarity_top_k=2, tokenizer=text_splitter.split_text)
+    vector_retriever = VectorIndexRetriever(index=vector_index, similarity_top_k=2)
+    hybrid_retriever = HybridRetriever(bm25_retriever=bm25_retriever, vector_retriever=vector_retriever)
+    chat_engine = ChatEngine(hybrid_retriever)
+    response = chat_engine.ask_question(question, llm)
+    return response
 
 def document_qa(file_upload, llm_choice, question_input):
-    query_engine = handle_file_upload(file_upload, llm_choice)
-    result = query_engine.query(question_input)
-    return str(result)
+    response = handle_file_upload(file_upload, llm_choice)
+    return response
 
 
 llm_choice = gr.Dropdown(choices=list(lm_list.values()), label="Choose LLM")