ABC

app.py

@@ -3,6 +3,7 @@ from pydantic import BaseModel
 from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer
 import torch
 import re
+import time
 
 app = FastAPI()
 
@@ -67,21 +68,33 @@ def translate_text(req: TranslateRequest):
     tokenizer.src_lang = req.source_lang
     text_chunks = split_by_words_and_dot(req.text, min_words=125, max_words=160, fallback_words=150)
     translated_chunks = []
-    for chunk in text_chunks:
+    timing_info = []
+
+    for idx, chunk in enumerate(text_chunks):
+        start_time = time.perf_counter()  # Bắt đầu đếm thời gian
+
         encoded = tokenizer(chunk, return_tensors="pt", truncation=True, max_length=256).to(device)
-        generated_tokens = model.generate(
-            **encoded,
-            forced_bos_token_id=tokenizer.get_lang_id(req.target_lang),
-            max_length=256,
-            num_beams=2,
-            no_repeat_ngram_size=3,
-        )
+        with torch.inference_mode():
+            generated_tokens = model.generate(
+                **encoded,
+                forced_bos_token_id=tokenizer.get_lang_id(req.target_lang),
+                max_length=256,
+                num_beams=2,
+                no_repeat_ngram_size=3,
+            )
         translated_text = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
         translated_chunks.append(translated_text)
+
+        end_time = time.perf_counter()  # Kết thúc đếm thời gian
+        elapsed = end_time - start_time
+        timing_info.append(f"Translated chunk {idx+1}/{len(text_chunks)} in {elapsed:.3f} seconds")
+
     full_translation = "\n".join(translated_chunks)
+
     return {
         "source_text": req.text,
         "translated_text": full_translation,
         "src_lang": req.source_lang,
-        "tgt_lang": req.target_lang
+        "tgt_lang": req.target_lang,
+        "timing": timing_info
     }