Update app.py

app.py

@@ -69,31 +69,46 @@ def load_model():
     model_name = "vietai/gpt-neo-1.3B-vietnamese-news"
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     model = AutoModelForCausalLM.from_pretrained(model_name)
+    
+    # Fix the padding token issue
+    tokenizer.pad_token = tokenizer.eos_token
+    model.config.pad_token_id = model.config.eos_token_id
+    
     return model, tokenizer
 
 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.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
-        )
-    
-    # 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
+    try:
+        # Prepare input
+        inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
+        
+        # Generate response
+        with torch.no_grad():
+            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.pad_token_id,
+                attention_mask=inputs.attention_mask
+            )
+        
+        # Decode response
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        # Remove the input prompt from the response
+        response = response[len(prompt):].strip()
+        
+        # If response is empty, return a default message
+        if not response:
+            return "Xin lỗi, tôi không thể tạo câu trả lời. Bạn có thể hỏi lại không?"
+            
+        return response
+        
+    except Exception as e:
+        st.error(f"Error generating response: {str(e)}")
+        return "Xin lỗi, đã có lỗi xảy ra. Vui lòng thử lại."
 
 def init_session_state():
     if 'messages' not in st.session_state:
@@ -118,65 +133,70 @@ def main():
     local_css()
     init_session_state()
     
-    # Load model
-    model, tokenizer = load_model()
-    
-    # Chat interface
-    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()
-    
-    # Input container
-    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"
-            )
-        with col2:
-            send_button = st.button("Gửi")
-    
-    if user_input and send_button:
-        # Add user message
-        st.session_state.messages.append({"role": "user", "content": user_input})
+    try:
+        # Load model
+        model, tokenizer = load_model()
         
-        # Show thinking animation
-        st.session_state.thinking = True
+        # Chat interface
+        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ì!")
         
-        # 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
-        ])
+        # Chat history container
+        chat_container = st.container()
         
-        # Generate response
-        prompt = f"{conversation_history}\nAssistant:"
-        bot_response = generate_response(prompt, model, tokenizer)
-        
-        # 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
+        # Input container
+        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"
+                )
+            with col2:
+                send_button = st.button("Gửi")
         
-        # Clear input and rerun
-        st.rerun()
-    
-    # Display chat history
-    with chat_container:
-        display_chat_history()
+        if user_input and send_button:
+            # Add user message
+            st.session_state.messages.append({"role": "user", "content": user_input})
+            
+            # Show thinking animation
+            st.session_state.thinking = True
+            
+            # 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
+            prompt = f"{conversation_history}\nAssistant:"
+            bot_response = generate_response(prompt, model, tokenizer)
+            
+            # 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 and rerun
+            st.rerun()
         
-        if st.session_state.thinking:
-            st.markdown("""
-                <div class="thinking-animation">
-                    <div class="dot"></div>
-                    <div class="dot"></div>
-                    <div class="dot"></div>
-                </div>
-            """, unsafe_allow_html=True)
+        # Display chat history
+        with chat_container:
+            display_chat_history()
+            
+            if st.session_state.thinking:
+                st.markdown("""
+                    <div class="thinking-animation">
+                        <div class="dot"></div>
+                        <div class="dot"></div>
+                        <div class="dot"></div>
+                    </div>
+                """, unsafe_allow_html=True)
+                
+    except Exception as e:
+        st.error(f"An error occurred: {str(e)}")
+        st.info("Please refresh the page to try again.")
 
 if __name__ == "__main__":
     main()