Spaces:
Runtime error
Runtime error
File size: 2,198 Bytes
7ffcd85 5c01bc1 7ffcd85 5c01bc1 7ffcd85 5c01bc1 7ffcd85 5c01bc1 7ffcd85 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 |
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
# Check for GPU availability
DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu"
# Global variables
conversation_retrieval_chain = None
chat_history = []
llm_hub = None
embeddings = None
def init_llm():
global llm_hub, embeddings
# Ensure API key is set in Hugging Face Spaces
hf_token = os.getenv("HUGGINGFACEHUB_API_TOKEN")
if not hf_token:
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})
embeddings = HuggingFaceInstructEmbeddings(
model_name="sentence-transformers/all-MiniLM-L6-v2", model_kwargs={"device": DEVICE}
)
def process_document(document_path):
global conversation_retrieval_chain
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)
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"
)
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({"question": prompt, "chat_history": chat_history})
answer = output["result"]
chat_history.append((prompt, answer))
return answer
init_llm()
|