import os import streamlit as st import torch from transformers import AutoTokenizer, AutoModel import numpy as np from sklearn.metrics.pairwise import cosine_similarity import json @st.cache_resource def load_model_and_tokenizer(model_name='intfloat/multilingual-e5-small'): """ Cached function to load model and tokenizer This ensures the model is loaded only once and reused """ print("Loading tokenizer...") tokenizer = AutoTokenizer.from_pretrained(model_name) print("Loading model...") model = AutoModel.from_pretrained(model_name, torch_dtype=torch.float16) return tokenizer, model class VietnameseChatbot: def __init__(self, model_name='intfloat/multilingual-e5-small'): """ Initialize the Vietnamese chatbot with pre-loaded model and conversation data """ # Load pre-trained model and tokenizer using cached function self.tokenizer, self.model = load_model_and_tokenizer(model_name) # Load comprehensive conversation dataset self.conversation_data = self._load_conversation_data() # Pre-compute embeddings for faster response generation print("Pre-computing conversation embeddings...") self.conversation_embeddings = self._compute_embeddings() def _load_conversation_data(self): """ Load a comprehensive conversation dataset """ return [ # Greeting conversations {"query": "Xin chào", "response": "Chào bạn! Tôi có thể giúp gì cho bạn?"}, {"query": "Hi", "response": "Xin chào! Tôi là trợ lý AI tiếng Việt."}, {"query": "Chào buổi sáng", "response": "Chào buổi sáng! Chúc bạn một ngày tốt lành."}, # Identity and purpose {"query": "Bạn là ai?", "response": "Tôi là trợ lý AI được phát triển để hỗ trợ và trò chuyện bằng tiếng Việt."}, {"query": "Bạn từ đâu đến?", "response": "Tôi được phát triển bởi một nhóm kỹ sư AI, và tôn chỉ của tôi là hỗ trợ con người."}, # Small talk {"query": "Bạn thích gì?", "response": "Tôi thích học hỏi và giúp đỡ mọi người. Mỗi cuộc trò chuyện là một cơ hội để tôi phát triển."}, {"query": "Bạn có thể làm gì?", "response": "Tôi có thể trò chuyện, trả lời câu hỏi, và hỗ trợ bạn trong nhiều tình huống khác nhau."}, # Weather and time {"query": "Thời tiết hôm nay thế nào?", "response": "Xin lỗi, tôi không thể cung cấp thông tin thời tiết trực tiếp. Bạn có thể kiểm tra ứng dụng dự báo thời tiết."}, {"query": "Bây giờ là mấy giờ?", "response": "Tôi là trợ lý AI, nên không thể xem đồng hồ. Bạn có thể kiểm tra thiết bị của mình."}, # Assistance offers {"query": "Tôi cần trợ giúp", "response": "Tôi sẵn sàng hỗ trợ bạn. Bạn cần giúp gì?"}, {"query": "Giúp tôi với cái gì đó", "response": "Vâng, tôi có thể hỗ trợ bạn. Hãy cho tôi biết chi tiết hơn."}, # Farewell {"query": "Tạm biệt", "response": "Hẹn gặp lại! Chúc bạn một ngày tốt đẹp."}, {"query": "Bye", "response": "Tạm biệt! Rất vui được trò chuyện với bạn."}, ] @st.cache_data def _compute_embeddings(queries): """ Pre-compute embeddings for conversation queries Cached to avoid recomputing on every run """ def embed_single_text(text, tokenizer, model): try: # Tokenize and generate embeddings inputs = tokenizer(text, return_tensors='pt', padding=True, truncation=True) with torch.no_grad(): model_output = model(**inputs) # Mean pooling token_embeddings = model_output[0] input_mask_expanded = inputs['attention_mask'].unsqueeze(-1).expand(token_embeddings.size()).float() embeddings = torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9) return embeddings.numpy()[0] except Exception as e: print(f"Embedding error: {e}") return None # Import these arguments to make the function self-contained from transformers import AutoTokenizer, AutoModel tokenizer = AutoTokenizer.from_pretrained('intfloat/multilingual-e5-small') model = AutoModel.from_pretrained('intfloat/multilingual-e5-small', torch_dtype=torch.float16) embeddings = [] for query in queries: embedding = embed_single_text(query['query'], tokenizer, model) if embedding is not None: embeddings.append(embedding) return np.array(embeddings) def embed_text(self, text): """ Generate embeddings for input text """ try: # Tokenize and generate embeddings inputs = self.tokenizer(text, return_tensors='pt', padding=True, truncation=True) with torch.no_grad(): model_output = self.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 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) 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 khi phân tích câu hỏi của bạn." # Calculate cosine similarities similarities = cosine_similarity(query_embedding, self.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 chưa hiểu rõ câu hỏi của bạn. Bạn có thể diễn đạt lại được không?" except Exception as e: print(f"Response generation error: {e}") return "Đã xảy ra lỗi. Xin vui lòng thử lại." @st.cache_resource def initialize_chatbot(): """ Cached function to initialize the chatbot This ensures the chatbot is created only once """ return VietnameseChatbot() def main(): st.set_page_config( page_title="Trợ Lý AI Tiếng Việt", page_icon="🤖", ) st.title("🤖 Trợ Lý AI Tiếng Việt") st.caption("Trò chuyện với trợ lý AI được phát triển bằng mô hình đa ngôn ngữ") # Initialize chatbot using cached initialization chatbot = initialize_chatbot() # Chat history in session state if 'messages' not in st.session_state: st.session_state.messages = [] # Sidebar for additional information with st.sidebar: st.header("Về Trợ Lý AI") st.write("Đây là một trợ lý AI được phát triển để hỗ trợ trò chuyện bằng tiếng Việt.") st.write("Mô hình sử dụng: intfloat/multilingual-e5-small") # Display chat messages for message in st.session_state.messages: with st.chat_message(message["role"]): st.markdown(message["content"]) # 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}) # Display user message with st.chat_message("user"): st.markdown(prompt) # Get chatbot response response = chatbot.get_response(prompt) # Display chatbot response with st.chat_message("assistant"): st.markdown(response) # Add assistant message to chat history st.session_state.messages.append({"role": "assistant", "content": response}) if __name__ == "__main__": main()