Transformers compatible Mamba
Collection
This release includes the `mamba` repositories compatible with the `transformers` library.
•
3 items
•
Updated
The following content is mostly from https://huggingface.co./state-spaces/mamba-2.8b-hf
This repository contains the transfromers
compatible mamba-2.8b-slimpj
. The checkpoints are untouched, but the full config.json
and tokenizer are pushed to this repo.
You need to install transformers
from main
until transformers=4.39.0
is released.
pip install git+https://github.com/huggingface/transformers@main
We also recommend you to install both causal_conv_1d
and mamba-ssm
using:
pip install causal-conv1d>=1.2.0
pip install mamba-ssm
If any of these two is not installed, the "eager" implementation will be used. Otherwise the more optimised cuda
kernels will be used.
You can use the classic generate
API:
>>> from transformers import MambaConfig, MambaForCausalLM, AutoTokenizer
>>> import torch
>>> tokenizer = AutoTokenizer.from_pretrained("han1997/mamba-2.8b-slimpj-hf")
>>> model = MambaForCausalLM.from_pretrained("han1997/mamba-2.8b-slimpj-hf")
>>> input_ids = tokenizer("Hey how are you doing?", return_tensors="pt")["input_ids"]
>>> out = model.generate(input_ids, max_new_tokens=10)
>>> print(tokenizer.batch_decode(out))
["Hey how are you doing?\n\nI'm doing great.\n\nI"]
In order to finetune using the peft
library, we recommend keeping the model in float32!
from datasets import load_dataset
from trl import SFTTrainer
from peft import LoraConfig
from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments
tokenizer = AutoTokenizer.from_pretrained("han1997/mamba-2.8b-slimpj-hf")
model = AutoModelForCausalLM.from_pretrained("han1997/mamba-2.8b-slimpj-hf")
dataset = load_dataset("Abirate/english_quotes", split="train")
training_args = TrainingArguments(
output_dir="./results",
num_train_epochs=3,
per_device_train_batch_size=4,
logging_dir='./logs',
logging_steps=10,
learning_rate=2e-3
)
lora_config = LoraConfig(
r=8,
target_modules=["x_proj", "embeddings", "in_proj", "out_proj"],
task_type="CAUSAL_LM",
bias="none"
)
trainer = SFTTrainer(
model=model,
tokenizer=tokenizer,
args=training_args,
peft_config=lora_config,
train_dataset=dataset,
dataset_text_field="quote",
)
trainer.train()