Update app.py

app.py

@@ -3,7 +3,8 @@ import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import spaces , os
 from quickmt import Translator
-from huggingface_hub import snapshot_download
+from quickmt.hub import hf_download, hf_list
+from pathlib import Path
 
 
 MODEL_ID = "bmiller22000/xyntrai-mistral-2.5-7b-chat-nsfw"
@@ -13,11 +14,28 @@ MODEL_NAME = "quickmt/quickmt-en-vi"
 model_path = snapshot_download(MODEL_NAME)
 translator = Translator(model_path, device="auto")
 
-def translate_text(text):
-    if not text.strip():
-        return ""
-    output = translator([text], beam_size=1)
-    return output[0]
+def translate_text(text,progress=gr.Progress(track_tqdm=True)):
+    global t
+    model_path = Path("/code/models") / MODEL_NAME
+    if not model_path.exists():
+        hf_download(
+            model_name="quickmt/" + MODEL_NAME,
+            output_dir=Path("/code/models") / MODEL_NAME,
+        )
+    try:
+        if t is None or str(input.model()) != str(Path(t.model_path).name):
+            print(f"Loading model {input.model()}")
+            t = Translator(
+                str(model_path),
+                device="auto",
+                inter_threads=2,
+            )
+            if not text.strip():
+                return text
+            output = t([text], beam_size=1)
+            return output[0]
+    except:
+        return text
 
 # Tải model và tokenizer 1 LẦN DUY NHẤT
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
@@ -30,7 +48,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):
+def chat_with_model(prompt, system_prompt, chatbot_display, internal_history,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).
@@ -77,7 +95,8 @@ def chat_with_model(prompt, system_prompt, chatbot_display, internal_history):
     
     # 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)
+    print(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})