datahandling/data API

dnallm.datahandling.data

DNA Dataset handling and processing utilities.

This module provides comprehensive tools for loading, processing, and managing DNA sequence datasets. It supports various file formats, data augmentation techniques, and statistical analysis.

Classes

DNADataset

DNADataset(ds, tokenizer=None, max_length=512)

A comprehensive wrapper for DNA sequence datasets with advanced processing capabilities.

This class provides methods for loading DNA datasets from various sources (local files, Hugging Face Hub, ModelScope), encoding sequences with tokenizers, data augmentation, statistical analysis, and more.

Attributes:

Name	Type	Description
dataset		The underlying Hugging Face Dataset or DatasetDict
tokenizer		Tokenizer for sequence encoding
max_length		Maximum sequence length for tokenization
sep	str | None	Separator for multi-label data
multi_label_sep	str | None	Separator for multi-label sequences
data_type	str | None	Type of the dataset (classification, regression, etc.)
stats	dict | None	Cached dataset statistics
stats_for_plot	DataFrame | None	Cached statistics for plotting

Initialize a DNADataset.

Parameters:

Name	Type	Description	Default
ds	Dataset | DatasetDict	A Hugging Face Dataset containing at least 'sequence' and 'label' fields	required
tokenizer	PreTrainedTokenizerBase | None	A Hugging Face tokenizer for encoding sequences	None
max_length	int	Maximum length for tokenization	512

Source code in dnallm/datahandling/data.py

def __init__(
    self,
    ds: Dataset | DatasetDict,
    tokenizer: PreTrainedTokenizerBase | None = None,
    max_length: int = 512,
) -> None:
    """Initialize a DNADataset.

    Args:
        ds: A Hugging Face Dataset containing at least 'sequence' and
            'label' fields
        tokenizer: A Hugging Face tokenizer for encoding sequences
        max_length: Maximum length for tokenization
    """
    if ds is None:
        raise TypeError("Dataset cannot be None")

    if max_length <= 0:
        raise ValueError("max_length must be positive")

    self.dataset = ds
    self.tokenizer = tokenizer
    self.max_length = max_length
    self.sep: str | None = None
    self.multi_label_sep: str | None = None
    self.data_type: str | None = None
    self.stats: dict | None = None
    self.stats_for_plot: pd.DataFrame | None = None
    self.__data_type__()  # Determine the data type of the dataset

Functions

__data_type__

__data_type__()

Get the data type of the dataset (classification, regression, etc.).

This method analyzes the labels to determine if the dataset is for: - classification (integer or string labels) - regression (float labels) - multi-label (multiple labels per sample) - multi-regression (multiple float values per sample)

Source code in dnallm/datahandling/data.py

def __data_type__(self) -> None:
    """Get the data type of the dataset (classification, regression, etc.).

    This method analyzes the labels to determine if the dataset is for:
    - classification (integer or string labels)
    - regression (float labels)
    - multi-label (multiple labels per sample)
    - multi-regression (multiple float values per sample)
    """
    labels = self._extract_labels()
    if labels is None:
        self.data_type = "unknown"
        return

    if not self._is_valid_labels(labels):
        self.data_type = "unknown"
        return

    first_label = self._get_first_label(labels)
    if first_label is None:
        self.data_type = "unknown"
        return

    self.data_type = self._determine_data_type(first_label)

__getitem__

__getitem__(idx)

Get an item from the dataset.

Parameters:

Name	Type	Description	Default
idx	int	Index of the item to retrieve	required

Returns:

Type	Description
	The item at the specified index

Raises:

Type	Description
ValueError	If dataset is a DatasetDict

Source code in dnallm/datahandling/data.py

def __getitem__(self, idx: int):
    """Get an item from the dataset.

    Args:
        idx: Index of the item to retrieve

    Returns:
        The item at the specified index

    Raises:
        ValueError: If dataset is a DatasetDict
    """
    if isinstance(self.dataset, DatasetDict):
        raise ValueError(
            "Dataset is a DatasetDict Object, please use "
            "`DNADataset.dataset[datatype].__getitem__(idx)` "
            "instead."
        )
    else:
        return self.dataset[idx]

__len__

__len__()

Return the length of the dataset.

Returns:

Type	Description
int	Length of the dataset or total length for DatasetDict

Source code in dnallm/datahandling/data.py

def __len__(self) -> int:
    """Return the length of the dataset.

    Returns:
        Length of the dataset or total length for DatasetDict
    """
    if isinstance(self.dataset, DatasetDict):
        # Return total length across all splits
        return sum(len(self.dataset[dt]) for dt in self.dataset)
    else:
        return len(self.dataset)

augment_reverse_complement

augment_reverse_complement(reverse=True, complement=True)

Augment the dataset by adding reverse complement sequences.

This method doubles the dataset size.

Parameters:

Name	Type	Description	Default
reverse	bool	Whether to do reverse	True
complement	bool	Whether to do complement	True

Source code in dnallm/datahandling/data.py

def augment_reverse_complement(
    self, reverse: bool = True, complement: bool = True
) -> None:
    """Augment the dataset by adding reverse complement sequences.

    This method doubles the dataset size.

    Args:
        reverse: Whether to do reverse
        complement: Whether to do complement
    """

    def process(ds, reverse, complement):
        # Create a dataset with an extra field for the reverse complement.
        def add_rc(example):
            example["rc_sequence"] = reverse_complement(
                example["sequence"], reverse=reverse, complement=complement
            )
            return example

        ds_with_rc = ds.map(add_rc, desc="Reverse complementary")
        # Build a new dataset where the reverse complement becomes the
        # 'sequence'
        rc_ds = ds_with_rc.map(
            lambda ex: {
                "sequence": ex["rc_sequence"],
                "labels": ex["labels"],
            },
            desc="Data augment",
        )
        ds = concatenate_datasets([ds, rc_ds])
        ds.remove_columns(["rc_sequence"])
        return ds

    if isinstance(self.dataset, DatasetDict):
        for dt in self.dataset:
            self.dataset[dt] = process(
                self.dataset[dt], reverse, complement
            )
    else:
        self.dataset = process(self.dataset, reverse, complement)

concat_reverse_complement

concat_reverse_complement(
    reverse=True, complement=True, sep=""
)

Augment each sample by concatenating the sequence with its reverse complement.

Parameters:

Name	Type	Description	Default
reverse	bool	Whether to do reverse	True
complement	bool	Whether to do complement	True
sep	str	Separator between the original and reverse complement sequences	''

Source code in dnallm/datahandling/data.py

def concat_reverse_complement(
    self, reverse: bool = True, complement: bool = True, sep: str = ""
) -> None:
    """Augment each sample by concatenating the sequence with its reverse
    complement.

    Args:
        reverse: Whether to do reverse
        complement: Whether to do complement
        sep: Separator between the original and reverse complement
            sequences
    """

    def process(ds, reverse, complement, sep):
        def concat_fn(example):
            rc = reverse_complement(
                example["sequence"], reverse=reverse, complement=complement
            )
            example["sequence"] = example["sequence"] + sep + rc
            return example

        ds = ds.map(concat_fn, desc="Data augment")
        return ds

    if isinstance(self.dataset, DatasetDict):
        for dt in self.dataset:
            self.dataset[dt] = process(
                self.dataset[dt], reverse, complement, sep
            )
    else:
        self.dataset = process(self.dataset, reverse, complement, sep)

encode_sequences

encode_sequences(
    padding="max_length",
    return_tensors="pt",
    remove_unused_columns=False,
    uppercase=False,
    lowercase=False,
    task="SequenceClassification",
    tokenizer=None,
)

Encode all sequences using the provided tokenizer.

The dataset is mapped to include tokenized fields along with the label, making it directly usable with Hugging Face Trainer.

Parameters:

Name	Type	Description	Default
padding	str	Padding strategy for sequences. Can be 'max_length' or 'longest'. Use 'longest' to pad to the length of the longest sequence in case of memory outage	'max_length'
return_tensors	str	Returned tensor types, can be 'pt', 'tf', 'np', or 'jax'	'pt'
remove_unused_columns	bool	Whether to remove the original 'sequence' and 'label' columns	False
uppercase	bool	Whether to convert sequences to uppercase	False
lowercase	bool	Whether to convert sequences to lowercase	False
task	str | None	Task type for the tokenizer. If not provided, defaults to 'SequenceClassification'	'SequenceClassification'
tokenizer	PreTrainedTokenizerBase | None	Tokenizer to use for encoding. If not provided, uses the instance's tokenizer	None

Raises:

Type	Description
ValueError	If tokenizer is not provided

Source code in dnallm/datahandling/data.py

def encode_sequences(
    self,
    padding: str = "max_length",
    return_tensors: str = "pt",
    remove_unused_columns: bool = False,
    uppercase: bool = False,
    lowercase: bool = False,
    task: str | None = "SequenceClassification",
    tokenizer: PreTrainedTokenizerBase | None = None,
) -> None:
    """Encode all sequences using the provided tokenizer.

    The dataset is mapped to include tokenized fields along with the
    label, making it directly usable with Hugging Face Trainer.

    Args:
        padding: Padding strategy for sequences. Can be 'max_length' or
            'longest'. Use 'longest' to pad to the length of the longest
            sequence in case of memory outage
        return_tensors: Returned tensor types, can be 'pt', 'tf', 'np', or
            'jax'
        remove_unused_columns: Whether to remove the original 'sequence'
            and 'label' columns
        uppercase: Whether to convert sequences to uppercase
        lowercase: Whether to convert sequences to lowercase
        task: Task type for the tokenizer. If not provided, defaults to
            'SequenceClassification'
        tokenizer: Tokenizer to use for encoding. If not provided, uses
            the instance's tokenizer

    Raises:
        ValueError: If tokenizer is not provided
    """
    if not self.tokenizer:
        if tokenizer:
            self.tokenizer = tokenizer
        else:
            raise ValueError("Tokenizer is required")

    # Get tokenizer configuration
    tokenizer_config = self._get_tokenizer_config()

    # Judge the task type and apply appropriate tokenization
    if task is None:
        task = "sequenceclassification"
    task = task.lower()

    if task in ["tokenclassification", "token", "ner"]:
        self._apply_token_classification_tokenization(
            tokenizer_config, padding, uppercase, lowercase
        )
    else:
        self._apply_sequence_classification_tokenization(
            tokenizer_config, padding, uppercase, lowercase
        )

    # Post-process dataset
    self._post_process_encoded_dataset(
        remove_unused_columns, return_tensors
    )

from_huggingface classmethod

from_huggingface(
    dataset_name,
    seq_col="sequence",
    label_col="labels",
    data_dir=None,
    tokenizer=None,
    max_length=512,
)

Load a dataset from the Hugging Face Hub.

Parameters:

Name	Type	Description	Default
dataset_name	str	Name of the dataset	required
seq_col	str	Column name for the DNA sequence	'sequence'
label_col	str	Column name for the label	'labels'
data_dir	str | None	Data directory in a dataset	None
tokenizer	PreTrainedTokenizerBase | None	Tokenizer for sequence encoding	None
max_length	int	Max token length	512

Returns:

Type	Description
DNADataset	An instance wrapping a datasets.Dataset

Source code in dnallm/datahandling/data.py

@classmethod
def from_huggingface(
    cls,
    dataset_name: str,
    seq_col: str = "sequence",
    label_col: str = "labels",
    data_dir: str | None = None,
    tokenizer: PreTrainedTokenizerBase | None = None,
    max_length: int = 512,
) -> "DNADataset":
    """Load a dataset from the Hugging Face Hub.

    Args:
        dataset_name: Name of the dataset
        seq_col: Column name for the DNA sequence
        label_col: Column name for the label
        data_dir: Data directory in a dataset
        tokenizer: Tokenizer for sequence encoding
        max_length: Max token length

    Returns:
        An instance wrapping a datasets.Dataset
    """
    if data_dir:
        ds = load_dataset(dataset_name, data_dir=data_dir)
    else:
        ds = load_dataset(dataset_name)
    # Rename columns if necessary
    if seq_col != "sequence":
        ds = ds.rename_column(seq_col, "sequence")
    if label_col != "labels":
        ds = ds.rename_column(label_col, "labels")
    return cls(ds, tokenizer=tokenizer, max_length=max_length)

from_modelscope classmethod

from_modelscope(
    dataset_name,
    seq_col="sequence",
    label_col="labels",
    data_dir=None,
    tokenizer=None,
    max_length=512,
)

Load a dataset from the ModelScope.

Parameters:

Name	Type	Description	Default
dataset_name	str	Name of the dataset	required
seq_col	str	Column name for the DNA sequence	'sequence'
label_col	str	Column name for the label	'labels'
data_dir	str | None	Data directory in a dataset	None
tokenizer	PreTrainedTokenizerBase | None	Tokenizer for sequence encoding	None
max_length	int	Max token length	512

Returns:

Type	Description
DNADataset	An instance wrapping a datasets.Dataset

Source code in dnallm/datahandling/data.py

@classmethod
def from_modelscope(
    cls,
    dataset_name: str,
    seq_col: str = "sequence",
    label_col: str = "labels",
    data_dir: str | None = None,
    tokenizer: PreTrainedTokenizerBase | None = None,
    max_length: int = 512,
) -> "DNADataset":
    """Load a dataset from the ModelScope.

    Args:
        dataset_name: Name of the dataset
        seq_col: Column name for the DNA sequence
        label_col: Column name for the label
        data_dir: Data directory in a dataset
        tokenizer: Tokenizer for sequence encoding
        max_length: Max token length

    Returns:
        An instance wrapping a datasets.Dataset
    """
    from modelscope import MsDataset

    if data_dir:
        ds = MsDataset.load(dataset_name, data_dir=data_dir)
    else:
        ds = MsDataset.load(dataset_name)
    # Rename columns if necessary
    if seq_col != "sequence":
        ds = ds.rename_column(seq_col, "sequence")
    if label_col != "labels":
        ds = ds.rename_column(label_col, "labels")
    return cls(ds, tokenizer=tokenizer, max_length=max_length)

get_split_lengths

get_split_lengths()

Get lengths of individual splits for DatasetDict.

Returns:

Type	Description
dict | None	Dictionary of split names and their lengths, or None for single
dict | None	dataset

Source code in dnallm/datahandling/data.py

def get_split_lengths(self) -> dict | None:
    """Get lengths of individual splits for DatasetDict.

    Returns:
        Dictionary of split names and their lengths, or None for single
        dataset
    """
    if isinstance(self.dataset, DatasetDict):
        return {dt: len(self.dataset[dt]) for dt in self.dataset}
    else:
        return None

head

head(head=10, show=False)

Fetch the head n data from the dataset.

Parameters:

Name	Type	Description	Default
head	int	Number of samples to fetch	10
show	bool	Whether to print the data or return it	False

Returns:

Type	Description
dict[Any, Any] | None	A dictionary containing the first n samples if show=False,
dict[Any, Any] | None	otherwise None

Source code in dnallm/datahandling/data.py

def head(
    self, head: int = 10, show: bool = False
) -> dict[Any, Any] | None:
    """Fetch the head n data from the dataset.

    Args:
        head: Number of samples to fetch
        show: Whether to print the data or return it

    Returns:
        A dictionary containing the first n samples if show=False,
        otherwise None
    """
    import pprint

    def format_convert(data):
        df: dict[Any, Any] = {}
        length = len(data["sequence"])
        for i in range(length):
            df[i] = {}
            for key in data.keys():
                df[i][key] = data[key][i]
        return df

    dataset = self.dataset
    if isinstance(dataset, DatasetDict):
        df = {}
        for dt in dataset.keys():
            data = dataset[dt][:head]
            if show:
                print(f"Dataset: {dt}")
                pprint.pp(format_convert(data))
            else:
                df[dt] = data
        return df if not show else None
    else:
        data = dataset[:head]
        if show:
            pprint.pp(format_convert(data))
            return None
        else:
            return dict(data)

iter_batches

iter_batches(batch_size)

Generator that yields batches of examples from the dataset.

Parameters:

Name	Type	Description	Default
batch_size	int	Size of each batch	required

Yields:

Type	Description
	A batch of examples

Raises:

Type	Description
ValueError	If dataset is a DatasetDict

Source code in dnallm/datahandling/data.py

def iter_batches(self, batch_size: int):
    """Generator that yields batches of examples from the dataset.

    Args:
        batch_size: Size of each batch

    Yields:
        A batch of examples

    Raises:
        ValueError: If dataset is a DatasetDict
    """
    if isinstance(self.dataset, DatasetDict):
        raise ValueError(
            "Dataset is a DatasetDict Object, please use "
            "`DNADataset.dataset[datatype].iter_batches(batch_size)` "
            "instead."
        )
    else:
        for i in range(0, len(self.dataset), batch_size):
            yield self.dataset[i : i + batch_size]

load_local_data classmethod

load_local_data(
    file_paths,
    seq_col="sequence",
    label_col="labels",
    sep=None,
    fasta_sep="|",
    multi_label_sep=None,
    tokenizer=None,
    max_length=512,
)

Load DNA sequence datasets from one or multiple local files.

Supports input formats: csv, tsv, json, parquet, arrow, dict, fasta, txt, pkl, pickle.

Parameters:

Name	Type	Description	Default
file_paths	str | list | dict	Single dataset: Provide one file path (e.g., "data.csv"). Pre-split datasets: Provide a dict like {"train": "train.csv", "test": "test.csv"}	required
seq_col	str	Column name for DNA sequences	'sequence'
label_col	str	Column name for labels	'labels'
sep	str | None	Delimiter for CSV, TSV, or TXT	None
fasta_sep	str	Delimiter for FASTA files	'|'
multi_label_sep	str | None	Delimiter for multi-label sequences	None
tokenizer	PreTrainedTokenizerBase | None	A tokenizer for sequence encoding	None
max_length	int	Max token length	512

Returns:

Type	Description
DNADataset	An instance wrapping a Dataset or DatasetDict

Raises:

Type	Description
ValueError	If file type is not supported

Source code in dnallm/datahandling/data.py

@classmethod
def load_local_data(
    cls,
    file_paths: str | list | dict,
    seq_col: str = "sequence",
    label_col: str = "labels",
    sep: str | None = None,
    fasta_sep: str = "|",
    multi_label_sep: str | None = None,
    tokenizer: PreTrainedTokenizerBase | None = None,
    max_length: int = 512,
) -> "DNADataset":
    """Load DNA sequence datasets from one or multiple local files.

    Supports input formats: csv, tsv, json, parquet, arrow, dict, fasta,
    txt, pkl, pickle.

    Args:
        file_paths: Single dataset: Provide one file path
            (e.g., "data.csv").
            Pre-split datasets: Provide a dict like
            {"train": "train.csv", "test": "test.csv"}
        seq_col: Column name for DNA sequences
        label_col: Column name for labels
        sep: Delimiter for CSV, TSV, or TXT
        fasta_sep: Delimiter for FASTA files
        multi_label_sep: Delimiter for multi-label sequences
        tokenizer: A tokenizer for sequence encoding
        max_length: Max token length

    Returns:
        An instance wrapping a Dataset or DatasetDict

    Raises:
        ValueError: If file type is not supported
    """
    # Set separators
    cls.sep = sep
    cls.multi_label_sep = multi_label_sep
    # Check if input is a list or dict
    if isinstance(
        file_paths, dict
    ):  # Handling multiple files (pre-split datasets)
        ds_dict = {}
        for split, path in file_paths.items():
            ds_dict[split] = cls._load_single_data(
                path, seq_col, label_col, sep, fasta_sep, multi_label_sep
            )
        dataset = DatasetDict(ds_dict)
    else:  # Handling a single file
        dataset = cls._load_single_data(
            file_paths, seq_col, label_col, sep, fasta_sep, multi_label_sep
        )
    dataset.stats = None  # Initialize stats as None

    return cls(dataset, tokenizer=tokenizer, max_length=max_length)

plot_statistics

plot_statistics(save_path=None)

Plot statistics of the dataset.

Includes sequence length distribution (histogram), GC content distribution (box plot) for each sequence. If dataset is a DatasetDict, length plots and GC content plots from different datasets will be concatenated into a single chart, respectively. Sequence length distribution is shown as a histogram, with min and max lengths for its' limit.

Parameters:

Name	Type	Description	Default
save_path	str | None	Path to save the plots. If None, plots will be shown interactively	None

Raises:

Type	Description
ValueError	If statistics have not been computed yet

Source code in dnallm/datahandling/data.py

def plot_statistics(self, save_path: str | None = None) -> None:
    """Plot statistics of the dataset.

    Includes sequence length distribution (histogram),
    GC content distribution (box plot) for each sequence.
    If dataset is a DatasetDict, length plots and GC content plots from
    different datasets will be concatenated into a single chart,
    respectively. Sequence length distribution is shown as a histogram,
    with min and max lengths for its' limit.

    Args:
        save_path: Path to save the plots. If None, plots will be shown
            interactively

    Raises:
        ValueError: If statistics have not been computed yet
    """
    import altair as alt

    alt.data_transformers.enable("vegafusion")

    if self.stats is None or self.stats_for_plot is None:
        raise ValueError(
            "Statistics have not been computed yet. Please call "
            "`statistics()` method first."
        )

    task_type = self.data_type or "unknown"
    df = self.stats_for_plot.copy()
    final = self._create_final_chart(df, task_type)
    self._display_or_save_chart(final, save_path)

process_missing_data

process_missing_data()

Filter out samples with missing or empty sequences or labels.

Source code in dnallm/datahandling/data.py

def process_missing_data(self) -> None:
    """Filter out samples with missing or empty sequences or labels."""

    def non_missing(example):
        return (
            example["sequence"]
            and example["labels"] is not None
            and example["sequence"].strip() != ""
        )

    self.dataset = self.dataset.filter(non_missing)

random_generate

random_generate(
    minl,
    maxl=0,
    samples=1,
    gc=(0, 1),
    n_ratio=0.0,
    padding_size=0,
    seed=None,
    label_func=None,
    append=False,
)

Replace the current dataset with randomly generated DNA sequences.

Parameters:

Name	Type	Description	Default
minl	int	Minimum length of the sequences	required
maxl	int	Maximum length of the sequences, default is the same as minl	0
samples	int	Number of sequences to generate, default 1	1
gc	tuple	GC content range, default (0,1)	(0, 1)
N_ratio		Include N base in the generated sequence, default 0.0	required
padding_size	int	Padding size for sequence length, default 0	0
seed	int | None	Random seed, default None	None
label_func	Callable | None	A function that generates a label from a sequence	None
append	bool	Append the random generated data to the existing dataset or use the data as a dataset	False

Source code in dnallm/datahandling/data.py

def random_generate(
    self,
    minl: int,
    maxl: int = 0,
    samples: int = 1,
    gc: tuple = (0, 1),
    n_ratio: float = 0.0,
    padding_size: int = 0,
    seed: int | None = None,
    label_func: Callable | None = None,
    append: bool = False,
) -> None:
    """Replace the current dataset with randomly generated DNA sequences.

    Args:
        minl: Minimum length of the sequences
        maxl: Maximum length of the sequences, default is the same as minl
        samples: Number of sequences to generate, default 1
        gc: GC content range, default (0,1)
        N_ratio: Include N base in the generated sequence, default 0.0
        padding_size: Padding size for sequence length, default 0
        seed: Random seed, default None
        label_func: A function that generates a label from a sequence
        append: Append the random generated data to the existing dataset
            or use the data as a dataset
    """

    def process(
        minl, maxl, number, gc, n_ratio, padding_size, seed, label_func
    ):
        sequences = random_generate_sequences(
            minl=minl,
            maxl=maxl,
            samples=number,
            gc=gc,
            n_ratio=n_ratio,
            padding_size=padding_size,
            seed=seed,
        )
        labels = []
        for seq in sequences:
            labels.append(label_func(seq) if label_func else 0)
        random_ds = Dataset.from_dict({
            "sequence": sequences,
            "labels": labels,
        })
        return random_ds

    if append:
        if isinstance(self.dataset, DatasetDict):
            for dt in self.dataset:
                total_length = sum(
                    len(self.dataset[split])
                    for split in self.dataset.keys()
                )
                number = round(
                    samples * len(self.dataset[dt]) / total_length
                )
                random_ds = process(
                    minl,
                    maxl,
                    number,
                    gc,
                    n_ratio,
                    padding_size,
                    seed,
                    label_func,
                )
                self.dataset[dt] = concatenate_datasets([
                    self.dataset[dt],
                    random_ds,
                ])
        else:
            random_ds = process(
                minl,
                maxl,
                samples,
                gc,
                n_ratio,
                padding_size,
                seed,
                label_func,
            )
            self.dataset = concatenate_datasets([self.dataset, random_ds])
    else:
        self.dataset = process(
            minl,
            maxl,
            samples,
            gc,
            n_ratio,
            padding_size,
            seed,
            label_func,
        )

raw_reverse_complement

raw_reverse_complement(ratio=0.5, seed=None)

Do reverse complement of sequences in the dataset.

Parameters:

Name	Type	Description	Default
ratio	float	Ratio of sequences to reverse complement	0.5
seed	int | None	Random seed for reproducibility	None

Source code in dnallm/datahandling/data.py

def raw_reverse_complement(
    self, ratio: float = 0.5, seed: int | None = None
) -> None:
    """Do reverse complement of sequences in the dataset.

    Args:
        ratio: Ratio of sequences to reverse complement
        seed: Random seed for reproducibility
    """

    def process(ds, ratio, seed):
        random.seed(seed)
        number = len(ds["sequence"])
        idxlist = set(random.sample(range(number), int(number * ratio)))

        def concat_fn(example, idx):
            rc = reverse_complement(example["sequence"])
            if idx in idxlist:
                example["sequence"] = rc
            return example

        # Create a dataset with random reverse complement.
        ds.map(concat_fn, with_indices=True, desc="Reverse complementary")
        return ds

    if isinstance(self.dataset, DatasetDict):
        for dt in self.dataset:
            self.dataset[dt] = process(self.dataset[dt], ratio, seed)
    else:
        self.dataset = process(self.dataset, ratio, seed)

sampling

sampling(ratio=1.0, seed=None, overwrite=False)

Randomly sample a fraction of the dataset.

Parameters:

Name	Type	Description	Default
ratio	float	Fraction of the dataset to sample. Default is 1.0 (no sampling)	1.0
seed	int | None	Random seed for reproducibility	None
overwrite	bool	Whether to overwrite the original dataset with the sampled one	False

Returns:

Type	Description
DNADataset	A DNADataset object with sampled data

Source code in dnallm/datahandling/data.py

def sampling(
    self,
    ratio: float = 1.0,
    seed: int | None = None,
    overwrite: bool = False,
) -> "DNADataset":
    """Randomly sample a fraction of the dataset.

    Args:
        ratio: Fraction of the dataset to sample. Default is 1.0
            (no sampling)
        seed: Random seed for reproducibility
        overwrite: Whether to overwrite the original dataset with the
            sampled one

    Returns:
        A DNADataset object with sampled data
    """
    if ratio <= 0 or ratio > 1:
        raise ValueError("ratio must be between 0 and 1")

    dataset = self.dataset
    if isinstance(dataset, DatasetDict):
        for dt in dataset.keys():
            random.seed(seed)
            random_idx = random.sample(
                range(len(dataset[dt])), int(len(dataset[dt]) * ratio)
            )
            dataset[dt] = dataset[dt].select(random_idx)
    else:
        random_idx = random.sample(
            range(len(dataset)), int(len(dataset) * ratio)
        )
        dataset = dataset.select(random_idx)

    if overwrite:
        self.dataset = dataset
        return self
    else:
        # Create a new DNADataset object with the sampled data
        return DNADataset(dataset, self.tokenizer, self.max_length)

show

show(head=10)

Display the dataset.

Parameters:

Name	Type	Description	Default
head	int	Number of samples to display	10

Source code in dnallm/datahandling/data.py

def show(self, head: int = 10) -> None:
    """Display the dataset.

    Args:
        head: Number of samples to display
    """
    self.head(head=head, show=True)

shuffle

shuffle(seed=None)

Shuffle the dataset.

Parameters:

Name	Type	Description	Default
seed	int | None	Random seed for reproducibility	None

Source code in dnallm/datahandling/data.py

def shuffle(self, seed: int | None = None) -> None:
    """Shuffle the dataset.

    Args:
        seed: Random seed for reproducibility
    """
    self.dataset.shuffle(seed=seed)

split_data

split_data(test_size=0.2, val_size=0.1, seed=None)

Split the dataset into train, test, and validation sets.

Parameters:

Name	Type	Description	Default
test_size	float	Proportion of the dataset to include in the test split	0.2
val_size	float	Proportion of the dataset to include in the validation split	0.1
seed	int | None	Random seed for reproducibility	None

Source code in dnallm/datahandling/data.py

def split_data(
    self,
    test_size: float = 0.2,
    val_size: float = 0.1,
    seed: int | None = None,
) -> None:
    """Split the dataset into train, test, and validation sets.

    Args:
        test_size: Proportion of the dataset to include in the test
            split
        val_size: Proportion of the dataset to include in the validation
            split
        seed: Random seed for reproducibility
    """
    # check if the dataset is already a DatasetDict
    if isinstance(self.dataset, DatasetDict):
        raise ValueError(
            "Dataset is already a DatasetDict, no need to split"
        )
    # First, split off test+validation from training data
    split_result = self.dataset.train_test_split(
        test_size=test_size + val_size, seed=seed
    )
    train_ds = split_result["train"]
    temp_ds = split_result["test"]
    # Further split temp_ds into test and validation sets
    if val_size > 0:
        rel_val_size = val_size / (test_size + val_size)
        temp_split = temp_ds.train_test_split(
            test_size=rel_val_size, seed=seed
        )
        test_ds = temp_split["train"]
        val_ds = temp_split["test"]
        self.dataset = DatasetDict({
            "train": train_ds,
            "test": test_ds,
            "val": val_ds,
        })
    else:
        self.dataset = DatasetDict({"train": train_ds, "test": temp_ds})

statistics

statistics()

Get statistics of the dataset.

Includes number of samples, sequence length (min, max, average, median), label distribution, GC content (by labels), nucleotide composition (by labels).

Returns:

Type	Description
dict	A dictionary containing statistics of the dataset

Raises:

Type	Description
ValueError	If statistics have not been computed yet

Source code in dnallm/datahandling/data.py

def statistics(self) -> dict:
    """Get statistics of the dataset.

    Includes number of samples, sequence length (min, max, average,
    median), label distribution, GC content (by labels), nucleotide
    composition (by labels).

    Returns:
        A dictionary containing statistics of the dataset

    Raises:
        ValueError: If statistics have not been computed yet
    """

    def prepare_dataframe(dataset) -> pd.DataFrame:
        """Convert a datasets.Dataset to pandas DataFrame if needed.

        If the input is already a pandas DataFrame, return a copy.
        """
        # avoid importing datasets at top-level to keep dependency optional
        try:
            from datasets import Dataset

            is_dataset = isinstance(dataset, Dataset)
        except Exception:
            is_dataset = False

        df: pd.DataFrame
        if is_dataset:
            df = dataset.to_pandas()
        elif isinstance(dataset, pd.DataFrame):
            df = dataset.copy()
        else:
            raise ValueError(
                "prepare_dataframe expects a datasets.Dataset or "
                "pandas.DataFrame"
            )
        return df

    def compute_basic_stats(
        df: pd.DataFrame, seq_col: str = "sequence"
    ) -> dict:
        """Compute number of samples and sequence length statistics."""
        seqs = df[seq_col].fillna("").astype(str)
        lens = seqs.str.len()
        return {
            "n_samples": len(lens),
            "min_len": int(lens.min()) if len(lens) > 0 else 0,
            "max_len": int(lens.max()) if len(lens) > 0 else 0,
            "mean_len": float(lens.mean())
            if len(lens) > 0
            else float("nan"),
            "median_len": float(lens.median())
            if len(lens) > 0
            else float("nan"),
        }

    stats = {}
    seq_col = "sequence"
    # label_col = "labels"  # Not used in current implementation
    if isinstance(self.dataset, DatasetDict):
        for split_name, split_ds in self.dataset.items():
            df = prepare_dataframe(split_ds)
            data_type = self.data_type
            basic = compute_basic_stats(df, seq_col)
            stats[split_name] = {"data_type": data_type, **basic}
    else:
        df = prepare_dataframe(self.dataset)
        data_type = self.data_type
        basic = compute_basic_stats(df, seq_col)
        stats["full"] = {"data_type": data_type, **basic}

    self.stats = stats  # Store stats in the instance for later use
    self.stats_for_plot = df

    return stats

validate_sequences

validate_sequences(
    minl=20, maxl=6000, gc=(0, 1), valid_chars="ACGTN"
)

Filter the dataset to keep sequences containing valid DNA bases or allowed length.

Parameters:

Name	Type	Description	Default
minl	int	Minimum length of the sequences	20
maxl	int	Maximum length of the sequences	6000
gc	tuple	GC content range between 0 and 1	(0, 1)
valid_chars	str	Allowed characters in the sequences	'ACGTN'

Source code in dnallm/datahandling/data.py

def validate_sequences(
    self,
    minl: int = 20,
    maxl: int = 6000,
    gc: tuple = (0, 1),
    valid_chars: str = "ACGTN",
) -> None:
    """Filter the dataset to keep sequences containing valid DNA bases or
    allowed length.

    Args:
        minl: Minimum length of the sequences
        maxl: Maximum length of the sequences
        gc: GC content range between 0 and 1
        valid_chars: Allowed characters in the sequences
    """
    self.dataset = self.dataset.filter(
        lambda example: check_sequence(
            example["sequence"], minl, maxl, gc, valid_chars
        )
    )

Functions

load_preset_dataset

load_preset_dataset(dataset_name, task=None)

Load a preset dataset from Hugging Face or ModelScope.

Parameters:

Name	Type	Description	Default
dataset_name	str	Name of the dataset	required
task	str | None	Task directory in a dataset	None

Returns:

Type	Description
DNADataset	An instance wrapping a datasets.Dataset

Raises:

Type	Description
ValueError	If dataset is not found in preset datasets

Source code in dnallm/datahandling/data.py

def load_preset_dataset(
    dataset_name: str, task: str | None = None
) -> "DNADataset":
    """Load a preset dataset from Hugging Face or ModelScope.

    Args:
        dataset_name: Name of the dataset
        task: Task directory in a dataset

    Returns:
        An instance wrapping a datasets.Dataset

    Raises:
        ValueError: If dataset is not found in preset datasets
    """
    from .dataset_auto import PRESET_DATASETS

    ds_info = _get_dataset_info(dataset_name, PRESET_DATASETS)
    ds = _load_dataset_from_modelscope(ds_info, task)
    ds = _standardize_column_names(ds)
    return _create_dna_dataset(ds, ds_info)

show_preset_dataset

show_preset_dataset()

Show all preset datasets available in Hugging Face or ModelScope.

Returns:

Type	Description
dict	A dictionary containing dataset names and their descriptions

Source code in dnallm/datahandling/data.py

def show_preset_dataset() -> dict:
    """Show all preset datasets available in Hugging Face or ModelScope.

    Returns:
        A dictionary containing dataset names and their descriptions
    """
    from .dataset_auto import PRESET_DATASETS

    return PRESET_DATASETS