Update app.py

app.py

@@ -2,7 +2,7 @@ import pandas as pd
 from sklearn.utils import shuffle
 from sklearn.model_selection import train_test_split
 from sklearn.metrics import accuracy_score, precision_recall_fscore_support
-from transformers import TrainingArguments, Trainer, AutoTokenizer, AutoModelForSequenceClassification
+from transformers import TrainingArguments, Trainer, AutoTokenizer, AutoModelForSequenceClassification, AutoModelForSeq2SeqLM
 import torch
 import gradio as gr
 from pathlib import Path
@@ -30,24 +30,36 @@ data['Label'] = data['Label'].map(groups)
 
 seeds = [1, 2, 3, 4]
 
+# โหลด model และ tokenizer ของ NLLB
+translation_model_name = "facebook/nllb-200-distilled-600M"
+translation_tokenizer = AutoTokenizer.from_pretrained(translation_model_name)
+translation_model = AutoModelForSeq2SeqLM.from_pretrained(translation_model_name)
+translation_model.eval()
+
 # Translate function
-async def translate_all(seed, texts, language):
-    dest_lang = "ja" if language == "Japanese" else "th"
-    semaphore = asyncio.Semaphore(12)
+async def translate_all(seed, texts, language, progress=gr.Progress(track_tqdm=True)):
+    dest_lang = "jpn_Jpan" if language == "Japanese" else "tha_Thai"
+    semaphore = asyncio.Semaphore(3)  # จำกัด concurrent request (GPU load)
+
+    translated = [None] * len(texts)
 
     async def sem_translate_task(text, idx):
         async with semaphore:
-            async with Translator() as translator:
-                result = await translator.translate(text, src='en', dest=dest_lang)
-                return result.text, idx
+            inputs = translation_tokenizer(text, return_tensors="pt", max_length=512, truncation=True)
+            if torch.cuda.is_available():
+                inputs = {k: v.to("cuda") for k, v in inputs.items()}
+            with torch.no_grad():
+                translated_tokens = translation_model.generate(
+                    **inputs,
+                    forced_bos_token_id=translation_tokenizer.lang_code_to_id[dest_lang]
+                )
+            translated_text = translation_tokenizer.decode(translated_tokens[0], skip_special_tokens=True)
+            translated[idx] = translated_text
+            progress(idx + 1, len(texts))
+            await asyncio.sleep(0)  # ให้ async loop ทำงาน
 
     tasks = [asyncio.create_task(sem_translate_task(text, idx)) for idx, text in enumerate(texts)]
-    translated = [None] * len(texts)
-
-    for coro in tqdm_asyncio.as_completed(tasks, total=len(tasks)):
-        result, index = await coro
-        translated[index] = result
-
+    await asyncio.gather(*tasks)
     return translated
 
 # Sample Dataset