reach-vb HF staff commited on
Commit
9a94062
1 Parent(s): 58aa497

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -3
app.py CHANGED
@@ -13,7 +13,7 @@ tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")
13
  model = MambaLMHeadModel.from_pretrained("state-spaces/mamba-2.8b", device=device, dtype=torch.float16)
14
  genlen = 200
15
 
16
- def 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)
@@ -26,7 +26,7 @@ def pred(text_in):
26
  output_scores=True,
27
  enable_timing=False,
28
  temperature=0.5,
29
- top_k=1,
30
  top_p=0.9,
31
  )
32
  out = fn()
@@ -35,7 +35,7 @@ def pred(text_in):
35
 
36
  demo = gr.Interface(
37
  title="Mamba: Selective State Space Model",
38
- description="A demo for [Mamba](https://github.com/state-spaces/mamba)",
39
  fn=pred, inputs="text", outputs="text")
40
 
41
  if __name__ == "__main__":
 
13
  model = MambaLMHeadModel.from_pretrained("state-spaces/mamba-2.8b", device=device, dtype=torch.float16)
14
  genlen = 200
15
 
16
+ def 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)
 
26
  output_scores=True,
27
  enable_timing=False,
28
  temperature=0.5,
29
+ top_k=10,
30
  top_p=0.9,
31
  )
32
  out = fn()
 
35
 
36
  demo = gr.Interface(
37
  title="Mamba: Selective State Space Model",
38
+ description="A demo for [Mamba](https://github.com/state-spaces/mamba) by Albert & Tri.",
39
  fn=pred, inputs="text", outputs="text")
40
 
41
  if __name__ == "__main__":