> ## Documentation Index
> Fetch the complete documentation index at: https://mintlify.com/YerevaNN/ChemLactica/llms.txt
> Use this file to discover all available pages before exploring further.

# Dataset Functions

> Dataset loading and preparation utilities

## Overview

ChemLactica provides utilities for loading and preparing datasets for both pretraining and supervised fine-tuning. The dataset system supports streaming JSONL files for efficient memory usage during pretraining and standard HuggingFace datasets for SFT.

## get\_dataset()

Main function for loading and preparing datasets based on training type.

### Function Signature

```python theme={null}
def get_dataset(
    train_type,
    training_data_dirs,
    valid_data_dir,
    dir_data_types,
    train_config,
    model_config,
    shared_jsonl_files,
    evaluate_only,
    slurm_eval,
    shuffle_buffer_size,
)
```

### Parameters

**train\_type** `str` *required*

Type of training. Options: `"pretrain"` or `"sft"`

**training\_data\_dirs** `list[str]` *required*

List of directories containing training data files. For pretrain: directories with JSONL files. For SFT: path to HuggingFace dataset

**valid\_data\_dir** `str` *required*

Directory containing validation data files (JSONL format for pretrain)

**dir\_data\_types** `list[str]` *required*

List of data types corresponding to each training directory. Must match length of `training_data_dirs`. Supported types defined in `DIR_DATA_TYPES`

**train\_config** `dict` *required*

Training configuration object containing training hyperparameters

**model\_config** `dict` *required*

Model configuration object containing model architecture parameters and tokenizer path

**shared\_jsonl\_files** `multiprocessing.Manager.dict` *required*

Shared dictionary for tracking JSONL file reading positions across processes (pretrain only). Pass `None` for SFT

**evaluate\_only** `bool` *required*

If True, only load validation dataset

**slurm\_eval** `bool` *required*

If True and not `evaluate_only`, skip loading validation dataset (evaluation runs separately via SLURM)

**shuffle\_buffer\_size** `int` *required*

Size of shuffle buffer for assay datasets during pretraining

### Returns

**dataset** `dict`

Dictionary containing `"train"` and `"validation"` datasets. For pretrain, returns iterable datasets. For SFT, returns standard HuggingFace dataset dict

### Pretrain Mode

When `train_type="pretrain"`, the function:

1. Validates that data types are supported
2. Loads JSONL files from each training directory
3. Creates iterable datasets using `samples_generator`
4. Processes each dataset with tokenization and formatting
5. Shuffles assay-type datasets with specified buffer size
6. Interleaves multiple datasets if provided
7. Loads validation dataset (unless `evaluate_only=True` or `slurm_eval=True`)

```python theme={null}
from chemlactica.get_dataset import get_dataset
import multiprocessing

with multiprocessing.Manager() as manager:
    shared_jsonl_files = manager.dict()
    
    dataset = get_dataset(
        train_type="pretrain",
        training_data_dirs=[
            "/data/molecules",
            "/data/assays"
        ],
        valid_data_dir="/data/valid",
        dir_data_types=["molecules", "assay_split"],
        train_config=train_config,
        model_config=model_config,
        shared_jsonl_files=shared_jsonl_files,
        evaluate_only=False,
        slurm_eval=False,
        shuffle_buffer_size=10000,
    )
    
    train_dataset = dataset["train"]
    eval_dataset = dataset["validation"]
```

### SFT Mode

When `train_type="sft"`, the function loads a standard HuggingFace dataset:

```python theme={null}
from chemlactica.get_dataset import get_dataset

dataset = get_dataset(
    train_type="sft",
    training_data_dirs=["username/dataset-name"],
    valid_data_dir=None,
    dir_data_types=None,
    train_config=train_config,
    model_config=model_config,
    shared_jsonl_files=None,
    evaluate_only=False,
    slurm_eval=False,
    shuffle_buffer_size=None,
)
```

### Data Types

Supported data types are defined in `DIR_DATA_TYPES`. Common types include:

* `"molecules"`: Molecular structure data
* `"assay_split"`: Assay data that should be shuffled
* Other domain-specific types

Datasets with "assay" in the type name are automatically shuffled using the specified `shuffle_buffer_size`.

## samples\_generator()

Generator function for streaming JSONL files in a distributed manner.

### Function Signature

```python theme={null}
def samples_generator(
    files: List[str],
    shared_jsonl_files,
    chunk_size=25000,
    return_line_info=False
)
```

### Parameters

**files** `list[str]` *required*

List of JSONL file paths to read

**shared\_jsonl\_files** `multiprocessing.Manager.dict` *required*

Shared dictionary for tracking file positions across processes, enabling checkpoint resumption

**chunk\_size** `int` (default: 25000)

Size of chunks for reading files (currently not used in line-by-line reading)

**return\_line\_info** `bool` (default: False)

Whether to return line position information with samples

### Yields

**sample** `dict`

Dictionary with `"text"` key containing the line content from the JSONL file

### Features

#### Distributed Reading

The generator distributes samples across multiple processes:

```python theme={null}
def should_yield_on_current_rank(i, num_processes, process_index):
    return i % num_processes == process_index
```

Each process only yields samples at indices matching its rank, ensuring no duplicate processing.

#### Checkpoint Resumption

The generator tracks file positions in `shared_jsonl_files`:

```python theme={null}
file_states = {f: {"position": 0, "line_number": 0} for f in files}
for file in file_states.keys():
    if shared_jsonl_files.get(file):
        jsonl_state = shared_jsonl_files[file]
        file_states[file] = jsonl_state
        print(f"loaded {file}: {jsonl_state['position']}")
```

When resuming from a checkpoint, the generator seeks to the saved position in each file.

#### Sample Format

Each line is formatted as:

```python theme={null}
def format_sample(line):
    sample = line.strip()
    ret = {"text": sample}
    return ret
```

### Usage Example

```python theme={null}
from chemlactica.jsonl_dataset import samples_generator
from datasets.iterable_dataset import IterableDataset
import multiprocessing

with multiprocessing.Manager() as manager:
    shared_jsonl_files = manager.dict()
    
    training_files = [
        "/data/train/molecules_001.jsonl",
        "/data/train/molecules_002.jsonl",
    ]
    
    dataset = IterableDataset.from_generator(
        samples_generator,
        gen_kwargs={
            "files": training_files,
            "shared_jsonl_files": shared_jsonl_files,
        },
    )
    
    for sample in dataset:
        print(sample["text"])
```

## Dataset Processing

All datasets loaded by `get_dataset()` are processed using `process_dataset` from `chemlactica.utils.dataset_utils`, which handles:

* Tokenization with the specified tokenizer
* Sequence formatting and truncation
* Batching and padding
* Special handling for assay vs. non-assay data

Processing parameters:

```python theme={null}
dataset = process_dataset(
    dataset=dataset,
    train_config=train_config,
    model_config=model_config,
    process_batch_sizes=(50, 50),
    is_eval=False,
    assay=is_assay_split,
)
```

## Source Reference

* `get_dataset()`: `chemlactica/get_dataset.py:10`
* `samples_generator()`: `chemlactica/jsonl_dataset.py:37`


This documentation is built and hosted on [Mintlify](https://mintlify.com), a developer documentation platform.