reach-vb/mamba

1import torch2import torch.nn.functional as F3 4from einops import rearrange5import gradio as gr6 7from transformers import AutoTokenizer, AutoModelForCausalLM8 9from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel10 11device = "cuda"12tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")13model = MambaLMHeadModel.from_pretrained("state-spaces/mamba-2.8b-slimpj", device=device, dtype=torch.float16)14genlen = 50015 16def pred(text_in,):17    tokens = tokenizer(text_in, return_tensors="pt")18    input_ids = tokens.input_ids.to(device=device)19    attn_mask = tokens.attention_mask.to(device=device)20    max_length = input_ids.shape[1] + genlen21    fn = lambda: model.generate(22        input_ids=input_ids,23        max_length=max_length,24        cg=True,25        return_dict_in_generate=True,26        output_scores=True,27        enable_timing=False,28        temperature=0.9,29        top_p=0.7,30    )31    out = fn()32    text_out = tokenizer.batch_decode(out.sequences.tolist())33    return text_out[0]34 35demo = gr.Interface(36    title="Mamba: Selective State Space Model",37    description="A demo for [Mamba](https://github.com/state-spaces/mamba) by Albert & Tri.",38    fn=pred, inputs="text", outputs="text")39    40if __name__ == "__main__":41    demo.launch()