Class: Transformers::Bert::BertPooler

Inherits:
Torch::NN::Module
  • Object
show all
Defined in:
lib/transformers/models/bert/modeling_bert.rb

Instance Method Summary collapse

Constructor Details

#initialize(config) ⇒ BertPooler

Returns a new instance of BertPooler.



441
442
443
444
445
 
# File 'lib/transformers/models/bert/modeling_bert.rb', line 441

def initialize(config)
  super()
  @dense = Torch::NN::Linear.new(config.hidden_size, config.hidden_size)
  @activation = Torch::NN::Tanh.new
end

Instance Method Details

#forward(hidden_states) ⇒ Object 



447
448
449
450
451
452
453
454
 
# File 'lib/transformers/models/bert/modeling_bert.rb', line 447

def forward(hidden_states)
  # We "pool" the model by simply taking the hidden state corresponding
  # to the first token.
  first_token_tensor = hidden_states[0.., 0]
  pooled_output = @dense.(first_token_tensor)
  pooled_output = @activation.(pooled_output)
  pooled_output
end