Update app.py

app.py

@@ -1,26 +1,59 @@
 import streamlit as st
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
 # Load model and tokenizer
-model_name = "gpt2"
+model_name = "gpt2-large"  # You can change this to a larger GPT model if needed
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name)
 
-def generate_blog_post(topic):
-    prompt = f"Write a detailed blog post about {topic}."
-    inputs = tokenizer.encode(prompt, return_tensors="pt")
-    outputs = model.generate(inputs, max_length=512, num_return_sequences=1)
-    blog_post = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return blog_post
+# Set up zero-shot classification pipeline
+classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")
+
+# Function to generate blog post
+def generate_blog_post(topic, max_length=500):
+    prompt = f"Write a blog post about {topic}:\n\n"
+    inputs = tokenizer(prompt, return_tensors="pt")
+    
+    outputs = model.generate(
+        inputs.input_ids,
+        max_length=max_length,
+        num_return_sequences=1,
+        no_repeat_ngram_size=2,
+        top_k=50,
+        top_p=0.95,
+        temperature=0.7,
+    )
+    
+    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return generated_text.replace(prompt, "")
+
+# Function to classify blog post
+def classify_blog_post(text, labels):
+    result = classifier(text, labels)
+    return result
 
 # Streamlit interface
 st.title("Blog Post Generator")
-st.write("Enter a topic to generate a detailed blog post.")
 
-topic = st.text_input("Topic", "")
-if st.button("Generate Blog Post"):
-    if topic:
+topic = st.text_input("Enter a topic for your blog post:")
+generate_button = st.button("Generate Blog Post")
+
+if generate_button and topic:
+    with st.spinner("Generating blog post..."):
         blog_post = generate_blog_post(topic)
+        st.subheader("Generated Blog Post")
         st.write(blog_post)
-    else:
-        st.write("Please enter a topic to generate a blog post.")
+        
+        # Classify the generated blog post
+        st.subheader("Blog Post Classification")
+        labels = ["Technology", "Travel", "Food", "Health", "Finance"]
+        classification = classify_blog_post(blog_post, labels)
+        
+        for label, score in zip(classification['labels'], classification['scores']):
+            st.write(f"{label}: {score:.2f}")
+
+st.sidebar.title("About")
+st.sidebar.info(
+    "This app generates a blog post on a given topic using a large GPT model. "
+    "It also classifies the generated post using zero-shot classification."
+)