Update app.py

app.py

@@ -8,53 +8,46 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 
 from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
 
-device = "cuda"
-tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")
-model = MambaLMHeadModel.from_pretrained("state-spaces/mamba-2.8b", device=device, dtype=torch.float16)
-genlen = 200
-
-tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")
 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="auto")
 
-def pred(text_in, d_model, n_layer, vocab_size, genlen, temperature, top_k, top_p):
-    model = MambaLMHeadModel(d_model=d_model, n_layer=n_layer, vocab_size=vocab_size).to(device)
-    model.eval() 
+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
-
-    output = model.generate(
+    max_length = input_ids.shape[1] + gen_length
+    out = model.generate(
         input_ids=input_ids,
         max_length=max_length,
-        cg=True,
-        do_sample = True,
-        return_dict_in_generate=True,
-        output_scores=True,
-        enable_timing=False,
+        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(output.sequences.tolist(), skip_special_tokens=True)
+    text_out = tokenizer.batch_decode(out.sequences.tolist(), skip_special_tokens=True)
     return text_out[0]
-    pass
+
 demo = gr.Interface(
     fn=pred,
     inputs=[
         gr.Textbox(label="Input Text"),
-        gr.Slider(minimum=128, maximum=1024, value=512, label="Model Dimension (d_model)"),
-        gr.Slider(minimum=1, maximum=24, value=12, label="Number of Layers (n_layer)"),
-        gr.Number(value=50257, label="Vocabulary Size (vocab_size)"),
-        gr.Slider(minimum=50, maximum=500, value=200, label="Generation Length (genlen)"),
         gr.Slider(minimum=0.1, maximum=1.0, value=0.5, label="Temperature"),
-        gr.Slider(minimum=1, maximum=50, value=10, label="Top K"),
+        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",
-    description="With this Demo, you can customize the model hyperparameters of [🐍Mamba](https://www.huggingface.co/state-spaces/mamba-2.8b). Every time you send a request it will instantiate the model accordingly, so please be patient."
+    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 on  Discord: [Discord](https://discord.gg/nXx5wbX9) On 🤗Huggingface: [TeamTonic](https://huggingface.co/TeamTonic) & [MultiTransformer](https://huggingface.co/MultiTransformer) On 🌐Github: [Polytonic](https://github.com/tonic-ai) & contribute to  🌟[PolyGPT](https://github.com/tonic-ai/polygpt-alpha)
+"""
 )
 
 if __name__ == "__main__":