Update modeling_motif.py

modeling_motif.py

@@ -1,36 +1,28 @@
 import math
+from dataclasses import dataclass
 from typing import List, Optional, Tuple, Union
 
 import torch
+import torch.nn.functional as F
 import torch.utils.checkpoint
 from torch import nn
 from torch.nn import CrossEntropyLoss
+from transformers.activations import ACT2CLS as _ACT2CLS
+from transformers.activations import ClassInstantier
 from transformers.cache_utils import Cache, DynamicCache, SlidingWindowCache, StaticCache
 from transformers.generation import GenerationMixin
 from transformers.modeling_attn_mask_utils import AttentionMaskConverter
 from transformers.modeling_flash_attention_utils import _flash_attention_forward
-from transformers.modeling_outputs import (
-    CausalLMOutputWithPast,
-    ModelOutput,
-)
+from transformers.modeling_outputs import CausalLMOutputWithPast, ModelOutput
 from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
 from transformers.modeling_utils import PreTrainedModel
 from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS
-from transformers.utils import (
-    add_start_docstrings,
-    add_start_docstrings_to_model_forward,
-    is_flash_attn_greater_or_equal_2_10,
-    is_flash_attn_2_available,
-    logging,
-    replace_return_docstrings,
-)
+from transformers.utils import (add_start_docstrings, add_start_docstrings_to_model_forward, is_flash_attn_2_available,
+                                is_flash_attn_greater_or_equal_2_10, logging, replace_return_docstrings)
+
 from .configuration_motif import MotifConfig
-from dataclasses import dataclass
 
-import torch.nn.functional as F
 
-from transformers.activations import ACT2CLS as _ACT2CLS
-from transformers.activations import ClassInstantier
 class PolyNorm(torch.nn.Module):
     """ 
     A trainable activation function introduced in https://arxiv.org/html/2411.03884v1.