104104 FluxTransformerModelPatcher ,
105105 Gemma2ModelPatcher ,
106106 Gemma3LMModelPatcher ,
107+ Gemma3nImageEmbeddingsModelPatcher ,
108+ Gemma3nLMModelPatcher ,
107109 Gemma4ImageEmbeddingsModelPatcher ,
108110 Gemma4LMModelPatcher ,
109111 Gemma4UnifiedImageEmbeddingsModelPatcher ,
@@ -1333,6 +1335,21 @@ class Gemma4UnifiedTextOpenVINOConfig(Gemma4TextOpenVINOConfig):
13331335 MAX_TRANSFORMERS_VERSION = "5.10.99"
13341336
13351337
1338+ @register_in_tasks_manager (
1339+ "gemma3n_text" ,
1340+ * [
1341+ "feature-extraction" ,
1342+ "feature-extraction-with-past" ,
1343+ "text-generation" ,
1344+ "text-generation-with-past" ,
1345+ "text-classification" ,
1346+ ],
1347+ library_name = "transformers" ,
1348+ )
1349+ class Gemma3nTextOpenVINOConfig (Gemma4TextOpenVINOConfig ):
1350+ MIN_TRANSFORMERS_VERSION = "5.0"
1351+
1352+
13361353@register_in_tasks_manager ("deci" , * ["text-generation" , "text-generation-with-past" ], library_name = "transformers" )
13371354class DeciOpenVINOConfig (TextDecoderWithPositionIdsOpenVINOConfig ):
13381355 MAX_TRANSFORMERS_VERSION = "4.57.6"
@@ -3925,7 +3942,7 @@ def __init__(
39253942 behavior = behavior ,
39263943 )
39273944 self ._behavior = behavior
3928- if self ._behavior == Gemma4ConfigBehavior .VISION_EMBEDDINGS :
3945+ if self ._behavior == Gemma4ConfigBehavior .VISION_EMBEDDINGS and config . model_type == "gemma4" :
39293946 self .DUMMY_INPUT_GENERATOR_CLASSES = (DummyGemma4VisionInputGenerator ,)
39303947 # Attach image_seq_length from preprocessor to normalized config so
39313948 # the dummy input generator can compute the correct number of patches.
@@ -3950,14 +3967,26 @@ def __init__(
39503967 break
39513968 if image_seq_length is not None :
39523969 self ._normalized_config .image_seq_length = image_seq_length
3953- elif self ._behavior in (
3954- Gemma4ConfigBehavior .TEXT_EMBEDDINGS ,
3955- Gemma4ConfigBehavior .TEXT_EMBEDDINGS_PER_LAYER ,
3970+ elif (
3971+ self ._behavior
3972+ in (
3973+ Gemma4ConfigBehavior .TEXT_EMBEDDINGS ,
3974+ Gemma4ConfigBehavior .TEXT_EMBEDDINGS_PER_LAYER ,
3975+ )
3976+ and config .model_type == "gemma4"
39563977 ):
39573978 self .DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator ,)
39583979 self ._config = config .text_config
39593980 self ._normalized_config = NormalizedTextConfig (self ._config )
39603981
3982+ @staticmethod
3983+ def _get_language_model (model ):
3984+ if hasattr (model , "language_model" ):
3985+ return model .language_model
3986+ if hasattr (model , "model" ) and hasattr (model .model , "language_model" ):
3987+ return model .model .language_model
3988+ raise AttributeError ("Gemma3n model does not expose language_model" )
3989+
39613990 def with_behavior (self , behavior : Union [str , Gemma4ConfigBehavior ]):
39623991 if isinstance (behavior , str ) and not isinstance (behavior , Gemma4ConfigBehavior ):
39633992 behavior = Gemma4ConfigBehavior (behavior )
@@ -4089,6 +4118,128 @@ def outputs(self) -> Dict[str, Dict[int, str]]:
40894118 return super ().outputs
40904119
40914120
4121+ @register_in_tasks_manager ("gemma3n" , * ["image-text-to-text" ], library_name = "transformers" )
4122+ class Gemma3nOpenVINOConfig (Gemma4OpenVINOConfig ):
4123+ def __init__ (
4124+ self ,
4125+ config : "PretrainedConfig" ,
4126+ task : str = "feature-extraction" ,
4127+ int_dtype : str = "int64" ,
4128+ float_dtype : str = "fp32" ,
4129+ behavior : Gemma4ConfigBehavior = Gemma4ConfigBehavior .VISION_EMBEDDINGS ,
4130+ preprocessors : Optional [List [Any ]] = None ,
4131+ ):
4132+ super ().__init__ (
4133+ config = config ,
4134+ task = task ,
4135+ int_dtype = int_dtype ,
4136+ float_dtype = float_dtype ,
4137+ behavior = behavior ,
4138+ preprocessors = preprocessors ,
4139+ )
4140+ # For VISION_EMBEDDINGS, the vision model contains a reshape that uses
4141+ # vision_soft_tokens_per_image as a fixed constant. The dummy input must
4142+ # produce spatial output matching that value, so we extract image size
4143+ # from the preprocessor to ensure the trace uses the correct dimensions.
4144+ if self ._behavior == Gemma4ConfigBehavior .VISION_EMBEDDINGS and config .model_type == "gemma3n" :
4145+ image_size = None
4146+ if preprocessors is not None :
4147+ for p in preprocessors :
4148+ ip = getattr (p , "image_processor" , p )
4149+ size = getattr (ip , "size" , None )
4150+ if size is not None :
4151+ if isinstance (size , dict ):
4152+ image_size = size .get ("height" , size .get ("shortest_edge" ))
4153+ elif isinstance (size , (int , float )):
4154+ image_size = int (size )
4155+ break
4156+ if image_size is not None :
4157+ self ._config .image_size = image_size
4158+
4159+ def with_behavior (self , behavior : Union [str , Gemma4ConfigBehavior ]):
4160+ if behavior == Gemma4ConfigBehavior .LANGUAGE :
4161+ return get_vlm_text_generation_config (
4162+ "gemma3n_text" ,
4163+ self ._orig_config .text_config ,
4164+ self .int_dtype ,
4165+ self .float_dtype ,
4166+ model_patcher = Gemma3nLMModelPatcher ,
4167+ inputs_update = {"per_layer_inputs" : {0 : "batch_size" , 1 : "sequence_length" , 2 : "num_hidden_layers" }},
4168+ )
4169+ return super ().with_behavior (behavior )
4170+
4171+ def get_model_for_behavior (self , model , behavior : Union [str , VLMConfigBehavior ]):
4172+ if behavior == Gemma4ConfigBehavior .TEXT_EMBEDDINGS_PER_LAYER :
4173+ import torch
4174+
4175+ class PerLayerInputsModule (torch .nn .Module ):
4176+ def __init__ (self , language_model , vocab_size_per_layer_input : int ):
4177+ super ().__init__ ()
4178+ self .language_model = language_model
4179+ self .vocab_size_per_layer_input = vocab_size_per_layer_input
4180+
4181+ # Original code: https://github.com/huggingface/transformers/blob/v5.0.0/src/transformers/models/gemma3n/modeling_gemma3n.py#L2081-L2089
4182+ # Patching is needed to move the preparation of per_layer_inputs outside the language model's forward,
4183+ # as this part requires input_ids to be passed to forward, while for VLM models we pass only input_embeds.
4184+ def forward (self , input_ids : torch .Tensor ):
4185+ per_layer_inputs_mask = torch .logical_and (
4186+ input_ids >= 0 , input_ids < self .vocab_size_per_layer_input
4187+ )
4188+ per_layer_inputs_tokens = torch .where (
4189+ per_layer_inputs_mask , input_ids , torch .zeros_like (input_ids )
4190+ )
4191+ per_layer_inputs = self .language_model .get_per_layer_inputs (per_layer_inputs_tokens )
4192+ return per_layer_inputs
4193+
4194+ model = PerLayerInputsModule (
4195+ self ._get_language_model (model ), model .config .text_config .vocab_size_per_layer_input
4196+ )
4197+ return model
4198+ if behavior == VLMConfigBehavior .TEXT_EMBEDDINGS :
4199+ import torch
4200+
4201+ class TextEmbeddingsModule (torch .nn .Module ):
4202+ def __init__ (self , model ):
4203+ super ().__init__ ()
4204+ self .model = model
4205+
4206+ def forward (self , input_ids : torch .Tensor ):
4207+ inputs_embeds = self .model .get_input_embeddings ()(input_ids )
4208+ vision_mask = torch .logical_and (
4209+ input_ids >= self .model .model .embed_vision .vocab_offset ,
4210+ input_ids < self .model .model .embed_audio .vocab_offset ,
4211+ )
4212+ dummy_vision_token_id = (
4213+ self .model .model .embed_vision .vocab_offset + self .model .model .embed_vision .vocab_size - 1
4214+ )
4215+ vision_input_ids = torch .where (vision_mask , input_ids , dummy_vision_token_id ).to (
4216+ inputs_embeds .device
4217+ )
4218+ vision_embeds = self .model .model .embed_vision (input_ids = vision_input_ids )
4219+ expanded_vision_mask = vision_mask .unsqueeze (- 1 ).expand_as (inputs_embeds )
4220+ inputs_embeds = torch .where (expanded_vision_mask , vision_embeds , inputs_embeds )
4221+
4222+ return inputs_embeds
4223+
4224+ text_embedding = TextEmbeddingsModule (model )
4225+ text_embedding .config = self ._get_language_model (model ).config
4226+ return text_embedding
4227+
4228+ return super ().get_model_for_behavior (model , behavior )
4229+
4230+ def patch_model_for_export (self , model : Union ["PreTrainedModel" ], model_kwargs : Optional [Dict [str , Any ]] = None ):
4231+ model_kwargs = model_kwargs or {}
4232+ if self ._behavior == VLMConfigBehavior .VISION_EMBEDDINGS :
4233+ return Gemma3nImageEmbeddingsModelPatcher (self , model , model_kwargs )
4234+ return super ().patch_model_for_export (model , model_kwargs )
4235+
4236+ @property
4237+ def inputs (self ) -> Dict [str , Dict [int , str ]]:
4238+ if self ._behavior == Gemma4ConfigBehavior .VISION_EMBEDDINGS :
4239+ return {"pixel_values" : {0 : "batch_size" , 2 : "height" , 3 : "width" }}
4240+ return super ().inputs
4241+
4242+
40924243@register_in_tasks_manager ("gemma4_unified" , * ["image-text-to-text" ], library_name = "transformers" )
40934244class Gemma4UnifiedOpenVINOConfig (Gemma3OpenVINOConfig ):
40944245 # gemma4_unified (e.g. google/gemma-4-12B) reuses the gemma3 VLM scaffolding but has an
0 commit comments