Update app.py

app.py

@@ -1,43 +1,29 @@
 import streamlit as st
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch
+from transformers import GPT2LMHeadModel, GPT2Tokenizer
 
-# Load the GPT-2 large model and tokenizer
-model_name = "gpt2-large"
-model = AutoModelForCausalLM.from_pretrained(model_name)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+@st.cache(allow_output_mutation=True)
+def load_model():
+    tokenizer = GPT2Tokenizer.from_pretrained("gpt2-large")
+    model = GPT2LMHeadModel.from_pretrained("gpt2-large")
+    return tokenizer, model
 
-def generate_blogpost(topic):
-    try:
-        # Prepare input text
-        input_text = f"Write a blog post about {topic}:"
-        inputs = tokenizer.encode(input_text, return_tensors="pt")
-        st.write(f"Input IDs: {inputs}")
+def generate_blog_post(topic, max_length=200):
+    tokenizer, model = load_model()
+    input_ids = tokenizer.encode(topic, return_tensors='pt')
+    output = model.generate(input_ids, max_length=max_length, num_return_sequences=1, no_repeat_ngram_size=2, pad_token_id=tokenizer.eos_token_id)
+    blog_post = tokenizer.decode(output[0], skip_special_tokens=True)
+    return blog_post
 
-        # Generate output
-        with torch.no_grad():
-            outputs = model.generate(inputs, max_length=500, num_return_sequences=1)
-        st.write(f"Output IDs: {outputs}")
-
-        # Decode the generated text
-        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        return generated_text
-
-    except Exception as e:
-        return f"An error occurred: {str(e)}"
-
-# Streamlit UI
 st.title("Blog Post Generator")
-st.write("Generate a blog post for a given topic using GPT-2 large.")
+st.write("Enter a topic to generate a blog post using GPT-2 large.")
 
-# Input for the topic
-topic = st.text_input("Enter the topic:")
+topic = st.text_input("Topic:", "")
+length = st.slider("Post Length (in tokens):", min_value=50, max_value=500, value=200)
 
-# Generate button
 if st.button("Generate"):
     if topic:
-        # Generate and display the blog post
-        blog_post = generate_blogpost(topic)
+        blog_post = generate_blog_post(topic, max_length=length)
+        st.subheader("Generated Blog Post")
         st.write(blog_post)
     else:
         st.write("Please enter a topic to generate a blog post.")