File size: 2,276 Bytes
fd3e507 | 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 | """BPE tokenizer for Fractus-1B (byte-level, GPT-2 compatible).
Wraps the HuggingFace `tokenizers` library. Builds or loads a byte-level BPE
tokenizer with vocab_size ~50k, suitable for multi-language code/math/text.
Usage:
tok = FractusTokenizer.build_or_load()
ids = tok.encode("def hello(): pass")
text = tok.decode(ids)
"""
import os
from typing import List, Optional
from tokenizers import Tokenizer
from tokenizers.models import BPE
from tokenizers.trainers import BpeTrainer
from tokenizers.pre_tokenizers import ByteLevel
from tokenizers.decoders import ByteLevel as ByteLevelDecoder
from tokenizers.processors import TemplateProcessing
DEFAULT_VOCAB_SIZE = 50257
DEFAULT_TOKENIZER_PATH = os.path.join(
os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))),
"data", "tokenizer", "fractus_bpe.json",
)
class FractusTokenizer:
"""Wrapper around a byte-level BPE tokenizer."""
def __init__(self, tokenizer: Tokenizer):
self.tokenizer = tokenizer
self.vocab_size = tokenizer.get_vocab_size()
@classmethod
def gpt2_compatible(cls) -> "FractusTokenizer":
"""Load a GPT-2-compatible tokenizer from HF (same vocab, same byte-level BPE).
This gives us a ready-to-use 50k vocab without training our own."""
from tokenizers import Tokenizer
tok = Tokenizer.from_pretrained("gpt2")
return cls(tok)
@classmethod
def build_or_load(cls, path: Optional[str] = None) -> "FractusTokenizer":
"""Load from path if exists, else use GPT-2 compatible."""
path = path or DEFAULT_TOKENIZER_PATH
if os.path.exists(path):
tok = Tokenizer.from_file(path)
return cls(tok)
return cls.gpt2_compatible()
def encode(self, text: str) -> List[int]:
return self.tokenizer.encode(text).ids
def encode_batch(self, texts: List[str]) -> List[List[int]]:
return [enc.ids for enc in self.tokenizer.encode_batch(texts)]
def decode(self, ids: List[int]) -> str:
return self.tokenizer.decode(ids)
def save(self, path: Optional[str] = None):
path = path or DEFAULT_TOKENIZER_PATH
os.makedirs(os.path.dirname(path), exist_ok=True)
self.tokenizer.save(path)
|