| |
| """Train SentencePiece model and save a Hugging Face–compatible fast tokenizer.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import sys |
| from pathlib import Path |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| sys.path.insert(0, str(ROOT / "src")) |
|
|
| from weather_llm.config_loader import repo_root_from |
| from weather_llm.tokenization.hf_tokenizer import load_llama_tokenizer_from_spm |
| from weather_llm.tokenization.spm_trainer import train_sentencepiece |
| from weather_llm.tokenization.validate import validate_tokenizer, write_validation_report |
|
|
|
|
| def main() -> int: |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--corpus", type=Path, required=True, help="JSONL with `text` field") |
| ap.add_argument("--out_dir", type=Path, default=None) |
| ap.add_argument("--vocab_size", type=int, default=32000) |
| args = ap.parse_args() |
| root = repo_root_from(ROOT) |
| corpus = args.corpus if args.corpus.is_absolute() else (root / args.corpus) |
| out_dir = args.out_dir or (root / "artifacts/tokenizer") |
| out_dir.mkdir(parents=True, exist_ok=True) |
| prefix = out_dir / "weather_spm" |
| model_file = train_sentencepiece( |
| corpus, |
| prefix, |
| vocab_size=args.vocab_size, |
| sample_docs=None, |
| ) |
| tok = load_llama_tokenizer_from_spm(model_file) |
| tok.save_pretrained(str(out_dir)) |
| report = validate_tokenizer(tok) |
| write_validation_report(out_dir / "validation_report.txt", report) |
| if report["issues"]: |
| print("Tokenizer validation warnings:") |
| for i in report["issues"]: |
| print(" -", i) |
| print("Saved tokenizer to", out_dir) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
