Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add MambaSSM as a library #802

Merged
merged 5 commits into from
Jul 16, 2024
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 22 additions & 0 deletions packages/tasks/src/model-libraries-snippets.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
8 changes: 8 additions & 0 deletions packages/tasks/src/model-libraries.ts
Original file line number Diff line number Diff line change
Expand Up @@ -298,6 +298,14 @@ export const MODEL_LIBRARIES_UI_ELEMENTS = {
repoName: "mindspore",
repoUrl: "https://github.com/mindspore-ai/mindspore",
},
mamba_ssm: {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

just mamba would be cleaner, branding-wise?

Suggested change
mamba_ssm: {
mamba: {

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

and it's pip install mamba-ssm on PyPI so maybe a better option too

I think i'd go with mamba personally

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We discussed this with @osanseviero. The problem of simply mamba is that it's a transformers architecture as well so it might create some conflicts. Also having mamba_ssm as library_name and mamba as tag helps with code snippets (to import Mamba class).

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We could have Mamba as prettyName but it might be misleading

Copy link
Member

@julien-c julien-c Jul 16, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ok but then dash rather than underscore no?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nvm let's not complicate stuff, LGTM as is :)

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

good point, addressed in e3e7e62 and state-spaces/mamba@961eccb

prettyLabel: "MambaSSM",
repoName: "MambaSSM",
repoUrl: "https://github.com/state-spaces/mamba",
filter: false,
countDownloads: `path:"config.json"`,
Wauplin marked this conversation as resolved.
Show resolved Hide resolved
snippets: snippets.mamba_ssm,
},
"mars5-tts": {
prettyLabel: "MARS5-TTS",
repoName: "MARS5-TTS",
Expand Down
Loading