app.py

import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import spaces , os
from quickmt import Translator
from quickmt.hub import hf_download, hf_list
from pathlib import Path


MODEL_ID = "bmiller22000/xyntrai-mistral-2.5-7b-chat-nsfw"
# Load model and tokenizer


# 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
)



# Thêm một ô system_prompt
@spaces.GPU(duration=60)
def chat_with_model(prompt, system_prompt, chatbot_display, internal_history,lang,gender,progress=gr.Progress(track_tqdm=True)):
    """
    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}]
    
    # 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)
    
    # 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:
            translated =  translate_text(response_text +"; Ms.",lang)

    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})
    
    # 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():
        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
            )

            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]
    )
    
    # 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()