Update app.py

app.py

@@ -1,181 +1,145 @@
+import os
 import streamlit as st
-from transformers import AutoModelForCausalLM, AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
-import time
+from transformers import AutoTokenizer, AutoModel
+import numpy as np
+from sklearn.metrics.pairwise import cosine_similarity
 
-# 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;
-        }
-        
-        .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)
+# Get the port from Heroku environment, default to 8501 for local development
+PORT = int(os.environ.get('PORT', 8501))
 
-# Load model and tokenizer
-@st.cache_resource
-def load_model():
-    # Using VietAI's Vietnamese GPT model
-    model_name = "google-t5/t5-base"
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
-    return model, tokenizer
+class LazyLoadModel:
+    def __init__(self, model_name='intfloat/multilingual-e5-small'):
+        self.model_name = model_name
+        self._tokenizer = None
+        self._model = None
 
-def generate_response(prompt, model, tokenizer, max_length=4169):
-    # 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
-
-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
+    @property
+    def tokenizer(self):
+        if self._tokenizer is None:
+            print("Loading tokenizer...")
+            self._tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+        return self._tokenizer
 
-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)
+    @property
+    def model(self):
+        if self._model is None:
+            print("Loading model...")
+            # Use float16 to reduce memory and potentially speed up loading
+            self._model = AutoModel.from_pretrained(self.model_name, torch_dtype=torch.float16)
+        return self._model
 
-def main():
-    st.set_page_config(
-        page_title="IOGPT",
-        page_icon="🤖",
-        layout="wide"
-    )
+class VietnameseChatbot:
+    def __init__(self):
+        """
+        Initialize the Vietnamese chatbot with lazy-loaded model
+        """
+        self.model_loader = LazyLoadModel()
+        
+        # Very minimal conversation data to reduce startup time
+        self.conversation_data = [
+            {"query": "Xin chào", "response": "Chào bạn!"},
+            {"query": "Bạn là ai?", "response": "Tôi là trợ lý AI."},
+        ]
+    
+    def embed_text(self, text):
+        """
+        Generate embeddings for input text
+        """
+        try:
+            # Tokenize and generate embeddings
+            inputs = self.model_loader.tokenizer(text, return_tensors='pt', padding=True, truncation=True)
+            
+            with torch.no_grad():
+                model_output = self.model_loader.model(**inputs)
+            
+            # Mean pooling
+            embeddings = self.mean_pooling(model_output, inputs['attention_mask'])
+            return embeddings.numpy()
+        except Exception as e:
+            print(f"Embedding error: {e}")
+            return None
     
-    local_css()
-    init_session_state()
+    def mean_pooling(self, model_output, attention_mask):
+        """
+        Perform mean pooling on model output
+        """
+        token_embeddings = model_output[0]
+        input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
+        return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
     
-    # Load model
-    model, tokenizer = load_model()
+    def get_response(self, user_query):
+        """
+        Find the most similar response from conversation data
+        """
+        try:
+            # Embed user query
+            query_embedding = self.embed_text(user_query)
+            
+            if query_embedding is None:
+                return "Xin lỗi, đã có lỗi xảy ra."
+            
+            # Embed conversation data
+            conversation_embeddings = np.array([
+                self.embed_text(item['query'])[0] for item in self.conversation_data
+            ])
+            
+            # Calculate cosine similarities
+            similarities = cosine_similarity(query_embedding, conversation_embeddings)[0]
+            
+            # Find most similar response
+            best_match_index = np.argmax(similarities)
+            
+            # Return response if similarity is above threshold
+            if similarities[best_match_index] > 0.5:
+                return self.conversation_data[best_match_index]['response']
+            
+            return "Xin lỗi, tôi không hiểu câu hỏi của bạn."
+        except Exception as e:
+            print(f"Response generation error: {e}")
+            return "Đã xảy ra lỗi. Xin vui lòng thử lại."
+
+def main():
+    # Server configuration to use Heroku-assigned port
+    if 'PORT' in os.environ:
+        #st.set_option('server.port', PORT)
+        print(f"Server starting on port {PORT}")
+
+    st.title("🤖 Trợ Lý AI Tiếng Việt")
     
-    # Chat interface
-    st.title("IOGPT 🤖")
-    st.markdown("Xin chào! Tôi là trợ lý IOGPT. Hãy hỏi tôi bất cứ điều gì!")
+    # Initialize chatbot
+    chatbot = VietnameseChatbot()
     
-    # Chat history container
-    chat_container = st.container()
+    # Chat history in session state
+    if 'messages' not in st.session_state:
+        st.session_state.messages = []
     
-    # 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")
+    # Display chat messages
+    for message in st.session_state.messages:
+        with st.chat_message(message["role"]):
+            st.markdown(message["content"])
     
-    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
+    # User input
+    if prompt := st.chat_input("Hãy nói gì đó..."):
+        # Add user message to chat history
+        st.session_state.messages.append({"role": "user", "content": prompt})
         
-        # 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
-        ])
+        # Display user message
+        with st.chat_message("user"):
+            st.markdown(prompt)
         
-        # Generate response
-        prompt = f"{conversation_history}\nAssistant:"
-        bot_response = generate_response(prompt, model, tokenizer)
+        # Get chatbot response
+        response = chatbot.get_response(prompt)
         
-        # Add bot response
-        st.session_state.messages.append({"role": "assistant", "content": bot_response})
-        st.session_state.thinking = False
+        # Display chatbot response
+        with st.chat_message("assistant"):
+            st.markdown(response)
         
-        # 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)
+        # Add assistant message to chat history
+        st.session_state.messages.append({"role": "assistant", "content": response})
+
+# Logging for Heroku diagnostics
+print("Chatbot application is initializing...")
 
 if __name__ == "__main__":
     main()