dela
584f7e9e73
Standalone tree split from LLMRL/projects/kda. Includes Triton dt_bias backward fix, train_k3 --preset 0.5b, SFT, Docker runtime, and tests.
52 lines
1.5 KiB
Python
52 lines
1.5 KiB
Python
"""Train a SentencePiece tokenizer on a Chinese Wikipedia subset.
|
|
|
|
用法:
|
|
uv run python kda/training/train_tokenizer.py \
|
|
--out data/spm_4k --vocab-size 4096 --limit 20000
|
|
|
|
产出:
|
|
data/spm_4k.model / data/spm_4k.vocab (BPE/unigram, 中文小语料)
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
|
|
import sentencepiece as spm
|
|
|
|
from .data import fetch_wiki_texts
|
|
|
|
|
|
def main() -> None:
|
|
p = argparse.ArgumentParser(description=__doc__)
|
|
p.add_argument("--out", default="data/spm_4k", help="输出前缀 (model/vocab 文件)")
|
|
p.add_argument("--vocab-size", type=int, default=8192)
|
|
p.add_argument("--limit", type=int, default=20000, help="用于训练的 wiki 文章数")
|
|
p.add_argument("--model-type", default="unigram", choices=["unigram", "bpe"])
|
|
p.add_argument("--character-coverage", type=float, default=0.9995)
|
|
args = p.parse_args()
|
|
|
|
texts = fetch_wiki_texts(args.limit)
|
|
corpus = "".join(texts)
|
|
tmp = args.out + ".corpus.txt"
|
|
with open(tmp, "w", encoding="utf-8") as f:
|
|
f.write(corpus)
|
|
print(f"corpus: {len(corpus):,} chars from {len(texts)} articles")
|
|
|
|
spm.SentencePieceTrainer.train(
|
|
input=tmp,
|
|
model_prefix=args.out,
|
|
vocab_size=args.vocab_size,
|
|
model_type=args.model_type,
|
|
character_coverage=args.character_coverage,
|
|
unk_id=0,
|
|
pad_id=1,
|
|
bos_id=-1,
|
|
eos_id=-1,
|
|
num_threads=4,
|
|
)
|
|
print(f"tokenizer saved: {args.out}.model / {args.out}.vocab")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|