File size: 3,065 Bytes
223d3da | 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 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 | #!/usr/bin/env python3
"""Train the Emese Hungarian SentencePiece tokenizer (32K unigram)."""
import argparse, json, os, sys, tempfile, time
import sentencepiece as spm
def build_corpus(path: str, limit: int) -> str:
"""Extract up to `limit` texts from JSONL into a temp file."""
if not os.path.exists(path):
sys.exit(f"Error: input file not found: {path}")
tmp = tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False, encoding="utf-8")
count = 0
with open(path) as f:
for line in f:
if text := json.loads(line).get("text"):
tmp.write(text + "\n")
count += 1
if count >= limit:
break
tmp.close()
print(f"Sampled {count:,} paragraphs from {path}")
return tmp.name
def train_tokenizer(corpus: str, prefix: str, vocab_size: int) -> None:
os.makedirs(os.path.dirname(prefix) or ".", exist_ok=True)
print(f"Training {vocab_size:,}-vocab tokenizer → {prefix}.model")
spm.SentencePieceTrainer.train(
input=corpus, model_prefix=prefix, model_type="unigram",
vocab_size=vocab_size, character_coverage=0.9999,
pad_id=-1, unk_id=0, bos_id=1, eos_id=2,
user_defined_symbols=["<eos>"],
shrinking_factor=0.85, num_sub_iterations=4,
max_sentencepiece_length=24, seed_sentencepiece_size=2_000_000,
split_digits=False, byte_fallback=False, num_threads=os.cpu_count(),
input_sentence_size=500_000, shuffle_input_sentence=True,
normalization_rule_name="nfkc",
)
def evaluate(model_path: str, corpus: str, limit: int = 50_000) -> None:
sp = spm.SentencePieceProcessor(model_file=f"{model_path}.model")
chars = tokens = words = 0
with open(corpus) as f:
for i, line in enumerate(f):
if i >= limit:
break
if text := line.strip():
ids = sp.encode(text, out_type=int)
chars += len(text); tokens += len(ids); words += len(text.split())
print(f" Chars/token: {chars/tokens:.2f} | Tokens/word: {tokens/words:.2f}")
def main():
p = argparse.ArgumentParser(description="Train Emese SentencePiece tokenizer")
p.add_argument("--input", default="data/clean/wiki_clean.jsonl", help="Input JSONL")
p.add_argument("--model-dir", default="tokenizer", help="Output directory")
p.add_argument("--model-name", default="emese-tokenizer", help="Model prefix name")
p.add_argument("--vocab-size", type=int, default=32_000, help="Vocabulary size")
p.add_argument("--sample-limit", type=int, default=500_000, help="Max paragraphs to sample")
args = p.parse_args()
t0 = time.time()
prefix = os.path.join(args.model_dir, args.model_name)
corpus = build_corpus(args.input, args.sample_limit)
try:
train_tokenizer(corpus, prefix, args.vocab_size)
evaluate(prefix, corpus)
finally:
os.unlink(corpus)
print(f"Done in {time.time()-t0:.1f}s → {prefix}.model")
if __name__ == "__main__":
main()
|