sync changes from github

modeling_aria.py

@@ -23,7 +23,7 @@ from typing import List, Optional, Tuple, Union
 import torch
 import torch.nn as nn
 from torch import nn
-from transformers import PreTrainedModel
+from transformers import GenerationMixin, PreTrainedModel
 from transformers.modeling_outputs import ModelOutput
 from transformers.utils import logging
 
@@ -122,7 +122,7 @@ def build_mm_projector(config: AriaConfig):
 
 
 # adapted from transformers.models.llava.modeling_llava.LlavaForConditionalGeneration
-class AriaForConditionalGeneration(AriaPretrainedModel):
+class AriaForConditionalGeneration(AriaPretrainedModel, GenerationMixin):
     """
     Aria model for conditional generation tasks.
 

moe_lm.py

@@ -25,7 +25,7 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from torch import nn
-from transformers import LlamaConfig
+from transformers import GenerationMixin, LlamaConfig
 from transformers.models.llama.modeling_llama import (
     ACT2FN,
     LLAMA_ATTENTION_CLASSES,
@@ -634,7 +634,7 @@ class AriaMoELMModel(LlamaModel):
         self.post_init()
 
 
-class AriaMoELMForCausalLM(LlamaForCausalLM):
+class AriaMoELMForCausalLM(LlamaForCausalLM, GenerationMixin):
     """
     AriaMoE model for causal language modeling tasks.