File size: 1,597 Bytes
5f5633f 038d9a4 5f5633f b6c0621 5f5633f |
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 |
import sys, os
os.system("!git clone https://github.com/karpathy/minGPT/")
sys.path.append('minGPT')
from transformers import GPT2Tokenizer, GPT2LMHeadModel
from mingpt.model import GPT
from minGPT.mingpt.utils import set_seed
use_mingpt = True # use minGPT or huggingface/transformers model?
model_type = 'gpt2-xl'
device = 'cuda'
if use_mingpt:
model = GPT.from_pretrained(model_type)
else:
model = GPT2LMHeadModel.from_pretrained(model_type)
model.config.pad_token_id = model.config.eos_token_id # suppress a warning
# ship model to device and set to eval mode
model.to(device)
model.eval();
def generate(prompt='', num_samples=10, steps=20, do_sample=True):
# tokenize the input prompt into integer input sequence
tokenizer = GPT2Tokenizer.from_pretrained(model_type)
if prompt == '': # to create unconditional samples we feed in the special start token
prompt = '<|endoftext|>'
encoded_input = tokenizer(prompt, return_tensors='pt').to(device)
x = encoded_input['input_ids']
x = x.expand(num_samples, -1)
# forward the model `steps` times to get samples, in a batch
y = model.generate(x, max_new_tokens=steps, do_sample=do_sample, top_k=40)
for i in range(num_samples):
out = tokenizer.decode(y[i].cpu().squeeze())
print('-'*80)
print(out)
def infer(input):
return generate(prompt=input, num_samples=3, steps=20)
import gradio as gr
gr.Interface(infer, "text", ["text", "text", "text"], examples=["I was commuting to my Silicon Valley job when I took a wrong turn. I"]) |