Base Utilities

Core data loading, saving, and splitting utilities for neural machine translation.

Overview

This module provides foundational functions for working with parallel corpora:

• Load data from TSV, CSV, JSON, and Parquet formats
• Convert parallel text files to DataFrames
• Save processed data
• Split datasets into train/val/test sets

API Reference

load_data

load_data(filepath: Union[Path, str], format: str = None) -> DataFrame

Load a single structured data file containing src/tgt columns.

Supports multiple file formats and automatically detects format from file extension if not explicitly specified. Single-file inputs only; for parallel .txt corpora, use parallel_txt_to_dataframe first.

Parameters:

Name	Type	Description	Default
filepath	Path | str	Path to the data file.	required
format	str	File format ('tsv', 'csv', 'parquet', 'json'). If None, inferred from file extension. Defaults to None.	None

Returns:

Type	Description
DataFrame	pd.DataFrame: Loaded data with src/tgt columns.

Raises:

Type	Description
ValueError	If format is 'txt' (must use parallel_txt_to_dataframe), or if format is unsupported.
FileNotFoundError	If filepath does not exist.

Examples:

>>> df = load_data('data.tsv')
>>> df = load_data('data.csv', format='csv')
>>> df = load_data('data.parquet')

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

def load_data(filepath: Union[Path, str], format: str = None) -> pd.DataFrame:
    """Load a single structured data file containing src/tgt columns.

    Supports multiple file formats and automatically detects format from file
    extension if not explicitly specified. Single-file inputs only; for
    parallel .txt corpora, use parallel_txt_to_dataframe first.

    Args:
        filepath (Path | str): Path to the data file.
        format (str, optional): File format ('tsv', 'csv', 'parquet', 'json').
            If None, inferred from file extension. Defaults to None.

    Returns:
        pd.DataFrame: Loaded data with src/tgt columns.

    Raises:
        ValueError: If format is 'txt' (must use parallel_txt_to_dataframe),
            or if format is unsupported.
        FileNotFoundError: If filepath does not exist.

    Examples:
        >>> df = load_data('data.tsv')
        >>> df = load_data('data.csv', format='csv')
        >>> df = load_data('data.parquet')
    """
    filepath = Path(filepath)
    if format is None:
        format = filepath.suffix.lstrip(".")
    if format == "tsv":
        return pd.read_csv(filepath, sep="\t")
    elif format == "csv":
        return pd.read_csv(filepath)
    elif format == "parquet":
        return pd.read_parquet(filepath)
    elif format == "json":
        return pd.read_json(filepath, lines=True)
    elif format == "txt":
        raise ValueError(
            "Single-file 'txt' format is not supported. Convert parallel txt files to a DataFrame first using parallel_txt_to_dataframe."
        )
    else:
        raise ValueError(f"Unsupported format: {format}")

save_data

save_data(df: DataFrame, filepath: Path, format: str = None)

Save a DataFrame to disk in the specified format.

Creates parent directories as needed. Automatically detects format from file extension if not explicitly specified. Prints a confirmation message with record count and output path.

Parameters:

Name	Type	Description	Default
df	DataFrame	DataFrame to save with src/tgt columns.	required
filepath	Path	Destination file path.	required
format	str	Output format ('tsv', 'csv', 'parquet', 'json'). If None, inferred from file extension. Defaults to None.	None

Raises:

Type	Description
ValueError	If format is unsupported.

Examples:

>>> save_data(df, Path('output.tsv'))
Saved 1000 records to output.tsv
>>> save_data(df, Path('output.json'), format='json')

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

def save_data(df: pd.DataFrame, filepath: Path, format: str = None):
    """Save a DataFrame to disk in the specified format.

    Creates parent directories as needed. Automatically detects format from
    file extension if not explicitly specified. Prints a confirmation message
    with record count and output path.

    Args:
        df (pd.DataFrame): DataFrame to save with src/tgt columns.
        filepath (Path): Destination file path.
        format (str, optional): Output format ('tsv', 'csv', 'parquet', 'json').
            If None, inferred from file extension. Defaults to None.

    Raises:
        ValueError: If format is unsupported.

    Examples:
        >>> save_data(df, Path('output.tsv'))
        Saved 1000 records to output.tsv
        >>> save_data(df, Path('output.json'), format='json')
    """
    if format is None:
        format = filepath.suffix.lstrip(".")
    filepath.parent.mkdir(parents=True, exist_ok=True)
    if format == "tsv":
        df.to_csv(filepath, sep="\t", index=False)
    elif format == "csv":
        df.to_csv(filepath, index=False)
    elif format == "parquet":
        df.to_parquet(filepath, index=False)
    elif format == "json":
        df.to_json(filepath, orient="records", lines=True)
    else:
        raise ValueError(f"Unsupported format: {format}")
    print(f"Saved {len(df)} records to {filepath}")

parallel_txt_to_dataframe

parallel_txt_to_dataframe(src_path: Path, tgt_path: Path, src_col: str = None, tgt_col: str = None, config: Config = None) -> DataFrame

Convert parallel .txt files into a single DataFrame with src/tgt columns.

Reads two aligned text files (source and target) and combines them into a single DataFrame with 'src' and 'tgt' columns. Files must have an equal number of lines. Each line is stripped of whitespace.

Parameters:

Name	Type	Description	Default
src_path	Path	Path to source language file (one sentence per line).	required
tgt_path	Path	Path to target language file (one sentence per line).	required
src_col	str	Column name for source text. Falls back to config.src_col if not provided. Defaults to None.	None
tgt_col	str	Column name for target text. Falls back to config.tgt_col if not provided. Defaults to None.	None
config	Config	Configuration object. If None, uses default config. Explicit parameters take priority over config values.	None

Returns:

Type	Description
DataFrame	pd.DataFrame: DataFrame with columns defined by src_col and tgt_col, containing one parallel pair per row.

Raises:

Type	Description
ValueError	If src_path and tgt_path have different number of lines.
FileNotFoundError	If either file does not exist.

Examples:

>>> df = parallel_txt_to_dataframe('en.txt', 'es.txt')
>>> print(df.columns)
Index(['src', 'tgt'], dtype='object')
>>> len(df)
1000

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

def parallel_txt_to_dataframe(
    src_path: Path,
    tgt_path: Path,
    src_col: str = None,
    tgt_col: str = None,
    config: Config = None,
) -> pd.DataFrame:
    """Convert parallel .txt files into a single DataFrame with src/tgt columns.

    Reads two aligned text files (source and target) and combines them into a
    single DataFrame with 'src' and 'tgt' columns. Files must have an equal
    number of lines. Each line is stripped of whitespace.

    Args:
        src_path (Path): Path to source language file (one sentence per line).
        tgt_path (Path): Path to target language file (one sentence per line).
        src_col (str, optional): Column name for source text. Falls back to
            config.src_col if not provided. Defaults to None.
        tgt_col (str, optional): Column name for target text. Falls back to
            config.tgt_col if not provided. Defaults to None.
        config (Config, optional): Configuration object. If None, uses default
            config. Explicit parameters take priority over config values.

    Returns:
        pd.DataFrame: DataFrame with columns defined by src_col and tgt_col,
            containing one parallel pair per row.

    Raises:
        ValueError: If src_path and tgt_path have different number of lines.
        FileNotFoundError: If either file does not exist.

    Examples:
        >>> df = parallel_txt_to_dataframe('en.txt', 'es.txt')
        >>> print(df.columns)
        Index(['src', 'tgt'], dtype='object')
        >>> len(df)
        1000
    """
    cfg = config if config is not None else get_default_config()
    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

    src_path = Path(src_path)
    tgt_path = Path(tgt_path)
    with open(src_path, "r", encoding="utf-8") as f_src:
        src_lines = [ln.rstrip("\n") for ln in f_src]
    with open(tgt_path, "r", encoding="utf-8") as f_tgt:
        tgt_lines = [ln.rstrip("\n") for ln in f_tgt]
    if len(src_lines) != len(tgt_lines):
        raise ValueError(
            f"Parallel files have different number of lines: {src_path} ({len(src_lines)}), {tgt_path} ({len(tgt_lines)})"
        )
    pairs = []
    for s, t in zip(src_lines, tgt_lines):
        pairs.append({src_col: s.strip(), tgt_col: t.strip()})
    return pd.DataFrame(pairs)

split_data

split_data(df: DataFrame, train_ratio: float = None, val_ratio: float = None, seed: int = None, config: Config = None) -> Tuple[DataFrame, DataFrame, DataFrame]

Split a DataFrame into train, validation, and test sets.

Shuffles the entire DataFrame using the provided seed for reproducibility, then partitions by ratio. Test set is the remainder after train and val.

Parameters:

Name	Type	Description	Default
df	DataFrame	Input DataFrame to split.	required
train_ratio	float	Fraction of data for training. Falls back to config.train_ratio if not provided. Defaults to None.	None
val_ratio	float	Fraction of data for validation. Falls back to config.val_ratio if not provided. Defaults to None. Remaining (1 - train_ratio - val_ratio) goes to test.	None
seed	int	Random seed for reproducibility. Falls back to config.seed if not provided. Defaults to None.	None
config	Config	Configuration object. If None, uses default config. Explicit parameters take priority over config values.	None

Returns:

Name	Type	Description
tuple	Tuple[DataFrame, DataFrame, DataFrame]	(train_df, val_df, test_df) as pandas DataFrames.

Examples:

>>> train, val, test = split_data(df, train_ratio=0.8, val_ratio=0.1)
>>> len(train), len(val), len(test)
(800, 100, 100)

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

def split_data(
    df: pd.DataFrame,
    train_ratio: float = None,
    val_ratio: float = None,
    seed: int = None,
    config: Config = None,
) -> Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
    """Split a DataFrame into train, validation, and test sets.

    Shuffles the entire DataFrame using the provided seed for reproducibility,
    then partitions by ratio. Test set is the remainder after train and val.

    Args:
        df (pd.DataFrame): Input DataFrame to split.
        train_ratio (float, optional): Fraction of data for training. Falls back
            to config.train_ratio if not provided. Defaults to None.
        val_ratio (float, optional): Fraction of data for validation. Falls back
            to config.val_ratio if not provided. Defaults to None.
            Remaining (1 - train_ratio - val_ratio) goes to test.
        seed (int, optional): Random seed for reproducibility. Falls back to
            config.seed if not provided. Defaults to None.
        config (Config, optional): Configuration object. If None, uses default
            config. Explicit parameters take priority over config values.

    Returns:
        tuple: (train_df, val_df, test_df) as pandas DataFrames.

    Examples:
        >>> train, val, test = split_data(df, train_ratio=0.8, val_ratio=0.1)
        >>> len(train), len(val), len(test)
        (800, 100, 100)
    """
    cfg = config if config is not None else get_default_config()
    train_ratio = train_ratio if train_ratio is not None else cfg.train_ratio
    val_ratio = val_ratio if val_ratio is not None else cfg.val_ratio
    seed = seed if seed is not None else cfg.seed

    df = df.sample(frac=1, random_state=seed).reset_index(drop=True)
    n = len(df)
    train_end = int(n * train_ratio)
    val_end = int(n * (train_ratio + val_ratio))
    train_df = df.iloc[:train_end]
    val_df = df.iloc[train_end:val_end]
    test_df = df.iloc[val_end:]
    return train_df, val_df, test_df

Examples

Loading Data

from torchlingo.preprocessing import load_data

# Auto-detect format from extension
df = load_data("data/train.tsv")

# Explicit format
df = load_data("data/train.csv", format="csv")

# Supported formats
df_tsv = load_data("data.tsv")        # Tab-separated
df_csv = load_data("data.csv")        # Comma-separated
df_json = load_data("data.json")      # JSON Lines
df_parquet = load_data("data.parquet") # Parquet

Converting Parallel Files

Many NMT datasets come as separate source and target files:

english.txt:
  Hello world
  Good morning

spanish.txt:
  Hola mundo
  Buenos días

Convert them to a DataFrame:

from torchlingo.preprocessing import parallel_txt_to_dataframe

df = parallel_txt_to_dataframe(
    src_path="english.txt",
    tgt_path="spanish.txt",
)

print(df)
#            src            tgt
# 0  Hello world     Hola mundo
# 1  Good morning  Buenos días

Custom Column Names

df = parallel_txt_to_dataframe(
    src_path="english.txt",
    tgt_path="spanish.txt",
    src_col="english",
    tgt_col="spanish",
)

Saving Data

from torchlingo.preprocessing import save_data

# Save as TSV (recommended)
save_data(df, "output/train.tsv")

# Save as other formats
save_data(df, "output/train.csv")
save_data(df, "output/train.json")
save_data(df, "output/train.parquet")

Splitting Data

from torchlingo.preprocessing import train_test_split

train_df, val_df, test_df = train_test_split(
    df,
    val_ratio=0.1,   # 10% validation
    test_ratio=0.1,  # 10% test
    shuffle=True,    # Shuffle before splitting
    random_state=42, # Reproducibility
)

print(f"Train: {len(train_df)}")  # 80%
print(f"Val: {len(val_df)}")      # 10%
print(f"Test: {len(test_df)}")    # 10%

Format Details

TSV (Tab-Separated Values)

Recommended for text data because tabs rarely appear in natural text.

src tgt
Hello world Hola mundo
How are you?    ¿Cómo estás?

CSV (Comma-Separated Values)

Requires quoting if text contains commas:

src,tgt
"Hello, world","Hola, mundo"
How are you?,¿Cómo estás?

JSON Lines

One JSON object per line, good for complex metadata:

{"src": "Hello world", "tgt": "Hola mundo"}
{"src": "Good morning", "tgt": "Buenos días"}

Parquet

Binary columnar format:

• Efficient compression
• Fast loading for large files
• Not human-readable

# Parquet is great for large datasets
save_data(df, "data/large_corpus.parquet")
df = load_data("data/large_corpus.parquet")  # Fast!

Error Handling

Mismatched Line Counts

# english.txt has 1000 lines, spanish.txt has 999 lines
df = parallel_txt_to_dataframe("english.txt", "spanish.txt")
# Raises: ValueError: Parallel files have different number of lines

Missing Columns

df = load_data("data_without_src_column.tsv")
# Later: ValueError when used with NMTDataset

Unsupported Format

df = load_data("data.xml")
# Raises: ValueError: Unsupported format: xml

Best Practices

1. Use TSV for text: Tabs are rare in natural language
2. Always shuffle before splitting: Ensures representative splits
3. Check data after loading: Verify row counts and sample content
4. Use Parquet for large data: Faster loading, smaller files

# Good workflow
from torchlingo.preprocessing import load_data

df = load_data("data/corpus.tsv")
print(f"Loaded {len(df)} rows")
print(df.head())  # Inspect samples
print(df.isnull().sum())  # Check for missing values