Update modeling_mmMamba_embedding.py

modeling_mmMamba_embedding.py

@@ -14,52 +14,30 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import math
-import queue
-import threading
-import warnings
-from typing import List, Optional, Tuple, Union
-from functools import partial
+from typing import Optional, Tuple
 
 import torch
 import torch.nn.functional as F
 import torch.utils.checkpoint
 from einops import rearrange
 from torch import nn
-from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
 from transformers.activations import ACT2FN
-from transformers.modeling_outputs import (
-    BaseModelOutputWithPast,
-    CausalLMOutputWithPast,
-    SequenceClassifierOutputWithPast,
-)
+
 from transformers.modeling_utils import PreTrainedModel
-from transformers.cache_utils import Cache
-from transformers.utils import (
-    add_start_docstrings,
-    add_start_docstrings_to_model_forward,
-    logging,
-    replace_return_docstrings,
-)
-from fla.modules import FusedRMSNormSwishGate, RMSNorm, ShortConvolution
-import copy
+from transformers.utils import logging
+
+from fused_norm_gate import FusedRMSNormSwishGate
+
 from mamba_ssm.ops.triton.ssd_combined import mamba_chunk_scan_combined
 from mamba_ssm.ops.triton.selective_state_update import selective_state_update
 from causal_conv1d import causal_conv1d_fn, causal_conv1d_update
-from transformers.cache_utils import Cache
-import time
+
 from timm.models.layers import DropPath
 
 compute_ARank = False # [ARank] Set this to True to compute attention rank
 
-try:
-    from transformers.generation.streamers import BaseStreamer
-except:  # noqa # pylint: disable=bare-except
-    BaseStreamer = None
-
 from .configuration_mmMamba_embedding import mmMambaEmbeddingConfig
 
-import time
-
 from .configuration_mmMamba import mmMambaConfig
 
 try: