|
import streamlit as st |
|
import torch |
|
from transformers import GPT2Tokenizer, GPT2LMHeadModel |
|
|
|
|
|
tokenizer = GPT2Tokenizer.from_pretrained('gpt2') |
|
model = GPT2LMHeadModel.from_pretrained('gpt2') |
|
|
|
|
|
def main(): |
|
st.title('GPT-2 Text Generation') |
|
st.write('Enter a prompt below to generate text:') |
|
|
|
|
|
prompt = st.text_area('Input Prompt') |
|
|
|
|
|
if st.button('Generate Text'): |
|
if prompt: |
|
|
|
input_ids = tokenizer.encode(prompt, return_tensors='pt') |
|
|
|
|
|
output = model.generate(input_ids, max_length=100, num_return_sequences=1, pad_token_id=tokenizer.eos_token_id) |
|
|
|
|
|
generated_text = tokenizer.decode(output[0], skip_special_tokens=True) |
|
st.write("Generated Text:") |
|
st.write(generated_text) |
|
else: |
|
st.warning('Please enter a prompt.') |
|
|
|
if __name__ == '__main__': |
|
main() |
|
|