Spaces:
Sleeping
Sleeping
import streamlit as st | |
from transformers import GPT2LMHeadModel, GPT2Tokenizer | |
# Load pre-trained model and tokenizer | |
model_name = "gpt2" | |
tokenizer = GPT2Tokenizer.from_pretrained(model_name) | |
model = GPT2LMHeadModel.from_pretrained(model_name) | |
def generate_text(prompt, max_length=50): | |
# Encode the input prompt | |
inputs = tokenizer.encode(prompt, return_tensors="pt") | |
# Generate text | |
outputs = model.generate(inputs, max_length=max_length, num_return_sequences=1) | |
# Decode the generated text | |
text = tokenizer.decode(outputs[0], skip_special_tokens=True) | |
return text | |
# Streamlit app | |
st.title("GPT-2 Text Generator") | |
prompt = st.text_area("Input", "Once upon a time...") | |
max_length = st.slider("Max Length", min_value=10, max_value=100, value=50) | |
if st.button("Generate"): | |
generated_text = generate_text(prompt, max_length) | |
st.subheader("Generated Text") | |
st.write(generated_text) |