Dataset

Now that we have split our data into (input, target) pairs, and remember the target is just the same tokens shifted one position to the right, we need an efficient way to feed them to the model. This is where the Dataset and DataLoader come in.

Dataset

The Dataset is a simple class that wraps our token data and turns it into chunks. Each chunk is a small object holding two fixed-size arrays, an input and its target, where the length of each is a predefined size we call the context size.

export interface Sample {
  input: number[];
  target: number[];
}

export class LlmDataset {
  private readonly samples: Sample[] = [];

  // maxLength is the context size: how many tokens each chunk holds.
  constructor(tokenIds: number[], maxLength: number, stride: number) {
    // Slide a fixed-size window across the token stream, stepping by `stride`.
    for (let i = 0; i + maxLength < tokenIds.length; i += stride) {
      this.samples.push({
        input: tokenIds.slice(i, i + maxLength), // the window
        target: tokenIds.slice(i + 1, i + maxLength + 1), // same window, shifted right by one
      });
    }
  }
}

The stride controls how far the window jumps between chunks: a stride equal to the context size gives non-overlapping chunks, while a smaller stride overlaps them.

Full implementation

The snippet above shows only the core sliding-window logic. The complete LlmDataset class, with length, at(), and iteration support, lives in @node-llm/core at packages/core/src/data/dataset.ts. You can run it through the chapter’s practice code at packages/chapters/src/ch02-tokenizer/04-dataset-dataloader.ts.

Tensors

We already know that looping through an array is slow, and once our dataset grows to thousands, or even hundreds of thousands, of words, that slowness adds up fast. To train at any reasonable speed, we need to take advantage of the parallelism that GPUs are built for.

This is where the tensor representation from @tensorflow/tfjs-core comes in. Instead of handing the model plain JavaScript arrays, we use tensors: multi-dimensional arrays held in a single contiguous, typed buffer. Because an entire batch lives in one tensor, the model can run its math on all of it at once, on CPU or GPU, rather than looping element by element. That is what makes data loading fast enough to train on.

PyTorch → JS mapping

Reference (Python)	Here (TypeScript)
torch.utils.data.Dataset	LlmDataset
__getitem__ / __len__	at(i) / length

Recap

• The Dataset turns a flat token stream into fixed-size (input, target) chunks using a sliding window and a stride.
• Each chunk holds an input and its target, where the target is the same tokens shifted right by one.
• Next, the DataLoader groups these chunks into batches and can hand them back as tensors.