Spaces:
Runtime error
Runtime error
File size: 1,761 Bytes
48f8345 82872e3 48f8345 6ee4952 8253e02 6ee4952 8253e02 3c50c94 17d1f44 48f8345 9224f39 48f8345 738800c 2988fa3 82872e3 738800c 48f8345 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 |
import torch
from transformers import AutoTokenizer
from palm_rlhf_pytorch import PaLM
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=num_tokens, dim=dim, depth=depth, dim_head=dim_head, heads=heads, flash_attn=flash_attn
# ).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")
encoded_text = tokenizer(prompt, return_tensors="pt")
output_tensor = model.generate(
seq_len=seq_len,
prompt=encoded_text["input_ids"].to(device),
temperature=temperature,
filter_thres=filter_thres,
pad_value=0.0,
eos_token=tokenizer.eos_token_id,
return_seq_without_prompt=False,
use_tqdm=True,
)
decoded_output = tokenizer.batch_decode(output_tensor, skip_special_tokens=True)
return decoded_output
iface = gr.Interface(
fn=generate,
title="PaLM",
description="Open-source PaLM demo.",
inputs="text",
outputs="text",
# seq_len=gr.Slider(minimum=1, maximum=8192, step=1, default=32, label="Sequence Length"),
# temperature=gr.Slider(minimum=0.0, maximum=1.0, step=0.01, default=0.8, label="Temperature"),
# filter_thres=gr.Slider(minimum=0.0, maximum=1.0, step=0.01, default=0.9, label="Filter Threshold"),
)
iface.launch() |