Spaces:
Runtime error
Runtime error
File size: 2,544 Bytes
7ffcd85 f8407d3 b7e8fd3 f8407d3 7ffcd85 5c01bc1 7ffcd85 f8407d3 7ffcd85 f8407d3 7ffcd85 f8407d3 7ffcd85 5c01bc1 7ffcd85 f8407d3 7ffcd85 b7e8fd3 7ffcd85 f8407d3 7ffcd85 b7e8fd3 f8407d3 b7e8fd3 7ffcd85 f8407d3 7ffcd85 f8407d3 7ffcd85 f8407d3 7ffcd85 5c01bc1 f8407d3 b7e8fd3 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 69 70 71 72 73 74 75 76 77 |
import os
import torch
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 HuggingFacePipeline
from transformers import pipeline
# Check for GPU availability
DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu"
# Global variables
conversation_retrieval_chain = None
chat_history = []
llm_pipeline = None
embeddings = None
def init_llm():
global llm_pipeline, 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"
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_pipeline 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)
# 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)
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({"question": prompt, "chat_history": chat_history})
answer = output["answer"]
chat_history.append((prompt, answer))
return answer
|