| """ |
| BoccaccioAI - BPE Tokenizer Training Script |
| |
| Trains a byte-level BPE tokenizer on a subset of CulturaX Italian data |
| using the HuggingFace tokenizers library. |
| |
| De Lauretis Tech |
| """ |
|
|
| import argparse |
| import logging |
| import os |
| import sys |
| from pathlib import Path |
|
|
| import yaml |
| from tokenizers import Tokenizer, models, trainers |
| from tokenizers.decoders import ByteLevel as ByteLevelDecoder |
| from tokenizers.normalizers import NFKC |
| from tokenizers.pre_tokenizers import ByteLevel |
| from tqdm import tqdm |
| from transformers import PreTrainedTokenizerFast |
|
|
| logging.basicConfig( |
| level=logging.INFO, |
| format="%(asctime)s [%(levelname)s] %(name)s - %(message)s", |
| datefmt="%Y-%m-%d %H:%M:%S", |
| ) |
| logger = logging.getLogger("boccaccio.tokenizer") |
|
|
| BYTES_PER_GB = 1024 ** 3 |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser( |
| description="Train a BPE tokenizer for BoccaccioAI on CulturaX Italian data." |
| ) |
| parser.add_argument( |
| "--config", |
| type=str, |
| default="configs/tokenizer.yaml", |
| help="Path to the tokenizer YAML configuration file.", |
| ) |
| parser.add_argument( |
| "--output-dir", |
| type=str, |
| default=None, |
| help="Override the output directory from config.", |
| ) |
| parser.add_argument( |
| "--corpus-size-gb", |
| type=float, |
| default=None, |
| help="Override the corpus size (in GB) used for training.", |
| ) |
| parser.add_argument( |
| "--vocab-size", |
| type=int, |
| default=None, |
| help="Override the vocabulary size from config.", |
| ) |
| return parser.parse_args() |
|
|
|
|
| def load_config(config_path: str) -> dict: |
| config_file = Path(config_path) |
| if not config_file.exists(): |
| logger.error("Config file not found: %s", config_path) |
| sys.exit(1) |
| with open(config_file, "r", encoding="utf-8") as f: |
| config = yaml.safe_load(f) |
| logger.info("Loaded config from %s", config_path) |
| return config |
|
|
|
|
| def corpus_iterator(corpus_size_gb: float, source: str, lang: str): |
| """Stream text samples from CulturaX until the target byte size is reached.""" |
| from datasets import load_dataset |
|
|
| target_bytes = int(corpus_size_gb * BYTES_PER_GB) |
| logger.info( |
| "Streaming %s (lang=%s), target size: %.2f GB (%s bytes)", |
| source, lang, corpus_size_gb, f"{target_bytes:,}", |
| ) |
|
|
| dataset = load_dataset(source, lang, split="train", streaming=True) |
|
|
| collected_bytes = 0 |
| pbar = tqdm( |
| total=target_bytes, |
| unit="B", |
| unit_scale=True, |
| desc="Collecting corpus", |
| ) |
|
|
| try: |
| for example in dataset: |
| text = example["text"] |
| text_bytes = len(text.encode("utf-8")) |
| collected_bytes += text_bytes |
| pbar.update(text_bytes) |
| yield text |
|
|
| if collected_bytes >= target_bytes: |
| break |
| finally: |
| pbar.close() |
|
|
| logger.info( |
| "Corpus collection complete: %.2f GB collected", |
| collected_bytes / BYTES_PER_GB, |
| ) |
|
|
|
|
| def build_tokenizer( |
| vocab_size: int, |
| special_tokens: list[str], |
| min_frequency: int, |
| show_progress: bool, |
| ) -> tuple[Tokenizer, trainers.BpeTrainer]: |
| """Construct a byte-level BPE tokenizer and its trainer.""" |
| tokenizer = Tokenizer(models.BPE()) |
| tokenizer.normalizer = NFKC() |
| tokenizer.pre_tokenizer = ByteLevel() |
| tokenizer.decoder = ByteLevelDecoder() |
|
|
| trainer = trainers.BpeTrainer( |
| vocab_size=vocab_size, |
| min_frequency=min_frequency, |
| special_tokens=special_tokens, |
| show_progress=show_progress, |
| ) |
|
|
| return tokenizer, trainer |
|
|
|
|
| def validate_tokenizer(tokenizer: Tokenizer, special_tokens: list[str]) -> None: |
| """Print validation stats and fertility metrics for the trained tokenizer.""" |
| vocab_size = tokenizer.get_vocab_size() |
| logger.info("Vocabulary size: %d", vocab_size) |
|
|
| for token in special_tokens: |
| token_id = tokenizer.token_to_id(token) |
| logger.info(" Special token: %-8s -> id %s", token, token_id) |
|
|
| test_sentences = [ |
| "L'intelligenza artificiale sta trasformando il mondo della tecnologia.", |
| "Il gatto nero dorme tranquillamente sul divano di casa.", |
| "La Repubblica Italiana e' una democrazia parlamentare fondata sul lavoro.", |
| ] |
|
|
| logger.info("--- Fertility analysis ---") |
| for sentence in test_sentences: |
| encoding = tokenizer.encode(sentence) |
| num_tokens = len(encoding.ids) |
| num_words = len(sentence.split()) |
| fertility = num_tokens / num_words if num_words > 0 else 0.0 |
|
|
| logger.info(" Text: %s", sentence) |
| logger.info(" Tokens: %d (words: %d, fertility: %.2f tok/word)", num_tokens, num_words, fertility) |
| logger.info(" Token ids: %s", encoding.ids) |
| logger.info(" Decoded: %s", tokenizer.decode(encoding.ids)) |
| logger.info("") |
|
|
|
|
| def main() -> None: |
| args = parse_args() |
| config = load_config(args.config) |
| tok_cfg = config["tokenizer"] |
|
|
| |
| vocab_size = args.vocab_size or tok_cfg["vocab_size"] |
| corpus_size_gb = args.corpus_size_gb if args.corpus_size_gb is not None else tok_cfg["training_corpus_size_gb"] |
| output_dir = args.output_dir or tok_cfg["output_dir"] |
| output_name = tok_cfg["output_name"] |
| source = tok_cfg["training_corpus_source"] |
| lang = tok_cfg["training_corpus_lang"] |
| min_frequency = tok_cfg.get("min_frequency", 2) |
| show_progress = tok_cfg.get("show_progress", True) |
|
|
| special_tokens_cfg = tok_cfg["special_tokens"] |
| special_tokens = [ |
| special_tokens_cfg["pad"], |
| special_tokens_cfg["bos"], |
| special_tokens_cfg["eos"], |
| special_tokens_cfg["unk"], |
| ] |
|
|
| logger.info("=== BoccaccioAI Tokenizer Training ===") |
| logger.info("Vocab size: %d", vocab_size) |
| logger.info("Corpus size: %.2f GB", corpus_size_gb) |
| logger.info("Output dir: %s", output_dir) |
| logger.info("Special tkns: %s", special_tokens) |
|
|
| |
| tokenizer, trainer = build_tokenizer( |
| vocab_size=vocab_size, |
| special_tokens=special_tokens, |
| min_frequency=min_frequency, |
| show_progress=show_progress, |
| ) |
|
|
| |
| logger.info("Starting tokenizer training...") |
| text_iterator = corpus_iterator(corpus_size_gb, source, lang) |
| tokenizer.train_from_iterator(text_iterator, trainer=trainer) |
| logger.info("Training complete.") |
|
|
| |
| output_path = Path(output_dir) |
| output_path.mkdir(parents=True, exist_ok=True) |
|
|
| json_path = output_path / f"{output_name}.json" |
| tokenizer.save(str(json_path)) |
| logger.info("Saved raw tokenizer to %s", json_path) |
|
|
| |
| hf_tokenizer = PreTrainedTokenizerFast( |
| tokenizer_object=tokenizer, |
| bos_token=special_tokens_cfg["bos"], |
| eos_token=special_tokens_cfg["eos"], |
| unk_token=special_tokens_cfg["unk"], |
| pad_token=special_tokens_cfg["pad"], |
| ) |
| hf_dir = output_path / f"{output_name}-hf" |
| hf_tokenizer.save_pretrained(str(hf_dir)) |
| logger.info("Saved HuggingFace tokenizer to %s", hf_dir) |
|
|
| |
| validate_tokenizer(tokenizer, special_tokens) |
| logger.info("Done.") |
|
|
|
|
| if __name__ == "__main__": |
| try: |
| main() |
| except KeyboardInterrupt: |
| logger.warning("Training interrupted by user.") |
| sys.exit(130) |
|
|
