Add GenerationMixin to LIMEForCausalLM

modeling_lime.py

@@ -1,6 +1,7 @@
 import torch
 from torch import nn
 from transformers import PreTrainedModel
+from transformers.generation import GenerationMixin
 from transformers.modeling_outputs import CausalLMOutputWithPast
 from typing import Optional, Tuple, Union
 from ukraine.research.transformer.transformer import Transformer
@@ -21,7 +22,7 @@ def make_norm(config: LIMEConfig):
     return nn.RMSNorm(config.d_model)
 
 
-class LIMEForCausalLM(PreTrainedModel):
+class LIMEForCausalLM(PreTrainedModel, GenerationMixin):
     config_class = LIMEConfig
     base_model_prefix = "lime"
     _tied_weights_keys = ["transformer.output_fc.weight"]