Enrico Shippole commited on
Commit
3c50c94
1 Parent(s): 8253e02
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -21,7 +21,7 @@ def generate(prompt, seq_len=128, temperature=0.8, filter_thres=0.9):
21
  num_tokens=50304, dim=1024, depth=24, dim_head=128, heads=8, flash_attn=False, qk_rmsnorm = False,
22
  ).to(device).eval()
23
 
24
- checkpoint = torch.load('./palm_1b_8k_v0.pt', map_location=device)
25
  model.load_state_dict(checkpoint)
26
 
27
  tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")
 
21
  num_tokens=50304, dim=1024, depth=24, dim_head=128, heads=8, flash_attn=False, qk_rmsnorm = False,
22
  ).to(device).eval()
23
 
24
+ checkpoint = torch.load('./palm_410m_8k_v0.pt', map_location=device)
25
  model.load_state_dict(checkpoint)
26
 
27
  tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")