Class: NanoGPT::DataLoader

Inherits:
Object
  • Object
show all
Defined in:
lib/nano_gpt/data_loader.rb

Overview

Loads batches from binary token files Memory-efficient: reads from file each batch (like Python's memmap recreation)

Constant Summary collapse

BYTES_PER_TOKEN =

Each token is stored as uint16 (2 bytes) in little-endian format

2

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(data_dir:, block_size:, batch_size:, device: "cpu") ⇒ DataLoader

Returns a new instance of DataLoader.



12
13
14
15
16
17
18
19
20
21
22
23
 
# File 'lib/nano_gpt/data_loader.rb', line 12

def initialize(data_dir:, block_size:, batch_size:, device: "cpu")
  @data_dir = data_dir
  @block_size = block_size
  @batch_size = batch_size
  @device = device

  @train_path = File.join(data_dir, "train.bin")
  @val_path = File.join(data_dir, "val.bin")

  @train_size = File.size(@train_path) / BYTES_PER_TOKEN
  @val_size = File.size(@val_path) / BYTES_PER_TOKEN
end

Instance Attribute Details

#batch_size ⇒ Object (readonly)

Returns the value of attribute batch_size.



10
11
12
 
# File 'lib/nano_gpt/data_loader.rb', line 10

def batch_size
  @batch_size
end

#block_size ⇒ Object (readonly)

Returns the value of attribute block_size.



10
11
12
 
# File 'lib/nano_gpt/data_loader.rb', line 10

def block_size
  @block_size
end

#train_size ⇒ Object (readonly)

Returns the value of attribute train_size.



10
11
12
 
# File 'lib/nano_gpt/data_loader.rb', line 10

def train_size
  @train_size
end

#val_size ⇒ Object (readonly)

Returns the value of attribute val_size.



10
11
12
 
# File 'lib/nano_gpt/data_loader.rb', line 10

def val_size
  @val_size
end

Instance Method Details

#get_batch(split) ⇒ Object

Get a batch of data Memory-efficient: reads only needed bytes per batch to avoid memory leak



27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
 
# File 'lib/nano_gpt/data_loader.rb', line 27

def get_batch(split)
  path = split == :train ? @train_path : @val_path
  data_size = split == :train ? @train_size : @val_size

  max_start = data_size - @block_size - 1
  indices = Array.new(@batch_size) { rand(0..max_start) }

  x_arrays = []
  y_arrays = []

  File.open(path, "rb") do |f|
    indices.each do |i|
      f.seek(i * BYTES_PER_TOKEN)
      x_bytes = f.read((@block_size + 1) * BYTES_PER_TOKEN)
      tokens = x_bytes.unpack("S<*") # uint16 little-endian

      x_arrays << tokens[0...@block_size]
      y_arrays << tokens[1..@block_size]
    end
  end

  x = Torch.tensor(x_arrays, dtype: :long)
  y = Torch.tensor(y_arrays, dtype: :long)

  if @device != "cpu"
    x = x.to(@device)
    y = y.to(@device)
  end

  [x, y]
end