|
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 retriever_from_chroma(docs, search_type, k): |
|
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/") |
|
retriever = vectordb.as_retriever(search_type=search_type, search_kwargs={"k": k}) |
|
return retriever |
|
|
|
|
|
|
|
|
|
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, 250, 20) |
|
|
|
retriever = retriever_from_chroma(docs,'mmr',7) |
|
|
|
|
|
|
|
|
|
|
|
|
|
def main(retriever): |
|
|
|
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:") |
|
|
|
with st.chat_message("Assistant"): |
|
st.write("Hello my name is Robert, how can i help you? ") |
|
user_question = st.text_input("Ask a question about your documents:") |
|
with st.chat_message("User"): |
|
st.write(user_question) |
|
if user_question: |
|
handle_userinput(user_question,vectorstore) |
|
|
|
|
|
def handle_userinput(user_question,retriever): |
|
docs = retriever.invoke(question) |
|
|
|
doc_txt = [doc.page_content for doc in docs] |
|
|
|
Rag_chain = create_conversational_rag_chain(retriever) |
|
response = rag_chain.invoke({"context": doc_txt, "question": question}) |
|
with st.chat_message("Assistant"): |
|
st.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=1, |
|
temperature=0.1, |
|
top_p=0.9, |
|
n_ctx=22000, |
|
max_tokens=200, |
|
repeat_penalty=1.7, |
|
|
|
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(vectorstore) |
