DataMine commited on
Commit
e8283e5
·
verified ·
1 Parent(s): 43d00d3

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +142 -143
app.py CHANGED
@@ -1,144 +1,143 @@
1
- import streamlit as st
2
- from langchain.document_loaders import PyPDFLoader
3
- from langchain.text_splitter import RecursiveCharacterTextSplitter
4
- from langchain.vectorstores import Chroma
5
- from langchain.embeddings import HuggingFaceEmbeddings
6
- from transformers import pipeline
7
- import time
8
-
9
-
10
- # Initialize session state for vector database
11
- if 'vector_db' not in st.session_state:
12
- # Load and process PDF only once
13
- @st.cache_resource
14
- def initialize_qa_system():
15
- # Load PDF
16
- pdf_path = "ai_buddy.pdf" # Update path accordingly
17
- data = PyPDFLoader(pdf_path).load()
18
-
19
- # Split documents
20
- text_splitter = RecursiveCharacterTextSplitter(chunk_size=750, chunk_overlap=150)
21
- splits = text_splitter.split_documents(data)
22
-
23
- # Initialize embeddings
24
- embeddings = HuggingFaceEmbeddings(
25
- model_name="sentence-transformers/all-MiniLM-L6-v2"
26
- )
27
-
28
- # Create vector store
29
- vector_db = Chroma.from_documents(
30
- documents=splits,
31
- embedding=embeddings,
32
- persist_directory="vector_db"
33
- )
34
-
35
- # Initialize QA pipeline
36
- qa_pipeline = pipeline("question-answering", model="deepset/roberta-base-squad2")
37
-
38
- return vector_db, qa_pipeline
39
-
40
- # Initialize the system
41
- st.session_state.vector_db, st.session_state.qa_pipeline = initialize_qa_system()
42
-
43
- # Streamlit interface setup
44
- st.set_page_config(page_title="AI-Buddy Assistant", page_icon="🤖", layout="centered")
45
- st.title("🤖 AI-Buddy Assistant")
46
-
47
- # Custom CSS
48
- st.markdown("""
49
- <style>
50
- .stChat {
51
- font-size: 1.3rem;
52
- }
53
- .stTextArea textarea {
54
- font-size: 1.2rem;
55
- }
56
- .stMarkdown {
57
- font-size: 1.2rem;
58
- }
59
- .stButton button {
60
- background-color: #ff4b4b;
61
- color: white;
62
- font-size: 1.2rem;
63
- padding: 10px 20px;
64
- border-radius: 5px;
65
- border: none;
66
- }
67
- .stButton button:hover {
68
- background-color: #ff6565;
69
- }
70
- </style>
71
- """, unsafe_allow_html=True)
72
-
73
- # Sidebar configuration
74
- st.sidebar.title("Want to know how AI helps in your profession?")
75
- professions = ["Software Engineer", "Data Scientist", "Marketing Specialist",
76
- "Financial Analyst", "Teacher", "Doctor", "Project Manager",
77
- "Consultant", "Business Analyst", "Other"]
78
- fields = ["IT", "Healthcare", "Education", "Finance", "Marketing",
79
- "Engineering", "Sales", "Human Resources", "Consulting", "Other"]
80
-
81
- profession = st.sidebar.selectbox("Choose Your Profession", professions)
82
- field = st.sidebar.selectbox("Choose Your Field/Domain", fields)
83
-
84
- if profession == "Other":
85
- profession = st.sidebar.text_input("Please specify your profession")
86
- if field == "Other":
87
- field = st.sidebar.text_input("Please specify your field")
88
-
89
- description = st.sidebar.text_area("About you",
90
- placeholder="Briefly describe your role")
91
-
92
- # Function to execute query with streaming response
93
- def execute_query_with_stream(question):
94
- retriever = st.session_state.vector_db.as_retriever()
95
- retrieved_docs = retriever.get_relevant_documents(question)
96
-
97
- if not retrieved_docs:
98
- return "No relevant information found."
99
-
100
- context = " ".join([doc.page_content.strip() for doc in retrieved_docs])
101
- response = st.session_state.qa_pipeline(question=question, context=context)
102
-
103
- return response['answer']
104
-
105
- # Initialize chat history
106
- if "messages" not in st.session_state:
107
- st.session_state.messages = [
108
- {"role": "assistant", "content": "Hello! How can I assist you today?"}
109
- ]
110
-
111
- # Handle sidebar submission
112
- if st.sidebar.button("Get AI Insights"):
113
- prompt = f"""
114
- My profession is {profession} in the {field} field.
115
- Here's about me: {description}.
116
- How can AI and AI-Buddy help me in my role?
117
- """
118
- st.session_state.messages.append({"role": "user", "content": prompt})
119
-
120
- # Chat interface
121
- if prompt := st.chat_input("Ask me anything about AI-Buddy"):
122
- st.session_state.messages.append({"role": "user", "content": prompt})
123
-
124
- # Display chat history and generate responses
125
- for message in st.session_state.messages:
126
- with st.chat_message(message["role"]):
127
- st.write(message["content"])
128
-
129
- if st.session_state.messages[-1]["role"] == "user":
130
- with st.chat_message("assistant"):
131
- with st.spinner("Thinking..."):
132
- response = execute_query_with_stream(st.session_state.messages[-1]["content"])
133
-
134
- # Simulate streaming effect
135
- placeholder = st.empty()
136
- displayed_response = ""
137
- for char in response:
138
- displayed_response += char
139
- placeholder.write(displayed_response)
140
- time.sleep(0.02)
141
-
142
- st.session_state.messages.append(
143
- {"role": "assistant", "content": response}
144
  )
 
1
+ import streamlit as st
2
+ from langchain.document_loaders import PyPDFLoader
3
+ from langchain.text_splitter import RecursiveCharacterTextSplitter
4
+ from langchain.vectorstores import Chroma
5
+ from langchain.embeddings import HuggingFaceEmbeddings
6
+ from transformers import pipeline
7
+ import time
8
+
9
+ # Set page config as the first Streamlit command
10
+ st.set_page_config(page_title="AI-Buddy Assistant", page_icon="🤖", layout="centered")
11
+
12
+ # Custom CSS
13
+ st.markdown("""
14
+ <style>
15
+ .stChat {
16
+ font-size: 1.3rem;
17
+ }
18
+ .stTextArea textarea {
19
+ font-size: 1.2rem;
20
+ }
21
+ .stMarkdown {
22
+ font-size: 1.2rem;
23
+ }
24
+ .stButton button {
25
+ background-color: #ff4b4b;
26
+ color: white;
27
+ font-size: 1.2rem;
28
+ padding: 10px 20px;
29
+ border-radius: 5px;
30
+ border: none;
31
+ }
32
+ .stButton button:hover {
33
+ background-color: #ff6565;
34
+ }
35
+ </style>
36
+ """, unsafe_allow_html=True)
37
+
38
+ # Cache the QA system initialization
39
+ @st.cache_resource
40
+ def initialize_qa_system():
41
+ # Load PDF
42
+ pdf_path = "ai_buddy.pdf" # Update path accordingly
43
+ data = PyPDFLoader(pdf_path).load()
44
+
45
+ # Split documents
46
+ text_splitter = RecursiveCharacterTextSplitter(chunk_size=750, chunk_overlap=150)
47
+ splits = text_splitter.split_documents(data)
48
+
49
+ # Initialize embeddings
50
+ embeddings = HuggingFaceEmbeddings(
51
+ model_name="sentence-transformers/all-MiniLM-L6-v2"
52
+ )
53
+
54
+ # Create vector store
55
+ vector_db = Chroma.from_documents(
56
+ documents=splits,
57
+ embedding=embeddings,
58
+ persist_directory="vector_db"
59
+ )
60
+
61
+ # Initialize QA pipeline
62
+ qa_pipeline = pipeline("question-answering", model="deepset/roberta-base-squad2")
63
+
64
+ return vector_db, qa_pipeline
65
+
66
+ # Initialize session state for vector database and QA pipeline
67
+ if 'vector_db' not in st.session_state:
68
+ st.session_state.vector_db, st.session_state.qa_pipeline = initialize_qa_system()
69
+
70
+ # Initialize chat history
71
+ if "messages" not in st.session_state:
72
+ st.session_state.messages = [
73
+ {"role": "assistant", "content": "Hello! How can I assist you today?"}
74
+ ]
75
+
76
+ st.title("🤖 AI-Buddy Assistant")
77
+
78
+ # Sidebar configuration
79
+ st.sidebar.title("Want to know how AI helps in your profession?")
80
+ professions = ["Software Engineer", "Data Scientist", "Marketing Specialist",
81
+ "Financial Analyst", "Teacher", "Doctor", "Project Manager",
82
+ "Consultant", "Business Analyst", "Other"]
83
+ fields = ["IT", "Healthcare", "Education", "Finance", "Marketing",
84
+ "Engineering", "Sales", "Human Resources", "Consulting", "Other"]
85
+
86
+ profession = st.sidebar.selectbox("Choose Your Profession", professions)
87
+ field = st.sidebar.selectbox("Choose Your Field/Domain", fields)
88
+
89
+ if profession == "Other":
90
+ profession = st.sidebar.text_input("Please specify your profession")
91
+ if field == "Other":
92
+ field = st.sidebar.text_input("Please specify your field")
93
+
94
+ description = st.sidebar.text_area("About you",
95
+ placeholder="Briefly describe your role")
96
+
97
+ # Function to execute query with streaming response
98
+ def execute_query_with_stream(question):
99
+ retriever = st.session_state.vector_db.as_retriever()
100
+ retrieved_docs = retriever.get_relevant_documents(question)
101
+
102
+ if not retrieved_docs:
103
+ return "No relevant information found."
104
+
105
+ context = " ".join([doc.page_content.strip() for doc in retrieved_docs])
106
+ response = st.session_state.qa_pipeline(question=question, context=context)
107
+
108
+ return response['answer']
109
+
110
+ # Handle sidebar submission
111
+ if st.sidebar.button("Get AI Insights"):
112
+ prompt = f"""
113
+ My profession is {profession} in the {field} field.
114
+ Here's about me: {description}.
115
+ How can AI and AI-Buddy help me in my role?
116
+ """
117
+ st.session_state.messages.append({"role": "user", "content": prompt})
118
+
119
+ # Chat interface
120
+ if prompt := st.chat_input("Ask me anything about AI-Buddy"):
121
+ st.session_state.messages.append({"role": "user", "content": prompt})
122
+
123
+ # Display chat history and generate responses
124
+ for message in st.session_state.messages:
125
+ with st.chat_message(message["role"]):
126
+ st.write(message["content"])
127
+
128
+ if st.session_state.messages[-1]["role"] == "user":
129
+ with st.chat_message("assistant"):
130
+ with st.spinner("Thinking..."):
131
+ response = execute_query_with_stream(st.session_state.messages[-1]["content"])
132
+
133
+ # Simulate streaming effect
134
+ placeholder = st.empty()
135
+ displayed_response = ""
136
+ for char in response:
137
+ displayed_response += char
138
+ placeholder.write(displayed_response)
139
+ time.sleep(0.02)
140
+
141
+ st.session_state.messages.append(
142
+ {"role": "assistant", "content": response}
 
143
  )