Add initial gradio setup

app.py

@@ -6,7 +6,12 @@ import gradio as gr
 def generate(prompt, seq_len, temperature, filter_thres, model):
     device = torch.device("cpu")
 
-    model = torch.hub.load("conceptofmind/PaLM", "palm_1b_8k_v0", map_location=device).to(device).eval()
+    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_410m_8k_v0.pt', map_location=device)
+    model.load_state_dict(checkpoint)
 
     tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")
 
@@ -25,7 +30,9 @@ def generate(prompt, seq_len, temperature, filter_thres, model):
 
     decoded_output = tokenizer.batch_decode(output_tensor, skip_special_tokens=True)
 
-    return decoded_output
+    return decoded_output[0]
+
+
 
 iface = gr.Interface(
     fn=generate,