Class: GRX::NN::Embedding

Inherits:
Module
  • Object
show all
Defined in:
lib/grx/nn.rb

Overview

================================================================

Embedding — Dense vector lookup table for token indices

Instance Attribute Summary collapse

Instance Method Summary collapse

Methods inherited from Module

#call, #load_weights, #parameters, #save_weights, #zero_grad

Constructor Details

#initialize(num_embeddings, embedding_dim) ⇒ Embedding

Returns a new instance of Embedding.



204
205
206
207
208
 
# File 'lib/grx/nn.rb', line 204

def initialize(num_embeddings, embedding_dim)
  @num_embeddings = num_embeddings
  @embedding_dim  = embedding_dim
  @weight         = Tensor.he_normal([num_embeddings, embedding_dim], requires_grad: true)
end

Instance Attribute Details

#embedding_dim ⇒ Object (readonly)

Returns the value of attribute embedding_dim.



202
203
204
 
# File 'lib/grx/nn.rb', line 202

def embedding_dim
  @embedding_dim
end

#num_embeddings ⇒ Object (readonly)

Returns the value of attribute num_embeddings.



202
203
204
 
# File 'lib/grx/nn.rb', line 202

def num_embeddings
  @num_embeddings
end

#weight ⇒ Object (readonly)

Returns the value of attribute weight.



202
203
204
 
# File 'lib/grx/nn.rb', line 202

def weight
  @weight
end

Instance Method Details

#forward(indices) ⇒ Object 



210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
 
# File 'lib/grx/nn.rb', line 210

def forward(indices)
  ids = indices.is_a?(Tensor) ? indices.to_a.map(&:to_i) : indices.map(&:to_i)
  batch_size = ids.size
  out_data = ids.flat_map do |id|
    raise IndexError, "Token index #{id} out of range [0, #{@num_embeddings})" if id < 0 || id >= @num_embeddings
    @weight.to_a.slice(id * @embedding_dim, @embedding_dim)
  end

  out = Tensor.create(out_data, [batch_size, @embedding_dim])
  if @weight.requires_grad
    out.requires_grad = true
    out._grafo_hijos << @weight
    w = @weight; dim = @embedding_dim; num_emb = @num_embeddings
    out.backward_fn = ->(g) {
      grad_w = Array.new(num_emb * dim, 0.0)
      g_data = g.to_a
      ids.each_with_index do |id, i|
        slice = g_data.slice(i * dim, dim)
        dim.times { |d| grad_w[id * dim + d] += slice[d] }
      end
      w.agregar_gradiente(Tensor.create(grad_w, w.shape))
    }
  end
  out
end

#to_s ⇒ Object 



236
237
238
 
# File 'lib/grx/nn.rb', line 236

def to_s
  "Embedding(#{@num_embeddings}, #{@embedding_dim})"
end