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.



1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
# File 'lib/informers/models.rb', line 1042

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