From 7ef74d1417e63aadbaae1d4f345093c4b5311df0 Mon Sep 17 00:00:00 2001 From: Wauplin Date: Tue, 16 Jul 2024 10:30:07 +0200 Subject: [PATCH 1/5] Add MambaSSM as a library --- .../tasks/src/model-libraries-snippets.ts | 22 +++++++++++++++++++ packages/tasks/src/model-libraries.ts | 8 +++++++ 2 files changed, 30 insertions(+) diff --git a/packages/tasks/src/model-libraries-snippets.ts b/packages/tasks/src/model-libraries-snippets.ts index c418b6b4cc..d7d14eef68 100644 --- a/packages/tasks/src/model-libraries-snippets.ts +++ b/packages/tasks/src/model-libraries-snippets.ts @@ -270,6 +270,28 @@ model = from_pretrained_keras("${model.id}") `, ]; +export const mamba_ssm = (model: ModelData): string[] => { + if (model.tags.includes("mamba2")) { + return [ + `from mamba_ssm import Mamba2 + +model = Mamba2.from_pretrained("${model.id}")`, + ]; + } else if (model.tags.includes("mamba2simple")) { + return [ + `from mamba_ssm.modules.mamba2_simple import Mamba2Simple + +model = Mamba2Simple.from_pretrained("${model.id}")`, + ]; + } else { + return [ + `from mamba_ssm import Mamba + +model = Mamba.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 97cb1e8753..3942269b92 100644 --- a/packages/tasks/src/model-libraries.ts +++ b/packages/tasks/src/model-libraries.ts @@ -298,6 +298,14 @@ 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, + countDownloads: `path:"config.json"`, + snippets: snippets.mamba_ssm, + }, "mars5-tts": { prettyLabel: "MARS5-TTS", repoName: "MARS5-TTS", From efaf43986ac117fb048848d30ea705f706e60aac Mon Sep 17 00:00:00 2001 From: Wauplin Date: Tue, 16 Jul 2024 11:02:10 +0200 Subject: [PATCH 2/5] no need for default download count rule --- packages/tasks/src/model-libraries.ts | 1 - 1 file changed, 1 deletion(-) diff --git a/packages/tasks/src/model-libraries.ts b/packages/tasks/src/model-libraries.ts index 3942269b92..1f7e72fdc5 100644 --- a/packages/tasks/src/model-libraries.ts +++ b/packages/tasks/src/model-libraries.ts @@ -303,7 +303,6 @@ export const MODEL_LIBRARIES_UI_ELEMENTS = { repoName: "MambaSSM", repoUrl: "https://github.com/state-spaces/mamba", filter: false, - countDownloads: `path:"config.json"`, snippets: snippets.mamba_ssm, }, "mars5-tts": { From e3e7e62c2bd3b0309e1e3a9a203abc430e17c98c Mon Sep 17 00:00:00 2001 From: Wauplin Date: Tue, 16 Jul 2024 13:37:54 +0200 Subject: [PATCH 3/5] mamba-ssm --- packages/tasks/src/model-libraries.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/tasks/src/model-libraries.ts b/packages/tasks/src/model-libraries.ts index 1f7e72fdc5..abb97a54e9 100644 --- a/packages/tasks/src/model-libraries.ts +++ b/packages/tasks/src/model-libraries.ts @@ -298,7 +298,7 @@ export const MODEL_LIBRARIES_UI_ELEMENTS = { repoName: "mindspore", repoUrl: "https://github.com/mindspore-ai/mindspore", }, - mamba_ssm: { + "mamba-ssm": { prettyLabel: "MambaSSM", repoName: "MambaSSM", repoUrl: "https://github.com/state-spaces/mamba", From 692581fcd53014ef9c12e527848ede7848814eb3 Mon Sep 17 00:00:00 2001 From: Wauplin Date: Tue, 16 Jul 2024 16:45:19 +0200 Subject: [PATCH 4/5] update snippet --- .../tasks/src/model-libraries-snippets.ts | 25 ++++--------------- 1 file changed, 5 insertions(+), 20 deletions(-) diff --git a/packages/tasks/src/model-libraries-snippets.ts b/packages/tasks/src/model-libraries-snippets.ts index d7d14eef68..902b36b05a 100644 --- a/packages/tasks/src/model-libraries-snippets.ts +++ b/packages/tasks/src/model-libraries-snippets.ts @@ -270,27 +270,12 @@ model = from_pretrained_keras("${model.id}") `, ]; -export const mamba_ssm = (model: ModelData): string[] => { - if (model.tags.includes("mamba2")) { - return [ - `from mamba_ssm import Mamba2 - -model = Mamba2.from_pretrained("${model.id}")`, - ]; - } else if (model.tags.includes("mamba2simple")) { - return [ - `from mamba_ssm.modules.mamba2_simple import Mamba2Simple +export const mamba_ssm = (model: ModelData): string[] => [ + `from mamba_ssm import MambaLMHeadModel -model = Mamba2Simple.from_pretrained("${model.id}")`, - ]; - } else { - return [ - `from mamba_ssm import Mamba - -model = Mamba.from_pretrained("${model.id}")`, - ]; - } -}; +model = MambaLMHeadModel.from_pretrained("${model.id}") +`, +]; export const mars5_tts = (model: ModelData): string[] => [ `# Install from https://github.com/Camb-ai/MARS5-TTS From db51e315039da05a59b0d0425effc531471eab32 Mon Sep 17 00:00:00 2001 From: Wauplin Date: Tue, 16 Jul 2024 16:46:22 +0200 Subject: [PATCH 5/5] lint --- packages/tasks/src/model-libraries-snippets.ts | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/packages/tasks/src/model-libraries-snippets.ts b/packages/tasks/src/model-libraries-snippets.ts index 902b36b05a..48097e3481 100644 --- a/packages/tasks/src/model-libraries-snippets.ts +++ b/packages/tasks/src/model-libraries-snippets.ts @@ -273,8 +273,7 @@ model = from_pretrained_keras("${model.id}") export const mamba_ssm = (model: ModelData): string[] => [ `from mamba_ssm import MambaLMHeadModel -model = MambaLMHeadModel.from_pretrained("${model.id}") -`, +model = MambaLMHeadModel.from_pretrained("${model.id}")`, ]; export const mars5_tts = (model: ModelData): string[] => [