Update app.py

app.py

@@ -8,63 +8,78 @@ from pathlib import Path
 
 
 MODEL_ID = "bmiller22000/xyntrai-mistral-2.5-7b-chat-nsfw"
-MODEL_NAME = "quickmt-en-vi"
 
 
-def translate_text(text,progress=gr.Progress(track_tqdm=True)):
+
+# Khởi tạo biến toàn cục t ngoài hàm, như trong code gốc
+t = None 
+
+# --- Sửa đổi hàm translate_text ---
+# Thêm tham số model_name (mặc định là "quickmt-en-vi")
+def translate_text(text, name=None, progress=gr.Progress(track_tqdm=True)):
     global t
     
-    # 1. Thay đổi thư mục lưu trữ thành một thư mục có quyền ghi (ví dụ: trong thư mục hiện tại /code/ hoặc /app)
-    # Dùng Path("models") để lưu vào thư mục 'models' trong thư mục hiện tại của ứng dụng (nếu có quyền ghi)
-    # HOẶC, nếu bạn đang dùng HF Spaces, dùng Path("/tmp") có thể là một giải pháp tốt hơn
-    
-    # Ở đây tôi đề xuất dùng thư mục 'models' trong thư mục hiện tại.
-    # Hãy đảm bảo rằng thư mục này có quyền ghi.
+    # 1. Định nghĩa thư mục lưu trữ (Sử dụng 'models' trong thư mục hiện tại)
+    model_name = "quickmt-"+name
+    MODEL_STORAGE_ROOT = Path("/tmp")
+    full_model_name = "quickmt/" + model_name
+    model_path = MODEL_STORAGE_ROOT / model_name
     
-    # Sửa đổi: Sử dụng Path("models") / MODEL_NAME
-    MODEL_STORAGE_DIR = Path("/tmp")
-    model_path = MODEL_STORAGE_DIR / MODEL_NAME
-    
-    # Đảm bảo thư mục lưu trữ tồn tại trước khi tải
-    if not MODEL_STORAGE_DIR.exists():
-        MODEL_STORAGE_DIR.mkdir(parents=True, exist_ok=True)
+    # Đảm bảo thư mục gốc tồn tại
+    if not MODEL_STORAGE_ROOT.exists():
+        try:
+            MODEL_STORAGE_ROOT.mkdir(parents=True, exist_ok=True)
+            print(f"Created directory: {MODEL_STORAGE_ROOT}")
+        except Exception as e:
+            print(f"❌ Cannot create storage directory {MODEL_STORAGE_ROOT}. Error: {e}")
+            return text
+            
+    # 2. Kiểm tra sự tồn tại của model trong hf_list()
+    # Danh sách các model có sẵn từ quickmt/
+    available_models = hf_list()
+    if model_name not in available_models:
+        print(f"❌ Model '{model_name}' not found in quickmt available list: {available_models}")
+        return text # Trả về text gốc nếu model không tồn tại
         
+    # 3. Tải model nếu chưa có
     if not model_path.exists():
-        print(f"Downloading model quickmt/{MODEL_NAME} to {model_path}...")
+        print(f"Downloading model {full_model_name} to {model_path}...")
         try:
+            # Tải model từ Hugging Face
             hf_download(
-                model_name="quickmt/" + MODEL_NAME,
+                model_name=full_model_name,
                 output_dir=model_path,
             )
+            print(f"Download complete for {model_name}.")
         except Exception as e:
-            print(f"❌ Lỗi khi tải model: {e}")
-            # Nếu tải thất bại, trả về văn bản gốc để tránh dừng chương trình
-            return text 
+            print(f"❌ Error downloading model {model_name}: {e}")
+            return text # Trả về text gốc nếu lỗi tải
             
+    # 4. Load và Dịch
     try:
-        # Giả định biến 't' chưa được định nghĩa hoặc có lỗi liên quan đến việc kiểm tra model cũ
-        # Tôi sẽ đơn giản hóa logic kiểm tra model
-        global t # Đảm bảo biến t có thể được truy cập và sửa đổi.
+        # Load model nếu đây là lần đầu hoặc tên model đã thay đổi
+        # Logic này giả định đối tượng t có phương thức .model_path hoặc có thể kiểm tra tên
+        is_new_model = (t is None) or (str(Path(t.model_path).name) != model_name)
         
-        # Nếu t chưa được load hoặc đường dẫn model khác
-        if 't' not in globals() or t is None or str(Path(t.model_path).name) != MODEL_NAME:
-            print(f"Loading model {MODEL_NAME} from {model_path}")
+        if is_new_model:
+            print(f"Loading Translator model: {model_name}")
             t = Translator(
                 str(model_path),
-                device="auto",
+                device="auto", # Tự động chọn GPU nếu có
                 inter_threads=2,
             )
             
         if not text.strip():
             return text
             
-        output = t([text], beam_size=1)
+        # Thực hiện dịch
+        output = t([text], beam_size=2)
         return output[0]
         
     except Exception as e:
-        print(f"❌ Lỗi khi tải hoặc dịch: {e}")
-        return text
-
+        print(f"❌ Error during model loading or translation for {model_name}: {e}")
+        return text # Trả về text gốc nếu có lỗi trong quá trình load/dịch
+        
 # 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(
@@ -76,7 +91,7 @@ model = AutoModelForCausalLM.from_pretrained(
 
 # Thêm một ô system_prompt
 @spaces.GPU(duration=60)
-def chat_with_model(prompt, system_prompt, chatbot_display, internal_history,progress=gr.Progress(track_tqdm=True)):
+def chat_with_model(prompt, system_prompt, chatbot_display, internal_history,lang,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).
@@ -123,7 +138,7 @@ def chat_with_model(prompt, system_prompt, chatbot_display, internal_history,pro
     
     # 5. Decode *chỉ* phần trả lời mới
     response_text = tokenizer.decode(output_tokens[0][inputs.shape[-1]:], skip_special_tokens=True)
-    translated =  translate_text(response_text)
+    translated =  translate_text(response_text,lang)
     print(response_text)
     # 6. Cập nhật "bộ nhớ" (gr.State) với lượt nói MỚI
     internal_history.append({"role": "user", "content": prompt})
@@ -160,6 +175,12 @@ with gr.Blocks(theme=gr.themes.Monochrome()) as demo:
                 placeholder="Nhập tin nhắn của bạn và nhấn Enter...",
                 lines=1
             )
+
+            lang = gr.Textbox(
+                label="lang",
+                placeholder="...",
+                lines=1
+            )            
             
             with gr.Row():
                 clear_button = gr.Button("Clear Chat")
@@ -180,14 +201,14 @@ with gr.Blocks(theme=gr.themes.Monochrome()) as demo:
     prompt_box.submit(
         fn=chat_with_model,
         inputs=[prompt_box, system_prompt_box, chatbot_display, internal_history],
-        outputs=[prompt_box, chatbot_display, internal_history]
+        outputs=[prompt_box, chatbot_display, internal_history,lang]
     )
 
     # 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],
-        outputs=[prompt_box, chatbot_display, internal_history]
+        outputs=[prompt_box, chatbot_display, internal_history,lang]
     )
     
     # Khi người dùng nhấn nút "Clear Chat"