diff --git a/packages/tasks/src/model-libraries-snippets.ts b/packages/tasks/src/model-libraries-snippets.ts index c418b6b4c..48097e348 100644 --- a/packages/tasks/src/model-libraries-snippets.ts +++ b/packages/tasks/src/model-libraries-snippets.ts @@ -270,6 +270,12 @@ model = from_pretrained_keras("${model.id}") `, ]; +export const mamba_ssm = (model: ModelData): string[] => [ + `from mamba_ssm import MambaLMHeadModel + +model = MambaLMHeadModel.from_pretrained("${model.id}")`, +]; + export const mars5_tts = (model: ModelData): string[] => [ `# Install from https://github.com/Camb-ai/MARS5-TTS diff --git a/packages/tasks/src/model-libraries.ts b/packages/tasks/src/model-libraries.ts index 97cb1e875..abb97a54e 100644 --- a/packages/tasks/src/model-libraries.ts +++ b/packages/tasks/src/model-libraries.ts @@ -298,6 +298,13 @@ export const MODEL_LIBRARIES_UI_ELEMENTS = { repoName: "mindspore", repoUrl: "https://github.com/mindspore-ai/mindspore", }, + "mamba-ssm": { + prettyLabel: "MambaSSM", + repoName: "MambaSSM", + repoUrl: "https://github.com/state-spaces/mamba", + filter: false, + snippets: snippets.mamba_ssm, + }, "mars5-tts": { prettyLabel: "MARS5-TTS", repoName: "MARS5-TTS",