Update worker.py

worker.py

@@ -1,11 +1,12 @@
 import os
 import torch
-from langchain.chains import RetrievalQA
+from langchain.chains import ConversationalRetrievalChain
 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
+from langchain_community.llms import HuggingFacePipeline
+from transformers import pipeline
 
 # Check for GPU availability
 DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu"
@@ -13,11 +14,12 @@ DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu"
 # Global variables
 conversation_retrieval_chain = None
 chat_history = []
-llm_hub = None
+llm_pipeline = None
 embeddings = None
 
+
 def init_llm():
-    global llm_hub, embeddings
+    global llm_pipeline, embeddings
 
     # Ensure API key is set in Hugging Face Spaces
     hf_token = os.getenv("HUGGINGFACEHUB_API_TOKEN")
@@ -25,17 +27,19 @@ def init_llm():
         raise ValueError("HUGGINGFACEHUB_API_TOKEN is not set in environment variables.")
 
     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})
+    hf_pipeline = pipeline("text-generation", model=model_id, device=DEVICE)
+    llm_pipeline = HuggingFacePipeline(pipeline=hf_pipeline)
 
     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:
+    if not llm_pipeline or not embeddings:
         init_llm()
 
     loader = PyPDFLoader(document_path)
@@ -44,22 +48,27 @@ def process_document(document_path):
     text_splitter = RecursiveCharacterTextSplitter(chunk_size=1024, chunk_overlap=64)
     texts = text_splitter.split_documents(documents)
     
-    db = Chroma.from_documents(texts, embedding=embeddings, persist_directory="./chroma_db")
+    # Load or create ChromaDB
+    persist_directory = "./chroma_db"
+    if os.path.exists(persist_directory):
+        db = Chroma(persist_directory=persist_directory, embedding_function=embeddings)
+    else:
+        db = Chroma.from_documents(texts, embedding=embeddings, persist_directory=persist_directory)
 
-    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
+    retriever = db.as_retriever(search_type="similarity", search_kwargs={'k': 6})
+    
+    conversation_retrieval_chain = ConversationalRetrievalChain.from_llm(
+        llm=llm_pipeline, retriever=retriever
     )
 
+
 def process_prompt(prompt):
     global conversation_retrieval_chain, chat_history
 
     if not conversation_retrieval_chain:
         return "No document has been processed yet. Please upload a PDF first."
 
-    output = conversation_retrieval_chain({"query": prompt, "chat_history": chat_history})
+    output = conversation_retrieval_chain({"question": prompt, "chat_history": chat_history})
     answer = output["answer"]
     
     chat_history.append((prompt, answer))