Add Mamba2

docs/source/en/_toctree.yml

@@ -430,6 +430,8 @@
         title: MADLAD-400
       - local: model_doc/mamba
         title: Mamba
+      - local: model_doc/mamba2
+        title: mamba2
       - local: model_doc/marian
         title: MarianMT
       - local: model_doc/markuplm

docs/source/en/model_doc/mamba2.md

@@ -0,0 +1,60 @@
+<!--Copyright 2024 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
+
+-->
+
+# mamba2
+
+# mamba2
+
+# mamba2
+
+## Overview
+
+The mamba2 model was proposed in [<INSERT PAPER NAME HERE>](<INSERT PAPER LINK HERE>) by <INSERT AUTHORS HERE>.
+<INSERT SHORT SUMMARY HERE>
+
+The abstract from the paper is the following:
+
+*<INSERT PAPER ABSTRACT HERE>*
+
+Tips:
+
+<INSERT TIPS ABOUT MODEL HERE>
+
+This model was contributed by [INSERT YOUR HF USERNAME HERE](https://huggingface.co/<INSERT YOUR HF USERNAME HERE>).
+The original code can be found [here](<INSERT LINK TO GITHUB REPO HERE>).
+
+
+## Mamba2Config
+
+[[autodoc]] Mamba2Config
+
+
+## Mamba2Model
+
+[[autodoc]] Mamba2Model
+    - forward
+
+
+## Mamba2ForCausalLM
+
+[[autodoc]] Mamba2ForCausalLM
+    - forward
+
+
+## Mamba2ForSequenceClassification
+
+[[autodoc]] transformers.Mamba2ForSequenceClassification
+    - forward

src/transformers/__init__.py

@@ -474,6 +474,7 @@
         "InstructBlipVisionConfig",
     ],
     "models.jamba": ["JambaConfig"],
+    "models.mamba2": ["Mamba2Config"],
     "models.jetmoe": ["JetMoeConfig"],
     "models.kosmos2": [
         "Kosmos2Config",
@@ -2311,6 +2312,14 @@
             "JambaPreTrainedModel",
         ]
     )
+    _import_structure["models.mamba2"].extend(
+        [
+            "Mamba2ForCausalLM",
+            "Mamba2ForSequenceClassification",
+            "Mamba2Model",
+            "Mamba2PreTrainedModel",
+        ]
+    )
     _import_structure["models.jetmoe"].extend(
         [
             "JetMoeForCausalLM",
@@ -5056,6 +5065,7 @@
         InstructBlipVisionConfig,
     )
     from .models.jamba import JambaConfig
+    from .models.mamba2 import Mamba2Config
     from .models.jetmoe import JetMoeConfig
     from .models.kosmos2 import (
         Kosmos2Config,
@@ -6738,6 +6748,12 @@
             JambaModel,
             JambaPreTrainedModel,
         )
+        from .models.mamba2 import (
+            Mamba2ForCausalLM,
+            Mamba2ForSequenceClassification,
+            Mamba2Model,
+            Mamba2PreTrainedModel,
+        )
         from .models.jetmoe import (
             JetMoeForCausalLM,
             JetMoeForSequenceClassification,

src/transformers/models/__init__.py

@@ -112,6 +112,7 @@
     informer,
     instructblip,
     jamba,
+    mamba2,
     jetmoe,
     kosmos2,
     layoutlm,

src/transformers/models/auto/configuration_auto.py

@@ -129,6 +129,7 @@
         ("informer", "InformerConfig"),
         ("instructblip", "InstructBlipConfig"),
         ("jamba", "JambaConfig"),
+        ("mamba2", "Mamba2Config"),
         ("jetmoe", "JetMoeConfig"),
         ("jukebox", "JukeboxConfig"),
         ("kosmos-2", "Kosmos2Config"),
@@ -403,6 +404,7 @@
         ("informer", "Informer"),
         ("instructblip", "InstructBLIP"),
         ("jamba", "Jamba"),
+        ("mamba2", "mamba2"),
         ("jetmoe", "JetMoe"),
         ("jukebox", "Jukebox"),
         ("kosmos-2", "KOSMOS-2"),

src/transformers/models/auto/modeling_auto.py

@@ -125,6 +125,7 @@
         ("imagegpt", "ImageGPTModel"),
         ("informer", "InformerModel"),
         ("jamba", "JambaModel"),
+        ("mamba2", "Mamba2Model"),
         ("jetmoe", "JetMoeModel"),
         ("jukebox", "JukeboxModel"),
         ("kosmos-2", "Kosmos2Model"),
@@ -461,6 +462,7 @@
         ("gpt_neox_japanese", "GPTNeoXJapaneseForCausalLM"),
         ("gptj", "GPTJForCausalLM"),
         ("jamba", "JambaForCausalLM"),
+        ("mamba2", "Mamba2ForCausalLM"),
         ("jetmoe", "JetMoeForCausalLM"),
         ("llama", "LlamaForCausalLM"),
         ("mamba", "MambaForCausalLM"),
@@ -866,6 +868,7 @@
         ("gptj", "GPTJForSequenceClassification"),
         ("ibert", "IBertForSequenceClassification"),
         ("jamba", "JambaForSequenceClassification"),
+        ("mamba2", "Mamba2ForSequenceClassification"),
         ("jetmoe", "JetMoeForSequenceClassification"),
         ("layoutlm", "LayoutLMForSequenceClassification"),
         ("layoutlmv2", "LayoutLMv2ForSequenceClassification"),

src/transformers/models/auto/tokenization_auto.py

@@ -212,6 +212,12 @@
                     "LlamaTokenizerFast" if is_tokenizers_available() else None,
                 ),
             ),
+            (
+                "mamba2",
+                (
+                    "LlamaTokenizer" if is_sentencepiece_available() else None,
+                    "LlamaTokenizerFast" if is_tokenizers_available() else None,
+                ),
             (
                 "jetmoe",
                 (

src/transformers/models/mamba2/__init__.py

@@ -0,0 +1,58 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import TYPE_CHECKING
+
+from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available
+
+
+_import_structure = {
+    "configuration_mamba2": ["Mamba2Config"],
+}
+
+
+try:
+    if not is_torch_available():
+        raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+    pass
+else:
+    _import_structure["modeling_mamba2"] = [
+        "Mamba2ForCausalLM",
+        "Mamba2ForSequenceClassification",
+        "Mamba2Model",
+        "Mamba2PreTrainedModel",
+    ]
+
+
+if TYPE_CHECKING:
+    from .configuration_mamba2 import Mamba2Config
+
+    try:
+        if not is_torch_available():
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        pass
+    else:
+        from .modeling_mamba2 import (
+            Mamba2ForCausalLM,
+            Mamba2ForSequenceClassification,
+            Mamba2Model,
+            Mamba2PreTrainedModel,
+        )
+
+
+else:
+    import sys
+
+    sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)