File size: 2,942 Bytes
ee7f83c 60d3f8f ee7f83c 60d3f8f ee7f83c 60d3f8f ee7f83c 60d3f8f ee7f83c 60d3f8f ee7f83c 60d3f8f ee7f83c 60d3f8f ee7f83c 60d3f8f ee7f83c 60d3f8f ee7f83c 60d3f8f ee7f83c 60d3f8f ee7f83c 60d3f8f ee7f83c 60d3f8f ee7f83c 60d3f8f ee7f83c 60d3f8f ee7f83c 60d3f8f | 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 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 | import json
import sentencepiece as spm
from src.utils.helpers import JSON_DATA, TOKENIZER_MODEL
def train_sentencepiece_model():
"""
Reads the tokenized JSONL training dataset, extracts the raw text,
and trains a Shared-Vocabulary SentencePiece (BPE) model.
"""
# Define working paths based on your config
jsonl_path = JSON_DATA
# Create a temporary text file in the same directory as the tokenizer model
tokenizer_dir = TOKENIZER_MODEL.parent
tokenizer_dir.mkdir(parents=True, exist_ok=True)
corpus_path = tokenizer_dir / "spm_training_corpus.txt"
# SentencePiece requires a prefix without the '.model' extension
spm_prefix = str(TOKENIZER_MODEL.with_suffix(""))
vocab_size = 4000
print(f"Building corpus from {jsonl_path}...")
line_count = 0
# Extract raw source + target text into a flat corpus file
with open(jsonl_path, "r", encoding="utf-8") as f_in, open(
corpus_path, "w", encoding="utf-8"
) as f_out:
for line in f_in:
item = json.loads(line)
# Prefer raw text; fall back to joined tokens
src = item.get("source_text") or " ".join(item.get("source_tokens", []))
tgt = item.get("target_text") or " ".join(item.get("target_tokens", []))
if src.strip():
f_out.write(src.strip() + "\n")
line_count += 1
if tgt.strip():
f_out.write(tgt.strip() + "\n")
line_count += 1
print(f"Corpus written: {line_count:,} lines → {corpus_path}")
# Train the SentencePiece BPE model
print(f"\nTraining SentencePiece BPE model (vocab_size={vocab_size})...")
spm.SentencePieceTrainer.train(
input=str(corpus_path),
model_prefix=spm_prefix,
vocab_size=vocab_size,
model_type="bpe",
# Match existing special token index layout exactly
pad_id=0,
pad_piece="<PAD>",
bos_id=1,
bos_piece="<SOS>",
eos_id=2,
eos_piece="<EOS>",
unk_id=3,
unk_piece="<UNK>",
character_coverage=1.0,
add_dummy_prefix=True,
input_sentence_size=500000,
shuffle_input_sentence=True,
)
print("\nDone! Saved:")
print(f" {spm_prefix}.model ← load this at inference time")
print(f" {spm_prefix}.vocab ← human-readable token list")
# Cleanup the temporary corpus text file to save space
if corpus_path.exists():
corpus_path.unlink()
print(f"Cleaned up temporary corpus file: {corpus_path}")
# Quick sanity check
sp = spm.SentencePieceProcessor()
sp.load(f"{spm_prefix}.model")
print(f"\nVocab Size: {sp.get_piece_size()}")
if __name__ == "__main__":
train_sentencepiece_model()
"""
To run:
python scripts/train_tokenizer.py
This is used for the transformer models to improve
performance by grouping tokens together.
"""
|