Deepakraj2006 commited on
Commit
b7e8fd3
·
verified ·
1 Parent(s): 829e296

Update worker.py

Browse files
Files changed (1) hide show
  1. worker.py +14 -14
worker.py CHANGED
@@ -1,12 +1,11 @@
1
  import os
2
  import torch
3
- from langchain import PromptTemplate
4
  from langchain.chains import RetrievalQA
5
- from langchain.embeddings import HuggingFaceInstructEmbeddings
6
- from langchain.document_loaders import PyPDFLoader
7
- from langchain.text_splitter import RecursiveCharacterTextSplitter
8
- from langchain.vectorstores import Chroma
9
- from langchain.llms import HuggingFaceHub
10
 
11
  # Check for GPU availability
12
  DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu"
@@ -28,27 +27,30 @@ def init_llm():
28
  model_id = "tiiuae/falcon-7b-instruct"
29
  llm_hub = HuggingFaceHub(repo_id=model_id, model_kwargs={"temperature": 0.1, "max_new_tokens": 600, "max_length": 600})
30
 
31
- embeddings = HuggingFaceInstructEmbeddings(
32
  model_name="sentence-transformers/all-MiniLM-L6-v2", model_kwargs={"device": DEVICE}
33
  )
34
 
35
  def process_document(document_path):
36
  global conversation_retrieval_chain
37
 
 
 
 
 
38
  loader = PyPDFLoader(document_path)
39
  documents = loader.load()
40
 
41
  text_splitter = RecursiveCharacterTextSplitter(chunk_size=1024, chunk_overlap=64)
42
  texts = text_splitter.split_documents(documents)
43
 
44
- db = Chroma.from_documents(texts, embedding=embeddings)
45
 
46
  conversation_retrieval_chain = RetrievalQA.from_chain_type(
47
  llm=llm_hub,
48
  chain_type="stuff",
49
  retriever=db.as_retriever(search_type="mmr", search_kwargs={'k': 6, 'lambda_mult': 0.25}),
50
- return_source_documents=False,
51
- input_key="question"
52
  )
53
 
54
  def process_prompt(prompt):
@@ -57,11 +59,9 @@ def process_prompt(prompt):
57
  if not conversation_retrieval_chain:
58
  return "No document has been processed yet. Please upload a PDF first."
59
 
60
- output = conversation_retrieval_chain({"question": prompt, "chat_history": chat_history})
61
- answer = output["result"]
62
 
63
  chat_history.append((prompt, answer))
64
 
65
  return answer
66
-
67
- init_llm()
 
1
  import os
2
  import torch
 
3
  from langchain.chains import RetrievalQA
4
+ from langchain_community.embeddings import HuggingFaceEmbeddings
5
+ from langchain_community.document_loaders import PyPDFLoader
6
+ from langchain_text_splitters import RecursiveCharacterTextSplitter
7
+ from langchain_community.vectorstores import Chroma
8
+ from langchain_community.llms import HuggingFaceHub
9
 
10
  # Check for GPU availability
11
  DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu"
 
27
  model_id = "tiiuae/falcon-7b-instruct"
28
  llm_hub = HuggingFaceHub(repo_id=model_id, model_kwargs={"temperature": 0.1, "max_new_tokens": 600, "max_length": 600})
29
 
30
+ embeddings = HuggingFaceEmbeddings(
31
  model_name="sentence-transformers/all-MiniLM-L6-v2", model_kwargs={"device": DEVICE}
32
  )
33
 
34
  def process_document(document_path):
35
  global conversation_retrieval_chain
36
 
37
+ # Ensure LLM and embeddings are initialized
38
+ if not llm_hub or not embeddings:
39
+ init_llm()
40
+
41
  loader = PyPDFLoader(document_path)
42
  documents = loader.load()
43
 
44
  text_splitter = RecursiveCharacterTextSplitter(chunk_size=1024, chunk_overlap=64)
45
  texts = text_splitter.split_documents(documents)
46
 
47
+ db = Chroma.from_documents(texts, embedding=embeddings, persist_directory="./chroma_db")
48
 
49
  conversation_retrieval_chain = RetrievalQA.from_chain_type(
50
  llm=llm_hub,
51
  chain_type="stuff",
52
  retriever=db.as_retriever(search_type="mmr", search_kwargs={'k': 6, 'lambda_mult': 0.25}),
53
+ return_source_documents=False
 
54
  )
55
 
56
  def process_prompt(prompt):
 
59
  if not conversation_retrieval_chain:
60
  return "No document has been processed yet. Please upload a PDF first."
61
 
62
+ output = conversation_retrieval_chain({"query": prompt, "chat_history": chat_history})
63
+ answer = output["answer"]
64
 
65
  chat_history.append((prompt, answer))
66
 
67
  return answer