File size: 5,460 Bytes
fe1fc2e 467c73a 3c4de7d bd26a11 8356c3c bd26a11 467c73a fe1fc2e 644332a 8356c3c dfb65c1 4f834c9 a76c41e 2390875 10edf7a 4b1699f 8e41c14 10edf7a 152b9b0 02d892a 152b9b0 852ae92 2bc8ef5 d1e59aa 152b9b0 d1e59aa 55e669b d1e59aa 687dbd6 55e669b 2bc8ef5 6a474c7 e709c16 6a474c7 261a9b2 3c62aaf 261a9b2 799176a ce39389 799176a 261a9b2 03d617b 261a9b2 9f3b8b8 4b1699f 687dbd6 3d410fb 8e41c14 2bc8ef5 922d7c8 3d410fb 687dbd6 9e61368 8e41c14 9e61368 687dbd6 9f3b8b8 fd0bd52 2bc8ef5 05d1ad8 fe1fc2e 3c4de7d c547536 8356c3c 9f3b8b8 ab2f443 0984971 8356c3c 8e29761 8356c3c 8e41c14 2bc8ef5 8356c3c fe1fc2e 9513e12 7d01f3a 6c3e7c4 8e41c14 2bc8ef5 fe1fc2e 2bc8ef5 fe1fc2e 81c159a fe8eb6d 5bd324a 815187b 10edf7a 815187b c547536 152b9b0 |
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 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 |
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='similarity', k=7, chunk_size=350, 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() |