irfantea's picture
Update app.py
1089abe
raw
history blame
1.02 kB
import torch
device = "cuda:0" if torch.cuda.is_available() else "cpu"
from transformers import GPT2LMHeadModel, GPT2Tokenizer
import streamlit as st
st.set_page_config(
page_title="GPT-2 Demo",
page_icon=":robot_face:",
layout="wide")
st.title("GPT-2 Text Generation Demo")
st.info("This is an GPT2 Text Generation Example using HuggingFace GPT2 Model")
pretrained = "gpt2-large"
tokenizer = GPT2Tokenizer.from_pretrained(pretrained)
model = GPT2LMHeadModel.from_pretrained(pretrained, pad_token_id=tokenizer.eos_token_id)
sentence = st.text_input('Input your sentence here:', value='My favorite ice cream flavor is ')
st.info("Max generated sentence: 100 words")
if (st.button("Generate")):
input_ids = tokenizer.encode(sentence, return_tensors='pt').to(device)
paragraph_generated = model.generate(input_ids, max_length=100, num_beams=5, no_repeat_ngram_size=2, early_stopping=True).to(device)
text = tokenizer.decode(paragraph_generated[0], skip_special_tokens=True)
st.write(text)