Update modeling_quasrav4.py

modeling_quasrav4.py

@@ -6,6 +6,7 @@ from typing import Optional, Tuple, List, Union
 
 from transformers import PreTrainedModel
 from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from transformers.generation import GenerationMixin
 from transformers.utils import logging
 
 from .configuration_quasrav4 import InfinityFormerConfig
@@ -213,7 +214,7 @@ class InfinityFormerModel(InfinityFormerPreTrainedModel):
             return (hidden_states,)
         return BaseModelOutputWithPast(last_hidden_state=hidden_states, past_key_values=None, hidden_states=None, attentions=None)
 
-class InfinityFormerForCausalLM(InfinityFormerPreTrainedModel):
+class InfinityFormerForCausalLM(GenerationMixin, InfinityFormerPreTrainedModel):
     _auto_class = "AutoModelForCausalLM"
     
     def __init__(self, config: InfinityFormerConfig):