Upload modeling_florence2.py

Florence2/modeling_florence2.py

@@ -29,6 +29,7 @@ from einops import rearrange
 from timm.layers import DropPath, trunc_normal_
 
 from transformers.modeling_utils import PreTrainedModel
+from transformers.generation import GenerationMixin
 from transformers.utils import (
     ModelOutput,
     add_start_docstrings,
@@ -1430,7 +1431,7 @@ class Florence2DecoderLayer(nn.Module):
 
 
 
-class Florence2LanguagePreTrainedModel(PreTrainedModel):
+class Florence2LanguagePreTrainedModel(PreTrainedModel, GenerationMixin ):
     config_class = Florence2LanguageConfig
     base_model_prefix = "model"
     supports_gradient_checkpointing = True
@@ -1462,7 +1463,7 @@ class Florence2LanguagePreTrainedModel(PreTrainedModel):
         return dummy_inputs
 
 
-class Florence2Encoder(Florence2LanguagePreTrainedModel):
+class Florence2Encoder(Florence2LanguagePreTrainedModel, GenerationMixin):
     """
     Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a
     [`Florence2EncoderLayer`].
@@ -1650,7 +1651,7 @@ class Florence2Encoder(Florence2LanguagePreTrainedModel):
         )
 
 
-class Florence2Decoder(Florence2LanguagePreTrainedModel):
+class Florence2Decoder(Florence2LanguagePreTrainedModel, GenerationMixin):
     """
     Transformer decoder consisting of *config.decoder_layers* layers. Each layer is a [`Florence2DecoderLayer`]
 
@@ -1937,7 +1938,7 @@ class Florence2Decoder(Florence2LanguagePreTrainedModel):
         )
 
 
-class Florence2LanguageModel(Florence2LanguagePreTrainedModel):
+class Florence2LanguageModel(Florence2LanguagePreTrainedModel, GenerationMixin):
     _tied_weights_keys = ["encoder.embed_tokens.weight", "decoder.embed_tokens.weight"]
 
     def __init__(self, config: Florence2LanguageConfig):
@@ -2058,7 +2059,6 @@ class Florence2LanguageModel(Florence2LanguagePreTrainedModel):
             encoder_attentions=encoder_outputs.attentions,
         )
 
-from transformers.generation.utils import GenerationMixin
 
 class Florence2LanguageForConditionalGeneration(Florence2LanguagePreTrainedModel, GenerationMixin):
     base_model_prefix = "model"
@@ -2327,7 +2327,7 @@ FLORENCE2_START_DOCSTRING = r"""
     "The bare Florence-2 Model outputting raw hidden-states without any specific head on top.",
     FLORENCE2_START_DOCSTRING,
 )
-class Florence2PreTrainedModel(PreTrainedModel):
+class Florence2PreTrainedModel(PreTrainedModel, GenerationMixin):
     config_class = Florence2Config
     base_model_prefix = "model"
     supports_gradient_checkpointing = True
@@ -2530,7 +2530,7 @@ class Florence2VisionModelWithProjection(Florence2PreTrainedModel):
     """The FLORENCE2 model which consists of a vision backbone and a language model.""",
     FLORENCE2_START_DOCSTRING,
 )
-class Florence2ForConditionalGeneration(Florence2PreTrainedModel):
+class Florence2ForConditionalGeneration(Florence2PreTrainedModel, GenerationMixin):
     def __init__(self, config: Florence2Config):
         super().__init__(config)
         assert config.vision_config.model_type == 'davit', 'only DaViT is supported for now'