Update app.py

app.py

@@ -1,202 +1,111 @@
 import streamlit as st
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
-import time
 
-# Custom CSS for the chat interface
-def local_css():
-    st.markdown("""
-        <style>
-        .chat-container {
-            padding: 10px;
-            border-radius: 5px;
-            margin-bottom: 10px;
-            display: flex;
-            flex-direction: column;
-        }
-        
-        .user-message {
-            background-color: #e3f2fd;
-            padding: 10px;
-            border-radius: 15px;
-            margin: 5px;
-            margin-left: 20%;
-            margin-right: 5px;
-            align-self: flex-end;
-            max-width: 70%;
-        }
-        
-        .bot-message {
-            background-color: #f5f5f5;
-            padding: 10px;
-            border-radius: 15px;
-            margin: 5px;
-            margin-right: 20%;
-            margin-left: 5px;
-            align-self: flex-start;
-            max-width: 70%;
-        }
-        
-        .thinking-animation {
-            display: flex;
-            align-items: center;
-            margin-left: 10px;
-        }
+class VietnameseChatbot:
+    def __init__(self, model_name="vinai/gpt-neo-1.3B-vietnamese-news"):
+        """
+        Initialize the Vietnamese chatbot with a pre-trained model
+        """
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.model = AutoModelForCausalLM.from_pretrained(model_name)
         
-        .dot {
-            width: 8px;
-            height: 8px;
-            margin: 0 3px;
-            background: #888;
-            border-radius: 50%;
-            animation: bounce 0.8s infinite;
-        }
-        
-        .dot:nth-child(2) { animation-delay: 0.2s; }
-        .dot:nth-child(3) { animation-delay: 0.4s; }
-        
-        @keyframes bounce {
-            0%, 100% { transform: translateY(0); }
-            50% { transform: translateY(-5px); }
-        }
-        </style>
-    """, unsafe_allow_html=True)
+        # Use GPU if available
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.model.to(self.device)
 
-# Load model and tokenizer
-@st.cache_resource
-def load_model():
-    # 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)
-    
-    # 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):
-    try:
-        # Prepare input
-        inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
+    def generate_response(self, conversation_history, max_length=200):
+        """
+        Generate a response based on conversation history
+        """
+        # Combine conversation history into a single prompt
+        prompt = "\n".join(conversation_history)
+        
+        # Tokenize input
+        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
         
         # 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
-            )
+        outputs = self.model.generate(
+            inputs.input_ids, 
+            max_length=max_length,
+            num_return_sequences=1,
+            no_repeat_ngram_size=2,
+            temperature=0.7,
+            top_k=50,
+            top_p=0.95
+        )
         
-        # Decode response
-        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        # Remove the input prompt from the response
+        # Decode the generated response
+        response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        # Extract only the new part of 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:
-        st.session_state.messages = []
-    if 'thinking' not in st.session_state:
-        st.session_state.thinking = False
-
-def display_chat_history():
-    for message in st.session_state.messages:
-        if message['role'] == 'user':
-            st.markdown(f'<div class="user-message">{message["content"]}</div>', unsafe_allow_html=True)
-        else:
-            st.markdown(f'<div class="bot-message">{message["content"]}</div>', unsafe_allow_html=True)
 
 def main():
     st.set_page_config(
-        page_title="AI Chatbot Tiếng Việt",
-        page_icon="🤖",
+            page_title="IOGPT", 
+        page_icon="🇻🇳",
         layout="wide"
     )
-    
-    local_css()
-    init_session_state()
-    
-    try:
-        # 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()
+
+    st.title("🇻🇳 Chat với IOGPT")
+    st.markdown("""
+    ### Trò chuyện với IOGPT
+    """)
+
+    # Initialize chatbot
+    if 'chatbot' not in st.session_state:
+        try:
+            st.session_state.chatbot = VietnameseChatbot()
+        except Exception as e:
+            st.error(f"Lỗi khởi tạo mô hình: {e}")
+            return
+
+    # Initialize conversation history
+    if 'conversation' not in st.session_state:
+        st.session_state.conversation = []
+
+    # Chat interface
+    with st.form(key='chat_form'):
+        user_input = st.text_input("Nhập tin nhắn của bạn:", placeholder="Viết tin nhắn...")
+        send_button = st.form_submit_button("Gửi")
+
+    # Process user input
+    if send_button and user_input:
+        # Add user message to conversation
+        st.session_state.conversation.append(f"User: {user_input}")
         
-        # 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"
+        try:
+            # Generate AI response
+            with st.spinner('Đang suy nghĩ...'):
+                ai_response = st.session_state.chatbot.generate_response(
+                    st.session_state.conversation
                 )
-            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})
-            
-            # 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()
+            # Add AI response to conversation
+            st.session_state.conversation.append(f"AI: {ai_response}")
         
-        # 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.")
+        except Exception as e:
+            st.error(f"Lỗi trong quá trình trò chuyện: {e}")
+
+    # Display conversation history
+    st.subheader("Lịch sử trò chuyện")
+    for msg in st.session_state.conversation:
+        if msg.startswith("User:"):
+            st.markdown(f"**{msg}**")
+        else:
+            st.markdown(f"*{msg}*")
+
+    # Model and usage information
+    st.sidebar.header("Thông tin mô hình")
+    st.sidebar.info("""
+    ### Mô hình AI Tiếng Việt
+    - Được huấn luyện trên dữ liệu tiếng Việt
+    - Hỗ trợ trò chuyện đa dạng
+    - Lưu ý: Chất lượng trả lời phụ thuộc vào mô hình
+    """)
 
 if __name__ == "__main__":
     main()