dataloader — Batched Dataset Iteration for Training

The dataloader package provides a DataLoader that draws random batches from a Dataset, returning input/label pairs as *variable.Variable tensors. Two dataset implementations are included: TokenDataset for pre-training on raw token sequences and AlpacaDataset for supervised fine-tuning (SFT) on instruction–response pairs.

Dataset interface

Any type that satisfies the Dataset interface can be used as a data source with DataLoader:

type Dataset interface {
    Len() int
    GetItem(i int) ([]int, []int)  // returns (x_ids, y_ids)
}

Method	Description
`Len() int`	Returns the total number of samples in the dataset
`GetItem(i int) ([]int, []int)`	Returns the input token IDs and corresponding label IDs for sample `i`

Both TokenDataset and AlpacaDataset satisfy this interface and can be swapped into a DataLoader without any other changes.

DataLoader struct

DataLoader manages iteration state and assembles individual samples into batched tensors. Unexported fields (indices, idx) track the current shuffle permutation and position within the epoch.

type DataLoader struct {
    BatchSize int
    Dataset   Dataset
    Shuffle   bool
}

BatchSize

int

required

Number of samples to include in each batch.

Dataset

required

The underlying dataset implementation. Accepts any type that satisfies the Dataset interface.

Shuffle

bool

When true, the sample index order is randomly permuted at the start of each epoch via Reset(). Defaults to false.

Batch()

func (l *DataLoader) Batch() (*variable.Variable, *variable.Variable)

Returns the next (x, y) pair of batched tensors, each of shape (BatchSize, ContextLen). Internally, Batch() calls Reset() automatically (which optionally shuffles indices) whenever the dataset is exhausted, enabling seamless multi-epoch training loops. For each sample in the batch, GetItem is called to obtain the raw token slices, which are converted to float64 tensors and stacked along dimension 0 using tensor.Stack. Usage:

loader := dataloader.DataLoader{
    BatchSize: 32,
    Shuffle:   true,
    Dataset: &dataloader.TokenDataset{
        Tokens:     dataloader.MustLoadTokens("testdata/tiny_codes.bin"),
        ContextLen: 256,
    },
}

x, y := loader.Batch()  // x, y: (32, 256)

Example with batch size 2 (no shuffle):

loader := dataloader.DataLoader{
    BatchSize: 2,
    Shuffle:   false,
    Dataset: &dataloader.TokenDataset{
        Tokens:     []int{0, 1, 2, 3, 4, 5},
        ContextLen: 2,
    },
}

x, y := loader.Batch()
// x: variable[2 2]([0 1 1 2])
// y: variable[2 2]([1 2 2 3])

When the epoch ends mid-batch, Reset() is called immediately and iteration continues from the beginning of the (re-shuffled) dataset.

TokenDataset

TokenDataset implements Dataset for pre-training on a flat token corpus. It produces next-token prediction pairs by sliding a window of length ContextLen across the token sequence.

type TokenDataset struct {
    Tokens     []int  // all token IDs from the corpus
    ContextLen int    // sequence length per sample
}

Tokens

[]int

required

The full corpus as a flat slice of integer token IDs, typically produced by a BPE tokenizer.

ContextLen

int

required

The length of each input and label sequence. Determines the window size for sliding-window sampling.

Method	Behaviour
`Len() int`	Returns `len(Tokens) - ContextLen` — the number of valid starting positions
`GetItem(i int) ([]int, []int)`	`x = Tokens[i : i+ContextLen]`, `y = Tokens[i+1 : i+ContextLen+1]`

The label sequence y is shifted one position to the right relative to x, implementing standard autoregressive next-token prediction.

MustLoadTokens()

func MustLoadTokens(path string) []int

Loads a gob-encoded []int token file from path and returns the decoded slice. Panics if the file cannot be opened or decoded. Use the non-panicking companion when you need explicit error handling:

func LoadTokens(path string) ([]int, error)

Tokens are serialised with encoding/gob. The typical workflow is to tokenise a corpus once, save the result with gob.Encoder, and load it at training time with MustLoadTokens.

AlpacaDataset

AlpacaDataset implements Dataset for supervised fine-tuning (SFT) on instruction–response data. Construct it with NewAlpacaDataset; the struct fields are unexported after construction.

func NewAlpacaDataset(alpaca []Alpaca, tokenizer Tokenizer, contextLen int) *AlpacaDataset

alpaca

[]Alpaca

required

Slice of Alpaca structs containing instruction–response pairs.

tokenizer

Tokenizer

required

Any type that satisfies Tokenizer (i.e. exposes Encode(inputText string) []int). This is typically the BPE tokenizer used during pre-training.

contextLen

int

required

Maximum sequence length. Samples shorter than contextLen are padded; samples longer are truncated.

Sample construction per Alpaca entry:

Prompt is formatted as AlpacaFormat(instruction).
Response is appended with the <|endoftext|> sentinel.
Both are encoded via tokenizer.Encode.
Input IDs = promptIDs + responseIDs (shifted: last token dropped).
Labels = -100 for every prompt token (ignored in cross-entropy) + responseIDs (shifted: first token dropped).
Sequences shorter than contextLen are right-padded with 0 IDs and -100 labels; longer sequences are truncated.

Method	Behaviour
`Len() int`	Returns the number of encoded samples
`GetItem(i int) ([]int, []int)`	Returns `(IDs, Labels)` for sample `i`

Alpaca struct

type Alpaca struct {
    Instruction string `json:"instruction"`
    Response    string `json:"response"`
}

Instruction

string

required

The instruction text shown to the model as the prompt.

Response

string

required

The expected model response. The <|endoftext|> sentinel is appended automatically by NewAlpacaDataset.

Helper functions

MustLoadAlpaca / LoadAlpaca

func MustLoadAlpaca(path string) []Alpaca
func LoadAlpaca(path string) ([]Alpaca, error)

MustLoadAlpaca reads a JSON file at path, unmarshals it as a []Alpaca slice, and panics on any error. Use LoadAlpaca when you need to handle errors explicitly.

data := dataloader.MustLoadAlpaca("alpaca_data.json")
dataset := dataloader.NewAlpacaDataset(data, tokenizer, 512)

AlpacaFormat

func AlpacaFormat(message string) string

Wraps an instruction string in the standard Alpaca prompt template. The formatted prompt is used as the model input; its tokens receive -100 labels so they do not contribute to the loss.

prompt := dataloader.AlpacaFormat("Write add function")
// Returns:
// ### Instruction:
// Write add function
//
// ### Response:

Model

Tokenizer

Layers

Data & Utilities

dataloader — Batched Dataset Iteration for Training

Dataset interface

DataLoader struct

Batch()

TokenDataset

MustLoadTokens()

AlpacaDataset

Alpaca struct

Helper functions

MustLoadAlpaca / LoadAlpaca

AlpacaFormat

Build docs developers (and LLMs) love

Model

Tokenizer

Layers

Data & Utilities

Documentation Index

​Dataset interface

​DataLoader struct

​Batch()

​TokenDataset

​MustLoadTokens()

​AlpacaDataset

​Alpaca struct

​Helper functions

​MustLoadAlpaca / LoadAlpaca

​AlpacaFormat

Build docs developers (and LLMs) love

Dataset interface

DataLoader struct

Batch()

TokenDataset

MustLoadTokens()

AlpacaDataset

Alpaca struct

Helper functions

MustLoadAlpaca / LoadAlpaca

AlpacaFormat