Update app.py

app.py

@@ -75,65 +75,72 @@ def translate_blocking(text, lang):
 
 # --- HÀM DỊCH STREAMING/WRAPPER ---
 def translate_text(text, lang=None, needStreaming=False, progress=gr.Progress(track_tqdm=True)):
+    """
+    Hàm dịch luôn trả về dạng generator (chỉ dùng yield).
+    needStreaming=True: Yield từng chunk.
+    needStreaming=False: Thực hiện Blocking, sau đó yield toàn bộ kết quả một lần.
+    """
     print("lang", lang)
+    
+    # TRƯỜNG HỢP KHÔNG CẦN DỊCH
     if lang is None:
-        if needStreaming:
-            yield text
-        else:
-            return text
+        yield text
+        return # return không giá trị là hợp lệ
 
     # Khối xây dựng thông điệp dịch
-    messages = build_translation_messages(text, lang)
+    messages = build_translation_messages(text, lang) 
     max_tokens = 2048
     temperature = 0.5
     top_p = 0.05
     
-    if needStreaming:
-        # Xây dựng tokenized_chat và generation_kwargs cho Streaming
-        tokenized_chat = tokenizer_trans.apply_chat_template(
-            messages,
-            tokenize=True,
-            add_generation_prompt=True,
-            return_tensors="pt"
-        ).to(model_trans.device)
+    # Xây dựng tokenized_chat và generation_kwargs
+    tokenized_chat = tokenizer_trans.apply_chat_template(
+        messages,
+        tokenize=True,
+        add_generation_prompt=True,
+        return_tensors="pt"
+    ).to(model_trans.device)
 
-        generation_kwargs = dict(
-            inputs=tokenized_chat,
-            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
-        )
+    generation_kwargs = dict(
+        inputs=tokenized_chat,
+        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
+    )
 
-        # --- Logic Streaming ---
+    if needStreaming:
+        # --- Logic Streaming (Yield từng chunk) ---
         streamer = TextIteratorStreamer(tokenizer_trans, skip_prompt=True, skip_special_tokens=True)
         generation_kwargs["streamer"] = streamer
         
-        # Chạy generate trong một thread riêng biệt
         thread = Thread(target=model_trans.generate, kwargs=generation_kwargs)
         thread.start()
         
         full_text = ""
         for new_text in streamer:
             full_text += new_text
-            # [FIX] Cần yield full text (đã được sửa trong code gốc)
-            yield full_text.strip() # Làm sạch khi stream (tùy chọn)
-            
+            yield full_text # Yield accumulated text
+        
     else:
-        # --- Logic Blocking (Gọi hàm đã tách) ---
-        response = translate_blocking(text, lang)
-        print("response (Blocking)", response)
-        return response
+        # --- Logic Blocking (Chạy Blocking, yield kết quả một lần) ---
+        # Ta sử dụng model.generate trực tiếp và đợi kết quả.
+        with torch.no_grad():
+            outputs = model_trans.generate(**generation_kwargs)
+        
+        # Decode và làm sạch chuỗi
+        response = tokenizer_trans.decode(outputs[0][tokenized_chat.shape[-1]:], skip_special_tokens=True).strip()
+        print("response (Yielded Blocking)", response)
+        
+        # Trả về kết quả cuối cùng dưới dạng Generator (yield một lần)
+        yield response
 
         
 # --- HÀM CHÍNH CHAT ---
 @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). Trả về dạng Streaming.
-    """
+    
     expected_key = os.environ.get("hf_key")
     if expected_key and expected_key not in prompt:
         print("❌ Invalid key.")
@@ -144,27 +151,24 @@ def chat_with_model(prompt, system_prompt, chatbot_display, internal_history, la
         prompt = prompt.replace(expected_key, "")
 
     isAuto = False
-    
-    # [FIX] Đảm bảo actual_prompt_for_model luôn được định nghĩa
-    actual_prompt_for_model = prompt
+    actual_prompt_for_model = prompt # Gán giá trị mặc định
 
     if "[AUTO]" in prompt:
         prompt = prompt.replace("[AUTO]", "")
         isAuto = True
-        # Gán lại prompt sau khi loại bỏ [AUTO]
         actual_prompt_for_model = prompt 
     else:
         # Dịch prompt input của user sang tiếng Anh (không cần streaming input này)
         if lang is not None:
-            # GỌI HÀM, NÓ VẪN TRẢ VỀ GENERATOR OBJECT, DÙ CHẠY BLOCKING LOGIC
-            generator_obj = translate_text(prompt, "English", needStreaming=False)
+            # GỌI HÀM, NÓ LUÔN TRẢ VỀ GENERATOR OBJECT
+            generator_obj = translate_text(prompt, "English", needStreaming=False) 
             print("prompt_translated (generator)", generator_obj) 
             
-            # [FIX CỐ ĐỊNH] Dùng next() để buộc Generator chạy và trích xuất chuỗi:
+            # [FIX] Dùng next() để buộc Generator chạy và trích xuất chuỗi đầu tiên (kết quả Blocking):
             try:
                 prompt_translated = next(generator_obj)
             except (StopIteration, TypeError):
-                # Xử lý trường hợp generator rỗng hoặc đã bị trích xuất
+                # Xử lý trường hợp generator rỗng
                 prompt_translated = "" 
 
             # Gán và làm sạch chuỗi
@@ -174,12 +178,13 @@ def chat_with_model(prompt, system_prompt, chatbot_display, internal_history, la
             
     print("prompt",prompt)
     print("actual_prompt_for_model",actual_prompt_for_model)
-            
     
-    # [FIX] Bỏ thao tác vô nghĩa "".join(list(...))
+    # [FIX] Bỏ thao tác vô nghĩa "".join(list(...)) và thực hiện phép cộng chuỗi
+    # Dòng này giờ sẽ hoạt động vì actual_prompt_for_model đã là một chuỗi (str)
     actual_prompt_for_model = actual_prompt_for_model + " [Detailed description of the physical actions and expressions.]"
     print("prompt for model: " + actual_prompt_for_model)
 
+    # ... (Phần còn lại của hàm chat_with_model giữ nguyên logic generate)
     if chatbot_display is None:
         chatbot_display = []
     if internal_history is None:
@@ -197,21 +202,20 @@ def chat_with_model(prompt, system_prompt, chatbot_display, internal_history, la
     ).to(model.device)
 
     # Chuẩn bị Chatbot Display Placeholder
-    # Append một list [user_msg, None] để bắt đầu streaming câu trả lời
     chatbot_display.append([prompt, ""]) 
     
     # --- LOGIC STREAMING CHÍNH ---
     
     if lang is not None:
         # TRƯỜNG HỢP CÓ DỊCH:
-        # 1. Generate tiếng Anh (nhanh/blocking) để lấy full context
+        # 1. Generate tiếng Anh (Blocking)
         output_tokens = model.generate(
             inputs,
             max_new_tokens=1024,
             do_sample=True,
             temperature=0.7,
             top_p=0.5,
-            pad_token_id=tokenizer.eos_token_id # Thêm pad token cho model chat
+            pad_token_id=tokenizer.eos_token_id
         )
         english_response = tokenizer.decode(output_tokens[0][inputs.shape[-1]:], skip_special_tokens=True)
         print("Eng response generated: ", english_response)
@@ -221,7 +225,6 @@ def chat_with_model(prompt, system_prompt, chatbot_display, internal_history, la
         
         partial_translation = ""
         for chunk in stream_translator:
-            # chunk ở đây là full text tích lũy từ hàm translate_text đã sửa
             partial_translation = chunk 
             
             # Cập nhật UI
@@ -229,7 +232,7 @@ def chat_with_model(prompt, system_prompt, chatbot_display, internal_history, la
             yield "", chatbot_display, internal_history, partial_translation, prompt
         
         final_response_text = english_response
-        final_translated = partial_translation.strip() # Làm sạch lần cuối
+        final_translated = partial_translation.strip()
 
     else:
         # TRƯỜNG HỢP KHÔNG DỊCH (Raw English):
@@ -242,7 +245,7 @@ def chat_with_model(prompt, system_prompt, chatbot_display, internal_history, la
             do_sample=True,
             temperature=0.7,
             top_p=0.5,
-            pad_token_id=tokenizer.eos_token_id # Thêm pad token cho model chat
+            pad_token_id=tokenizer.eos_token_id
         )
         
         thread = Thread(target=model.generate, kwargs=generation_kwargs)
@@ -255,13 +258,12 @@ def chat_with_model(prompt, system_prompt, chatbot_display, internal_history, la
             yield "", chatbot_display, internal_history, partial_text, prompt
         
         final_response_text = partial_text
-        final_translated = partial_text # Giống nhau vì không dịch
+        final_translated = partial_text 
 
     # 6. Cập nhật "bộ nhớ" (gr.State) sau khi hoàn tất
     internal_history.append({"role": "user", "content": actual_prompt_for_model})
     internal_history.append({"role": "assistant", "content": final_response_text})
     
-    # Yield lần cuối để đảm bảo state được lưu
     yield "", chatbot_display, internal_history, final_translated, prompt
 
 def clear_chat():