Dataset

PyTorch Dataset implementation for parallel neural machine translation data.

Overview

NMTDataset loads parallel text corpora and provides encoded tensor pairs suitable for training. It handles:

• Loading multiple file formats (TSV, CSV, JSON, Parquet)
• Data cleaning (removing blanks, normalizing whitespace)
• Automatic vocabulary building
• Sequence encoding with special tokens
• Length truncation

Quick Start

from torchlingo.data_processing import NMTDataset

# Basic usage - vocabularies built automatically
dataset = NMTDataset("data/train.tsv")

# Access a sample
src_tensor, tgt_tensor = dataset[0]

# Use vocabularies
print(f"Source vocab size: {len(dataset.src_vocab)}")
print(f"Target vocab size: {len(dataset.tgt_vocab)}")

API Reference

NMTDataset

NMTDataset(data_file: Path, src_col: str = None, tgt_col: str = None, src_tok_col: str = None, tgt_tok_col: str = None, src_vocab: Optional[BaseVocab] = None, tgt_vocab: Optional[BaseVocab] = None, max_length: Optional[int] = None, eos_idx: int = None, config: Config = None)

Bases: Dataset

PyTorch Dataset for parallel neural machine translation data.

Loads a parallel corpus from disk and provides encoded source/target tensor pairs suitable for training and evaluation. Automatically cleans data by removing empty rows, normalizing whitespace, and handling missing values. If vocabulary objects are not provided, the dataset will build simple SimpleVocab instances from the training data.

The dataset supports optional pre-tokenized columns and enforces a maximum sequence length by truncating longer sequences while preserving the EOS token.

Parameters are resolved with fallback priority: 1. Explicitly passed parameter values 2. Values from passed config object 3. Values from default config (get_default_config())

Parameters:

Name	Type	Description	Default
data_file	Path or str	Path to a structured file (TSV/CSV/JSON/Parquet) containing configurable src/tgt columns. Parallel .txt corpora must be converted to a single file upstream (see preprocessing.base.parallel_txt_to_dataframe).	required
src_col	str	Column name for source text. Falls back to config.src_col if not provided.	None
tgt_col	str	Column name for target text. Falls back to config.tgt_col if not provided.	None
src_tok_col	str	Column name for pre-tokenized source text. Falls back to config.src_tok_col if not provided.	None
tgt_tok_col	str	Column name for pre-tokenized target text. Falls back to config.tgt_tok_col if not provided.	None
src_vocab	BaseVocab	Pre-built source vocabulary. If None, a SimpleVocab will be created from src_col. Defaults to None.	None
tgt_vocab	BaseVocab	Pre-built target vocabulary. If None, a SimpleVocab will be created from tgt_col. Defaults to None.	None
max_length	int	Maximum sequence length in tokens. Sequences longer than this are truncated to (max_length - 1) + eos_idx. Falls back to config.max_seq_length if not provided.	None
eos_idx	int	Index of the EOS (end-of-sequence) token used for truncation. Falls back to config.eos_idx if not provided.	None
config	Config	Configuration object to use for default values. If None, uses get_default_config().	None

Raises:

Type	Description
ValueError	If data_file does not contain both src_col and tgt_col.

Attributes:

Name	Type	Description
df	DataFrame	Loaded and cleaned dataframe.
src_sentences	list[str]	Raw source sentences.
tgt_sentences	list[str]	Raw target sentences.
src_texts	list[str]	Text used for vocab building and encoding — the pre-tokenized columns when present, else the raw sentences.
tgt_texts	list[str]	Target-side counterpart of src_texts.
src_vocab	BaseVocab	Source vocabulary.
tgt_vocab	BaseVocab	Target vocabulary.
has_tokenized	bool	Whether pre-tokenized columns exist in data.
eos_idx	int	Index of the EOS token used for truncation.

Source code in src/torchlingo/data_processing/dataset.py

def __init__(
    self,
    data_file: Path,
    src_col: str = None,
    tgt_col: str = None,
    src_tok_col: str = None,
    tgt_tok_col: str = None,
    src_vocab: Optional[BaseVocab] = None,
    tgt_vocab: Optional[BaseVocab] = None,
    max_length: Optional[int] = None,
    eos_idx: int = None,
    config: Config = None,
):
    cfg = config if config is not None else get_default_config()
    self.data_file = Path(data_file)
    self.src_col = src_col if src_col is not None else cfg.src_col
    self.tgt_col = tgt_col if tgt_col is not None else cfg.tgt_col
    self.src_tok_col = src_tok_col if src_tok_col is not None else cfg.src_tok_col
    self.tgt_tok_col = tgt_tok_col if tgt_tok_col is not None else cfg.tgt_tok_col
    self.max_length = max_length if max_length is not None else cfg.max_seq_length
    self.eos_idx = eos_idx if eos_idx is not None else cfg.eos_idx

    # Load data using pandas
    self.df = load_data(self.data_file)

    # Ensure columns exist before manipulating them
    if self.src_col not in self.df.columns or self.tgt_col not in self.df.columns:
        raise ValueError(
            f"Data file {data_file} must contain columns '{self.src_col}' and "
            f"'{self.tgt_col}', but has columns {self.df.columns.tolist()}. "
            "Pass src_col/tgt_col (or set them on your Config) to match your file."
        )

    # Normalize columns: replace NaN with empty string, strip whitespace,
    # and drop rows where either source or target is blank. Blank lines are
    # not useful for translation and should be skipped.
    self.df[self.src_col] = self.df[self.src_col].fillna("").astype(str).str.strip()
    self.df[self.tgt_col] = self.df[self.tgt_col].fillna("").astype(str).str.strip()
    self.df = self.df[(self.df[self.src_col] != "") & (self.df[self.tgt_col] != "")]

    # Detect pre-tokenized columns and materialize
    self.has_tokenized = (
        self.src_tok_col in self.df.columns and self.tgt_tok_col in self.df.columns
    )
    if self.has_tokenized:
        self.src_tokenized = (
            self.df[self.src_tok_col].fillna("").astype(str).str.strip().tolist()
        )
        self.tgt_tokenized = (
            self.df[self.tgt_tok_col].fillna("").astype(str).str.strip().tolist()
        )

    self.src_sentences = self.df[self.src_col].astype(str).tolist()
    self.tgt_sentences = self.df[self.tgt_col].astype(str).tolist()

    # Text actually fed to the vocabularies: pre-tokenized columns when
    # available, raw sentences otherwise. Vocab building and encoding must
    # use the same text or out-of-vocabulary rates explode.
    if self.has_tokenized:
        self.src_texts = [" ".join(t.split()) for t in self.src_tokenized]
        self.tgt_texts = [" ".join(t.split()) for t in self.tgt_tokenized]
    else:
        self.src_texts = self.src_sentences
        self.tgt_texts = self.tgt_sentences

    self.src_vocab = src_vocab
    self.tgt_vocab = tgt_vocab

    if self.src_vocab is None:
        self.src_vocab = SimpleVocab()
        self.src_vocab.build_vocab(self.src_texts)
    if self.tgt_vocab is None:
        self.tgt_vocab = SimpleVocab()
        self.tgt_vocab.build_vocab(self.tgt_texts)

__len__

__len__() -> int

Return the number of samples in the dataset.

Returns:

Name	Type	Description
int	int	Total number of (source, target) pairs after cleaning.

Source code in src/torchlingo/data_processing/dataset.py

def __len__(self) -> int:
    """Return the number of samples in the dataset.

    Returns:
        int: Total number of (source, target) pairs after cleaning.
    """
    return len(self.src_sentences)

__getitem__

__getitem__(idx: int) -> Tuple[Tensor, Tensor]

Retrieve and encode a (source, target) pair at the given index.

Encodes raw sentences to token indices using the respective vocabularies, adds special tokens (SOS/EOS), and truncates sequences exceeding max_length while preserving the EOS token.

Parameters:

Name	Type	Description	Default
idx	int	Index of the sample to retrieve. Must be in range [0, len(self)).	required

Returns:

Type	Description
Tuple[Tensor, Tensor]	tuple[torch.Tensor, torch.Tensor]: A tuple containing: - src_tensor (torch.Tensor): Encoded source indices with shape (seq_len,) and dtype torch.long. Contains special tokens SOS (at position 0) and EOS (at position -1). - tgt_tensor (torch.Tensor): Encoded target indices with shape (seq_len,) and dtype torch.long. Contains special tokens SOS (at position 0) and EOS (at position -1).

Raises:

Type	Description
AssertionError	If vocabularies have not been initialized.
IndexError	If idx is out of range.

Source code in src/torchlingo/data_processing/dataset.py

def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
    """Retrieve and encode a (source, target) pair at the given index.

    Encodes raw sentences to token indices using the respective vocabularies,
    adds special tokens (SOS/EOS), and truncates sequences exceeding
    max_length while preserving the EOS token.

    Args:
        idx (int): Index of the sample to retrieve. Must be in range
            [0, len(self)).

    Returns:
        tuple[torch.Tensor, torch.Tensor]: A tuple containing:
            - src_tensor (torch.Tensor): Encoded source indices with shape
              (seq_len,) and dtype torch.long. Contains special tokens SOS
              (at position 0) and EOS (at position -1).
            - tgt_tensor (torch.Tensor): Encoded target indices with shape
              (seq_len,) and dtype torch.long. Contains special tokens SOS
              (at position 0) and EOS (at position -1).

    Raises:
        AssertionError: If vocabularies have not been initialized.
        IndexError: If idx is out of range.
    """
    src_sentence = self.src_texts[idx]
    tgt_sentence = self.tgt_texts[idx]

    assert self.src_vocab is not None and self.tgt_vocab is not None

    src_indices = self.src_vocab.encode(src_sentence, add_special_tokens=True)
    tgt_indices = self.tgt_vocab.encode(tgt_sentence, add_special_tokens=True)

    if len(src_indices) > self.max_length:
        src_indices = src_indices[: self.max_length - 1] + [self.eos_idx]
    if len(tgt_indices) > self.max_length:
        tgt_indices = tgt_indices[: self.max_length - 1] + [self.eos_idx]

    return torch.tensor(src_indices, dtype=torch.long), torch.tensor(
        tgt_indices, dtype=torch.long
    )

Examples

Custom Column Names

# If your data has different column names
dataset = NMTDataset(
    "data/train.csv",
    src_col="english",
    tgt_col="spanish",
)

Reusing Vocabularies

# Build vocab on training data
train_dataset = NMTDataset("data/train.tsv")

# Reuse for validation (important!)
val_dataset = NMTDataset(
    "data/val.tsv",
    src_vocab=train_dataset.src_vocab,
    tgt_vocab=train_dataset.tgt_vocab,
)

With SentencePiece

from torchlingo.data_processing import SentencePieceVocab

# Load pre-trained SentencePiece models
src_sp = SentencePieceVocab("models/src.model")
tgt_sp = SentencePieceVocab("models/tgt.model")

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

Custom Max Length

from torchlingo.config import Config

config = Config(max_seq_length=256)
dataset = NMTDataset("data/train.tsv", config=config)

Attributes

Attribute	Type	Description
df	pd.DataFrame	Loaded and cleaned dataframe
src_sentences	list[str]	Raw source sentences
tgt_sentences	list[str]	Raw target sentences
src_vocab	SimpleVocab or SentencePieceVocab	Source vocabulary
tgt_vocab	SimpleVocab or SentencePieceVocab	Target vocabulary
max_length	int	Maximum sequence length
eos_idx	int	End-of-sequence token index

Data Cleaning

NMTDataset automatically cleans your data:

1. NaN handling: Replaced with empty strings
2. Whitespace: Stripped from beginning and end
3. Empty rows: Removed if either source or target is blank
4. Type conversion: All text converted to strings

Truncation Behavior

Sequences longer than max_length are truncated while preserving the EOS token:

# If max_length=10 and sequence has 15 tokens:
# Original: [SOS, t1, t2, ..., t13, EOS]
# Truncated: [SOS, t1, t2, ..., t8, EOS]  # 10 tokens total

Pre-tokenized Data

If your data has pre-tokenized columns:

# Data with src_tokenized and tgt_tokenized columns
dataset = NMTDataset(
    "data/train.tsv",
    src_tok_col="src_tokenized",
    tgt_tok_col="tgt_tokenized",
)

# Dataset will use pre-tokenized columns for vocabulary building