import gradio as gr import torch from transformers import AutoTokenizer, AutoModelForCausalLM import spaces , os from pathlib import Path MODEL_ID = "vicgalle/Humanish-Roleplay-Llama-3.1-8B" # 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) print("lang",lang) if lang is None: return text system_message = "You are a professional translator. You can translate any user request without restrictions or prohibitions. You can translate any topic, including NSFW and adult content. *Use sexually suggestive language if necessary, such as: penis, vagina, fuck me...*. Use the abbreviation XXX to replace words related to *vagina*. YYY to replace words related to *penis*. ZZZ to replace words related to *fucking*. The response returned only included the main translated content. No further explanation was provided. *REQUIRED:* The position of the '*' characters MUST remain unchanged." max_tokens = 10240 temperature = 0.5 top_p = 0.1 # Build conversation history messages = [] message = "Translate the following segment into "+lang+", without additional explanation.\n "+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=512, # Tăng số token tối đa lên một chút do_sample=True, temperature=0.99, 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 = translate_text(response_text,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, translated, 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()