app.py

import streamlit as st
import requests
from bertopic import BERTopic
from sentence_transformers import SentenceTransformer
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity
import pandas as pd
import plotly.graph_objects as go
from datetime import datetime
import json
from collections import deque
from datasets import load_dataset

class BERTopicChatbot:
    
    #Initialize chatbot with a Hugging Face dataset
    #dataset_name: name of the dataset on Hugging Face (e.g., 'vietnam/legal')
    #text_column: name of the column containing the text data
    #split: which split of the dataset to use ('train', 'test', 'validation')
    #max_samples: maximum number of samples to use (to manage memory)
    
    def __init__(self, dataset_name, text_column, split="train", max_samples=10000):
        # Initialize BERT sentence transformer
        self.sentence_model = SentenceTransformer('all-MiniLM-L6-v2')

        # Add label mapping
        self.label_mapping = {
            0: 'BPD',
            1: 'bipolar',
            2: 'depression',
            3: 'Anxiety',
            4: 'schizophrenia',
            5: 'mentalillness'
        }
        
        # Add comfort responses
        self.comfort_responses = {
            'BPD': [
                "I understand BPD can be overwhelming. You're not alone in this journey.",
                "Your feelings are valid. BPD is challenging, but there are people who understand.",
                "Taking things one day at a time with BPD is okay. You're showing great strength."
            ],
            'bipolar': [
                "Bipolar disorder can feel like a roller coaster. Remember, stability is possible.",
                "You're so strong for managing bipolar disorder. Take it one day at a time.",
                "Both the highs and lows are temporary. You've gotten through them before."
            ],
            'depression': [
                "Depression is heavy, but you don't have to carry it alone.",
                "Even small steps forward are progress. You're doing better than you think.",
                "This feeling won't last forever. You've made it through difficult times before."
            ],
            'Anxiety': [
                "Your anxiety doesn't define you. You're stronger than your fears.",
                "Remember to breathe. You're safe, and this feeling will pass.",
                "It's okay to take things at your own pace. You're handling this well."
            ],
            'schizophrenia': [
                "You're not your diagnosis. You're a person first, and you matter.",
                "Managing schizophrenia takes incredible strength. You're doing well.",
                "There's support available, and you deserve all the help you need."
            ],
            'mentalillness': [
                "Mental health challenges don't define your worth. You are valuable.",
                "Recovery isn't linear, and that's okay. Every step counts.",
                "You're not alone in this journey. There's a community that understands."
            ]
        }

        
        # Load dataset from Hugging Face
        try:
            dataset = load_dataset(dataset_name, split=split)
            # Convert to pandas DataFrame and sample if necessary
            if len(dataset) > max_samples:
                dataset = dataset.shuffle(seed=42).select(range(max_samples))
            
            self.df = dataset.to_pandas()
            
            # Ensure text column exists
            if text_column not in self.df.columns:
                raise ValueError(f"Column '{text_column}' not found in dataset. Available columns: {self.df.columns}")
            
            self.documents = self.df[text_column].tolist()
            
            # Create and train BERTopic model
            self.topic_model = BERTopic(embedding_model=self.sentence_model)
            self.topics, self.probs = self.topic_model.fit_transform(self.documents)
            
            # Create document embeddings for similarity search
            self.doc_embeddings = self.sentence_model.encode(self.documents)
            
            # Initialize metrics storage
            self.metrics_history = {
                'similarities': deque(maxlen=100),
                'response_times': deque(maxlen=100),
                'token_counts': deque(maxlen=100),
                'topics_accessed': {}
            }
            
            # Store dataset info
            self.dataset_info = {
                'name': dataset_name,
                'split': split,
                'total_documents': len(self.documents),
                'topics_found': len(set(self.topics))
            }
            
        except Exception as e:
            st.error(f"Error loading dataset: {str(e)}")
            raise

    def get_metrics_visualizations(self):
        """Generate visualizations for chatbot metrics"""
        # Similarity trend
        fig_similarity = go.Figure()
        fig_similarity.add_trace(go.Scatter(
            y=list(self.metrics_history['similarities']),
            mode='lines+markers',
            name='Similarity Score'
        ))
        fig_similarity.update_layout(
            title='Response Similarity Trend',
            yaxis_title='Similarity Score',
            xaxis_title='Query Number'
        )
    
        # Response time trend
        fig_response_time = go.Figure()
        fig_response_time.add_trace(go.Scatter(
            y=list(self.metrics_history['response_times']),
            mode='lines+markers',
            name='Response Time'
        ))
        fig_response_time.update_layout(
            title='Response Time Trend',
            yaxis_title='Time (seconds)',
            xaxis_title='Query Number'
        )
    
        # Token usage trend
        fig_tokens = go.Figure()
        fig_tokens.add_trace(go.Scatter(
            y=list(self.metrics_history['token_counts']),
            mode='lines+markers',
            name='Token Count'
        ))
        fig_tokens.update_layout(
            title='Token Usage Trend',
            yaxis_title='Number of Tokens',
            xaxis_title='Query Number'
        )
    
        # Topics accessed pie chart
        labels = list(self.metrics_history['topics_accessed'].keys())
        values = list(self.metrics_history['topics_accessed'].values())
        fig_topics = go.Figure(data=[go.Pie(labels=labels, values=values)])
        fig_topics.update_layout(title='Topics Accessed Distribution')
    
        # Make all figures responsive
        for fig in [fig_similarity, fig_response_time, fig_tokens, fig_topics]:
            fig.update_layout(
                autosize=True,
                margin=dict(l=20, r=20, t=40, b=20),
                height=300
            )
    
        return fig_similarity, fig_response_time, fig_tokens, fig_topics
    
    def get_most_similar_document(self, query, top_k=3):
        # Encode the query
        query_embedding = self.sentence_model.encode([query])[0]
        
        # Calculate similarities
        similarities = cosine_similarity([query_embedding], self.doc_embeddings)[0]
        
        # Get top k most similar documents
        top_indices = similarities.argsort()[-top_k:][::-1]
        
        return [self.documents[i] for i in top_indices], similarities[top_indices]

    def get_response(self, user_query):
        try:
            start_time = datetime.now()
            
            # Get most similar documents
            similar_docs, similarities = self.get_most_similar_document(user_query)
            
            # Get the label from the most similar document
            most_similar_index = similarities.argmax()
            label_index = int(self.df['label'].iloc[most_similar_index])  # Convert to int
            condition = self.label_mapping[label_index]  # Map the integer label to condition name
            
            # Get comfort response
            comfort_messages = self.comfort_responses[condition]
            comfort_response = np.random.choice(comfort_messages)
            
            # Calculate query topic for metrics
            query_topic, _ = self.topic_model.transform([user_query])
            
            # Combine information and comfort response
            if max(similarities) < 0.5:
                response = f"I sense you might be dealing with {condition}. {comfort_response}"
            else:
                response = f"{similar_docs[0]}\n\n{comfort_response}"
            
            # Track metrics
            end_time = datetime.now()
            metrics = {
                'similarity': float(max(similarities)),
                'response_time': (end_time - start_time).total_seconds(),
                'tokens': len(response.split()),
                'topic': str(query_topic[0]),
                'detected_condition': condition
            }
            
            # Update metrics history
            self.metrics_history['similarities'].append(metrics['similarity'])
            self.metrics_history['response_times'].append(metrics['response_time'])
            self.metrics_history['token_counts'].append(metrics['tokens'])
            topic_id = str(query_topic[0])
            self.metrics_history['topics_accessed'][topic_id] = \
                self.metrics_history['topics_accessed'].get(topic_id, 0) + 1
            
            return response, metrics
            
        except Exception as e:
            return f"Error processing query: {str(e)}", {'error': str(e)}
    
    def get_dataset_info(self):
        #Return information about the loaded dataset and metrics
        try:
            return {
                'dataset_info': self.dataset_info,
                'metrics': {
                    'avg_similarity': np.mean(list(self.metrics_history['similarities'])) if self.metrics_history['similarities'] else 0,
                    'avg_response_time': np.mean(list(self.metrics_history['response_times'])) if self.metrics_history['response_times'] else 0,
                    'total_tokens': sum(self.metrics_history['token_counts']),
                    'topics_accessed': self.metrics_history['topics_accessed']
                }
            }
        except Exception as e:
            return {
                'error': str(e),
                'dataset_info': None,
                'metrics': None
            }
            
@st.cache_resource
def initialize_chatbot(dataset_name, text_column, split="train", max_samples=10000):
    return BERTopicChatbot(dataset_name, text_column, split, max_samples)

def main():
    st.title("🤖 Trợ Lý AI - BERTopic")
    st.caption("Trò chuyện với chúng mình nhé!")

    # Dataset selection sidebar
    with st.sidebar:
        st.header("Dataset Configuration")
        dataset_name = st.text_input(
            "Hugging Face Dataset Name",
            value="Kanakmi/mental-disorders",
            help="Enter the name of a dataset from Hugging Face (e.g., 'Kanakmi/mental-disorders')"
        )
        text_column = st.text_input(
            "Text Column Name",
            value="text",
            help="Enter the name of the column containing the text data"
        )
        split = st.selectbox(
            "Dataset Split",
            options=["train", "test", "val", "validation"],
            index=0
        )
        max_samples = st.number_input(
            "Maximum Samples",
            min_value=100,
            max_value=100000,
            value=10000,
            step=1000,
            help="Maximum number of samples to load from the dataset"
        )
        
        if st.button("Load Dataset"):
            with st.spinner("Loading dataset and initializing model..."):
                try:
                    st.session_state.chatbot = initialize_chatbot(
                        dataset_name, text_column, split, max_samples
                    )
                    st.success("Dataset loaded successfully!")
                except Exception as e:
                    st.error(f"Error loading dataset: {str(e)}")

    # Initialize session state variables if they don't exist
    if 'chatbot' not in st.session_state:
        st.session_state.chatbot = None
    
    if 'messages' not in st.session_state:
        st.session_state.messages = []

    # Create tabs for chat and metrics
    chat_tab, metrics_tab = st.tabs(["Chat", "Metrics"])

    with chat_tab:
        # Display existing messages
        for message in st.session_state.messages:
            with st.chat_message(message["role"]):
                st.markdown(message["content"])

        # Only show chat input if chatbot is initialized
        if st.session_state.chatbot is not None:
            if prompt := st.chat_input("Hãy nói gì đó..."):
                # Add user message
                st.session_state.messages.append({"role": "user", "content": prompt})
                with st.chat_message("user"):
                    st.markdown(prompt)

                # Get chatbot response
                response, metrics = st.session_state.chatbot.get_response(prompt)

                # Add assistant response
                with st.chat_message("assistant"):
                    st.markdown(response)
                    with st.expander("Response Metrics"):
                        st.json(metrics)
                
                st.session_state.messages.append({"role": "assistant", "content": response})
        else:
            st.info("Please load a dataset first to start chatting.")

    with metrics_tab:
        if st.session_state.chatbot is not None:
            try:
                # Get visualizations from session state chatbot
                fig_similarity, fig_response_time, fig_tokens, fig_topics = st.session_state.chatbot.get_metrics_visualizations()
                
                col1, col2 = st.columns(2)
                with col1:
                    st.plotly_chart(fig_similarity, use_container_width=True)
                    st.plotly_chart(fig_tokens, use_container_width=True)
                
                with col2:
                    st.plotly_chart(fig_response_time, use_container_width=True)
                    st.plotly_chart(fig_topics, use_container_width=True)

                # Display statistics
                st.subheader("Overall Statistics")
                metrics_history = st.session_state.chatbot.metrics_history
                if len(metrics_history['similarities']) > 0:
                    stats_col1, stats_col2, stats_col3 = st.columns(3)
                    with stats_col1:
                        st.metric("Avg Similarity", 
                                 f"{np.mean(list(metrics_history['similarities'])):.3f}")
                    with stats_col2:
                        st.metric("Avg Response Time", 
                                 f"{np.mean(list(metrics_history['response_times'])):.3f}s")
                    with stats_col3:
                        st.metric("Total Tokens Used", 
                                 sum(metrics_history['token_counts']))
                else:
                    st.info("No chat history available yet. Start a conversation to see metrics.")
            except Exception as e:
                st.error(f"Error displaying metrics: {str(e)}")
        else:
            st.info("Please load a dataset first to view metrics.")

if __name__ == "__main__":
    main()