ahmadrocks's picture
Update app.py
58c6588 verified
import streamlit as st
from transformers import GPT2LMHeadModel, GPT2Tokenizer
import torch
# Load the GPT-2 model and tokenizer
model_name = 'gpt2-large'
model = GPT2LMHeadModel.from_pretrained(model_name)
tokenizer = GPT2Tokenizer.from_pretrained(model_name)
st.title("AI Article Generator")
# Input for the article title
title = st.text_input("Enter the title of the article")
# Parameters for text generation
max_length = st.slider("Maximum length of the article", min_value=50, max_value=1000, value=500, step=50)
temperature = st.slider("Temperature (creativity level)", min_value=0.7, max_value=1.5, value=1.0, step=0.1)
# Generate the article
if st.button("Generate Article"):
if title:
with st.spinner("Generating article..."):
try:
input_ids = tokenizer.encode(title, return_tensors='pt')
output = model.generate(
input_ids,
max_length=max_length,
temperature=temperature,
num_return_sequences=1,
no_repeat_ngram_size=2,
early_stopping=True
)
article = tokenizer.decode(output[0], skip_special_tokens=True)
st.success("Article generated successfully!")
st.write(article)
except Exception as e:
st.error(f"An error occurred: {e}")
else:
st.warning("Please enter a title to generate an article")
st.markdown("""
<style>
.reportview-container {
flex-direction: column;
align-items: center;
}
.css-1kyxreq {
width: 100%;
max-width: 700px;
margin: auto;
}
</style>
""", unsafe_allow_html=True)