Update app.py

app.py

@@ -1,20 +1,118 @@
-from fastapi import FastAPI, Request
+# from fastapi import FastAPI, Request
+# from pydantic import BaseModel
+# from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer
+# import torch
+# import re
+# import time
+
+# app = FastAPI()
+
+# def split_by_words_and_dot(text, min_words=125, max_words=160, fallback_words=150):
+#     import re
+#     words = re.findall(r'\S+|\n', text)  # giữ nguyên \n như một "từ"
+#     chunks = []
+#     start = 0
+#     while start < len(words):
+#         end = min(start + max_words, len(words))
+#         # Tìm dấu chấm trong khoảng min_words đến max_words
+#         dot_idx = -1
+#         for i in range(start + min_words, min(start + max_words, len(words))):
+#             if words[i] == '.' or (words[i].endswith('.') and words[i] != '\n'):
+#                 dot_idx = i
+#         if dot_idx != -1:
+#             chunk_end = dot_idx + 1
+#         elif end - start > fallback_words:
+#             chunk_end = start + fallback_words
+#         else:
+#             chunk_end = end
+#         chunk = ' '.join([w if w != '\n' else '\n' for w in words[start:chunk_end]]).replace(' \n ', '\n').replace(' \n', '\n').replace('\n ', '\n')
+#         chunks.append(chunk.strip())
+#         start = chunk_end
+#     return chunks
+
+# # Load model
+# model_path = "longvnhue1/facebook-m2m100_418M-fine_tuning"
+# tokenizer = M2M100Tokenizer.from_pretrained(model_path)
+# model = M2M100ForConditionalGeneration.from_pretrained(model_path)
+# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# model.to(device)
+
+# class TranslateRequest(BaseModel):
+#     text: str
+#     source_lang: str
+#     target_lang: str
+
+
+# @app.post("/translate")
+# 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 = []
+#     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)
+#         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)
+#     print(timing_info)
+
+#     return {
+#         "source_text": req.text,
+#         "translated_text": full_translation,
+#         "src_lang": req.source_lang,
+#         "tgt_lang": req.target_lang,
+#     }
+
+
+from fastapi import FastAPI
 from pydantic import BaseModel
 from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer
 import torch
 import re
 import time
 
+# Giới hạn số luồng CPU nếu không dùng GPU
+torch.set_num_threads(1)
+
 app = FastAPI()
 
+@app.on_event("startup")
+def startup_event():
+    print("🔁 Loading model...")
+
+    global tokenizer, model, device
+
+    model_path = "longvnhue1/facebook-m2m100_418M-fine_tuning"
+    tokenizer = M2M100Tokenizer.from_pretrained(model_path)
+    model = M2M100ForConditionalGeneration.from_pretrained(model_path)
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.to(device)
+
+    print("✅ Model loaded and ready.")
+
 def split_by_words_and_dot(text, min_words=125, max_words=160, fallback_words=150):
-    import re
     words = re.findall(r'\S+|\n', text)  # giữ nguyên \n như một "từ"
     chunks = []
     start = 0
     while start < len(words):
         end = min(start + max_words, len(words))
-        # Tìm dấu chấm trong khoảng min_words đến max_words
         dot_idx = -1
         for i in range(start + min_words, min(start + max_words, len(words))):
             if words[i] == '.' or (words[i].endswith('.') and words[i] != '\n'):
@@ -30,71 +128,45 @@ def split_by_words_and_dot(text, min_words=125, max_words=160, fallback_words=15
         start = chunk_end
     return chunks
 
-# Load model
-model_path = "longvnhue1/facebook-m2m100_418M-fine_tuning"
-tokenizer = M2M100Tokenizer.from_pretrained(model_path)
-model = M2M100ForConditionalGeneration.from_pretrained(model_path)
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model.to(device)
-
 class TranslateRequest(BaseModel):
     text: str
     source_lang: str
     target_lang: str
 
-# @app.post("/translate")
-# def translate_text(req: TranslateRequest):
-#     tokenizer.src_lang = req.source_lang
-#     encoded = tokenizer(req.text, return_tensors="pt", truncation=True, max_length=512).to(device)
-#     generated_tokens = model.generate(
-#         **encoded,
-#         forced_bos_token_id=tokenizer.get_lang_id(req.target_lang),
-#         max_length=512,      # tăng lên nếu cần dịch đoạn dài, nhưng không nên quá lớn
-#         num_beams=2,         # giảm beam search để nhanh hơn
-#         no_repeat_ngram_size=3,
-#         early_stopping=True
-#     )
-#     translated_text = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
-#     return {
-#         "source_text": req.text,
-#         "translated_text": translated_text,
-#         "src_lang": req.source_lang,
-#         "tgt_lang": req.target_lang
-#       }
-
-
 @app.post("/translate")
 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)
+    text_chunks = split_by_words_and_dot(req.text)
     translated_chunks = []
     timing_info = []
 
+    global_start = time.perf_counter()
+
     for idx, chunk in enumerate(text_chunks):
-        start_time = time.perf_counter()  # Bắt đầu đếm thời gian
+        start_time = time.perf_counter()
 
         encoded = tokenizer(chunk, return_tensors="pt", truncation=True, max_length=256).to(device)
         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,
+                max_length=256,              # <-- Giữ mức vừa phải
+                num_beams=2,                 # <-- Đơn giản beam search
                 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")
+        end_time = time.perf_counter()
+        timing_info.append(f"Translated chunk {idx+1}/{len(text_chunks)} in {end_time - start_time:.3f} seconds")
 
-    full_translation = "\n".join(translated_chunks)
+    global_end = time.perf_counter()
+    print(f"⚡️ Total translation time: {global_end - global_start:.3f} seconds")
     print(timing_info)
 
     return {
         "source_text": req.text,
-        "translated_text": full_translation,
+        "translated_text": "\n".join(translated_chunks),
         "src_lang": req.source_lang,
         "tgt_lang": req.target_lang,
-    }
+    }