mirror of https://github.com/razor-ai/soup.git
248 lines
8.6 KiB
Python
248 lines
8.6 KiB
Python
"""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: <pad>, <s>, </s>, <unk>.",
|
|
),
|
|
) -> 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 = ["<pad>", "<s>", "</s>", "<unk>"]
|
|
# 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="<unk>"))
|
|
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}[/]"
|
|
)
|