|
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 |
|
|
|
|
|
|
|
|
|
def get_vectorstore(text_chunks): |
|
model_name = "sentence-transformers/all-mpnet-base-v2" |
|
model_kwargs = {'device': 'cpu'} |
|
encode_kwargs = {'normalize_embeddings': True} |
|
embeddings = HuggingFaceEmbeddings( |
|
model_name=model_name, |
|
model_kwargs=model_kwargs, |
|
encode_kwargs=encode_kwargs |
|
) |
|
vectorstore_path = "docs/chroma/" |
|
if not os.path.exists(vectorstore_path): |
|
os.makedirs(vectorstore_path) |
|
vectorstore = Chroma.from_documents( |
|
documents=text_chunks, embedding=embeddings, persist_directory="docs/chroma/") |
|
return vectorstore |
|
|
|
|
|
|
|
data_path = "data" |
|
|
|
|
|
|
|
documents = [] |
|
|
|
for filename in os.listdir(data_path): |
|
|
|
if filename.endswith('.txt'): |
|
|
|
file_path = os.path.join(data_path, filename) |
|
|
|
documents = TextLoader(file_path).load() |
|
|
|
documents.extend(documents) |
|
|
|
|
|
|
|
|
|
docs = split_docs(documents, 350, 40) |
|
|
|
vectorstore = get_vectorstore(docs) |
|
|
|
|
|
|
|
|
|
|
|
|
|
def main(): |
|
|
|
st.set_page_config(page_title="Chat with multiple PDFs", |
|
page_icon=":books:") |
|
st.write(css, unsafe_allow_html=True) |
|
|
|
if "conversation" not in st.session_state: |
|
st.session_state.conversation = None |
|
if "chat_history" not in st.session_state: |
|
st.session_state.chat_history = None |
|
|
|
st.header("Chat with multiple PDFs :books:") |
|
user_question = st.text_input("Ask a question about your documents:") |
|
if user_question: |
|
handle_userinput(user_question) |
|
|
|
|
|
def handle_userinput(user_question,vectorstore): |
|
Rag_chain = create_conversational_rag_chain() |
|
response = Rag_chain.invoke( |
|
{"input": user_question "}, |
|
|
|
config={ |
|
"configurable": {"session_id": "k"} |
|
}, |
|
)["answer"] |
|
st.session_state.chat_history = response['chat_history'] |
|
|
|
for i, message in enumerate(st.session_state.chat_history): |
|
if i % 2 == 0: |
|
st.write(user_template.replace( |
|
"{{MSG}}", message.content), unsafe_allow_html=True) |
|
else: |
|
st.write(bot_template.replace( |
|
"{{MSG}}", message.content), unsafe_allow_html=True) |
|
|
|
if 'retrieved_documents' in response: |
|
st.subheader("Retrieved Documents") |
|
for doc in response['source_documents']: |
|
st.write(f"Document: {doc.metadata['source']}") |
|
st.write(doc.page_content) |
|
|
|
|
|
|
|
def create_conversational_rag_chain(vectorstore): |
|
|
|
model_path = ('qwen2-0_5b-instruct-q4_0.gguf') |
|
|
|
callback_manager = CallbackManager([StreamingStdOutCallbackHandler()]) |
|
|
|
llm = llamacpp.LlamaCpp( |
|
model_path=model_path, |
|
n_gpu_layers=1, |
|
temperature=0.1, |
|
top_p=0.9, |
|
n_ctx=22000, |
|
max_tokens=200, |
|
repeat_penalty=1.7, |
|
callback_manager=callback_manager, |
|
verbose=False, |
|
) |
|
|
|
contextualize_q_system_prompt = """Given a context, chat history and the latest user question |
|
which maybe reference context in the chat history, formulate a standalone question |
|
which can be understood without the chat history. Do NOT answer the question, |
|
just reformulate it if needed and otherwise return it as is.""" |
|
|
|
ha_retriever = history_aware_retriever(llm, vectorstore.as_retriever(), contextualize_q_system_prompt) |
|
|
|
qa_system_prompt = """You are an assistant for question-answering tasks. Use the following pieces of retrieved context to answer the question. If you don't know the answer, just say that you don't know. Be as informative as possible, be polite and formal.\n{context}""" |
|
|
|
qa_prompt = ChatPromptTemplate.from_messages( |
|
[ |
|
("system", qa_system_prompt), |
|
MessagesPlaceholder("chat_history"), |
|
("human", "{input}"), |
|
] |
|
) |
|
|
|
question_answer_chain = create_stuff_documents_chain(llm, qa_prompt) |
|
|
|
rag_chain = create_retrieval_chain(ha_retriever, question_answer_chain) |
|
msgs = StreamlitChatMessageHistory(key="special_app_key") |
|
|
|
return rag_chain |
|
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
main() |