Fix missing generate method by inheriting from GenerationMixin

#2
Files changed (1) hide show
  1. modeling_florence2.py +2 -1
modeling_florence2.py CHANGED
@@ -30,6 +30,7 @@ from timm.models.layers import DropPath, trunc_normal_
30
  import torch.nn.init as init
31
 
32
  from transformers.modeling_utils import PreTrainedModel
 
33
  from transformers.utils import (
34
  ModelOutput,
35
  add_start_docstrings,
@@ -2064,7 +2065,7 @@ class Florence2LanguageModel(Florence2LanguagePreTrainedModel):
2064
  )
2065
 
2066
 
2067
- class Florence2LanguageForConditionalGeneration(Florence2LanguagePreTrainedModel):
2068
  base_model_prefix = "model"
2069
  _tied_weights_keys = ["encoder.embed_tokens.weight", "decoder.embed_tokens.weight", "lm_head.weight"]
2070
  _keys_to_ignore_on_load_missing = ["final_logits_bias"]
 
30
  import torch.nn.init as init
31
 
32
  from transformers.modeling_utils import PreTrainedModel
33
+ from transformers.generation.utils import GenerationMixin
34
  from transformers.utils import (
35
  ModelOutput,
36
  add_start_docstrings,
 
2065
  )
2066
 
2067
 
2068
+ class Florence2LanguageForConditionalGeneration(Florence2LanguagePreTrainedModel, GenerationMixin):
2069
  base_model_prefix = "model"
2070
  _tied_weights_keys = ["encoder.embed_tokens.weight", "decoder.embed_tokens.weight", "lm_head.weight"]
2071
  _keys_to_ignore_on_load_missing = ["final_logits_bias"]