From 18c21287ed1b1f003dd60a678afc4b186c48df4a Mon Sep 17 00:00:00 2001 From: ai-modelscope Date: Tue, 5 Nov 2024 06:05:36 +0800 Subject: [PATCH] Inherit from GenerationMixin (#80) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add GenerationMixin as parent class (c719442dd7d237edafe787824ebe245dcfa997d4) Co-authored-by: Marvin Gazibarić --- modeling_florence2.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/modeling_florence2.py b/modeling_florence2.py index e5ee651..ccca154 100644 --- a/modeling_florence2.py +++ b/modeling_florence2.py @@ -29,6 +29,7 @@ from einops import rearrange from timm.models.layers import DropPath, trunc_normal_ from transformers.modeling_utils import PreTrainedModel +from transformers.generation.utils import GenerationMixin from transformers.utils import ( ModelOutput, add_start_docstrings, @@ -2059,7 +2060,7 @@ class Florence2LanguageModel(Florence2LanguagePreTrainedModel): ) -class Florence2LanguageForConditionalGeneration(Florence2LanguagePreTrainedModel): +class Florence2LanguageForConditionalGeneration(Florence2LanguagePreTrainedModel, GenerationMixin): base_model_prefix = "model" _tied_weights_keys = ["encoder.embed_tokens.weight", "decoder.embed_tokens.weight", "lm_head.weight"] _keys_to_ignore_on_load_missing = ["final_logits_bias"]