Update app.py

app.py

@@ -13,30 +13,50 @@ tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")
 model = MambaLMHeadModel.from_pretrained("state-spaces/mamba-2.8b", device=device, dtype=torch.float16)
 genlen = 200
 
-def pred(text_in,):
+tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+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() 
     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
-    fn = lambda: model.generate(
+
+    output = 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,
-        temperature=0.5,
-        top_k=10,
-        top_p=0.9,
+        temperature=temperature,
+        top_k=top_k,
+        top_p=top_p,
     )
-    out = fn()
-    text_out = tokenizer.batch_decode(out.sequences.tolist())
+
+    text_out = tokenizer.batch_decode(output.sequences.tolist(), skip_special_tokens=True)
     return text_out[0]
 
+# Define Gradio interface
 demo = gr.Interface(
-    title="Mamba: Selective State Space Model",
-    description="A demo for [Mamba](https://github.com/state-spaces/mamba) by Albert & Tri.",
-    fn=pred, inputs="text", outputs="text")
-    
+    fn=pred,
+    inputs=[
+        gr.inputs.Textbox(label="Input Text"),
+        gr.inputs.Slider(minimum=128, maximum=1024, default=512, label="Model Dimension (d_model)"),
+        gr.inputs.Slider(minimum=1, maximum=24, default=12, label="Number of Layers (n_layer)"),
+        gr.inputs.Number(default=50257, label="Vocabulary Size (vocab_size)"),
+        gr.inputs.Slider(minimum=50, maximum=500, default=200, label="Generation Length (genlen)"),
+        gr.inputs.Slider(minimum=0.1, maximum=1.0, default=0.5, label="Temperature"),
+        gr.inputs.Slider(minimum=1, maximum=50, default=10, label="Top K"),
+        gr.inputs.Slider(minimum=0.1, maximum=1.0, default=0.9, label="Top P"),
+    ],
+    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) . Everytime you send a request it will instantiate the model accordingly, so please be patient."
+)
+
 if __name__ == "__main__":
-    demo.launch() 
+    demo.launch()