"""soup tokenizer — BPE tokenizer training (v0.53.9 #15).""" from __future__ import annotations import json import os import stat from pathlib import Path from typing import Optional import typer from rich.console import Console from rich.panel import Panel from soup_cli.utils.paths import is_under_cwd console = Console() app = typer.Typer( help="Tokenizer tools: train a BPE tokenizer from a JSONL corpus.", no_args_is_help=True, ) _MAX_VOCAB_SIZE = 200_000 _MIN_VOCAB_SIZE = 256 _MAX_CORPUS_BYTES = 50 * 1024 * 1024 # 50 MB total cap _MAX_LINE_BYTES = 8 * 1024 # 8 KB per-line cap def _validate_under_cwd(path: str, *, label: str) -> str: if not isinstance(path, str) or not path: raise typer.BadParameter(f"{label} must be a non-empty path") if "\x00" in path: raise typer.BadParameter(f"{label} contains NUL byte") if not is_under_cwd(path): raise typer.BadParameter( f"{label} must stay under the current working directory: {path}" ) # TOCTOU defence: lstat the RAW user-supplied path BEFORE realpath # resolution (matches v0.53.7 #106 policy). `os.path.realpath` would # follow the link and `lstat` on the resolved target would see a # regular file even when the user-supplied entry is a symlink. try: st = os.lstat(path) except OSError: # Missing-file path is handled by the caller's `is_file` check. return os.path.realpath(path) if stat.S_ISLNK(st.st_mode): raise typer.BadParameter( f"{label} target must not be a symlink" ) return os.path.realpath(path) def _extract_texts(input_path: Path) -> list[str]: """Read JSONL or plaintext lines into a list of training strings. Capped at 50 MB total + 8 KB per line to bound RAM on adversarial input. """ texts: list[str] = [] total_bytes = 0 if input_path.suffix.lower() == ".jsonl": with open(input_path, encoding="utf-8") as fh: for raw in fh: if len(raw.encode("utf-8", errors="replace")) > _MAX_LINE_BYTES: continue # silently skip oversized line total_bytes += len(raw) if total_bytes > _MAX_CORPUS_BYTES: break text = raw.strip() if not text: continue try: row = json.loads(text) except json.JSONDecodeError: continue if isinstance(row, dict): for key in ("text", "content", "prompt"): value = row.get(key) if isinstance(value, str) and value: texts.append(value) break else: # ShareGPT-style messages list msgs = row.get("messages") if isinstance(msgs, list): for msg in msgs: if isinstance(msg, dict): content = msg.get("content") if isinstance(content, str) and content: texts.append(content) elif isinstance(row, str): texts.append(row) else: with open(input_path, encoding="utf-8") as fh: for raw in fh: if len(raw.encode("utf-8", errors="replace")) > _MAX_LINE_BYTES: continue total_bytes += len(raw) if total_bytes > _MAX_CORPUS_BYTES: break text = raw.rstrip("\n\r") if text: texts.append(text) return texts @app.command(name="train") def train( input_path: str = typer.Option( ..., "--input", "-i", help="Path to JSONL or .txt corpus (must stay under cwd).", ), vocab_size: int = typer.Option( 32_000, "--vocab-size", "-v", help=f"Target vocabulary size ({_MIN_VOCAB_SIZE}-{_MAX_VOCAB_SIZE}).", ), output: str = typer.Option( "tokenizer_out", "--output", "-o", help="Output directory for tokenizer.json + vocab.json (under cwd).", ), min_frequency: int = typer.Option( 2, "--min-frequency", help="Minimum pair frequency (>=1).", ), special_tokens: Optional[list[str]] = typer.Option( None, "--special-token", help="Add special token (repeatable). Defaults: , , , .", ), ) -> None: """Train a BPE tokenizer from a local corpus. Outputs `tokenizer.json` (Hugging Face format) and `vocab.json` to `--output`. Requires the `tokenizers` library (bundled with `transformers`). """ if isinstance(vocab_size, bool) or not isinstance(vocab_size, int): raise typer.BadParameter("--vocab-size must be int") if not (_MIN_VOCAB_SIZE <= vocab_size <= _MAX_VOCAB_SIZE): raise typer.BadParameter( f"--vocab-size must be in [{_MIN_VOCAB_SIZE}, {_MAX_VOCAB_SIZE}]" ) if isinstance(min_frequency, bool) or not isinstance(min_frequency, int): raise typer.BadParameter("--min-frequency must be int") if min_frequency < 1: raise typer.BadParameter("--min-frequency must be >= 1") input_real = _validate_under_cwd(input_path, label="--input") output_real = _validate_under_cwd(output, label="--output") resolved_input = Path(input_real) if not resolved_input.is_file(): console.print(f"[red]Input file not found:[/] {input_path}") raise typer.Exit(1) output_path = Path(output_real) try: from tokenizers import Tokenizer, models, pre_tokenizers, trainers except ImportError as exc: console.print( "[red]The 'tokenizers' package is required.[/]\n" "Install with: [bold]pip install tokenizers[/]" ) raise typer.Exit(1) from exc texts = _extract_texts(resolved_input) if not texts: console.print( f"[red]No usable training text found in[/] {resolved_input.name}" ) raise typer.Exit(1) tokens = list(special_tokens or []) if not tokens: tokens = ["", "", "", ""] # Sanitise: reject NUL / oversize entries; dedup preserves first-seen order. cleaned: list[str] = [] seen: set[str] = set() for tok in tokens: if not isinstance(tok, str) or not tok or "\x00" in tok or len(tok) > 128: raise typer.BadParameter( f"--special-token entry must be a non-empty NUL-free str <=128 chars: {tok!r}" ) if tok in seen: continue seen.add(tok) cleaned.append(tok) console.print( Panel( f"Input: [bold]{resolved_input.name}[/]\n" f"Texts: [bold]{len(texts)}[/]\n" f"Vocab size: [bold]{vocab_size}[/]\n" f"Output dir: [bold]{output_path.name}[/]", title="BPE Tokenizer Training", ) ) tokenizer = Tokenizer(models.BPE(unk_token="")) tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False) trainer = trainers.BpeTrainer( vocab_size=vocab_size, min_frequency=min_frequency, special_tokens=cleaned, show_progress=False, ) tokenizer.train_from_iterator(texts, trainer=trainer) output_path.mkdir(parents=True, exist_ok=True) # Post-mkdir re-check: defends against a symlink planted between the # containment check and the write. Matches v0.43.0 Part D `copy_bundle_to` # TOCTOU policy. try: post_st = os.lstat(output_path) except OSError as exc: console.print(f"[red]Output directory not accessible:[/] {exc}") raise typer.Exit(1) from exc if stat.S_ISLNK(post_st.st_mode): console.print( "[red]Output directory is a symlink — refusing to write.[/]" ) raise typer.Exit(1) if not is_under_cwd(str(output_path)): console.print("[red]Output directory escaped cwd after mkdir.[/]") raise typer.Exit(1) tokenizer_file = output_path / "tokenizer.json" tokenizer.save(str(tokenizer_file)) vocab_file = output_path / "vocab.json" vocab = tokenizer.get_vocab() vocab_file.write_text( json.dumps(vocab, ensure_ascii=False, indent=2), encoding="utf-8", ) console.print( f"[green]Trained tokenizer with[/] [bold]{len(vocab)}[/] tokens " f"-> [bold]{output_path.name}/{tokenizer_file.name}[/]" )