import streamlit as st
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# Load GPT-2 model and tokenizer
@st.cache(allow_output_mutation=True)
def load_model():
    tokenizer = AutoTokenizer.from_pretrained("gpt2")
    model = AutoModelForCausalLM.from_pretrained("gpt2")
    return tokenizer, model

tokenizer, model = load_model()

st.title("Blog Post Generator")
st.write("Generate a blog post for a given topic using GPT-2.")

# User input for the blog post topic
topic = st.text_input("Enter the topic for your blog post:")

# Generate blog post button
if st.button("Generate Blog Post"):
    if topic:
        # Prepare the input for the model
        input_text = f"Write a blog post about {topic}."
        inputs = tokenizer.encode(input_text, return_tensors="pt")

        # Generate the blog post using GPT-2
        outputs = model.generate(inputs, max_length=500, num_return_sequences=1, no_repeat_ngram_size=2, early_stopping=True)

        # Decode the generated text
        blog_post = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        st.write("### Generated Blog Post:")
        st.write(blog_post)
    else:
        st.write("Please enter a topic to generate a blog post.")