Spaces:
Sleeping
Sleeping
File size: 4,918 Bytes
e8283e5 d29989b |
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 |
import streamlit as st
from langchain.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import Chroma
from langchain.embeddings import HuggingFaceEmbeddings
from transformers import pipeline
import time
# Set page config as the first Streamlit command
st.set_page_config(page_title="AI-Buddy Assistant", page_icon="🤖", layout="centered")
# Custom CSS
st.markdown("""
<style>
.stChat {
font-size: 1.3rem;
}
.stTextArea textarea {
font-size: 1.2rem;
}
.stMarkdown {
font-size: 1.2rem;
}
.stButton button {
background-color: #ff4b4b;
color: white;
font-size: 1.2rem;
padding: 10px 20px;
border-radius: 5px;
border: none;
}
.stButton button:hover {
background-color: #ff6565;
}
</style>
""", unsafe_allow_html=True)
# Cache the QA system initialization
@st.cache_resource
def initialize_qa_system():
# Load PDF
pdf_path = "ai_buddy.pdf" # Update path accordingly
data = PyPDFLoader(pdf_path).load()
# Split documents
text_splitter = RecursiveCharacterTextSplitter(chunk_size=750, chunk_overlap=150)
splits = text_splitter.split_documents(data)
# Initialize embeddings
embeddings = HuggingFaceEmbeddings(
model_name="sentence-transformers/all-MiniLM-L6-v2"
)
# Create vector store
vector_db = Chroma.from_documents(
documents=splits,
embedding=embeddings,
persist_directory="vector_db"
)
# Initialize QA pipeline
qa_pipeline = pipeline("question-answering", model="deepset/roberta-base-squad2")
return vector_db, qa_pipeline
# Initialize session state for vector database and QA pipeline
if 'vector_db' not in st.session_state:
st.session_state.vector_db, st.session_state.qa_pipeline = initialize_qa_system()
# Initialize chat history
if "messages" not in st.session_state:
st.session_state.messages = [
{"role": "assistant", "content": "Hello! How can I assist you today?"}
]
st.title("🤖 AI-Buddy Assistant")
# Sidebar configuration
st.sidebar.title("Want to know how AI helps in your profession?")
professions = ["Software Engineer", "Data Scientist", "Marketing Specialist",
"Financial Analyst", "Teacher", "Doctor", "Project Manager",
"Consultant", "Business Analyst", "Other"]
fields = ["IT", "Healthcare", "Education", "Finance", "Marketing",
"Engineering", "Sales", "Human Resources", "Consulting", "Other"]
profession = st.sidebar.selectbox("Choose Your Profession", professions)
field = st.sidebar.selectbox("Choose Your Field/Domain", fields)
if profession == "Other":
profession = st.sidebar.text_input("Please specify your profession")
if field == "Other":
field = st.sidebar.text_input("Please specify your field")
description = st.sidebar.text_area("About you",
placeholder="Briefly describe your role")
# Function to execute query with streaming response
def execute_query_with_stream(question):
retriever = st.session_state.vector_db.as_retriever()
retrieved_docs = retriever.get_relevant_documents(question)
if not retrieved_docs:
return "No relevant information found."
context = " ".join([doc.page_content.strip() for doc in retrieved_docs])
response = st.session_state.qa_pipeline(question=question, context=context)
return response['answer']
# Handle sidebar submission
if st.sidebar.button("Get AI Insights"):
prompt = f"""
My profession is {profession} in the {field} field.
Here's about me: {description}.
How can AI and AI-Buddy help me in my role?
"""
st.session_state.messages.append({"role": "user", "content": prompt})
# Chat interface
if prompt := st.chat_input("Ask me anything about AI-Buddy"):
st.session_state.messages.append({"role": "user", "content": prompt})
# Display chat history and generate responses
for message in st.session_state.messages:
with st.chat_message(message["role"]):
st.write(message["content"])
if st.session_state.messages[-1]["role"] == "user":
with st.chat_message("assistant"):
with st.spinner("Thinking..."):
response = execute_query_with_stream(st.session_state.messages[-1]["content"])
# Simulate streaming effect
placeholder = st.empty()
displayed_response = ""
for char in response:
displayed_response += char
placeholder.write(displayed_response)
time.sleep(0.02)
st.session_state.messages.append(
{"role": "assistant", "content": response}
) |