import streamlit as st
from transformers import GPT2LMHeadModel, GPT2Tokenizer
import torch

# Title and description
st.title("🧠 AI Blog Writer")
st.write("""
This AI-powered blog generator uses **GPT-2** from Hugging Face Transformers to generate creative, 
context-aware blog content from a short user prompt.  
Simply enter a topic or a starting line, and watch it create a complete paragraph.
""")

# Load model and tokenizer with caching (to avoid reloading on every run)
@st.cache_resource
def load_model():
    tokenizer = GPT2Tokenizer.from_pretrained("gpt2-large")
    model = GPT2LMHeadModel.from_pretrained("gpt2-large", pad_token_id=tokenizer.eos_token_id)
    return tokenizer, model

tokenizer, model = load_model()

# User input
user_input = st.text_area("✍️ Enter your blog topic or starting line:", "Love is")

# Generation settings
st.sidebar.header("Generation Settings")
max_length = st.sidebar.slider("Max Length", 50, 500, 200)
temperature = st.sidebar.slider("Creativity (Temperature)", 0.5, 1.5, 0.9)
top_k = st.sidebar.slider("Top-K Sampling", 20, 100, 50)

# Generate button
if st.button("Generate Blog"):
    with st.spinner("Generating text... please wait ⏳"):
        input_ids = tokenizer.encode(user_input, return_tensors="pt")
        output = model.generate(
            input_ids,
            max_length=max_length,
            temperature=temperature,
            top_k=top_k,
            no_repeat_ngram_size=2,
            early_stopping=True
        )
        result = tokenizer.decode(output[0], skip_special_tokens=True)
        st.success("✅ Blog generated successfully!")
        st.write(result)

        # Option to download
        st.download_button("📥 Download Blog Text", result, file_name="generated_blog.txt")

st.caption("Built with 🤗 Hugging Face Transformers & Streamlit | Developed by Anshul")