Spaces:
Runtime error
Runtime error
| import torch | |
| import torch.nn.functional as F | |
| import einops | |
| from einops import rearrange | |
| import gradio as gr | |
| from transformers import AutoTokenizer, AutoModelForCausalLM | |
| from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b") | |
| model = MambaLMHeadModel.from_pretrained("state-spaces/mamba-2.8b", device=device, dtype=torch.float16) | |
| def pred(text_in, temperature, top_k, top_p, gen_length, cg, return_dict_in_generate, output_scores, enable_timing): | |
| tokens = tokenizer(text_in, return_tensors="pt") | |
| input_ids = tokens.input_ids.to(device=device) | |
| attn_mask = tokens.attention_mask.to(device=device) | |
| max_length = input_ids.shape[1] + genlen | |
| out = model.generate( | |
| input_ids=input_ids, | |
| max_length=max_length, | |
| cg=cg, | |
| return_dict_in_generate=return_dict_in_generate, | |
| output_scores=output_scores, | |
| enable_timing=enable_timing, | |
| temperature=temperature, | |
| top_k=top_k, | |
| top_p=top_p, | |
| ) | |
| text_out = tokenizer.batch_decode(out.sequences.tolist(), skip_special_tokens=True) | |
| return text_out[0] | |
| demo = gr.Interface( | |
| fn=pred, | |
| inputs=[ | |
| gr.Textbox(label="Input Text"), | |
| gr.Slider(minimum=0.1, maximum=1.0, value=0.5, label="Temperature"), | |
| gr.Slider(minimum=1, maximum=10, value=10, label="Top K"), | |
| gr.Slider(minimum=0.1, maximum=1.0, value=0.9, label="Top P"), | |
| gr.Slider(minimum=50, maximum=650, value=200, label="Generation Length (gen_length)"), | |
| gr.Checkbox(value=True, label="Cache Graph (cg)"), | |
| gr.Checkbox(value=True, label="Return Dict in Generate"), | |
| gr.Checkbox(value=True, label="Output Scores"), | |
| gr.Checkbox(value=False, label="Enable Timing"), | |
| ], | |
| outputs="text", | |
| title="Welcome👋🏻to🌟Tonic's🐍Mamba 2.8B! 🚀", | |
| description="""🐍Mamba is quite special because it uses a unique model architecture, has reasonable🏆performance, and a👌🏻tiny size. You can use this Space to test out the current model 🐍[state-spaces/mamba-2.8b](https://huggingface.co/state-spaces/mamba-2.8b) You can also use 🐍mamba-2.8b by cloning this space. Simply click here: [Duplicate Space](https://huggingface.co/spaces/Tonic1/VLChat?duplicate=true) | |
| Join us : 🌟TeamTonic🌟 is always making cool demos! Join our active builder's🛠️community 👻 [](https://discord.gg/GWpVpekp) On 🤗Huggingface: [TeamTonic](https://huggingface.co/TeamTonic) & [MultiTransformer](https://huggingface.co/MultiTransformer) On 🌐Github: [Tonic-AI](https://github.com/tonic-ai) & contribute to 🌟 [DataTonic](https://github.com/Tonic-AI/DataTonic) 🤗Big thanks to Yuvi Sharma and all the folks at huggingface for the community grant 🤗 | |
| """ | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() |