import os import streamlit as st import torch from PIL import Image from transformers import BlipProcessor, BlipForConditionalGeneration from transformers import AutoModelForCausalLM, AutoTokenizer # Cấu hình giao diện Web st.set_page_config(page_title="Lala AI Multimodal Chatbot", layout="wide") # ========================================== # 1. KHỞI TẠO MÔ HÌNH (CACHE TRÊN CPU) # ========================================== @st.cache_resource def load_models(): device = torch.device("cpu") # Model đọc ảnh vision_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base") vision_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base").to(device) # Model ngôn ngữ (Qwen 0.5B tối ưu cho CPU) model_name = "Qwen/Qwen2.5-0.5B-Instruct" chat_tokenizer = AutoTokenizer.from_pretrained(model_name) chat_model = AutoModelForCausalLM.from_pretrained(model_name).to(device) return vision_processor, vision_model, chat_tokenizer, chat_model, device vision_processor, vision_model, chat_tokenizer, chat_model, device = load_models() # ========================================== # 2. HÀM QUÉT THƯ MỤC TỰ ĐỘNG # ========================================== def scan_and_process_directory(folder_path): if not folder_path or not os.path.exists(folder_path): return "Thư mục không tồn tại.", 0, 0 combined_context = "" txt_count = 0 jpg_count = 0 for filename in os.listdir(folder_path): file_path = os.path.join(folder_path, filename) # Đọc file .txt if filename.lower().endswith('.txt'): try: with open(file_path, 'r', encoding='utf-8') as f: content = f.read().strip() combined_context += f"\n[Tài liệu từ file {filename}]:\n{content}\n" txt_count += 1 except Exception as e: combined_context += f"\n[Lỗi đọc file {filename}: {str(e)}]\n" # Phân tích file ảnh .jpg elif filename.lower().endswith(('.jpg', '.jpeg')): try: raw_image = Image.open(file_path).convert('RGB') inputs = vision_processor(raw_image, return_tensors="pt").to(device) out = vision_model.generate(**inputs, max_new_tokens=40) caption = vision_processor.decode(out, skip_special_tokens=True) combined_context += f"\n[Thông tin từ ảnh {filename}]: {caption}\n" jpg_count += 1 except Exception as e: combined_context += f"\n[Lỗi đọc ảnh {filename}: {str(e)}]\n" return combined_context, txt_count, jpg_count # ========================================== # 3. HÀM XỬ LÝ CHATBOX CÓ NHỚ LỊCH SỬ # ========================================== def generate_chat_response(history_messages, folder_context): # Khởi tạo Prompt hệ thống nghiêm ngặt messages = [ { "role": "system", "content": ( "Bạn là trợ lý AI tên Lala. Hãy dựa vào 'Dữ liệu từ thư mục' được cung cấp dưới đây để trả lời câu hỏi. " "Nếu câu hỏi trùng với dữ liệu dạng Q&A trong tài liệu, hãy trả về chính xác nội dung phần A tương ứng. " "Tuyệt đối không tự bịa đặt câu trả lời nằm ngoài dữ liệu." ) }, {"role": "user", "content": f"Dữ liệu từ thư mục:\n{folder_context}\n\nHãy ghi nhớ dữ liệu này để trò chuyện."} ] # Nối thêm toàn bộ lịch sử trò chuyện trước đó vào để AI nhớ ngữ cảnh for msg in history_messages: messages.append({"role": msg["role"], "content": msg["content"]}) # Định dạng prompt text = chat_tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) model_inputs = chat_tokenizer([text], return_tensors="pt").to(device) # Chạy trên CPU với độ chính xác tuyệt đối (temperature=0.0) generated_ids = chat_model.generate( **model_inputs, max_new_tokens=256, temperature=0.0, do_sample=False ) generated_ids = [ output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids) ] return chat_tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0] # ========================================== # 4. GIAO DIỆN WEB STREAMLIT # ========================================== st.title("🤖 Chatbot AI Đa Phương Thức (Lala Voice)") st.caption("Ứng dụng chạy thực tế: Đọc hiểu tệp tin văn bản & Hình ảnh cục bộ.") # Thanh Sidebar cấu hình thư mục dữ liệu with st.sidebar: st.header("📂 Cấu hình Dữ liệu") folder_input = st.text_input("Đường dẫn thư mục nguồn:", value="./data") btn_scan = st.button("🔄 Quét dữ liệu thư mục") if btn_scan or "folder_context" not in st.session_state: with st.spinner("Đang quét dữ liệu..."): context, txt_c, jpg_c = scan_and_process_directory(folder_input) st.session_state["folder_context"] = context st.session_state["txt_count"] = txt_c st.session_state["jpg_count"] = jpg_c st.success("Cập nhật dữ liệu thành công!") st.metric("Số file .txt tìm thấy", st.session_state.get("txt_count", 0)) st.metric("Số file .jpg tìm thấy", st.session_state.get("jpg_count", 0)) # Khởi tạo bộ nhớ lưu lịch sử chat chat nếu chưa có if "messages" not in st.session_state: st.session_state.messages = [] # Nút xóa lịch sử chat để trò chuyện lại từ đầu if st.button("🗑️ Xóa lịch sử trò chuyện"): st.session_state.messages = [] st.rerun() # Hiển thị các tin nhắn cũ trong phiên làm việc for message in st.session_state.messages: with st.chat_message(message["role"]): st.markdown(message["content"]) # Ô nhập chat của người dùng if user_query := st.chat_input("Nhập câu hỏi tại đây..."): # 1. Hiển thị tin nhắn người dùng with st.chat_message("user"): st.markdown(user_query) st.session_state.messages.append({"role": "user", "content": user_query}) # 2. Gọi AI xử lý dựa trên lịch sử câu thoại và file đã quét with st.chat_message("assistant"): with st.spinner("Lala đang xử lý..."): context_data = st.session_state.get("folder_context", "") ai_response = generate_chat_response(st.session_state.messages, context_data) st.markdown(ai_response) # 3. Lưu câu trả lời của AI vào lịch sử bộ nhớ st.session_state.messages.append({"role": "assistant", "content": ai_response})