Update app.py

app.py

@@ -1,21 +1,20 @@
 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="auto")
+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)
-    max_length = input_ids.shape[1] + gen_length
+    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,