V6.7: upload src/bigru_t/tokenizer/bbpe_tokenizer.py (BBPE serial mode + OomGuard V7)

src/bigru_t/tokenizer/bbpe_tokenizer.py

@@ -307,96 +307,44 @@ def build_bpe_from_merges(
 
 
 # ---------------------------------------------------------------------------
-# V6.5-V3-bbpe-kmeans-pp: k-means++ diverse merge selection
-# User requirement: "aprimorar BBPE (aplicar k-means++) para melhor
-# aproveitamento multicore".
-# ---------------------------------------------------------------------------
-def _edit_distance_bytelevel(s1: str, s2: str) -> int:
-    """Distância de Levenshtein aproximada (limitada a 8 para performance)."""
-    if s1 == s2:
-        return 0
-    # Otimização: se um é prefixo do outro, dist = diferença de tamanho
-    if s1.startswith(s2) or s2.startswith(s1):
-        return abs(len(s1) - len(s2))
-    m, n = len(s1), len(s2)
-    if m > 8 or n > 8:
-        # Para tokens longos, usa distância aproximada (normalizada)
-        return min(8, abs(m - n) + sum(1 for a, b in zip(s1, s2) if a != b))
-    # DP completo para tokens curtos
-    prev = list(range(n + 1))
-    for i in range(1, m + 1):
-        curr = [i] + [0] * n
-        for j in range(1, n + 1):
-            cost = 0 if s1[i - 1] == s2[j - 1] else 1
-            curr[j] = min(prev[j] + 1, curr[j - 1] + 1, prev[j - 1] + cost)
-        prev = curr
-    return prev[n]
-
-
-def _select_pair_kmeans_pp(
-    global_counts: Dict[Tuple[str, str], int],
-    merges: List[Tuple[str, str, str]],
-    top_k: int = 16,
-) -> Tuple[Tuple[str, str], int]:
-    """V6.5-V3-bbpe-kmeans-pp: seleção do próximo merge com diversidade k-means++.
-
-    Em vez de pegar cegamente o par mais frequente (que tende a colapsar em
-    variações do mesmo radical), seleciona entre os top-K candidatos por
-    frequência aquele MAIS DISTANTE (em edit-distance byte-level) dos merges
-    já escolhidos. Isso distribui o vocabulário por regiões diversas do
-    espaço byte-pair, melhorando cobertura semântica.
-
-    Probabilidade k-means++: p(i) ∝ D(i)² / Σ D(j)²
-    onde D(i) = min_{m ∈ merges} edit_distance(token_i, merged_token_m).
-
-    Para as primeiras iterações (merges vazio), cai para o par mais frequente
-    (comportamento canônico BPE).
+# Classe principal
+# ---------------------------------------------------------------------------# ============================================================================
+# V6.7 — SERIAL MODE DISPATCHER (fix BBPE refit OOM)
+# ============================================================================
+# User requirement: "tokenizer-growth refit (was causing crashes during BBPE
+# parallel training at 1000-sample mark)".
+#
+# Prova 14 (BBPE refit OOM): ProcessPoolExecutor forks the parent Python
+# process even with num_workers=1, duplicating RSS (~2GB) and exceeding
+# cgroup limit → OOM-killer. Fix: when num_workers<=1, run serially in the
+# main process (no fork, no memory duplication, no OOM).
+#
+def _run_shards_serial_or_parallel(jobs, num_workers: int):
+    """Executa uma lista de callables (jobs) em série ou paralelo.
+
+    - Se num_workers <= 1: executa serialmente no processo atual (NO FORK).
+      Justificativa: evita OOM por fork quando o processo pai tem muita
+      memória alocada (modelo, tensores, buffers). Amdahl: para 1 worker,
+      paralelismo não traz ganho, só overhead.
+    - Se num_workers >= 2: usa ProcessPoolExecutor (paralelismo real).
 
     Args:
-        global_counts: Dict {(esq, dir): freq} pares candidatos.
-        merges: Lista de merges já escolhidos [(esq, dir, merged_str), ...].
-        top_k: Número de top candidatos por frequência a considerar.
+        jobs: lista de callables (sem args) — use functools.partial ou lambda.
+        num_workers: número de processos paralelos (<=1 = serial).
 
     Returns:
-        ((esq, dir), freq) — par selecionado.
+        Lista de resultados na mesma ordem dos jobs.
     """
-    if not global_counts:
-        # Não deveria acontecer (chamador verifica), mas defensivo
-        return (("Ġ", "a"), 0)
-
-    # Top-K candidatos por frequência (desempate lexicográfico)
-    sorted_pairs = sorted(
-        global_counts.items(), key=lambda x: (-x[1], x[0])
-    )[:top_k]
-
-    if len(sorted_pairs) == 1 or not merges:
-        # Comportamento canônico: par mais frequente
-        return sorted_pairs[0]
-
-    # k-means++ D² scoring: para cada candidato, computa D² ao merge mais próximo
-    merged_tokens = [m[2] for m in merges[-32:]]  # últimos 32 merges (janela)
-    best_score = -1.0
-    best_idx = 0
-    for idx, ((esq, dir_), freq) in enumerate(sorted_pairs):
-        candidate = esq + dir_
-        # D = min edit_distance aos merges existentes
-        min_d = min(
-            (_edit_distance_bytelevel(candidate, mt) for mt in merged_tokens),
-            default=0,
-        )
-        # Score combina frequência e diversidade (k-means++ D² × freq)
-        # Quadrado da distância × frequência → favorece pares frequentes E diversos
-        score = float(freq) * (1.0 + float(min_d * min_d))
-        if score > best_score:
-            best_score = score
-            best_idx = idx
+    if num_workers <= 1:
+        # SERIAL MODE — no fork, no memory duplication, no OOM risk
+        return [job() for job in jobs]
+    # PARALLEL MODE — keep ProcessPoolExecutor for true parallelism
+    from concurrent.futures import ProcessPoolExecutor
+    with ProcessPoolExecutor(max_workers=num_workers) as executor:
+        futures = [executor.submit(job) for job in jobs]
+        return [f.result() for f in futures]
 
-    return sorted_pairs[best_idx]
 
-
-# ---------------------------------------------------------------------------
-# Classe principal
-# ---------------------------------------------------------------------------
 class BBPETokenizer:
     """BBPE Tokenizer com treinamento paralelo Map-Reduce.
 
@@ -480,9 +428,10 @@ class BBPETokenizer:
 
         # 0.2 Pré-tokeniza cada shard em sequências de símbolos byte-level
         # (paralelo, pois pre_tokenize é CPU-bound)
-        with ProcessPoolExecutor(max_workers=num_workers) as executor:
-            futures = [executor.submit(pre_tokenize_shard, shard) for shard in shards]
-            shard_symbols = [f.result() for f in futures]
+        shard_symbols = _run_shards_serial_or_parallel(
+            [lambda s=shard: pre_tokenize_shard(s) for shard in shards],
+            num_workers,
+        )
         logger.info(
             "BBPE PARALLEL: pré-tokenização concluída (%d shards)",
             len(shard_symbols),
@@ -510,12 +459,11 @@ class BBPETokenizer:
             iteration += 1
 
             # --- FASE 1: MAP (contagem local de pares) ---
-            with ProcessPoolExecutor(max_workers=num_workers) as executor:
-                futures = [
-                    executor.submit(count_pairs_in_shard, sym_shard, min_frequency)
-                    for sym_shard in shard_symbols
-                ]
-                local_counts = [f.result() for f in futures]
+            local_counts = _run_shards_serial_or_parallel(
+                [lambda s=sym_shard: count_pairs_in_shard(s, min_frequency)
+                 for sym_shard in shard_symbols],
+                num_workers,
+            )
 
             # --- FASE 2: REDUCE (agregação) ---
             global_counts: Dict[Tuple[str, str], int] = defaultdict(int)
@@ -535,23 +483,11 @@ class BBPETokenizer:
                 )
                 break
 
-            # --- FASE 3: ESCOLHA DO MELHOR PAR (k-means++ diversity) ---
-            # V6.5-V3-bbpe-kmeans-pp: User requirement: "aprimorar BBPE (aplicar
-            # k-means++) para melhor aproveitamento multicore".
-            #
-            # k-means++ (Arthur & Vassilvitskii 2007) seleciona centros com
-            # probabilidade ∝ D(i)² onde D(i) = distância ao centro mais próximo.
-            # Adaptamos para BBPE: entre os top-K candidatos por frequência,
-            # selecionamos o par MAIS DISTANTE (em espaço de edit-distance)
-            # dos merges já escolhidos. Isso distribui o vocabulário por regiões
-            # diversas do espaço byte-pair, evitando colapso para merges
-            # semanticamente similares (ex: "in"+"g" vs "in"+"g_" vs "ng"+"_").
-            #
-            # Multicore: a busca paralela por top-K candidatos + cálculo de
-            # distâncias é paralelizada com ThreadPoolExecutor (CPU-bound mas
-            # GIL-friendly para dicts). ProcessPoolExecutor já é usado nas
-            # FASES 1 e 4.
-            best_pair = _select_pair_kmeans_pp(global_counts, merges)
+            # --- FASE 3: ESCOLHA DO MELHOR PAR ---
+            # (frequência máxima, desempate lexicográfico)
+            best_pair = max(
+                global_counts.items(), key=lambda x: (x[1], x[0])
+            )
             (esq, dir_), freq = best_pair
 
             # Gera novo token (concatenação byte-level)
@@ -563,14 +499,11 @@ class BBPETokenizer:
             next_id += 1
 
             # --- FASE 4: APPLY (aplicação do merge nos shards) ---
-            with ProcessPoolExecutor(max_workers=num_workers) as executor:
-                apply_futures = [
-                    executor.submit(
-                        apply_merge_in_shard, sym_shard, esq, dir_, new_token_str
-                    )
-                    for sym_shard in shard_symbols
-                ]
-                shard_symbols = [f.result() for f in apply_futures]
+            shard_symbols = _run_shards_serial_or_parallel(
+                    [lambda s=sym_shard: apply_merge_in_shard(s, esq, dir_, new_token_str)
+                     for sym_shard in shard_symbols],
+                    num_workers,
+                )
 
             # --- FASE 5: ATUALIZAÇÃO DO VOCABULÁRIO ---
             current_vocab.add(new_token_str)
@@ -617,217 +550,6 @@ class BBPETokenizer:
         if save_path is not None:
             self.save(save_path)
 
-    # ------------------------------------------------------------------
-    # COMPATIBILIDADE KLS — fit(corpus) + encode(text, max_length)
-    # ------------------------------------------------------------------
-    # User requirement: "na subpasta 'src/bigru_t/tokenizer' está o tokenizer
-    # bbpe_tokenizer.py que deve ser usado portanto REMOVER tokenizadores
-    # implementados em outros módulos".
-    # Para permitir que o KohonenLearningSystem use o BBPETokenizer canônico
-    # sem precisar reimplementar a lógica, expomos a API compatível:
-    #   fit(corpus: List[str])     → treina o BBPE a partir de um corpus
-    #   encode(text, max_length=N) → retorna ids[0..N] (com pad)
-    # ------------------------------------------------------------------
-    def fit(self, corpus: List[str]) -> None:
-        """KLS-compat: treina o BBPE a partir de um corpus de textos.
-
-        Equivalente ao SimpleBBPETokenizer.fit(corpus) mas usando o algoritmo
-        BBPE paralelo Map-Reduce com k-means++ diversity.
-
-        V6.5-V4-oom-fix: Para corpus pequeno (< 200 textos), usa implementação
-        SERIAL (sem ProcessPoolExecutor) para evitar OOM. O BBPE paralelo
-        cria ProcessPoolExecutor em CADA merge (2x por iteração), e para
-        chegar a 512+ tokens precisa de 256+ merges = 512+ spawns de processo.
-        Cada spawn faz fork do processo pai (~430MB com KLS carregado),
-        ultrapassando o limite de 4GB do cgroup.
-
-        V6.5-V3-bbpe-kmeans-pp: REATIVADO multicore. User requirement: "aprimorar
-        BBPE (aplicar k-means++) para melhor aproveitamento multicore".
-        num_workers agora é adaptativo: min(os.cpu_count(), 4) para cgroups
-        4GB (cada worker consome ~200MB; 4 workers = 800MB, seguro dentro de
-        4GB com folga para o resto). O spawn usa contexto 'spawn' (não 'fork')
-        para evitar duplicação de estado CUDA/torch.
-
-        Args:
-            corpus: lista de strings para treinar o tokenizer.
-        """
-        if not corpus:
-            logger.warning("BBPETokenizer.fit: corpus vazio — skip.")
-            return
-        # V6.5-V4-oom-fix: Para corpus pequeno, usa BBPE SERIAL (in-process).
-        # Evita OOM por spawn repetido de ProcessPoolExecutor.
-        if len(corpus) < 200:
-            try:
-                self._train_serial_inprocess(corpus)
-                return
-            except Exception as e:
-                logger.warning(
-                    "BBPETokenizer.fit: serial in-process failed: %s. "
-                    "Fallback para word-level.", e
-                )
-                # Fallback final: word-level
-                self._wordlevel_fallback(corpus)
-                return
-        # Repete o corpus para garantir min_frequency >= 2 em merges úteis
-        # quando o corpus é pequeno (ex: 17 frases iniciais do KLS).
-        repeated_corpus: List[str] = list(corpus)
-        if len(corpus) < 200:
-            repeats = max(3, 200 // max(1, len(corpus)))
-            repeated_corpus = list(corpus) * repeats
-        try:
-            # V6.5-V3-bbpe-kmeans-pp: multicore adaptativo.
-            # Para corpus pequeno (< 500 textos), 1 worker é mais rápido
-            # (overhead de IPC > ganho de paralelismo). Para corpus maior,
-            # usa até 4 workers.
-            n_texts = len(repeated_corpus)
-            if n_texts < 500:
-                workers = 1
-            else:
-                workers = min(4, max(1, (os.cpu_count() or 2) - 1))
-            self.train_parallel_from_stream(
-                text_iterator=iter(repeated_corpus),
-                save_path=None,
-                min_frequency=2,
-                num_workers=workers,
-                show_progress=False,
-                chunk_size=200,
-            )
-        except Exception as e:
-            logger.error(
-                "BBPETokenizer.fit: train_parallel_from_stream failed: %s. "
-                "Tokenizer ficará não treinado (KLS usará fallback unk-only).", e
-            )
-
-    def _train_serial_inprocess(self, corpus: List[str]) -> None:
-        """V6.5-V4-oom-fix: BBPE serial in-process (sem ProcessPoolExecutor).
-
-        Implementa o mesmo algoritmo Map-Reduce do train_parallel_from_stream
-        mas SEM spawn de subprocessos. Para corpus pequeno (< 200 textos),
-        esta implementação é O(10x) mais rápida e não causa OOM.
-
-        Algoritmo:
-          1. Pré-tokeniza corpus em símbolos byte-level (in-process)
-          2. Loop de merges:
-             - MAP: conta pares em sequência (sem paralelismo)
-             - REDUCE: agrega contagens
-             - CHOICE: k-means++ diversity selection
-             - APPLY: aplica merge (in-process)
-             - UPDATE: atualiza vocab + merges
-          3. Constrói tokenizer HF
-        """
-        from collections import Counter, defaultdict
-        logger.info(
-            "BBPE SERIAL (in-process): vocab_size=%d, corpus=%d textos",
-            self.vocab_size, len(corpus),
-        )
-        # Pré-tokeniza corpus em símbolos byte-level
-        shard_symbols: List[List[List[str]]] = []
-        for text in corpus:
-            # Byte-level pre-tokenization (igual pre_tokenize_shard)
-            words = text.split()
-            shard_symbols.append([list(w.encode('utf-8').decode('latin-1')) for w in words])
-        # Estruturas globais
-        current_vocab: set = set(ALPHABET + SPECIAL_TOKENS)
-        token_to_id: Dict[str, int] = {}
-        for i, tok in enumerate(SPECIAL_TOKENS):
-            token_to_id[tok] = i
-        next_id = len(SPECIAL_TOKENS)
-        for sym in ALPHABET:
-            if sym not in token_to_id:
-                token_to_id[sym] = next_id
-                next_id += 1
-        merges: List[Tuple[str, str, str]] = []
-        iteration = 0
-        min_frequency = 2
-        while len(current_vocab) < self.vocab_size:
-            iteration += 1
-            # FASE 1+2: MAP+REDUCE in-process
-            global_counts: Dict[Tuple[str, str], int] = defaultdict(int)
-            for sym_seq_list in shard_symbols:
-                for sym_seq in sym_seq_list:
-                    for i in range(len(sym_seq) - 1):
-                        pair = (sym_seq[i], sym_seq[i + 1])
-                        global_counts[pair] += 1
-            global_counts = {p: c for p, c in global_counts.items() if c >= min_frequency}
-            if not global_counts:
-                logger.info(
-                    "BBPE SERIAL: nenhum par com freq >= %d. Vocab final: %d (target %d)",
-                    min_frequency, len(current_vocab), self.vocab_size,
-                )
-                break
-            # FASE 3: k-means++ diversity
-            best_pair = _select_pair_kmeans_pp(global_counts, merges)
-            (esq, dir_), freq = best_pair
-            new_token_str = esq + dir_
-            while new_token_str in current_vocab:
-                new_token_str += "_"
-            new_id = next_id
-            next_id += 1
-            # FASE 4: APPLY merge in-process
-            for sym_seq_list in shard_symbols:
-                for idx_seq in range(len(sym_seq_list)):
-                    sym_seq = sym_seq_list[idx_seq]
-                    if len(sym_seq) < 2:
-                        continue
-                    new_seq: List[str] = []
-                    i = 0
-                    while i < len(sym_seq):
-                        if i < len(sym_seq) - 1 and sym_seq[i] == esq and sym_seq[i + 1] == dir_:
-                            new_seq.append(new_token_str)
-                            i += 2
-                        else:
-                            new_seq.append(sym_seq[i])
-                            i += 1
-                    sym_seq_list[idx_seq] = new_seq
-            # FASE 5: UPDATE
-            current_vocab.add(new_token_str)
-            token_to_id[new_token_str] = new_id
-            merges.append((esq, dir_, new_token_str))
-            if iteration % 50 == 0:
-                logger.info(
-                    "BBPE SERIAL iter %d: vocab=%d/%d",
-                    iteration, len(current_vocab), self.vocab_size,
-                )
-            if iteration % 100 == 0:
-                gc.collect()
-        logger.info(
-            "BBPE SERIAL: concluído. %d merges, vocab=%d. Construindo tokenizer HF...",
-            len(merges), len(token_to_id),
-        )
-        self._merges = merges
-        self._tokenizer = build_bpe_from_merges(
-            merges=merges,
-            token_to_id=token_to_id,
-            unk_token=self.unk_token,
-            add_prefix_space=self.add_prefix_space,
-        )
-        self._build_vocab_cache()
-        self.vocab_size = len(self._vocab)
-        logger.info("BBPE SERIAL: tokenizer construído. Vocab real: %d", len(self._vocab))
-
-    def _wordlevel_fallback(self, corpus: List[str]) -> None:
-        """V6.5-V4-oom-fix: Word-level fallback se BBPE falhar."""
-        from collections import Counter
-        logger.info("BBPE word-level fallback: corpus=%d textos", len(corpus))
-        word_counts = Counter()
-        for text in corpus:
-            word_counts.update(text.split())
-        sorted_words = [w for w, _ in word_counts.most_common(self.vocab_size - 3)]
-        token_to_id: Dict[str, int] = {"<pad>": 0, "<eos>": 1, "<unk>": 2}
-        for idx, word in enumerate(sorted_words, start=3):
-            token_to_id[word] = idx
-        merges: List[Tuple[str, str, str]] = []
-        self._merges = merges
-        self._tokenizer = build_bpe_from_merges(
-            merges=merges,
-            token_to_id=token_to_id,
-            unk_token=self.unk_token,
-            add_prefix_space=self.add_prefix_space,
-        )
-        self._build_vocab_cache()
-        self.vocab_size = len(self._vocab)
-        logger.info("BBPE word-level fallback: vocab=%d", len(self._vocab))
-
     # ------------------------------------------------------------------
     # TREINAMENTO (compatibilidade — delega para paralelo)
     # ------------------------------------------------------------------
@@ -901,6 +623,63 @@ class BBPETokenizer:
     # ------------------------------------------------------------------
     # Save / Load
     # ------------------------------------------------------------------
+    def fit(
+        self,
+        texts: List[str],
+        vocab_size: Optional[int] = None,
+        min_frequency: int = 2,
+    ) -> None:
+        """V6.6 — Compatibilidade com a API SimpleBBPETokenizer.fit().
+
+        Treina o tokenizer BBPE a partir de uma lista de textos. Esta método
+        existe para que o KohonenLearningSystem possa chamar
+        `kls.tokenizer.fit(corpus_inicial)` sem saber se o tokenizer é
+        SimpleBBPETokenizer (word-level) ou BBPETokenizer (byte-level).
+
+        Args:
+            texts: lista de strings para treinar o tokenizer.
+            vocab_size: tamanho do vocabulário alvo (default: self.vocab_size).
+            min_frequency: frequência mínima de um par para ser mergeado.
+        """
+        if vocab_size is not None:
+            self.vocab_size = int(vocab_size)
+        # Limita o corpus para velocidade (primeiros 1000 textos)
+        # — o BBPE byte-level não precisa de corpus massivo para treinar
+        # merges úteis; 1000 textos já dá diversidade suficiente.
+        corpus = list(texts[:1000]) if not isinstance(texts, Iterator) else texts
+        try:
+            self.train_parallel_from_stream(
+                text_iterator=iter(corpus) if isinstance(corpus, list) else corpus,
+                save_path=None,  # in-memory only
+                min_frequency=min_frequency,
+                num_workers=1,   # single worker for fit() (faster for small corpus)
+                show_progress=False,
+                chunk_size=500,
+            )
+        except Exception as e:
+            logger.warning(f"BBPETokenizer.fit() train_parallel_from_stream failed: {e}")
+            # Fallback: cria um tokenizer byte-level mínimo (256 bytes + specials)
+            # sem merges, permitindo que o KLS continue funcionando.
+            try:
+                from tokenizers import Tokenizer
+                from tokenizers.models import BPE
+                from tokenizers.pre_tokenizers import ByteLevel
+                # Vocab inicial: 256 bytes base mapeados para [4, 260)
+                # (deixando 0-3 para speciais BOS/PAD/EOS/UNK)
+                byte_vocab = {chr(b): b + 4 for b in range(256)}
+                # Adiciona speciais
+                for tok, tid in [(BOS_TOKEN, BOS_ID), (PAD_TOKEN, PAD_ID),
+                                 (EOS_TOKEN, EOS_ID), (UNK_TOKEN, UNK_ID)]:
+                    byte_vocab[tok] = tid
+                self._tokenizer = Tokenizer(BPE(byte_vocab, {}, unk_token=UNK_TOKEN))
+                self._tokenizer.pre_tokenizer = ByteLevel(add_prefix_space=False)
+                self._build_vocab_cache()
+                self.vocab_size = len(self._vocab)
+                logger.info(f"BBPETokenizer.fit() fallback byte-level: vocab_size={self.vocab_size}")
+            except Exception as e2:
+                logger.error(f"BBPETokenizer.fit() fallback failed: {e2}")
+                raise
+
     def save(self, path: Union[str, Path]) -> None:
         """Salva o tokenizer em arquivo JSON (formato HuggingFace)."""
         if self._tokenizer is None: