DeciLMForCausalLM(DeciLMPreTrainedModel, GenerationMixin) for v4.50

modeling_decilm.py

@@ -25,7 +25,7 @@ import torch.utils.checkpoint
 from torch import nn
 from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
 from transformers import GenerationConfig
-from transformers.generation.utils import NEED_SETUP_CACHE_CLASSES_MAPPING
+from transformers.generation.utils import GenerationMixin, NEED_SETUP_CACHE_CLASSES_MAPPING
 from transformers.modeling_utils import PreTrainedModel
 from transformers.utils import (
     add_start_docstrings,
@@ -1131,7 +1131,7 @@ class DeciLMModel(DeciLMPreTrainedModel):
         return causal_mask
 
 
-class DeciLMForCausalLM(DeciLMPreTrainedModel):
+class DeciLMForCausalLM(DeciLMPreTrainedModel, GenerationMixin):
     _tied_weights_keys = ["lm_head.weight"]
 
     def __init__(self, config):