app.py

import numpy as np
import gradio as gr
from transformers import GPT2Tokenizer, TrainingArguments, Trainer, GPT2LMHeadModel

tokenizer = GPT2Tokenizer.from_pretrained('mindwrapped/gpt2-lotr-fellowship', bos_token='<|startoftext|>',
                                          eos_token='<|endoftext|>', pad_token='<|pad|>')
model = GPT2LMHeadModel.from_pretrained('mindwrapped/gpt2-lotr-fellowship').cuda()


def generate_text(text, temperature):
  generated = tokenizer("<|startoftext|> " + text, return_tensors="pt").input_ids.cuda()
  sample_outputs = model.generate(generated, do_sample=True, top_k=50, 
                                  max_length=300, top_p=0.95, temperature=float(temperature), num_return_sequences=1)

  for i, sample_output in enumerate(sample_outputs):
    out = tokenizer.decode(sample_output, skip_special_tokens=True)

  return out


demo = gr.Interface(
  fn=generate_text, 
  inputs=[gr.Text(),gr.Slider(minimum=0.0, maximum=5.0, value=1.0, step=0.1)],
  outputs='text',
  examples=[['', 1.9],['Frodo and Sam moved quietly through the night.', 2.0], ['Frodo and Sam went to the pub. ', 3.0]],
  title='LOTR Generator',
  description='This space uses GPT2 model fine-tuned on the "The Fellowship of the Rings" to generate text. Try inputting no text to the model and messing around with the temperature.',
  article=' ![visitor badge](https://visitor-badge.glitch.me/badge?page_id=mindwrapped.gpt2-lotr-fellowship-generator-space)',
  live=False,
  )

demo.launch(debug=True)