Class: Informers::VisionEncoderDecoderModel

Inherits:
PreTrainedModel show all
Defined in:
lib/informers/models.rb

Constant Summary collapse

MAIN_INPUT_NAME = 
:pixel_values

Instance Attribute Summary

Attributes inherited from PreTrainedModel

#config

Instance Method Summary collapse

Methods inherited from PreTrainedModel

#call, construct_session, from_pretrained, #generate

Constructor Details

#initialize(config, session, decoder_merged_session, generation_config) ⇒ VisionEncoderDecoderModel

Returns a new instance of VisionEncoderDecoderModel.



1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
 
# File 'lib/informers/models.rb', line 1019

def initialize(config, session, decoder_merged_session, generation_config)
  super(config, session)
  @decoder_merged_session = decoder_merged_session
  @generation_config = generation_config

  # Extract configs
  encoder_config = @config["encoder"]
  decoder_config = @config["decoder"]

  # Validate encoder
  encoder_model_type = encoder_config["model_type"]
  encoder_model = MODEL_MAPPING_NAMES_ENCODER_ONLY[encoder_model_type] || MODEL_MAPPING_NAMES_ENCODER_DECODER[encoder_model_type]
  if !encoder_model
    warn "Model type for encoder '#{encoder_model_type}' not found, assuming encoder-only architecture. Please report this."
  end

  # Validate decoder
  decoder_model = MODEL_WITH_LM_HEAD_MAPPING_NAMES[decoder_config["model_type"]]
  if !decoder_model
    raise Error, "Unable to construct `VisionEncoderDecoder` due to unsupported decoder: \"#{decoder_config["model_type"]}\""
  end

  decoder_model_class = decoder_model[1]
  decoder = decoder_model_class.new(decoder_config, decoder_merged_session, generation_config)

  @add_encoder_pkv = decoder.respond_to?(:num_decoder_layers)
  if @add_encoder_pkv
    # Decoder is part of an encoder-decoder model
    @num_decoder_layers = decoder.num_decoder_layers
    @num_decoder_heads = decoder.num_decoder_heads
    @decoder_dim_kv = decoder.decoder_dim_kv

    @num_encoder_layers = decoder.num_encoder_layers
    @num_encoder_heads = decoder.num_encoder_heads
    @encoder_dim_kv = decoder.encoder_dim_kv
  else
    # Decoder is a decoder-only model
    @num_layers = decoder.num_layers
    @num_heads = decoder.num_heads
    @dim_kv = decoder.dim_kv
  end
end