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}
            )