Harshdhsvguyt commited on
Commit
487bbdc
·
verified ·
1 Parent(s): 0bc0c18

Create src/streamlit_app.py

Browse files
Files changed (1) hide show
  1. src/streamlit_app.py +135 -0
src/streamlit_app.py ADDED
@@ -0,0 +1,135 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+
3
+ # --- Fix for Hugging Face permission issue ---
4
+ os.environ["STREAMLIT_BROWSER_GATHERUSAGESTATS"] = "false"
5
+ os.environ["STREAMLIT_HOME"] = "/tmp"
6
+ os.makedirs("/tmp/.streamlit", exist_ok=True)
7
+ with open("/tmp/.streamlit/config.toml", "w") as f:
8
+ f.write("[browser]\ngatherUsageStats = false\n[server]\nheadless = true\n")
9
+
10
+ #-----------------------------------------------------------------------------------------------------------------------------------------------------------------------
11
+ ## RAG Q&A Conversation With PDF Including Chat History
12
+ import streamlit as st
13
+ from langchain.chains import create_history_aware_retriever, create_retrieval_chain
14
+ from langchain.chains.combine_documents import create_stuff_documents_chain
15
+ from langchain_chroma import Chroma
16
+ from langchain_community.chat_message_histories import ChatMessageHistory
17
+ from langchain_core.chat_history import BaseChatMessageHistory
18
+ from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
19
+ from langchain_groq import ChatGroq
20
+ from langchain_core.runnables.history import RunnableWithMessageHistory
21
+ from langchain_huggingface import HuggingFaceEmbeddings
22
+ from langchain_text_splitters import RecursiveCharacterTextSplitter
23
+ from langchain_community.document_loaders import PyPDFLoader
24
+ import os
25
+
26
+ from dotenv import load_dotenv
27
+ load_dotenv()
28
+
29
+ os.environ['HF_TOKEN'] = os.getenv("HF_TOKEN")
30
+ embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
31
+
32
+ ## Set up Streamlit
33
+ st.title("Conversational RAG With PDF uploads and chat history")
34
+ st.write("Upload PDFs and chat with their content")
35
+
36
+ ## Input the Groq API Key
37
+ api_key = st.text_input("Enter your Groq API key:", type="password")
38
+
39
+ ## Check if Groq API key is provided
40
+ if api_key:
41
+ llm = ChatGroq(groq_api_key=api_key, model_name="Gemma2-9b-It")
42
+
43
+ ## Chat interface
44
+ session_id = st.text_input("Session ID", value="default_session")
45
+
46
+ ## Statefully manage chat history
47
+ if 'store' not in st.session_state:
48
+ st.session_state.store = {}
49
+
50
+ uploaded_files = st.file_uploader("Choose a PDF file", type="pdf", accept_multiple_files=True)
51
+
52
+ ## Process uploaded PDFs
53
+ if uploaded_files:
54
+ documents = []
55
+ for uploaded_file in uploaded_files:
56
+ temppdf = f"./temp.pdf"
57
+ with open(temppdf, "wb") as file:
58
+ file.write(uploaded_file.getvalue())
59
+ file_name = uploaded_file.name
60
+
61
+ loader = PyPDFLoader(temppdf)
62
+ docs = loader.load()
63
+ documents.extend(docs)
64
+
65
+ # Split and create embeddings for the documents
66
+ text_splitter = RecursiveCharacterTextSplitter(chunk_size=5000, chunk_overlap=500)
67
+ splits = text_splitter.split_documents(documents)
68
+ vectorstore = Chroma.from_documents(documents=splits, embedding=embeddings)
69
+ retriever = vectorstore.as_retriever()
70
+
71
+ contextualize_q_system_prompt = (
72
+ "Given a chat history and the latest user question "
73
+ "which might reference context in the chat history, "
74
+ "formulate a standalone question which can be understood "
75
+ "without the chat history. Do NOT answer the question, "
76
+ "just reformulate it if needed and otherwise return it as is."
77
+ )
78
+ contextualize_q_prompt = ChatPromptTemplate.from_messages(
79
+ [
80
+ ("system", contextualize_q_system_prompt),
81
+ MessagesPlaceholder("chat_history"),
82
+ ("human", "{input}"),
83
+ ]
84
+ )
85
+
86
+ history_aware_retriever = create_history_aware_retriever(llm, retriever, contextualize_q_prompt)
87
+
88
+ ## Answer question
89
+ system_prompt = (
90
+ "You are an assistant for question-answering tasks. "
91
+ "Use the following pieces of retrieved context to answer "
92
+ "the question. If you don't know the answer, say that you "
93
+ "don't know. Use three sentences maximum and keep the "
94
+ "answer concise."
95
+ "\n\n"
96
+ "{context}"
97
+ )
98
+ qa_prompt = ChatPromptTemplate.from_messages(
99
+ [
100
+ ("system", system_prompt),
101
+ MessagesPlaceholder("chat_history"),
102
+ ("human", "{input}"),
103
+ ]
104
+ )
105
+
106
+ question_answer_chain = create_stuff_documents_chain(llm, qa_prompt)
107
+ rag_chain = create_retrieval_chain(history_aware_retriever, question_answer_chain)
108
+
109
+ def get_session_history(session: str) -> BaseChatMessageHistory:
110
+ if session_id not in st.session_state.store:
111
+ st.session_state.store[session_id] = ChatMessageHistory()
112
+ return st.session_state.store[session_id]
113
+
114
+ conversational_rag_chain = RunnableWithMessageHistory(
115
+ rag_chain,
116
+ get_session_history,
117
+ input_messages_key="input",
118
+ history_messages_key="chat_history",
119
+ output_messages_key="answer"
120
+ )
121
+
122
+ user_input = st.text_input("Your question:")
123
+ if user_input:
124
+ session_history = get_session_history(session_id)
125
+ response = conversational_rag_chain.invoke(
126
+ {"input": user_input},
127
+ config={
128
+ "configurable": {"session_id": session_id}
129
+ },
130
+ )
131
+ st.write(st.session_state.store)
132
+ st.write("Assistant:", response['answer'])
133
+ st.write("Chat History:", session_history.messages)
134
+ else:
135
+ st.warning("Please enter the Groq API Key")