File size: 5,690 Bytes
49b8205 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 | """bbpe_tokenizer.py — BBPE (Byte-Level BPE) Tokenizer simplificado.
Inspirado em PowerMachine/BiGRU_T_version/src/bigru_t/tokenizer/bbpe_tokenizer.py,
mas simplificado para o projeto CNN-BiGRU. Usa a biblioteca `tokenizers` da
HuggingFace com modelo BPE byte-level.
Cobertura universal UTF-8 (qualquer string é tokenizável sem <unk>).
"""
from __future__ import annotations
import json
import logging
import os
from pathlib import Path
from typing import List, Optional, Sequence
from tokenizers import Tokenizer
from tokenizers.models import BPE
from tokenizers.pre_tokenizers import ByteLevel
from tokenizers.processors import ByteLevel as ByteLevelProcessor
from tokenizers.decoders import ByteLevel as ByteLevelDecoder
from tokenizers.trainers import BpeTrainer
logger = logging.getLogger(__name__)
BOS_TOKEN = "<s>"
PAD_TOKEN = "<pad>"
EOS_TOKEN = "</s>"
UNK_TOKEN = "<unk>"
SPECIAL_TOKENS = [BOS_TOKEN, PAD_TOKEN, EOS_TOKEN, UNK_TOKEN]
BOS_ID = 0
PAD_ID = 1
EOS_ID = 2
UNK_ID = 3
class BBPETokenizer:
"""Wrapper sobre Tokenizer (HuggingFace) com API conveniente.
Garante: dec(enc(s)) = s para todo s UTF-8.
"""
def __init__(self, tokenizer: Optional[Tokenizer] = None):
self._tok = tokenizer
@property
def tokenizer(self) -> Tokenizer:
if self._tok is None:
raise RuntimeError("Tokenizer não inicializado. Use train_from_* ou load.")
return self._tok
@property
def vocab_size(self) -> int:
return self.tokenizer.get_vocab_size()
@property
def bos_id(self) -> int:
return self.tokenizer.token_to_id(BOS_TOKEN) or BOS_ID
@property
def pad_id(self) -> int:
return self.tokenizer.token_to_id(PAD_TOKEN) or PAD_ID
@property
def eos_id(self) -> int:
return self.tokenizer.token_to_id(EOS_TOKEN) or EOS_ID
@property
def unk_id(self) -> int:
return self.tokenizer.token_to_id(UNK_TOKEN) or UNK_ID
def encode(self, text: str, add_special: bool = True) -> List[int]:
if add_special:
text = f"{BOS_TOKEN}{text}{EOS_TOKEN}"
return self.tokenizer.encode(text).ids
def encode_batch(self, texts: Sequence[str], add_special: bool = True) -> List[List[int]]:
if add_special:
texts = [f"{BOS_TOKEN}{t}{EOS_TOKEN}" for t in texts]
return [e.ids for e in self.tokenizer.encode_batch(list(texts))]
def decode(self, ids: List[int], skip_special: bool = True) -> str:
if skip_special:
ids = [i for i in ids if i not in (self.bos_id, self.pad_id, self.eos_id, self.unk_id)]
return self.tokenizer.decode(ids)
def save(self, path: str) -> None:
Path(path).parent.mkdir(parents=True, exist_ok=True)
self.tokenizer.save(path)
logger.info("Tokenizer salvo em %s", path)
@classmethod
def load(cls, path: str) -> "BBPETokenizer":
tok = Tokenizer.from_file(path)
return cls(tok)
@classmethod
def train_from_files(
cls,
files: Sequence[str],
vocab_size: int = 8000,
min_frequency: int = 2,
) -> "BBPETokenizer":
# Desativa paralelismo de tokenizers para evitar problemas de GIL no shutdown
os.environ["TOKENIZERS_PARALLELISM"] = "false"
tokenizer = Tokenizer(BPE(unk_token=UNK_TOKEN))
tokenizer.pre_tokenizer = ByteLevel(add_prefix_space=False)
trainer = BpeTrainer(
vocab_size=vocab_size,
min_frequency=min_frequency,
special_tokens=SPECIAL_TOKENS,
show_progress=False,
)
valid_files = [str(Path(p).resolve()) for p in files if Path(p).exists()]
if not valid_files:
raise FileNotFoundError(f"Nenhum arquivo válido em: {files}")
tokenizer.train(valid_files, trainer)
tokenizer.post_processor = ByteLevelProcessor(trim_offsets=False)
# Decoder byte-level: reverte a codificação byte-level para UTF-8
tokenizer.decoder = ByteLevelDecoder()
logger.info("BBPE treinado com vocab_size=%d", tokenizer.get_vocab_size())
return cls(tokenizer)
@classmethod
def train_from_texts(
cls,
texts: Sequence[str],
vocab_size: int = 8000,
min_frequency: int = 1,
) -> "BBPETokenizer":
"""Treina a partir de uma lista de textos em memória."""
# Salva temporariamente
import tempfile
with tempfile.NamedTemporaryFile(
mode="w", suffix=".txt", delete=False, encoding="utf-8"
) as f:
for line in texts:
line = line.strip()
if line:
f.write(line + "\n")
tmp_path = f.name
try:
return cls.train_from_files([tmp_path], vocab_size=vocab_size,
min_frequency=min_frequency)
finally:
try:
os.unlink(tmp_path)
except OSError:
pass
def validate_roundtrip(self, test_texts: Sequence[str]) -> float:
"""Valida que dec(enc(s)) == s. Retorna fração de sucessos."""
if not test_texts:
return 1.0
ok = 0
for s in test_texts:
try:
ids = self.encode(s, add_special=False)
decoded = self.decode(ids, skip_special=False)
if decoded == s:
ok += 1
except Exception:
pass
return ok / len(test_texts)
__all__ = [
"BBPETokenizer",
"BOS_TOKEN", "PAD_TOKEN", "EOS_TOKEN", "UNK_TOKEN",
"BOS_ID", "PAD_ID", "EOS_ID", "UNK_ID",
"SPECIAL_TOKENS",
]
|