import os import streamlit as st from dotenv import load_dotenv from langchain_community.embeddings import HuggingFaceEmbeddings from langchain_community.llms import llamacpp from langchain_core.runnables.history import RunnableWithMessageHistory from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder from langchain_core.callbacks import CallbackManager, StreamingStdOutCallbackHandler from langchain.chains import create_history_aware_retriever, create_retrieval_chain, ConversationalRetrievalChain from langchain.document_loaders import TextLoader from langchain.chains.combine_documents import create_stuff_documents_chain from langchain_community.chat_message_histories.streamlit import StreamlitChatMessageHistory from langchain.prompts import PromptTemplate from langchain.vectorstores import Chroma from utills import load_txt_documents, split_docs, load_uploaded_documents, retriever_from_chroma from langchain.text_splitter import TokenTextSplitter, RecursiveCharacterTextSplitter from langchain_community.document_loaders.directory import DirectoryLoader from HTML_templates import css, bot_template, user_template from langchain_core.output_parsers import StrOutputParser from langchain_core.runnables import RunnablePassthrough def create_retriever_from_chroma(vectorstore_path="docs/chroma/", search_type='mmr', k=7, chunk_size=250, chunk_overlap=20): data_path = "data" model_name = "Alibaba-NLP/gte-base-en-v1.5" model_kwargs = {'device': 'cpu', "trust_remote_code" : 'True'} encode_kwargs = {'normalize_embeddings': True} embeddings = HuggingFaceEmbeddings( model_name=model_name, model_kwargs=model_kwargs, encode_kwargs=encode_kwargs ) # Check if vectorstore exists if os.path.exists(vectorstore_path) and os.listdir(vectorstore_path): # Load the existing vectorstore vectorstore = Chroma(persist_directory=vectorstore_path,embedding_function=embeddings) else: # Load documents from the specified data path documents = [] for filename in os.listdir(data_path): if filename.endswith('.txt'): file_path = os.path.join(data_path, filename) loaded_docs = TextLoader(file_path).load() documents.extend(loaded_docs) # Split documents into chunks text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap) split_docs = text_splitter.split_documents(documents) # Ensure the directory for storing vectorstore exists if not os.path.exists(vectorstore_path): os.makedirs(vectorstore_path) # Create the vectorstore vectorstore = Chroma.from_documents( documents=split_docs, embedding=embeddings, persist_directory=vectorstore_path ) # Create and return the retriever retriever = vectorstore.as_retriever(search_type=search_type, search_kwargs={"k": k}) return retriever def main(): st.set_page_config(page_title="Chat with multiple PDFs", page_icon=":books:") st.write(css, unsafe_allow_html=True) st.header("Chat with multiple PDFs :books:") retriever = create_retriever_from_chroma(vectorstore_path="docs/chroma/", search_type='mmr', k=9, chunk_size=250, chunk_overlap=20) user_question = st.text_input("Ask a question about your documents:") if "messages" not in st.session_state: st.session_state["messages"] = [ {"role": "assistant", "content": "Hi, I'm a chatbot who is based on lithuanian law documents. How can I help you?"} ] for msg in st.session_state.messages: st.chat_message(msg["role"]).write(msg["content"]) if user_question: handle_userinput(user_question,retriever) def handle_userinput(user_question,retriever): st.session_state.messages.append({"role": "user", "content": user_question}) st.chat_message("user").write(user_question) docs = retriever.invoke(user_question) with st.sidebar: st.subheader("Your documents") with st.spinner("Processing"): for doc in docs: st.write(f"Document: {doc}") doc_txt = [doc.page_content for doc in docs] rag_chain = create_conversational_rag_chain(retriever) response = rag_chain.invoke({"context": doc_txt, "question": user_question}) st.session_state.messages.append({"role": "assistant", "content": response}) st.chat_message("assistant").write(response) def create_conversational_rag_chain(retriever): model_path = ('qwen2-0_5b-instruct-q4_0.gguf') callback_manager = CallbackManager([StreamingStdOutCallbackHandler()]) llm = llamacpp.LlamaCpp( model_path=model_path, n_gpu_layers=0, temperature=0.0, top_p=0.9, n_ctx=22000, n_batch=2000, max_tokens=200, repeat_penalty=1.7, #last_n_tokens_size = 200, # callback_manager=callback_manager, verbose=False, ) template = """Answer the question based only on the following context: {context} Question: {question} """ prompt = ChatPromptTemplate.from_template(template) rag_chain = prompt | llm | StrOutputParser() return rag_chain if __name__ == "__main__": main()