Spaces:
Runtime error
Runtime error
Update worker.py
Browse files
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
|
6 |
-
from
|
7 |
-
from
|
8 |
-
from
|
9 |
-
from
|
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 =
|
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({"
|
61 |
-
answer = output["
|
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
|
|
|
|