Updated BERT models.

app.py

@@ -2,9 +2,9 @@ import streamlit as st
 from transformers import BertForSequenceClassification, BertTokenizer
 import torch
 import time
-import streamlit.components.v1 as components
+import random
 
-# Custom CSS for chat interface
+# [Previous CSS styles remain the same]
 def local_css():
     st.markdown("""
         <style>
@@ -71,7 +71,6 @@ def local_css():
         </style>
     """, unsafe_allow_html=True)
 
-# Load model and tokenizer
 @st.cache_resource
 def load_model():
     model = BertForSequenceClassification.from_pretrained("trituenhantaoio/bert-base-vietnamese-uncased")
@@ -85,19 +84,56 @@ def predict(text, model, tokenizer):
         outputs = model(**inputs)
         predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
         predicted_class = torch.argmax(predictions, dim=1).item()
+        confidence = predictions[0][predicted_class].item()
     
-    return predicted_class, predictions[0]
+    return predicted_class, confidence
 
-def get_bot_response(predicted_class, confidence):
-    # Customize these responses based on your model's classes
+def get_bot_response(text, predicted_class, confidence):
+    # Define response templates based on classes and confidence levels
     responses = {
-        0: "I understand this is about [Class 0]. Let me help you with that.",
-        1: "This seems to be related to [Class 1]. Here's what I can tell you.",
-        # Add more responses for your classes
+        0: {  # Example for class 0 (positive sentiment)
+            'high_conf': [
+                "Tôi cảm nhận được sự tích cực trong câu nói của bạn. Xin chia sẻ thêm nhé!",
+                "Thật vui khi nghe điều đó. Bạn có thể kể thêm không?",
+                "Tuyệt vời! Tôi rất đồng ý với bạn về điều này."
+            ],
+            'low_conf': [
+                "Có vẻ như đây là điều tích cực. Đúng không nhỉ?",
+                "Tôi nghĩ đây là một góc nhìn thú vị đấy.",
+                "Nghe có vẻ tốt đấy, bạn nghĩ sao?"
+            ]
+        },
+        1: {  # Example for class 1 (negative sentiment)
+            'high_conf': [
+                "Tôi hiểu đây là điều khó khăn với bạn. Hãy chia sẻ thêm nhé.",
+                "Tôi rất tiếc khi nghe điều này. Bạn cần tôi giúp gì không?",
+                "Đúng là một tình huống khó khăn. Chúng ta cùng tìm giải pháp nhé."
+            ],
+            'low_conf': [
+                "Có vẻ như bạn đang gặp khó khăn. Tôi có hiểu đúng không?",
+                "Tôi không chắc mình hiểu hết, bạn có thể giải thích thêm được không?",
+                "Hãy chia sẻ thêm để tôi có thể hiểu rõ hơn nhé."
+            ]
+        }
     }
+
+    # Add more classes based on your model's output
+    
+    # Determine confidence level
+    confidence_threshold = 0.8
+    conf_level = 'high_conf' if confidence > confidence_threshold else 'low_conf'
     
-    default_response = "I'm not quite sure about that. Could you please rephrase?"
-    return responses.get(predicted_class, default_response)
+    # Get appropriate response list
+    try:
+        response_list = responses[predicted_class][conf_level]
+        response = random.choice(response_list)
+    except KeyError:
+        response = "Xin lỗi, tôi không chắc chắn về điều này. Bạn có thể giải thích rõ hơn được không?"
+
+    # Add context from user's input
+    context_response = f"{response}"
+    
+    return context_response
 
 def init_session_state():
     if 'messages' not in st.session_state:
@@ -113,7 +149,7 @@ def display_chat_history():
             st.markdown(f'<div class="bot-message">{message["content"]}</div>', unsafe_allow_html=True)
 
 def main():
-    st.set_page_config(page_title="AI Chatbot", page_icon="🤖", layout="wide")
+    st.set_page_config(page_title="Vietnamese Chatbot", page_icon="🤖", layout="wide")
     local_css()
     init_session_state()
     
@@ -121,46 +157,45 @@ def main():
     model, tokenizer = load_model()
     
     # Chat interface
-    st.title("AI Chatbot 🤖")
-    st.markdown("Welcome! I'm here to help answer your questions in Vietnamese.")
+    st.title("Chatbot Tiếng Việt 🤖")
+    st.markdown("Xin chào! Tôi có thể giúp gì cho bạn?")
     
     # Chat history container
     chat_container = st.container()
     
-    # Input container at the bottom
+    # Input container
     with st.container():
         col1, col2 = st.columns([6, 1])
         with col1:
-            user_input = st.text_input("Type your message...", key="user_input", label_visibility="hidden")
+            user_input = st.text_input("Nhập tin nhắn của bạn...", key="user_input", label_visibility="hidden")
         with col2:
-            send_button = st.button("Send")
+            send_button = st.button("Gửi")
     
     if user_input and send_button:
-        # Add user message to chat
+        # Add user message
         st.session_state.messages.append({"role": "user", "content": user_input})
         
         # Show thinking animation
         st.session_state.thinking = True
         
-        # Get model prediction
-        predicted_class, probabilities = predict(user_input, model, tokenizer)
+        # Get prediction
+        predicted_class, confidence = predict(user_input, model, tokenizer)
         
-        # Get bot response
-        bot_response = get_bot_response(predicted_class, probabilities)
+        # Generate response
+        bot_response = get_bot_response(user_input, predicted_class, confidence)
         
-        # Add bot response to chat
-        time.sleep(1)  # Simulate processing time
+        # Add bot response
+        time.sleep(0.5)  # Brief delay for natural feeling
         st.session_state.messages.append({"role": "assistant", "content": bot_response})
         st.session_state.thinking = False
         
-        # Clear input
+        # Clear input and rerun
         st.rerun()
     
     # Display chat history
     with chat_container:
         display_chat_history()
         
-        # Show thinking animation
         if st.session_state.thinking:
             st.markdown("""
                 <div class="thinking-animation">