|
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 |
|
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 |
|
|
|
|
|
model.to(device) |
|
model.eval(); |
|
|
|
|
|
def generate(prompt='', num_samples=10, steps=20, do_sample=True): |
|
|
|
|
|
tokenizer = GPT2Tokenizer.from_pretrained(model_type) |
|
if prompt == '': |
|
prompt = '<|endoftext|>' |
|
encoded_input = tokenizer(prompt, return_tensors='pt').to(device) |
|
x = encoded_input['input_ids'] |
|
x = x.expand(num_samples, -1) |
|
|
|
|
|
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"]) |