Update tokenizer.py

tokenizer.py

@@ -5,8 +5,11 @@ import logging
 from pathlib import Path
 from typing import Sequence, List, Optional
 
+import numpy as np
 import sentencepiece as spm
 
+DEFAULT_CHUNK_CHARS = 2_000_000
+
 logger = logging.getLogger(__name__)
 
 # 32K is the well-established baseline vocab size for BPE/SentencePiece
@@ -199,6 +202,7 @@ class TokenizerWrapper:
         model_path = str(Path(model_path))
         if not Path(model_path).exists():
             raise FileNotFoundError(model_path)
+        self.model_path = model_path
         self.sp = spm.SentencePieceProcessor(model_file=model_path)
         self.vocab_size = int(self.sp.vocab_size())
         self.pad_id = self.sp.pad_id()
@@ -228,53 +232,52 @@ class TokenizerWrapper:
         text: str,
         add_bos: bool = True,
         add_eos: bool = False,
-        chunk_chars: int = 2_000_000,
+        chunk_chars: int = DEFAULT_CHUNK_CHARS,
     ) -> List[int]:
-        """OOM-safe encode for large (multi-MB+) strings, e.g. a whole
-        training corpus file read in one go. `encode()` above calls
-        self.sp.encode() on the ENTIRE string in one call -- even though
-        SentencePiece's C++ core is fast (not a slow Python BPE loop),
-        the RETURN VALUE is still built as one giant Python list[int],
-        and Python's own string handling means a large UTF-8 file can
-        already be 2-4x its byte size once decoded into a `str` object.
-        For a real, confirmed example: a 300MB single-file corpus OOM'd
-        a ~12-13GB Colab/Kaggle box within ~30 seconds calling encode()
-        on the whole file at once -- the failure was in Python-level
-        memory (str + list[int] materialization), not inside
-        SentencePiece's own encoding step.
-
-        This method processes `text` in fixed-size character chunks,
-        converting each chunk's ids to the accumulator and discarding
-        the chunk's own Python list before moving to the next chunk, so
-        peak memory is bounded by `chunk_chars` instead of len(text).
-        Chunk boundaries are placed at whitespace (never mid-word/
-        mid-token), so the output is IDENTICAL to what encode(text,
-        add_bos, add_eos) would produce on the whole string at once --
-        this is not an approximation, it's the same encoding, just
-        computed incrementally. Verified by direct comparison across
-        many chunk sizes and edge cases (empty string, no-whitespace
-        text, exact chunk-boundary alignment) against the whole-string
-        encode() path.
-
-        Use this instead of encode() for anything that might be large
-        (a whole file's contents) -- keep encode() for genuinely small
-        strings (a single sentence/prompt at inference time) where the
-        chunking overhead isn't worth paying.
-        """
+        """Backward-compatible wrapper that now delegates to the chunk iterator
+        so encode_large_text and iter_encode_chunks share exactly one loop."""
         if text is None:
             raise ValueError('encode_large_text() received None')
+        parts = [
+            arr for arr in self.iter_encode_chunks(
+                text,
+                add_bos=add_bos,
+                add_eos=add_eos,
+                chunk_chars=chunk_chars,
+            )
+        ]
+        if not parts:
+            return []
+        return np.concatenate(parts, axis=0).astype(np.int32).tolist()
 
+    def iter_encode_chunks(
+        self,
+        text: str,
+        add_bos: bool = True,
+        add_eos: bool = True,
+        chunk_chars: int = DEFAULT_CHUNK_CHARS,
+    ):
+        """Yield per-chunk int32 arrays while preserving the first/last chunk
+        BOS/EOS semantics of the existing whole-file encode path.
+
+        This iterator mirrors `encode_large_text()` chunk placement logic,
+        but streams one chunk's token ids out as a NumPy array instead of
+        materializing a Python list for the whole file.
+        """
+        if text is None:
+            raise ValueError('iter_encode_chunks() received None')
         if text == '':
             ids: List[int] = []
             if add_bos:
                 ids = [self.bos_id] + ids
             if add_eos:
                 ids = ids + [self.eos_id]
-            return ids
+            yield np.asarray(ids, dtype=np.int32)
+            return
 
-        all_ids: List[int] = []
         pos = 0
         n = len(text)
+        first_chunk = True
         while pos < n:
             end = min(pos + chunk_chars, n)
             is_eof = (end >= n)
@@ -283,52 +286,26 @@ class TokenizerWrapper:
                 piece = window
                 pos = end
             else:
-                # Cut at the LAST whitespace boundary inside this
-                # window, so no token is split across two encode()
-                # calls (which would silently tokenize the seam
-                # differently than encoding the whole string at once).
                 cut = max(window.rfind(' '), window.rfind('\n'))
                 if cut <= 0:
-                    # cut == -1: no whitespace anywhere in this window
-                    # (pathological -- one token/URL longer than
-                    # chunk_chars). cut == 0: the window's FIRST
-                    # character is whitespace, which would make
-                    # `piece = window[:0]` empty and, worse, silently
-                    # drop that whitespace character (it's before the
-                    # cut, so it's never included in this piece OR the
-                    # next one, since pos would advance past it without
-                    # encoding it). Both cases: take the whole window
-                    # as one piece and advance past it -- correctness
-                    # at this one seam matters less than a silently
-                    # dropped character or a stalled loop, and this
-                    # only ever triggers when chunk_chars is smaller
-                    # than a typical word (never true at the real
-                    # ~2MB default; only reachable with a deliberately
-                    # tiny chunk_chars in tests).
                     piece = window
                     pos = end
                 else:
                     piece = window[:cut]
                     pos += cut
             if not piece:
-                # Should be unreachable now that cut<=0 is folded into
-                # the "take whole window" branch above, but keep this
-                # as a hard backstop against any future edge case that
-                # produces an empty piece with is_eof=False -- without
-                # it, pos would never advance and the loop would spin
-                # forever.
                 if not is_eof:
                     pos += 1
                 continue
+
             piece_ids = self.sp.encode(piece, out_type=int)
-            all_ids.extend(piece_ids)
-            del piece_ids
+            if add_bos and first_chunk:
+                piece_ids = [self.bos_id] + piece_ids
+            if add_eos and is_eof:
+                piece_ids = piece_ids + [self.eos_id]
 
-        if add_bos:
-            all_ids = [self.bos_id] + all_ids
-        if add_eos:
-            all_ids = all_ids + [self.eos_id]
-        return all_ids
+            yield np.asarray(piece_ids, dtype=np.int32)
+            first_chunk = False
 
     def encode_batch(
         self,