SentencePiece

SentencePiece tokenization utilities for subword segmentation.

Overview

SentencePiece is a language-independent subword tokenizer that can handle any language without pre-tokenization. This module provides utilities for training SentencePiece models and using them with TorchLingo.

Why SentencePiece?

Problem	Solution
Unknown words	Split into known subwords
Large vocabulary	Control vocabulary size
Language-agnostic	Works for any language
Rare words	Shared subword representations

Example

"unfortunately" → ["▁un", "fortun", "ately"]
"preprocessing" → ["▁pre", "process", "ing"]

The ▁ character marks the start of a word.

Quick Start

Training a Model

from torchlingo.preprocessing import train_sentencepiece

train_sentencepiece(
    input_files=["data/train.txt"],
    model_prefix="models/sp",
    vocab_size=8000,
    model_type="bpe",
)
# Creates: models/sp.model, models/sp.vocab

Using a Model

from torchlingo.data_processing import SentencePieceVocab

sp = SentencePieceVocab("models/sp.model")

# Tokenize
tokens = sp.tokenize("Hello world")
print(tokens)  # ['▁Hello', '▁world']

# Encode to indices
indices = sp.encode("Hello world", add_special_tokens=True)
print(indices)  # [2, 145, 892, 3]

# Decode back
text = sp.decode(indices, skip_special_tokens=True)
print(text)  # "Hello world"

API Reference

train_sentencepiece

train_sentencepiece(input_files: List[Path], model_prefix: str, vocab_size: int = None, columns: Optional[List[str]] = None, src_col: str = None, tgt_col: str = None, model_type: str = None, character_coverage: float = None, normalization_rule_name: str = None, pad_idx: int = None, unk_idx: int = None, sos_idx: int = None, eos_idx: int = None, pad_token: str = None, unk_token: str = None, sos_token: str = None, eos_token: str = None, user_defined_symbols: Optional[List[str]] = None, config: Config = None)

Train a SentencePiece tokenization model on raw text data.

Reads all input files, extracts selected columns (src/tgt by default), concatenates them into a temporary file, and trains a SentencePiece model. Model files are saved with the specified prefix (e.g., 'model.model' and 'model.vocab').

Model training respects configuration for normalization, special tokens, and character coverage.

Parameters:

Name	Type	Description	Default
input_files	List[Path]	Data files to use for training.	required
model_prefix	str	Output prefix for .model and .vocab files.	required
vocab_size	int	Target vocabulary size. Falls back to config.vocab_size.	None
columns	List[str]	Columns to use for training data. Defaults to [src_col, tgt_col].	None
src_col	str	Source column name. Falls back to config.src_col.	None
tgt_col	str	Target column name. Falls back to config.tgt_col.	None
model_type	str	SentencePiece model type. Falls back to config.sp_model_type.	None
character_coverage	float	Character coverage. Falls back to config.sp_character_coverage.	None
normalization_rule_name	str	Normalization rule. Falls back to config.sp_normalization_rule_name.	None
pad_idx	int	Padding token index. Falls back to config.pad_idx.	None
unk_idx	int	Unknown token index. Falls back to config.unk_idx.	None
sos_idx	int	Start-of-sequence token index. Falls back to config.sos_idx.	None
eos_idx	int	End-of-sequence token index. Falls back to config.eos_idx.	None
pad_token	str	Padding token string. Falls back to config.pad_token.	None
unk_token	str	Unknown token string. Falls back to config.unk_token.	None
sos_token	str	Start-of-sequence token string. Falls back to config.sos_token.	None
eos_token	str	End-of-sequence token string. Falls back to config.eos_token.	None
user_defined_symbols	List[str]	Symbols to register as whole vocabulary pieces that are never split during tokenization. Useful for multilingual language tags (e.g., ['<2es>', '<2fr>']).	None
config	Config	Configuration object. Falls back to get_default_config().	None

Side Effects

• Creates temporary file during training (automatically cleaned up).
• Writes {model_prefix}.model and {model_prefix}.vocab to disk.
• Prints confirmation message with output path.

Examples:

>>> train_sentencepiece([Path('train.tsv')], 'models/sp', vocab_size=16000)
SentencePiece model saved to models/sp.model

Source code in src/torchlingo/preprocessing/sentencepiece.py

def train_sentencepiece(
    input_files: List[Path],
    model_prefix: str,
    vocab_size: int = None,
    columns: Optional[List[str]] = None,
    src_col: str = None,
    tgt_col: str = None,
    model_type: str = None,
    character_coverage: float = None,
    normalization_rule_name: str = None,
    pad_idx: int = None,
    unk_idx: int = None,
    sos_idx: int = None,
    eos_idx: int = None,
    pad_token: str = None,
    unk_token: str = None,
    sos_token: str = None,
    eos_token: str = None,
    user_defined_symbols: Optional[List[str]] = None,
    config: Config = None,
):
    """Train a SentencePiece tokenization model on raw text data.

    Reads all input files, extracts selected columns (src/tgt by default),
    concatenates them into a temporary file, and trains a SentencePiece model. Model files
    are saved with the specified prefix (e.g., 'model.model' and 'model.vocab').

    Model training respects configuration for normalization, special tokens,
    and character coverage.

    Args:
        input_files (List[Path]): Data files to use for training.
        model_prefix (str): Output prefix for .model and .vocab files.
        vocab_size (int, optional): Target vocabulary size. Falls back to config.vocab_size.
        columns (List[str], optional): Columns to use for training data. Defaults
            to [src_col, tgt_col].
        src_col (str, optional): Source column name. Falls back to config.src_col.
        tgt_col (str, optional): Target column name. Falls back to config.tgt_col.
        model_type (str, optional): SentencePiece model type. Falls back to config.sp_model_type.
        character_coverage (float, optional): Character coverage. Falls back to config.sp_character_coverage.
        normalization_rule_name (str, optional): Normalization rule. Falls back to config.sp_normalization_rule_name.
        pad_idx (int, optional): Padding token index. Falls back to config.pad_idx.
        unk_idx (int, optional): Unknown token index. Falls back to config.unk_idx.
        sos_idx (int, optional): Start-of-sequence token index. Falls back to config.sos_idx.
        eos_idx (int, optional): End-of-sequence token index. Falls back to config.eos_idx.
        pad_token (str, optional): Padding token string. Falls back to config.pad_token.
        unk_token (str, optional): Unknown token string. Falls back to config.unk_token.
        sos_token (str, optional): Start-of-sequence token string. Falls back to config.sos_token.
        eos_token (str, optional): End-of-sequence token string. Falls back to config.eos_token.
        user_defined_symbols (List[str], optional): Symbols to register as whole
            vocabulary pieces that are never split during tokenization. Useful
            for multilingual language tags (e.g., ['<2es>', '<2fr>']).
        config (Config, optional): Configuration object. Falls back to get_default_config().

    Side Effects:
        - Creates temporary file during training (automatically cleaned up).
        - Writes {model_prefix}.model and {model_prefix}.vocab to disk.
        - Prints confirmation message with output path.

    Examples:
        >>> train_sentencepiece([Path('train.tsv')], 'models/sp', vocab_size=16000)
        SentencePiece model saved to models/sp.model
    """
    import tempfile

    cfg = config if config is not None else get_default_config()
    vocab_size = vocab_size if vocab_size is not None else cfg.vocab_size
    src_col = src_col if src_col is not None else cfg.src_col
    tgt_col = tgt_col if tgt_col is not None else cfg.tgt_col
    model_type = model_type if model_type is not None else cfg.sp_model_type
    character_coverage = (
        character_coverage
        if character_coverage is not None
        else cfg.sp_character_coverage
    )
    normalization_rule_name = (
        normalization_rule_name
        if normalization_rule_name is not None
        else cfg.sp_normalization_rule_name
    )
    pad_idx = pad_idx if pad_idx is not None else cfg.pad_idx
    unk_idx = unk_idx if unk_idx is not None else cfg.unk_idx
    sos_idx = sos_idx if sos_idx is not None else cfg.sos_idx
    eos_idx = eos_idx if eos_idx is not None else cfg.eos_idx
    pad_token = pad_token if pad_token is not None else cfg.pad_token
    unk_token = unk_token if unk_token is not None else cfg.unk_token
    sos_token = sos_token if sos_token is not None else cfg.sos_token
    eos_token = eos_token if eos_token is not None else cfg.eos_token

    cols = columns or [src_col, tgt_col]

    with tempfile.NamedTemporaryFile(mode="w", encoding="utf-8", delete=False) as tmp:
        for f in input_files:
            df = load_data(f)
            for col in cols:
                if col in df.columns:
                    for line in df[col].astype(str):
                        tmp.write(line + "\n")
        tmp_path = tmp.name

    try:
        spm.SentencePieceTrainer.train(
            input=tmp_path,
            model_prefix=model_prefix,
            vocab_size=vocab_size,
            model_type=model_type,
            character_coverage=character_coverage,
            normalization_rule_name=normalization_rule_name,
            pad_id=pad_idx,
            unk_id=unk_idx,
            bos_id=sos_idx,
            eos_id=eos_idx,
            pad_piece=pad_token,
            unk_piece=unk_token,
            bos_piece=sos_token,
            eos_piece=eos_token,
            user_defined_symbols=user_defined_symbols or [],
        )
        print(f"SentencePiece model saved to {model_prefix}.model")
    finally:
        Path(tmp_path).unlink()

Training Options

Model Types

Type	Description	Use Case
bpe	Byte Pair Encoding	General purpose, most common
unigram	Unigram LM	Better for some Asian languages
char	Character-level	When subwords aren't needed
word	Word-level	Pre-tokenized data

# BPE (default, recommended)
train_sentencepiece(input_files, model_prefix, model_type="bpe")

# Unigram
train_sentencepiece(input_files, model_prefix, model_type="unigram")

Vocabulary Size

# Small vocab (faster, less coverage)
train_sentencepiece(input_files, model_prefix, vocab_size=4000)

# Medium vocab (balanced)
train_sentencepiece(input_files, model_prefix, vocab_size=16000)

# Large vocab (better coverage, more parameters)
train_sentencepiece(input_files, model_prefix, vocab_size=32000)

Rule of thumb:

• Small datasets: 4K-8K
• Medium datasets: 8K-16K
• Large datasets: 16K-32K

Character Coverage

# Full coverage (all characters in vocabulary)
train_sentencepiece(input_files, model_prefix, character_coverage=1.0)

# Allow rare characters to be UNK (useful for noisy data)
train_sentencepiece(input_files, model_prefix, character_coverage=0.9995)

For CJK languages (Chinese, Japanese, Korean), use 0.9995 since they have many characters.

Normalization

# NMT-NFKC (default, good for translation)
train_sentencepiece(
    input_files, model_prefix,
    normalization_rule_name="nmt_nfkc"
)

# Other options
train_sentencepiece(input_files, model_prefix, normalization_rule_name="nfc")
train_sentencepiece(input_files, model_prefix, normalization_rule_name="identity")

Examples

Complete Training Pipeline

from pathlib import Path
from torchlingo.preprocessing import (
    load_data,
    save_data,
    train_sentencepiece,
)

# Prepare training text (combine src and tgt for shared vocab)
df = load_data("data/train.tsv")
train_text = df['src'].tolist() + df['tgt'].tolist()

# Save as text file for SentencePiece
text_path = Path("data/train_combined.txt")
with open(text_path, "w") as f:
    f.write("\n".join(train_text))

# Train SentencePiece
train_sentencepiece(
    input_files=[str(text_path)],
    model_prefix="models/shared_sp",
    vocab_size=16000,
    model_type="bpe",
)

Separate Source/Target Models

# Source language model
with open("data/train_src.txt", "w") as f:
    f.write("\n".join(df['src']))

train_sentencepiece(
    ["data/train_src.txt"],
    "models/sp_en",
    vocab_size=8000,
)

# Target language model
with open("data/train_tgt.txt", "w") as f:
    f.write("\n".join(df['tgt']))

train_sentencepiece(
    ["data/train_tgt.txt"],
    "models/sp_es",
    vocab_size=8000,
)

Using with NMTDataset

from torchlingo.data_processing import NMTDataset, SentencePieceVocab

# Load trained models
src_sp = SentencePieceVocab("models/sp_en.model")
tgt_sp = SentencePieceVocab("models/sp_es.model")

# Use with dataset
dataset = NMTDataset(
    "data/train.tsv",
    src_vocab=src_sp,
    tgt_vocab=tgt_sp,
)

Configuration via Config

from torchlingo.config import Config

config = Config(
    use_sentencepiece=True,
    sentencepiece_model_prefix="models/sp",
    vocab_size=16000,
    sp_model_type="bpe",
    sp_character_coverage=1.0,
    sp_normalization_rule_name="nmt_nfkc",
)

Best Practices

1. Train on Training Data Only

# ✅ Good: Train on training data
train_sentencepiece(["data/train.txt"], model_prefix)

# ❌ Bad: Including test data
train_sentencepiece(["data/train.txt", "data/test.txt"], model_prefix)

2. Combine Source and Target for Shared Vocab

For similar languages (e.g., Spanish-Portuguese), a shared vocabulary works well:

combined = df['src'].tolist() + df['tgt'].tolist()

3. Use Appropriate Vocab Size

Dataset Size	Recommended Vocab
< 100K sentences	4K-8K
100K-1M sentences	8K-16K
> 1M sentences	16K-32K

4. Inspect the Model

import sentencepiece as spm

sp = spm.SentencePieceProcessor()
sp.load("models/sp.model")

# Check vocabulary
print(f"Vocab size: {sp.get_piece_size()}")

# See how text is tokenized
text = "The quick brown fox"
pieces = sp.encode_as_pieces(text)
print(f"Pieces: {pieces}")

ids = sp.encode_as_ids(text)
print(f"IDs: {ids}")

Troubleshooting

Model produces too many UNK tokens

Increase character_coverage:

train_sentencepiece(..., character_coverage=1.0)

Vocabulary is too small/large

Adjust vocab_size:

train_sentencepiece(..., vocab_size=16000)

Training is slow

For large datasets, use sampling:

# SentencePiece automatically samples if data is large
train_sentencepiece(..., input_sentence_size=10000000)