scripts/train_tokenizer.py

"""Command line tokenizer trainer for Delta Ultra Mini."""

from __future__ import annotations

import argparse
import logging
import os
import sys
from pathlib import Path

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from delta.tokenizer import train_tokenizer

logging.basicConfig(level=os.getenv("DELTA_LOG_LEVEL", "INFO").upper())
logger = logging.getLogger(__name__)


def parse_args() -> argparse.Namespace:
    """Parse command line arguments."""

    parser = argparse.ArgumentParser(description="Train Delta BPE tokenizer.")
    parser.add_argument("--corpus_files", nargs="+", required=True, help="Input .txt corpus files.")
    parser.add_argument("--output_path", required=True, help="Output tokenizer JSON path.")
    return parser.parse_args()


def main() -> None:
    """Train and save a tokenizer."""

    args = parse_args()
    train_tokenizer([Path(path) for path in args.corpus_files], Path(args.output_path))
    logger.info("Tokenizer training complete.")


if __name__ == "__main__":
    main()