Add benchmark phase profile and paragraph chunk fix

hachimimt-local.zip

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:aa891b8cf49e96a8ce2d9abdcd54dec984f4748aed6ee608ff347d97de070b92
-size 96397
+oid sha256:9d639cba50e5ca3824f08f5026276012522bdc6cb1319ee7f73d049d0f997c51
+size 95965

src/benchmark_file.py

@@ -0,0 +1,114 @@
+"""Benchmark local file translation throughput."""
+
+from __future__ import annotations
+
+import argparse
+import time
+from pathlib import Path
+
+from app import read_text_file
+from hardware import detect_hardware_profile
+from text_preprocess import (
+    NORMALIZE_AUTO,
+    NORMALIZE_MODES,
+    normalization_message,
+    normalize_chinese_text,
+)
+from translator import Backend, HachimiTranslator
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("path", type=Path, help="Input .txt file")
+    parser.add_argument("--model", default="HachimiMT-60")
+    parser.add_argument("--backend", choices=[Backend.CT2.value, Backend.TRANSFORMERS.value], default=Backend.CT2.value)
+    parser.add_argument("--beam", type=int, default=2)
+    parser.add_argument("--chunk-mode", choices=["sentence", "paragraph"], default="sentence")
+    parser.add_argument("--normalize", choices=sorted(NORMALIZE_MODES), default=NORMALIZE_AUTO)
+    parser.add_argument("--progress-seconds", type=float, default=15.0)
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    total_start = time.perf_counter()
+
+    profile = detect_hardware_profile()
+    print(f"BENCH_START file={args.path}", flush=True)
+    print(f"PROFILE {profile.summary}", flush=True)
+
+    read_start = time.perf_counter()
+    text = read_text_file(args.path)
+    read_s = time.perf_counter() - read_start
+    original_chars = len(text)
+    normalize_start = time.perf_counter()
+    normalized_text = normalize_chinese_text(text, args.normalize)
+    normalize_s = time.perf_counter() - normalize_start
+    normalize_msg = normalization_message(text, normalized_text, args.normalize)
+    text = normalized_text
+    print(
+        f"INPUT chars={len(text)} original_chars={original_chars} "
+        f"lines={text.count(chr(10)) + 1 if text else 0} read_s={read_s:.3f} "
+        f"normalize_s={normalize_s:.3f} normalize={args.normalize} "
+        f"normalize_msg={normalize_msg}",
+        flush=True,
+    )
+
+    translator = HachimiTranslator(profile)
+    load_start = time.perf_counter()
+    status = translator.load(args.model, backend=args.backend)
+    load_s = time.perf_counter() - load_start
+    print(f"LOAD seconds={load_s:.3f} status={status}", flush=True)
+
+    translate_start = time.perf_counter()
+    last_print = 0.0
+    rows = []
+    full_text = ""
+
+    for done, total, message, result_rows, result_text in translator.translate_text_iter(
+        text,
+        chunk_mode=args.chunk_mode,
+        beam_size=args.beam,
+    ):
+        now = time.perf_counter()
+        if result_rows is not None and result_text is not None:
+            rows = result_rows
+            full_text = result_text
+
+        if done == 0 or done == total or now - last_print >= args.progress_seconds:
+            elapsed = now - translate_start
+            rate = done / elapsed if elapsed > 0 else 0.0
+            eta = (total - done) / rate if rate > 0 else 0.0
+            pct = done / total * 100 if total else 100.0
+            print(
+                "PROGRESS "
+                f"done={done} total={total} pct={pct:.2f} elapsed_s={elapsed:.1f} "
+                f"rate_chunks_s={rate:.2f} eta_s={eta:.1f} message={message}",
+                flush=True,
+            )
+            last_print = now
+
+    translate_s = time.perf_counter() - translate_start
+    total_s = time.perf_counter() - total_start
+    profile = translator.last_profile
+    if profile:
+        profile_parts = []
+        for key in sorted(profile):
+            value = profile[key]
+            if isinstance(value, float):
+                profile_parts.append(f"{key}={value:.6f}")
+            else:
+                profile_parts.append(f"{key}={value}")
+        print("BENCH_PROFILE " + " ".join(profile_parts), flush=True)
+    print(
+        "BENCH_DONE "
+        f"total_s={total_s:.3f} translate_s={translate_s:.3f} minutes={total_s / 60:.3f} "
+        f"chunks={len(rows)} chars_in={len(text)} chars_out={len(full_text)} "
+        f"chunks_s={(len(rows) / translate_s) if translate_s > 0 else 0:.3f} "
+        f"chars_s={(len(text) / translate_s) if translate_s > 0 else 0:.1f}",
+        flush=True,
+    )
+
+
+if __name__ == "__main__":
+    main()

src/translator.py

@@ -3,6 +3,8 @@
 from __future__ import annotations
 
 import os
+import time
+from collections import OrderedDict
 from concurrent.futures import Future, ThreadPoolExecutor
 from dataclasses import dataclass
 from enum import Enum
@@ -309,6 +311,40 @@ class CT2SentencePieceTokenizer:
     def __init__(self, model_path: Path) -> None:
         self._source_sp = spm.SentencePieceProcessor(model_file=str(model_path / "source.spm"))
         self._target_sp = spm.SentencePieceProcessor(model_file=str(model_path / "target.spm"))
+        self._encode_cache_max = _env_int(
+            "HACHIMIMT_TOKEN_CACHE_SIZE",
+            0,
+            min_value=0,
+            max_value=500_000,
+        )
+        self._encode_cache: OrderedDict[str, list[int]] = OrderedDict()
+        self._cache_hits = 0
+        self._cache_misses = 0
+
+    def _full_encode_one(self, text: str) -> list[int]:
+        if self._encode_cache_max > 0:
+            cached = self._encode_cache.get(text)
+            if cached is not None:
+                self._encode_cache.move_to_end(text)
+                self._cache_hits += 1
+                return list(cached)
+
+        self._cache_misses += 1
+        token_ids = list(self._source_sp.encode(text, out_type=int))
+        token_ids.append(EOS_TOKEN_ID)
+        if self._encode_cache_max > 0:
+            self._encode_cache[text] = list(token_ids)
+            self._encode_cache.move_to_end(text)
+            while len(self._encode_cache) > self._encode_cache_max:
+                self._encode_cache.popitem(last=False)
+        return token_ids
+
+    def cache_stats(self) -> dict[str, int]:
+        return {
+            "token_cache_entries": len(self._encode_cache),
+            "token_cache_hits": self._cache_hits,
+            "token_cache_misses": self._cache_misses,
+        }
 
     def _encode_one(
         self,
@@ -317,8 +353,7 @@ class CT2SentencePieceTokenizer:
         truncation: bool = False,
         max_length: int | None = None,
     ) -> list[int]:
-        token_ids = list(self._source_sp.encode(text, out_type=int))
-        token_ids.append(EOS_TOKEN_ID)
+        token_ids = self._full_encode_one(text)
         if truncation and max_length is not None and len(token_ids) > max_length:
             token_ids = token_ids[:max_length]
             if token_ids:
@@ -388,6 +423,13 @@ class CT2SentencePieceTokenizer:
             decoded.append(self._target_sp.decode(pieces))
         return decoded
 
+    def decode_tokens_batch(self, tokens_batch: list[list[str]]) -> list[str]:
+        decoded: list[str] = []
+        for tokens in tokens_batch:
+            pieces = [token for token in tokens if token not in SPECIAL_TOKEN_TO_ID]
+            decoded.append(self._target_sp.decode(pieces).strip())
+        return decoded
+
 
 class CT2FastTokenizer:
     """Minimal tokenizer.json wrapper for CTranslate2 inference."""
@@ -403,6 +445,39 @@ class CT2FastTokenizer:
         self._tokenizer = Tokenizer.from_file(str(model_path / "tokenizer.json"))
         self.pad_token_id = self._tokenizer.token_to_id("<pad>")
         self._eos_token_id = self._tokenizer.token_to_id("</s>")
+        self._encode_cache_max = _env_int(
+            "HACHIMIMT_TOKEN_CACHE_SIZE",
+            0,
+            min_value=0,
+            max_value=500_000,
+        )
+        self._encode_cache: OrderedDict[str, list[int]] = OrderedDict()
+        self._cache_hits = 0
+        self._cache_misses = 0
+
+    def _full_encode_one(self, text: str) -> list[int]:
+        if self._encode_cache_max > 0:
+            cached = self._encode_cache.get(text)
+            if cached is not None:
+                self._encode_cache.move_to_end(text)
+                self._cache_hits += 1
+                return list(cached)
+
+        self._cache_misses += 1
+        token_ids = list(self._tokenizer.encode(text).ids)
+        if self._encode_cache_max > 0:
+            self._encode_cache[text] = list(token_ids)
+            self._encode_cache.move_to_end(text)
+            while len(self._encode_cache) > self._encode_cache_max:
+                self._encode_cache.popitem(last=False)
+        return token_ids
+
+    def cache_stats(self) -> dict[str, int]:
+        return {
+            "token_cache_entries": len(self._encode_cache),
+            "token_cache_hits": self._cache_hits,
+            "token_cache_misses": self._cache_misses,
+        }
 
     def _encode_one(
         self,
@@ -411,7 +486,7 @@ class CT2FastTokenizer:
         truncation: bool = False,
         max_length: int | None = None,
     ) -> list[int]:
-        token_ids = list(self._tokenizer.encode(text).ids)
+        token_ids = self._full_encode_one(text)
         if truncation and max_length is not None and len(token_ids) > max_length:
             token_ids = token_ids[:max_length]
             if token_ids and self._eos_token_id is not None:
@@ -474,6 +549,13 @@ class CT2FastTokenizer:
             for token_ids in token_ids_batch
         ]
 
+    def decode_tokens_batch(self, tokens_batch: list[list[str]]) -> list[str]:
+        token_ids_batch = [self.convert_tokens_to_ids(tokens) for tokens in tokens_batch]
+        return [
+            text.strip()
+            for text in self.batch_decode(token_ids_batch, skip_special_tokens=True)
+        ]
+
 
 def _load_ct2_tokenizer(model_path: Path):
     if (model_path / "source.spm").exists() and (model_path / "target.spm").exists():
@@ -584,6 +666,7 @@ class HachimiTranslator:
         self._batch_size = self._profile.batch_size
         self._tokenize_workers = self._profile.tokenize_workers
         self._tokenize_pool: ThreadPoolExecutor | None = None
+        self._last_profile: dict[str, float | int | str] = {}
 
     @property
     def hardware_profile(self) -> HardwareProfile:
@@ -593,6 +676,19 @@ class HachimiTranslator:
     def batch_size(self) -> int:
         return self._batch_size
 
+    @property
+    def last_profile(self) -> dict[str, float | int | str]:
+        return dict(self._last_profile)
+
+    def _reset_profile(self) -> None:
+        self._last_profile = {}
+
+    def _profile_add(self, key: str, seconds: float) -> None:
+        self._last_profile[key] = float(self._last_profile.get(key, 0.0)) + seconds
+
+    def _profile_set(self, key: str, value: float | int | str) -> None:
+        self._last_profile[key] = value
+
     def set_batch_size(self, batch_size: int) -> None:
         self._batch_size = max(4, min(128, int(batch_size)))
 
@@ -749,12 +845,20 @@ class HachimiTranslator:
         return [tokens for job in jobs for tokens in job.result()]
 
     def _decode_ct2_results(self, results) -> list[str]:
+        start = time.perf_counter()
         hypotheses = [result.hypotheses[0] for result in results]
+        decode_tokens_batch = getattr(self._tokenizer, "decode_tokens_batch", None)
+        if callable(decode_tokens_batch):
+            decoded = decode_tokens_batch(hypotheses)
+            self._profile_add("decode_s", time.perf_counter() - start)
+            return decoded
         token_ids = [self._tokenizer.convert_tokens_to_ids(tokens) for tokens in hypotheses]
-        return [
+        decoded = [
             text.strip()
             for text in self._tokenizer.batch_decode(token_ids, skip_special_tokens=True)
         ]
+        self._profile_add("decode_s", time.perf_counter() - start)
+        return decoded
 
     def _load_ct2(self, config: ModelConfig) -> None:
         model_path = ensure_model_files(config, Backend.CT2)
@@ -978,7 +1082,10 @@ class HachimiTranslator:
     ) -> list[str]:
         config = MODELS[self._model_key]
         if source_batches is None:
+            tokenize_start = time.perf_counter()
             source_batches = self._tokenize_chunks_parallel(chunks)
+            self._profile_add("tokenize_s", time.perf_counter() - tokenize_start)
+        infer_start = time.perf_counter()
         results = self._ct2_model.translate_batch(
             source_batches,
             max_batch_size=self._ct2_max_batch_size(config),
@@ -987,6 +1094,7 @@ class HachimiTranslator:
             max_decoding_length=config.ct2_max_output_tokens,
             **self._ct2_repetition_kwargs(config),
         )
+        self._profile_add("ct2_infer_s", time.perf_counter() - infer_start)
         return self._decode_ct2_results(results)
 
     def _translate_ct2_batch_pipelined(
@@ -997,7 +1105,9 @@ class HachimiTranslator:
         prefetched_tokens: SourceTokenJobs | None,
     ) -> list[str]:
         if prefetched_tokens is not None:
+            wait_start = time.perf_counter()
             source_batches = self._collect_tokenize_jobs(prefetched_tokens)
+            self._profile_add("tokenize_wait_s", time.perf_counter() - wait_start)
         else:
             source_batches = None
         return self._translate_ct2_batch(
@@ -1024,8 +1134,14 @@ class HachimiTranslator:
     ) -> Iterator[tuple[int, int, str, list[tuple[int, str, str]] | None, str | None]]:
         """Yield (done, total, message, rows_or_none, full_text_or_none) sau mỗi batch."""
         beam_size = self.clamp_beam(beam_size)
+        self._reset_profile()
+        self._profile_set("backend", self._backend.value if self._backend else "")
+        self._profile_set("beam", beam_size)
+        chunk_start = time.perf_counter()
         chunks = self._chunk_text(text, chunk_mode)
+        self._profile_set("chunk_s", time.perf_counter() - chunk_start)
         total = len(chunks)
+        self._profile_set("chunks", total)
 
         yield 0, total, f"Đã chia {total} chunk, chuẩn bị dịch...", None, None
 
@@ -1035,7 +1151,9 @@ class HachimiTranslator:
         next_tokens: SourceTokenJobs | None = None
         if self._backend == Backend.CT2 and total:
             first_end = min(window_size, total)
+            submit_start = time.perf_counter()
             next_tokens = self._submit_tokenize_jobs(chunks[:first_end])
+            self._profile_add("tokenize_submit_s", time.perf_counter() - submit_start)
 
         for start in range(0, total, window_size):
             end = min(start + window_size, total)
@@ -1065,7 +1183,9 @@ class HachimiTranslator:
                 next_end = min(next_start + window_size, total)
                 if next_start < total:
                     next_batch = chunks[next_start:next_end]
+                    submit_start = time.perf_counter()
                     next_tokens = self._submit_tokenize_jobs(next_batch)
+                    self._profile_add("tokenize_submit_s", time.perf_counter() - submit_start)
 
                 translations.extend(
                     self._translate_ct2_batch_pipelined(
@@ -1079,11 +1199,17 @@ class HachimiTranslator:
 
             yield end, total, f"Đã xong {end}/{total} chunk", None, None
 
+        assemble_start = time.perf_counter()
         rows = [
             (index, chunk, translated)
             for index, (chunk, translated) in enumerate(zip(chunks, translations), start=1)
         ]
         full_text = "\n".join(translations)
+        self._profile_add("assemble_s", time.perf_counter() - assemble_start)
+        cache_stats = getattr(self._tokenizer, "cache_stats", None)
+        if callable(cache_stats):
+            for key, value in cache_stats().items():
+                self._profile_set(key, value)
         yield total, total, "Hoàn tất dịch.", rows, full_text
 
     def translate_text(