src/easytranslate/utils/seed.py

"""Global RNG seeding for reproducible runs."""

import os
import random

import numpy as np
import torch


def set_seed(seed: int = 42, *, deterministic_cudnn: bool = True) -> None:
    """
    Set Python / NumPy / Torch seeds for reproducibility.

    Args:
        seed: Integer seed.
        deterministic_cudnn: If True, use deterministic cuDNN (slower, more reproducible).
    """
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)
    if deterministic_cudnn:
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False
    os.environ["PYTHONHASHSEED"] = str(seed)