File size: 5,126 Bytes
fe1fc2e
467c73a
3c4de7d
bd26a11
 
 
 
 
8356c3c
bd26a11
 
 
467c73a
fe1fc2e
644332a
8356c3c
dfb65c1
4f834c9
a76c41e
 
2390875
10edf7a
4b1699f
 
10edf7a
 
 
152b9b0
 
10edf7a
 
 
 
 
152b9b0
 
 
 
02d892a
152b9b0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
852ae92
2bc8ef5
d1e59aa
 
 
 
 
 
152b9b0
d1e59aa
55e669b
 
 
d1e59aa
687dbd6
55e669b
2bc8ef5
687dbd6
c694eae
687dbd6
 
03d617b
 
ed519b4
03d617b
687dbd6
55e669b
4b1699f
a9e2b3b
687dbd6
55e669b
03d617b
9f3b8b8
 
4b1699f
687dbd6
 
3d410fb
2bc8ef5
 
 
922d7c8
3d410fb
687dbd6
 
 
9f3b8b8
fd0bd52
2bc8ef5
05d1ad8
 
fe1fc2e
3c4de7d
c547536
8356c3c
9f3b8b8
8356c3c
 
 
 
 
 
2bc8ef5
8356c3c
 
fe1fc2e
2bc8ef5
 
6c3e7c4
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
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 = "sentence-transformers/all-mpnet-base-v2"
    model_kwargs = {'device': 'cpu'}
    encode_kwargs = {'normalize_embeddings': True}

    # Initialize embeddings
    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:")

    if "messages" not in st.session_state:
        st.session_state["messages"] = [
        {"role": "assistant", "content": "Hi, I'm a chatbot who can search the web. How can I help you?"}
    ]

    for msg in st.session_state.messages:
        st.chat_message(msg["role"]).write(msg["content"])


    user_question = st.text_input("Ask a question about your documents:")

    retriever = create_retriever_from_chroma(vectorstore_path="docs/chroma/", search_type='mmr', k=7, chunk_size=250, chunk_overlap=20)
 
    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)
    
    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=1,
        temperature=0.1,
        top_p=0.9,
        n_ctx=22000,
        max_tokens=200,
        repeat_penalty=1.7,
        # 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()