Use 1b

app.py

@@ -6,11 +6,22 @@ import gradio as gr
 def generate(prompt, seq_len=128, temperature=0.8, filter_thres=0.9):
     device = torch.device("cpu")
 
+    num_tokens = 50304
+    dim = 2048
+    depth = 16
+    dim_head = 128
+    heads = 8
+    flash_attn = True
+
     model = PaLM(
-        num_tokens=50304, dim=1024, depth=24, dim_head=128, heads=8, flash_attn=False, qk_rmsnorm = False,
+        num_tokens=num_tokens, dim=dim, depth=depth, dim_head=dim_head, heads=heads, flash_attn=flash_attn
     ).to(device).eval()
 
-    checkpoint = torch.load('./palm_410m_8k_v0.pt', map_location=device)
+    # model = PaLM(
+    #     num_tokens=50304, dim=1024, depth=24, dim_head=128, heads=8, flash_attn=False, qk_rmsnorm = False,
+    # ).to(device).eval()
+
+    checkpoint = torch.load('./palm_1b_8k_v0.pt', map_location=device)
     model.load_state_dict(checkpoint)
 
     tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")