This tutorial shows how to prepare parallel data for training a machine translation (MT) model with PyTorch. We filter the data, train SentencePiece subword models, convert the text to IDs, split it into training, validation and test sets, and build padded batches. The code uses only SentencePiece and plain PyTorch.
Updated October 2026. An earlier version of this tutorial used Torchtext. Torchtext development has stopped, and version 0.18 (April 2024) was the last stable release. This version replaces the Torchtext functions with SentencePiece and PyTorch, and the steps stay the same. The code was tested with Python 3.11, PyTorch 2.9 and SentencePiece 0.2.
Table of Contents:
- Installation and Importing
- Filtering Data
- Subword Tokenization
- Data Preparation Function
- Vocabulary
- Build Batches
- Next Steps
Installation and Importing
Let’s first install PyTorch and SentencePiece:
pip3 install torch sentencepiece
Then, import what we need.
import random
import sentencepiece as spm
import torch
from torch.utils.data import DataLoader
from torch.nn.utils.rnn import pad_sequence
Filtering Data
It is highly recommended you filter your data before training an MT model. The purpose is mainly to remove empty rows, duplicates, and too-long source/target rows. You can either use one of Moses scripts (clean-corpus-n.perl), or one of my MT Preparation scripts (filter.py).
Subword Tokenization
Subwording is a type of text tokenization. Dividing words into subword units help decrease the size of vocabulary and hence minimize out-of-vocabulary situations. Among subwording techniques are Byte-Pair Encoding (BPE) (Sennrich et al., 2016) and unigram language model (Kudo, 2018), supported by several tools. SentencePiece is one of the popular tools used for subwording text.
Training SentencePiece Models
Unless you use a shared vocabulary, you have to train one model for the source text and another model for the target text. The following code creates source.spm.model and source.spm.vocab for the source, and the same two files for the target.
The IDs of the special tokens are set explicitly, so we know them later: <unk> is 0, <pad> is 1, <s> (start of a sentence) is 2, and </s> (end of a sentence) is 3.
source_file = "source.txt"
target_file = "target.txt"
vocab_size = 32000
for name, path in [("source", source_file), ("target", target_file)]:
spm.SentencePieceTrainer.train(
input=path,
model_prefix=f"{name}.spm",
vocab_size=vocab_size,
model_type="unigram",
unk_id=0, pad_id=1, bos_id=2, eos_id=3,
)
print("Done! Training SentencePiece models completed.")
Training prints a long log. You can choose model_type="bpe" instead of "unigram" for Byte-Pair Encoding.
💡Tip: What to do if “Vocabulary size too high”?
If you are experimenting with a very small file, a vocabulary size of 32000 might be too big. SentencePiece stops with an error such as “Vocabulary size too high (32000). Please set it to a value <= 664.” The message tells you the largest size your data supports, so use a value up to that number. Obviously, it would be better to add more data.
💡Tip: What if the corpus is very large?
SentencePiece does not need every sentence to train a good model. For a corpus of tens of millions of sentences, train on a random sample by adding input_sentence_size=10000000 and shuffle_input_sentence=True to the arguments of SentencePieceTrainer.train. For languages with a large character set, such as Chinese or Japanese, keep the default character_coverage=0.9995. For languages with a small character set, you can set it to 1.0.
💡Tip: How to rather build a shared vocabulary?
To train one model on both languages, give SentencePiece both files, separated by a comma. This works well for languages that share a script or many words.
spm.SentencePieceTrainer.train(
input=f"{source_file},{target_file}",
model_prefix="shared.spm",
vocab_size=vocab_size,
model_type="unigram",
unk_id=0, pad_id=1, bos_id=2, eos_id=3,
)
Then use the same model for the source and the target in the next steps.
Subwording Data with the SentencePiece Models
Load the trained models.
source_sp = spm.SentencePieceProcessor(model_file="source.spm.model")
target_sp = spm.SentencePieceProcessor(model_file="target.spm.model")
The encode method converts text into subword pieces or into IDs, depending on out_type. Both of the following accept a single sentence or a list of sentences.
print(source_sp.encode("This is an example.", out_type=str))
print(source_sp.encode("This is an example.", out_type=int))
For training, we use the IDs.
Data Preparation Function
This function reads the source and target files, optionally lowercases them, converts the sentences to IDs, and splits the data into training, development and test sets. The pairs are shuffled together before splitting, so the source and target sentences stay aligned.
def prepare(source_file, target_file, dev_size, test_size, source_sp, target_sp, lower=False, seed=1):
with open(source_file) as source, open(target_file) as target:
source_lines = [line.strip() for line in source]
target_lines = [line.strip() for line in target]
if len(source_lines) != len(target_lines):
raise ValueError("Length of source and target lines must be the same!")
if lower:
source_lines = [line.lower() for line in source_lines]
target_lines = [line.lower() for line in target_lines]
# Subword the sentences, i.e. convert them to lists of IDs
source_ids = source_sp.encode(source_lines, out_type=int)
target_ids = target_sp.encode(target_lines, out_type=int)
# Shuffle the pairs together, then split into dev, test, and train
data = list(zip(source_ids, target_ids))
random.Random(seed).shuffle(data)
dev_dataset = data[:dev_size]
test_dataset = data[dev_size:dev_size + test_size]
train_dataset = data[dev_size + test_size:]
return train_dataset, dev_dataset, test_dataset
Now, we can use this function to prepare the data.
dev_size = 2000
test_size = 2000
train_dataset, dev_dataset, test_dataset = prepare(source_file, target_file, dev_size, test_size,
source_sp, target_sp, lower=False)
Let’s print some data statistics.
print("Training dataset size: %s" % len(train_dataset),
"Dev dataset size: %s" % len(dev_dataset),
"Test dataset size: %s" % len(test_dataset),
sep=" | ")
Now, let’s print a few lines of the test dataset.
print(*test_dataset[:5], sep="\n")
Vocabulary
With SentencePiece, the vocabulary is the model itself, so there is no separate step to build one. Each model already maps every piece to an ID and back.
print("Source vocabulary size:", source_sp.get_piece_size())
print("Target vocabulary size:", target_sp.get_piece_size())
# IDs of the special tokens
print("pad, unk, bos, eos:", target_sp.pad_id(), target_sp.unk_id(), target_sp.bos_id(), target_sp.eos_id())
# Convert IDs back to pieces
print(source_sp.id_to_piece([2, 15, 16, 3]))
Keep the .model files together with your trained model, because you need the same ones at translation time. Unknown pieces map to the ID of <unk> automatically.
At translation time, encode the source with the source model, and decode the IDs that the model generates with the target model.
print(target_sp.decode([2, 15, 16, 3]))
Special tokens such as <s> and </s> are left out of the decoded text.
Build Batches
As training large datasets is time consuming, we build batches from the training data. First, let’s determine the device, the batch size, and the maximum length of the source and target sentences.
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
batch_size = 128 # examples
max_len = 100 # tokens in source or target
pad_idx = target_sp.pad_id()
Very long or empty sentences make training slower and less stable, so we keep only training pairs of a reasonable length.
train_dataset = [(source, target) for source, target in train_dataset
if 1 < len(source) < max_len and 1 < len(target) < max_len]
The collate function builds one batch. It adds the start and end tokens, which help the model know when to stop, and converts each sentence to a tensor. Sentences are padded to the longest sentence in the batch, so all of them have the same length.
def collate_batch(batch):
sources, targets = [], []
for source, target in batch:
source = [source_sp.bos_id()] + source + [source_sp.eos_id()]
target = [target_sp.bos_id()] + target + [target_sp.eos_id()]
sources.append(torch.tensor(source, dtype=torch.int64))
targets.append(torch.tensor(target, dtype=torch.int64))
sources = pad_sequence(sources, padding_value=pad_idx)
targets = pad_sequence(targets, padding_value=pad_idx)
return sources, targets
Finally, let’s create the data loaders, which yield batches. Each takes the collate_batch function that we have just created as an argument. We shuffle the training data only.
train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_batch)
valid_dataloader = DataLoader(dev_dataset, batch_size=batch_size, shuffle=False, collate_fn=collate_batch)
test_dataloader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, collate_fn=collate_batch)
If you want to test a data loader, you can use the following code. In your training loop, move each batch to the device.
for x, y in test_dataloader:
x, y = x.to(device), y.to(device)
print(x.shape, y.shape)
Each shape is (sequence length, batch size), because pad_sequence puts the sequence first. If your model expects the batch first, pass batch_first=True to pad_sequence.
Bucketing [optional]
Sentences in a random batch have very different lengths, so a large part of each batch can be padding. Bucketing groups sentences of similar lengths in the same batch, which reduces the padding and speeds up training. The following function shuffles the data, sorts chunks of it by source length, cuts the chunks into batches, and shuffles the batches.
def bucket_batches(dataset, batch_size, chunk_size=100, seed=1):
rng = random.Random(seed)
data = list(dataset)
rng.shuffle(data)
batches = []
for i in range(0, len(data), batch_size * chunk_size):
chunk = sorted(data[i:i + batch_size * chunk_size], key=lambda pair: len(pair[0]))
batches += [chunk[j:j + batch_size] for j in range(0, len(chunk), batch_size)]
rng.shuffle(batches)
return batches
Since the batches are already formed, set batch_size=None in the data loader. Create the data loader again for each epoch, with a different seed, to get new batches.
train_dataloader = DataLoader(bucket_batches(train_dataset, batch_size),
batch_size=None, collate_fn=collate_batch)
Next Steps
You can now feed these batches to a PyTorch model, such as torch.nn.Transformer. If you prefer to train with a toolkit, see How to Build and Deploy Machine Translation Models.