chore: import upstream snapshot with attribution

2026-07-13 13:34:58 +08:00
commit a203934033
1368 changed files with 175001 additions and 0 deletions
@@ -0,0 +1,40 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from transformers import PreTrainedModel
+
+from swift.template import TemplateType
+from swift.utils import get_logger
+from ..constant import LLMModelType
+from ..model_meta import Model, ModelGroup, ModelMeta
+from ..register import ModelLoader, register_model
+
+logger = get_logger()
+
+
+class MambaLoader(ModelLoader):
+
+    def get_model(self, model_dir: str, *args, **kwargs) -> PreTrainedModel:
+        logger.info(
+            '[IMPORTANT] Remember installing causal-conv1d>=1.2.0 and mamba-ssm, or you training and inference will'
+            'be really slow!')
+        return super().get_model(model_dir, *args, **kwargs)
+
+
+register_model(
+    ModelMeta(
+        LLMModelType.mamba,
+        [
+            ModelGroup([
+                Model('AI-ModelScope/mamba-130m-hf', 'state-spaces/mamba-130m-hf'),
+                Model('AI-ModelScope/mamba-370m-hf', 'state-spaces/mamba-370m-hf'),
+                Model('AI-ModelScope/mamba-390m-hf', 'state-spaces/mamba-390m-hf'),
+                Model('AI-ModelScope/mamba-790m-hf', 'state-spaces/mamba-790m-hf'),
+                Model('AI-ModelScope/mamba-1.4b-hf', 'state-spaces/mamba-1.4b-hf'),
+                Model('AI-ModelScope/mamba-2.8b-hf', 'state-spaces/mamba-2.8b-hf'),
+            ])
+        ],
+        MambaLoader,
+        template=TemplateType.default,
+        architectures=['MambaForCausalLM'],
+        model_arch=None,
+        requires=['transformers>=4.39.0'],
+    ))