Add GenerationMixin to ConditionalGeneration class

#34
Files changed (1) hide show
  1. modeling_florence2.py +2 -1
modeling_florence2.py CHANGED
@@ -28,6 +28,7 @@ from collections import OrderedDict
28
  from einops import rearrange
29
  from timm.models.layers import DropPath, trunc_normal_
30
 
 
31
  from transformers.modeling_utils import PreTrainedModel
32
  from transformers.utils import (
33
  ModelOutput,
@@ -2529,7 +2530,7 @@ class Florence2VisionModelWithProjection(Florence2PreTrainedModel):
2529
  """The FLORENCE2 model which consists of a vision backbone and a language model.""",
2530
  FLORENCE2_START_DOCSTRING,
2531
  )
2532
- class Florence2ForConditionalGeneration(Florence2PreTrainedModel):
2533
  def __init__(self, config: Florence2Config):
2534
  super().__init__(config)
2535
  assert config.vision_config.model_type == 'davit', 'only DaViT is supported for now'
 
28
  from einops import rearrange
29
  from timm.models.layers import DropPath, trunc_normal_
30
 
31
+ from transformers import GenerationMixin
32
  from transformers.modeling_utils import PreTrainedModel
33
  from transformers.utils import (
34
  ModelOutput,
 
2530
  """The FLORENCE2 model which consists of a vision backbone and a language model.""",
2531
  FLORENCE2_START_DOCSTRING,
2532
  )
2533
+ class Florence2ForConditionalGeneration(Florence2PreTrainedModel, GenerationMixin):
2534
  def __init__(self, config: Florence2Config):
2535
  super().__init__(config)
2536
  assert config.vision_config.model_type == 'davit', 'only DaViT is supported for now'