Update app.py

app.py

@@ -1,135 +1,264 @@
-# app.py (Sử dụng model.generate())
-
 import gradio as gr
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-import warnings
-import spaces
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import spaces , os
+from quickmt import Translator
+from quickmt.hub import hf_download, hf_list
+from pathlib import Path
 
-# Tắt cảnh báo về việc chuyển đổi tensor
-warnings.filterwarnings("ignore", category=UserWarning, module="transformers")
 
-# --- Cấu hình Mô hình và Tokenizer ---
 MODEL_ID = "vicgalle/Humanish-Roleplay-Llama-3.1-8B"
+# Load model and tokenizer
 
-# Tải Tokenizer
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 
-# Tải Mô hình: Vẫn giữ nguyên device_map="auto" để accelerate quản lý GPU/CPU
-try:
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_ID,
-        torch_dtype=torch.bfloat16,
-        device_map="auto",
-    )
-except:
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_ID,
-        torch_dtype=torch.float16,
-        device_map="auto",
+# Khởi tạo biến toàn cục t ngoài hàm, như trong code gốc
+t = None 
+
+model_name_or_path = "tencent/Hunyuan-MT-7B"
+print("Loading model... This may take a few minutes.")
+
+tokenizer_trans = AutoTokenizer.from_pretrained(model_name_or_path)
+model_trans = AutoModelForCausalLM.from_pretrained(
+    model_name_or_path,
+    torch_dtype=torch.bfloat16,
+    device_map="auto"
+)
+
+# --- Sửa đổi hàm translate_text ---
+# Thêm tham số model_name (mặc định là "quickmt-en-vi")
+@spaces.GPU(duration=60)
+def translate_text(text, lang=None, progress=gr.Progress(track_tqdm=True)):
+  # Set default values if None (happens during example caching)
+    if lang is None:
+        return text
+    system_message = "You are a helpful AI assistant."
+    max_tokens = 1024
+    temperature = 0.1
+    top_p = 0.05
+    
+    # Build conversation history
+    messages = []
+    message = "Translate to "+lang+": "+text
+    print("message : "+message)
+    # Add system message if provided
+    if system_message:
+        messages.append({"role": "system", "content": system_message})
+    
+    # Add current message
+    messages.append({"role": "user", "content": message})
+    
+    # Tokenize the conversation
+    tokenized_chat = tokenizer_trans.apply_chat_template(
+        messages,
+        tokenize=True,
+        add_generation_prompt=True,
+        return_tensors="pt"
     )
+    
+    # Generate response
+    with torch.no_grad():
+        outputs = model_trans.generate(
+            tokenized_chat.to(model.device),
+            max_new_tokens=max_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            do_sample=True if temperature > 0 else False,
+            pad_token_id=tokenizer_trans.eos_token_id
+        )
+    
+    # Decode only the new tokens
+    response = tokenizer_trans.decode(outputs[0][tokenized_chat.shape[-1]:], skip_special_tokens=True)
+    
+    return response
+        
+# Tải model và tokenizer 1 LẦN DUY NHẤT
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    torch_dtype=torch.float16,
+    device_map="auto", # Tự động dùng GPU cố định
+    trust_remote_code=True
+)
 
-# --- Hàm Tạo Văn bản MỚI với model.generate() ---
+
+
+# Thêm một ô system_prompt
 @spaces.GPU(duration=60)
-def generate_text(system_prompt, prompt, max_new_tokens, temperature, top_k, top_p):
+def chat_with_model(prompt, system_prompt, chatbot_display, internal_history,lang,gender,progress=gr.Progress(track_tqdm=True)):
     """
-    Sử dụng model.generate() để tạo văn bản dựa trên prompt và các tham số.
+    Hàm này nhận prompt mới, system_prompt, lịch sử hiển thị (của gr.Chatbot)
+    và lịch sử nội bộ (của gr.State).
     """
+    expected_key = os.environ.get("hf_key")
+    if expected_key not in prompt:
+        print("❌ Invalid key.")
+        return "", chatbot_display, internal_history
+    prompt = prompt.replace(expected_key, "")
+    isAuto = False
+    if "[AUTO]" in prompt:
+        prompt = prompt.replace("[AUTO]", "")
+        isAuto = True
+    else:
+        if lang != None:
+            prompt = translate_text(prompt,"English")
+
+    prompt = prompt +" [Detailed description of the physical actions and expressions.]"
+    print("prompt : "+prompt)
+    # 1. Khởi tạo nếu đây là lần chạy đầu tiên
+    # chatbot_display là [[user_msg, ai_msg], ...]
+    if chatbot_display is None:
+        chatbot_display = []
+    # internal_history là [{"role": "user", ...}, {"role": "assistant", ...}]
+    if internal_history is None:
+        internal_history = []
+
+    # 2. Xây dựng toàn bộ lịch sử để đưa cho model
+    # Bắt đầu với System Prompt (luôn lấy cái mới nhất từ Textbox)
+    messages_for_model = [{"role": "system", "content": system_prompt}]
     
-    # 1. Định dạng tin nhắn (System + User)
-    messages = [
-        {"role": "system", "content": system_prompt}, 
-        {"role": "user", "content": prompt}
-    ]
+    # Thêm toàn bộ các lượt nói cũ (user/assistant) từ "bộ nhớ" gr.State
+    messages_for_model.extend(internal_history)
     
-    # 2. Áp dụng template và chuyển thành chuỗi prompt
-    prompt_str = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-
-    # 3. Mã hóa (Encoding): Chuyển chuỗi thành Tensor và đặt nó trên thiết bị của mô hình
-    # Sử dụng model.device để đảm bảo tensor nằm trên cùng thiết bị (GPU/CPU)
-    input_ids = tokenizer.encode(prompt_str, return_tensors="pt").to(model.device)
-
-    # Cấu hình các tham số tạo văn bản
-    generation_args = {
-        "max_new_tokens": max_new_tokens,
-        "temperature": temperature,
-        "top_k": top_k,
-        "top_p": top_p,
-        "do_sample": True,
-        "pad_token_id": tokenizer.eos_token_id,
-        "eos_token_id": tokenizer.eos_token_id,
-    }
-
-    try:
-        # 4. Tạo Văn bản (Generation): Dùng model.generate()
-        output_ids = model.generate(
-            input_ids,
-            **generation_args
-        )
-        
-        # 5. Giải mã (Decoding): Chuyển Tensor đầu ra thành chuỗi văn bản
-        generated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-        
-        # 6. Hậu xử lý: Loại bỏ prompt ban đầu khỏi kết quả
-        # Llama template bao gồm cả prompt trong output_ids, nên cần cắt bỏ
-        if prompt_str in generated_text:
-            # Tìm vị trí của chuỗi prompt và chỉ lấy phần sau đó
-            response = generated_text.split(prompt_str)[-1].strip()
+    # Thêm prompt MỚI của người dùng
+    messages_for_model.append({"role": "user", "content": prompt})
+
+    # 3. Áp dụng Chat Template
+    inputs = tokenizer.apply_chat_template(
+        messages_for_model,
+        tokenize=True,
+        add_generation_prompt=True,
+        return_tensors="pt"
+    ).to(model.device)
+
+    # 4. Generate
+    output_tokens = model.generate(
+        inputs,
+        max_new_tokens=5120, # Tăng số token tối đa lên một chút
+        do_sample=True,
+        temperature=0.7,
+        top_p=0.9
+    )
+    
+    # 5. Decode *chỉ* phần trả lời mới
+    response_text = tokenizer.decode(output_tokens[0][inputs.shape[-1]:], skip_special_tokens=True)
+    print("response_text : "+response_text)
+    translated = response_text
+    if lang != None :
+        if gender :
+            translated =  translate_text(response_text+"; Mr.",lang)
         else:
-            response = generated_text # Trường hợp dự phòng
+            translated =  translate_text(response_text +"; Ms.",lang)
 
-        return response
+    print("translated : "+translated)    
+    # 6. Cập nhật "bộ nhớ" (gr.State) với lượt nói MỚI
+    internal_history.append({"role": "user", "content": prompt})
+    internal_history.append({"role": "assistant", "content": response_text})
     
-    except Exception as e:
-        # Trong trường hợp có lỗi (ví dụ: OOM - Out of Memory trên GPU nhỏ)
-        return f"Đã xảy ra lỗi trong quá trình tạo văn bản: {e}"
-
-# --- Giao diện Gradio (Giữ nguyên) ---
-
-# Định nghĩa các tham số điều khiển
-MAX_TOKENS = gr.Slider(minimum=10, maximum=2048, value=512, step=10, label="Max New Tokens (Độ dài câu trả lời)")
-TEMPERATURE = gr.Slider(minimum=0.1, maximum=1.5, value=0.7, step=0.05, label="Temperature (Độ ngẫu nhiên)")
-TOP_K = gr.Slider(minimum=0, maximum=100, value=50, step=1, label="Top K")
-TOP_P = gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.05, label="Top P")
-
-# Tạo giao diện chính
-with gr.Blocks() as demo:
-    gr.Markdown(
-        f"""
-        # 🎭 Humanish-Roleplay-Llama-3.1-8B (vicgalle)
-        **Mô hình Llama 3.1 8B Fine-tuned cho Roleplay/Chat.**
-        """)
-
-    system_input = gr.Textbox(
-        lines=2, 
-        label="System Prompt (Thiết lập vai trò/Ngữ cảnh)",
-        value="You are a helpful and human-like roleplay assistant. Maintain a friendly and immersive tone and stick to the character's role.",
-        placeholder="Thiết lập vai trò, tính cách, hoặc quy tắc cho mô hình."
-    )
+    # 7. Cập nhật lịch sử hiển thị (gr.Chatbot)
+    chatbot_display.append([prompt, translated])
+
+    # 8. Trả về cả hai để cập nhật UI
+    # (chuỗi rỗng "" để xóa nội dung trong ô prompt_box)
+    return "", chatbot_display, internal_history, response_text, prompt
+
+def clear_chat():
+    """Xóa lịch sử."""
+    return None, None
 
+# --- 4. Xây dựng giao diện Gradio Blocks ---
+with gr.Blocks(theme=gr.themes.Monochrome()) as demo:
+    # "Bộ nhớ" ẩn để lưu lịch sử ChatML (list of dicts)
+    internal_history = gr.State()
+    
     with gr.Row():
-        prompt_input = gr.Textbox(
-            lines=5, 
-            label="User Prompt (Lời thoại hoặc kịch bản Roleplay)",
-            placeholder="Ví dụ: 'Chúng ta đang ở trong một khu rừng tối. Tôi nghe thấy một tiếng động lớn từ phía sau. Bạn làm gì?'"
-        )
-        output_text = gr.Textbox(lines=5, label="Phản hồi của Mô hình")
+        with gr.Column(scale=3):
+            # Khung chat chính
+            chatbot_display = gr.Chatbot(
+                label="Chat History",
+                bubble_full_width=False,
+                height=500
+            )
+            
+            # Ô nhập prompt
+
+            lang = gr.Textbox(
+                label="lang",
+                placeholder="Nhập tin nhắn của bạn....",
+                lines=1
+            )      
+            
+            prompt_box = gr.Textbox(
+                label="Your Message",
+                placeholder="Nhập tin nhắn của bạn và nhấn Enter...",
+                lines=1
+            )
+
+            gender = gr.Checkbox(
+                label="Gender", 
+                value=True, # Mặc định KHÔNG được chọn
+                interactive=True # Cho phép người dùng tương tác
+            )
+
+            prompt = gr.Textbox(
+                label="",
+                placeholder="",
+                lines=1
+            )
 
-    generate_button = gr.Button("🔥 Tạo Phản Hồi 🔥")
+            response = gr.Textbox(
+                label="",
+                placeholder="",
+                lines=1
+            )   
+
+            text_translate = gr.Textbox(
+                label="",
+                placeholder="",
+                lines=1
+            )                 
+            
+            with gr.Row():
+                clear_button = gr.Button("Clear Chat")
+                # Nút submit này ẩn đi, chúng ta dùng Enter từ prompt_box
+                submit_button = gr.Button("Send") 
+
+        with gr.Column(scale=1):
+            # Ô System Prompt
+            system_prompt_box = gr.Textbox(
+                label="System Prompt (AI's Role & Rules)",
+                value="",
+                lines=30
+            )
+
+    # --- 5. Kết nối các hành động ---
+
+    # Khi người dùng nhấn Enter trong `prompt_box`
+    prompt_box.submit(
+        fn=chat_with_model,
+        inputs=[prompt_box, system_prompt_box, chatbot_display, internal_history,lang,gender],
+        outputs=[prompt_box, chatbot_display, internal_history, response, prompt]
+    )
+
+    text_translate.submit(
+        fn=translate_text,
+        inputs=[text_translate,lang],
+        outputs=[prompt]
+    )
+
+    # Khi người dùng nhấn nút "Send" (ẩn)
+    submit_button.click(
+        fn=chat_with_model,
+        inputs=[prompt_box, system_prompt_box, chatbot_display, internal_history,lang,gender],
+        outputs=[prompt_box, chatbot_display, internal_history, response, prompt]
+    )
     
-    with gr.Accordion("⚙️ Cấu hình Tham số Tạo Văn bản", open=False):
-        with gr.Row():
-            MAX_TOKENS.render()
-            TEMPERATURE.render()
-        with gr.Row():
-            TOP_K.render()
-            TOP_P.render()
-
-    generate_button.click(
-        fn=generate_text,
-        inputs=[system_input, prompt_input, MAX_TOKENS, TEMPERATURE, TOP_K, TOP_P],
-        outputs=output_text
+    # Khi người dùng nhấn nút "Clear Chat"
+    clear_button.click(
+        fn=clear_chat,
+        inputs=None,
+        outputs=[chatbot_display, internal_history]
     )
 
+
 if __name__ == "__main__":
     demo.launch()