Skip to content

Add MambaSSM as a library #802

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.

Sign up for GitHub

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

Merged
merged 5 commits into from
Jul 16, 2024
Merged

Add MambaSSM as a library #802

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
7ef74d1
Add MambaSSM as a library
Wauplin Jul 16, 2024
efaf439
no need for default download count rule
Wauplin Jul 16, 2024
e3e7e62
mamba-ssm
Wauplin Jul 16, 2024
692581f
update snippet
Wauplin Jul 16, 2024
db51e31
lint
Wauplin Jul 16, 2024
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
View file
Open in desktop
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
7 changes: 7 additions & 0 deletions packages/tasks/src/model-libraries.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -298,6 +298,13 @@ 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 •
edited
Loading

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,
snippets: snippets.mamba_ssm,
},
"mars5-tts": {
prettyLabel: "MARS5-TTS",
repoName: "MARS5-TTS",
Expand Down