Spaces:
Running
Running
import torch | |
from transformers import AutoTokenizer | |
from palm_rlhf_pytorch import PaLM | |
import gradio as gr | |
def generate(prompt, seq_len, temperature, filter_thres): | |
device = torch.device("cpu") | |
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" | |
[ | |
gr.Slider(minimum=1, maximum=512, step=1, default=128, label="Sequence Length"), | |
gr.Slider(minimum=0.0, maximum=1.0, step=0.01, default=0.9, label="Temperature"), | |
gr.Slider(minimum=0.0, maximum=1.0, step=0.01, default=0.9, label="Filter Threshold"), | |
], | |
) | |
iface.launch() |