Update worker.py

worker.py

@@ -1,12 +1,11 @@
 import os
 import torch
-from langchain import PromptTemplate
 from langchain.chains import RetrievalQA
-from langchain.embeddings import HuggingFaceInstructEmbeddings
-from langchain.document_loaders import PyPDFLoader
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.vectorstores import Chroma
-from langchain.llms import HuggingFaceHub
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_community.document_loaders import PyPDFLoader
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain_community.vectorstores import Chroma
+from langchain_community.llms import HuggingFaceHub
 
 # Check for GPU availability
 DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu"
@@ -28,27 +27,30 @@ def init_llm():
     model_id = "tiiuae/falcon-7b-instruct"
     llm_hub = HuggingFaceHub(repo_id=model_id, model_kwargs={"temperature": 0.1, "max_new_tokens": 600, "max_length": 600})
 
-    embeddings = HuggingFaceInstructEmbeddings(
+    embeddings = HuggingFaceEmbeddings(
         model_name="sentence-transformers/all-MiniLM-L6-v2", model_kwargs={"device": DEVICE}
     )
 
 def process_document(document_path):
     global conversation_retrieval_chain
 
+    # Ensure LLM and embeddings are initialized
+    if not llm_hub or not embeddings:
+        init_llm()
+
     loader = PyPDFLoader(document_path)
     documents = loader.load()
     
     text_splitter = RecursiveCharacterTextSplitter(chunk_size=1024, chunk_overlap=64)
     texts = text_splitter.split_documents(documents)
     
-    db = Chroma.from_documents(texts, embedding=embeddings)
+    db = Chroma.from_documents(texts, embedding=embeddings, persist_directory="./chroma_db")
 
     conversation_retrieval_chain = RetrievalQA.from_chain_type(
         llm=llm_hub,
         chain_type="stuff",
         retriever=db.as_retriever(search_type="mmr", search_kwargs={'k': 6, 'lambda_mult': 0.25}),
-        return_source_documents=False,
-        input_key="question"
+        return_source_documents=False
     )
 
 def process_prompt(prompt):
@@ -57,11 +59,9 @@ def process_prompt(prompt):
     if not conversation_retrieval_chain:
         return "No document has been processed yet. Please upload a PDF first."
 
-    output = conversation_retrieval_chain({"question": prompt, "chat_history": chat_history})
-    answer = output["result"]
+    output = conversation_retrieval_chain({"query": prompt, "chat_history": chat_history})
+    answer = output["answer"]
     
     chat_history.append((prompt, answer))
     
     return answer
-
-init_llm()