File size: 1,704 Bytes
78dea75 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 | #!/usr/bin/env python3
"""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())
|