Update app.py

app.py

@@ -1,10 +1,9 @@
 import streamlit as st
-from transformers import BertForSequenceClassification, BertTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 import time
-import random
 
-# [Previous CSS styles remain the same]
+# Custom CSS for the chat interface
 def local_css():
     st.markdown("""
         <style>
@@ -38,14 +37,6 @@ def local_css():
             max-width: 70%;
         }
         
-        .chat-input {
-            position: fixed;
-            bottom: 0;
-            width: 100%;
-            padding: 20px;
-            background-color: white;
-        }
-        
         .thinking-animation {
             display: flex;
             align-items: center;
@@ -71,69 +62,38 @@ 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")
-    tokenizer = BertTokenizer.from_pretrained("trituenhantaoio/bert-base-vietnamese-uncased")
+    # Using VietAI's Vietnamese GPT model
+    model_name = "vietai/gpt-neo-1.3B-vietnamese-news"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name)
     return model, tokenizer
 
-def predict(text, model, tokenizer):
-    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
+def generate_response(prompt, model, tokenizer, max_length=100):
+    # Prepare input
+    inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
     
+    # Generate response
     with torch.no_grad():
-        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, confidence
-
-def get_bot_response(text, predicted_class, confidence):
-    # Define response templates based on classes and confidence levels
-    responses = {
-        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'
+        outputs = model.generate(
+            inputs.input_ids,
+            max_length=max_length,
+            num_return_sequences=1,
+            temperature=0.7,
+            top_k=50,
+            top_p=0.95,
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id,
+            attention_mask=inputs.attention_mask
+        )
     
-    # 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
+    # Decode response
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    # Remove the input prompt from the response
+    response = response[len(prompt):].strip()
+    return response
 
 def init_session_state():
     if 'messages' not in st.session_state:
@@ -149,7 +109,12 @@ 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="Vietnamese Chatbot", page_icon="🤖", layout="wide")
+    st.set_page_config(
+        page_title="AI Chatbot Tiếng Việt",
+        page_icon="🤖",
+        layout="wide"
+    )
+    
     local_css()
     init_session_state()
     
@@ -157,8 +122,8 @@ def main():
     model, tokenizer = load_model()
     
     # Chat interface
-    st.title("Chatbot Tiếng Việt 🤖")
-    st.markdown("Xin chào! Tôi có thể giúp gì cho bạn?")
+    st.title("AI Chatbot Tiếng Việt 🤖")
+    st.markdown("Xin chào! Tôi là trợ lý AI có thể trò chuyện bằng tiếng Việt. Hãy hỏi tôi bất cứ điều gì!")
     
     # Chat history container
     chat_container = st.container()
@@ -167,7 +132,11 @@ def main():
     with st.container():
         col1, col2 = st.columns([6, 1])
         with col1:
-            user_input = st.text_input("Nhập tin nhắn của bạn...", 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("Gửi")
     
@@ -178,11 +147,15 @@ def main():
         # Show thinking animation
         st.session_state.thinking = True
         
-        # Get prediction
-        predicted_class, confidence = predict(user_input, model, tokenizer)
+        # Prepare conversation history
+        conversation_history = "\n".join([
+            f"{'User: ' if msg['role'] == 'user' else 'Assistant: '}{msg['content']}"
+            for msg in st.session_state.messages[-3:]  # Last 3 messages for context
+        ])
         
         # Generate response
-        bot_response = get_bot_response(user_input, predicted_class, confidence)
+        prompt = f"{conversation_history}\nAssistant:"
+        bot_response = generate_response(prompt, model, tokenizer)
         
         # Add bot response
         time.sleep(0.5)  # Brief delay for natural feeling