| """SentencePiece tokenizer training for SAGE.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| from pathlib import Path |
| from typing import Iterable, Iterator |
|
|
| DEFAULT_SPECIAL_TOKENS = ("<bos>", "<eos>", "<pad>", "<unk>", "[INST]", "[/INST]") |
|
|
|
|
| def iter_training_text(corpus_paths: Iterable[str], text_key: str = "text") -> Iterator[str]: |
| """Yield training lines from plain-text or JSONL corpus files.""" |
| for path in corpus_paths: |
| source = Path(path) |
| suffix = source.suffix.lower() |
| with source.open("r", encoding="utf-8") as handle: |
| if suffix == ".jsonl": |
| for raw_line in handle: |
| raw_line = raw_line.strip() |
| if not raw_line: |
| continue |
| payload = json.loads(raw_line) |
| text = payload.get(text_key) |
| if isinstance(text, str) and text.strip(): |
| yield text.strip() |
| continue |
| for raw_line in handle: |
| text = raw_line.strip() |
| if text: |
| yield text |
|
|
|
|
| def write_training_text(corpus_paths: Iterable[str], output_path: str, text_key: str = "text") -> str: |
| """Concatenate corpus text into a plain-text file for SentencePiece.""" |
| output = Path(output_path) |
| output.parent.mkdir(parents=True, exist_ok=True) |
| with output.open("w", encoding="utf-8") as sink: |
| for line in iter_training_text(corpus_paths, text_key=text_key): |
| sink.write(line) |
| sink.write("\n") |
| return str(output) |
|
|
|
|
| def train_sentencepiece(input_path: str, model_prefix: str, vocab_size: int = 50_000) -> None: |
| """Train a byte-fallback SentencePiece BPE model.""" |
| import sentencepiece as spm |
|
|
| spm.SentencePieceTrainer.train( |
| input=input_path, |
| model_prefix=model_prefix, |
| model_type="bpe", |
| vocab_size=vocab_size, |
| character_coverage=0.9995, |
| byte_fallback=True, |
| bos_id=0, |
| eos_id=1, |
| pad_id=2, |
| unk_id=3, |
| user_defined_symbols=list(DEFAULT_SPECIAL_TOKENS[4:]), |
| split_digits=False, |
| split_by_unicode_script=False, |
| remove_extra_whitespaces=False, |
| normalization_rule_name="identity", |
| hard_vocab_limit=False, |
| ) |
|
|
|
|
| def build_argparser() -> argparse.ArgumentParser: |
| """Build the CLI parser.""" |
| parser = argparse.ArgumentParser(description="Train the SAGE SentencePiece tokenizer.") |
| parser.add_argument("--input", nargs="+", required=True, help="Plain-text or JSONL corpus files.") |
| parser.add_argument("--model-prefix", default="tokenizer/tokenizer", help="SentencePiece model prefix.") |
| parser.add_argument("--vocab-size", type=int, default=50_000, help="Tokenizer vocabulary size.") |
| parser.add_argument("--training-text", default="tokenizer/training_corpus.txt", help="Temporary combined text file.") |
| parser.add_argument("--text-key", default="text", help="JSONL field to read when --input contains .jsonl files.") |
| return parser |
|
|
|
|
| def main() -> None: |
| """Train the tokenizer from CLI arguments.""" |
| args = build_argparser().parse_args() |
| training_text = write_training_text(args.input, args.training_text, text_key=args.text_key) |
| train_sentencepiece(training_text, args.model_prefix, args.vocab_size) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
