Spaces:
Runtime error
Runtime error
import os | |
import pickle | |
import gradio as gr | |
import torch | |
from model import GPT, GPTConfig | |
ckpt_path = 'model/ckpt.pt' | |
meta_path = 'model/meta.pkl' | |
seed = 1337 | |
device = 'cpu' | |
torch.manual_seed(seed) | |
# Load the model and meta data | |
checkpoint = torch.load(ckpt_path, map_location=device) | |
gptconf = GPTConfig(**checkpoint['model_args']) | |
model = GPT(gptconf) | |
state_dict = checkpoint['model'] | |
unwanted_prefix = '_orig_mod.' | |
for k, v in list(state_dict.items()): | |
if k.startswith(unwanted_prefix): | |
state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k) | |
model.load_state_dict(state_dict) | |
model.eval() | |
model.to(device) | |
with open(meta_path, 'rb') as f: | |
meta = pickle.load(f) | |
stoi, itos = meta['stoi'], meta['itos'] | |
encode = lambda s: [stoi[c] for c in s] | |
decode = lambda l: ''.join([itos[i] for i in l]) | |
# Define the function for generating text | |
def generate_text(start, temperature, max_new_tokens): | |
start_ids = encode(start) | |
x = (torch.tensor(start_ids, dtype=torch.long, device=device)[None, ...]) | |
# Generate text | |
with torch.no_grad(): | |
y = model.generate(x, max_new_tokens, temperature=temperature) | |
generated_text = decode(y[0].tolist()) | |
return generated_text | |
# Create a Gradio interface with sliders | |
examples = [['sport', 0.7, 200], ['lord', 1.2, 300]] | |
iface = gr.Interface( | |
fn=generate_text, | |
inputs=[ | |
gr.Textbox(label="Starting Prompt"), | |
gr.Slider(minimum=0.1, maximum=4, step=0.1, label="Temperature"), | |
gr.Slider(minimum=100, maximum=1000, step=50, label="Max New Tokens"), | |
], | |
outputs=gr.Textbox(label="Generated Text"), | |
examples = examples | |
) | |
iface.launch() |