Skip to content

Commit c73abf9

Browse files
authored
Add MambaSSM as a library (#802)
This PR adds https://github.com/state-spaces/mamba as an official library on the Hub. Let's wait for Mamba integration to be merged (coming soon) + having a few first models uploaded on the Hub. cc @osanseviero @NielsRogge @tridao Related PR on Mamba side: state-spaces/mamba#471
1 parent 8cba743 commit c73abf9

File tree

2 files changed

+13
-0
lines changed

2 files changed

+13
-0
lines changed

packages/tasks/src/model-libraries-snippets.ts

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -270,6 +270,12 @@ model = from_pretrained_keras("${model.id}")
270270
`,
271271
];
272272

273+
export const mamba_ssm = (model: ModelData): string[] => [
274+
`from mamba_ssm import MambaLMHeadModel
275+
276+
model = MambaLMHeadModel.from_pretrained("${model.id}")`,
277+
];
278+
273279
export const mars5_tts = (model: ModelData): string[] => [
274280
`# Install from https://github.com/Camb-ai/MARS5-TTS
275281

packages/tasks/src/model-libraries.ts

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -298,6 +298,13 @@ export const MODEL_LIBRARIES_UI_ELEMENTS = {
298298
repoName: "mindspore",
299299
repoUrl: "https://github.com/mindspore-ai/mindspore",
300300
},
301+
"mamba-ssm": {
302+
prettyLabel: "MambaSSM",
303+
repoName: "MambaSSM",
304+
repoUrl: "https://github.com/state-spaces/mamba",
305+
filter: false,
306+
snippets: snippets.mamba_ssm,
307+
},
301308
"mars5-tts": {
302309
prettyLabel: "MARS5-TTS",
303310
repoName: "MARS5-TTS",

0 commit comments

Comments
 (0)