Update app.py

app.py

@@ -1,109 +1,182 @@
 import streamlit as st
-import os
-from transformers import AutoModelForCausalLM, AutoTokenizer, AutoModel
-from huggingface_hub import login
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
+import time
 
-class VietnameseChatbot:
-    def __init__(self, model_name="tamgrnguyen/Gemma-2-2b-it-Vietnamese-Aesthetic"):
-        """
-        Initialize the Vietnamese chatbot with a pre-trained model
-        """
-        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
-        self.model = AutoModelForCausalLM.from_pretrained(model_name)
+# 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;
+        }
         
-        # Use GPU if available
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        self.model.to(self.device)
-
-    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)
+        .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%;
+        }
         
-        # Tokenize input
-        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
+        .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%;
+        }
         
-        # Generate response
-        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
-        )
+        .thinking-animation {
+            display: flex;
+            align-items: center;
+            margin-left: 10px;
+        }
         
-        # Decode the generated response
-        response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        .dot {
+            width: 8px;
+            height: 8px;
+            margin: 0 3px;
+            background: #888;
+            border-radius: 50%;
+            animation: bounce 0.8s infinite;
+        }
         
-        # Extract only the new part of the response
-        response = response[len(prompt):].strip()
+        .dot:nth-child(2) { animation-delay: 0.2s; }
+        .dot:nth-child(3) { animation-delay: 0.4s; }
         
-        return response
+        @keyframes bounce {
+            0%, 100% { transform: translateY(0); }
+            50% { transform: translateY(-5px); }
+        }
+        </style>
+    """, unsafe_allow_html=True)
 
-def main():
-    st.set_page_config(page_title="IOGPT", layout="wide")
+# Load model and tokenizer
+@st.cache_resource
+def load_model():
+    # Using VietAI's Vietnamese GPT model
+    model_name = "tamgrnguyen/Gemma-2-2b-it-Vietnamese-Aesthetic"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name)
+    return model, tokenizer
 
-    st.title("Chat với IOGPT")
-    st.markdown("""
-    ### Trò chuyện với IOGPT
-    """)
+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
 
-    # 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
+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
 
-    # Initialize conversation history
-    if 'conversation' not in st.session_state:
-        st.session_state.conversation = []
+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="🤖",
+        layout="wide"
+    )
+    
+    local_css()
+    init_session_state()
+    
+    # Load model
+    model, tokenizer = load_model()
+    
     # 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}")
+    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:
-            # Generate AI response
-            with st.spinner('Đang suy nghĩ...'):
-                ai_response = st.session_state.chatbot.generate_response(
-                    st.session_state.conversation
-                )
-            
-            # Add AI response to conversation
-            st.session_state.conversation.append(f"AI: {ai_response}")
+        # Show thinking animation
+        st.session_state.thinking = True
         
-        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
-    """)
+        # 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()
+    
+    # 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)
 
 if __name__ == "__main__":
     main()