File size: 6,715 Bytes
1f41ee0
dd53874
fe1fc2e
 
 
 
2353baf
3f669b6
 
fe1fc2e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
929b727
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
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184


import streamlit as st
import os
import sys
import shutil
from langchain_community.text_splitter import TokenTextSplitter,RecursiveCharacterTextSplitter
from langchain.document_loaders import PyPDFLoader
from langchain.document_loaders.pdf import PyPDFDirectoryLoader
from langchain_community.embeddings import HuggingFaceEmbeddings
from transformers import pipeline
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig,AutoModelForSeq2SeqLM
from langchain.retrievers.self_query.base import SelfQueryRetriever
from langchain.chains.query_constructor.base import AttributeInfo
from langchain.retrievers import ContextualCompressionRetriever
from langchain.retrievers.document_compressors import LLMChainExtractor
from langchain_huggingface.llms import HuggingFacePipeline
from langchain.text_splitter import CharacterTextSplitter, RecursiveCharacterTextSplitter
from langchain.vectorstores import DocArrayInMemorySearch
from langchain.document_loaders import TextLoader
from langchain.chains import RetrievalQA,  ConversationalRetrievalChain
from langchain.memory import ConversationBufferMemory
from transformers import DistilBertTokenizer, DistilBertForQuestionAnswering
from langchain.chains import create_history_aware_retriever, create_retrieval_chain
from langchain.chains.combine_documents import create_stuff_documents_chain
import panel as pn
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_core.chat_history import BaseChatMessageHistory
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_community.llms import Aphrodite
from typing import Callable, Dict, List, Optional, Union
from langchain.vectorstores import Chroma
import re
import streamlit as st
from langchain_community.llms import llamacpp
from utills import split_docs, retriever_from_chroma, history_aware_retriever,chroma_db
from langchain_community.chat_message_histories.streamlit import  StreamlitChatMessageHistory






script_dir = os.path.dirname(os.path.abspath(__file__))
data_path = os.path.join(script_dir, "data")
model_path = os.path.join(script_dir, '/mistral-7b-v0.1-layla-v4-Q4_K_M.gguf.2')
store = {}

model_name = "sentence-transformers/all-mpnet-base-v2"
model_kwargs = {'device': 'cpu'}
encode_kwargs = {'normalize_embeddings': True}
hf = HuggingFaceEmbeddings(
    model_name=model_name,
    model_kwargs=model_kwargs,
    encode_kwargs=encode_kwargs)




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, 450, 20)
chroma_db = chroma_db(docs,hf)
retriever = retriever_from_chroma(chroma_db, "mmr", 6)


model_name = "sentence-transformers/all-mpnet-base-v2"
model_kwargs = {'device': 'cpu'}
encode_kwargs = {'normalize_embeddings': True}
hf = HuggingFaceEmbeddings(
    model_name=model_name,
    model_kwargs=model_kwargs,
    encode_kwargs=encode_kwargs
)


callback_manager = CallbackManager([StreamingStdOutCallbackHandler()])

llm = llamacpp.LlamaCpp(
    model_path= model_path,
    n_gpu_layers=0,
    temperature=0.1,
    top_p=0.5,
    n_ctx=7000,
    max_tokens=350,
    repeat_penalty=1.7,
    stop=["", "Instruction:", "### Instruction:", "###<user>", "</user>"],
    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, 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")

conversational_rag_chain = RunnableWithMessageHistory(
    rag_chain,
    lambda session_id: msgs,
    input_messages_key="input",
    history_messages_key="chat_history",
    output_messages_key="answer",
)






def display_chat_history(chat_history):
    """Displays the chat history in Streamlit."""
    for msg in chat_history.messages:
        st.chat_message(msg.type).write(msg.content)

def display_documents(docs, on_click=None):
    """Displays retrieved documents with optional click action."""
    if docs:  # Check if documents exist before displaying
        for i, document in enumerate(docs):  # Iterate over docs, not documents
            st.write(f"**Docs {i+1}**")
            st.markdown(document, unsafe_allow_html=True)  # Allow HTML formatting
            if on_click:
                if st.button(f"Expand Article {i+1}"):
                    on_click(i)  # Call the user-defined click function

def main(conversational_rag_chain):
    """Main function for the Streamlit app."""
    msgs = st.session_state.get("chat_history", StreamlitChatMessageHistory())  # Initialize chat history
    chain_with_history =conversational_rag_chain

    st.title("Conversational RAG Chatbot")

    # Display chat history
    display_chat_history(msgs)

    if prompt := st.chat_input():
        st.chat_message("human").write(prompt)

        # Process user input
        config = {"configurable": {"session_id": "any"}}
        response = chain_with_history.invoke({"question": prompt}, config)
        st.chat_message("ai").write(response.content)

        # Display retrieved documents (if any and present in response)
        if "docs" in response and response["documents"]:
            docs = response["documents"]
            def expand_document(index):
                # Implement your document expansion logic here (e.g., show extra details)
                st.write(f"Expanding document {index+1}...")
            display_documents(docs, expand_document)  # Pass click function

    st.session_state["chat_history"] = msgs  # Update chat history in session state

if __name__ == "__main__":
    main()