Update app.py

app.py

@@ -2,128 +2,98 @@ import streamlit as st
 from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification, AutoModelForCausalLM
 import torch
 
-# -------------------------------
+# ----------------------------------
 # Page Configuration
-# -------------------------------
-st.set_page_config(
-    page_title="AI Sentiment Essay Generator",
-    page_icon="🧠",
-    layout="wide"
-)
+# ----------------------------------
+st.set_page_config(page_title="AI Sentiment Essay Generator", page_icon="🧠", layout="wide")
 
 st.markdown("""
-<style>
-body {
-    background: radial-gradient(circle at top, #1f1f2e, #0f0f16);
-    color: #f0f0f0;
-}
-.main {
-    background-color: rgba(20, 20, 30, 0.8);
-    padding: 30px;
-    border-radius: 20px;
-    box-shadow: 0px 0px 15px rgba(100,100,255,0.2);
-}
-textarea {
-    border-radius: 10px !important;
-}
-</style>
+    <style>
+    body {
+        background-color: #0e1117;
+        color: white;
+    }
+    .positive {background: linear-gradient(90deg, #3fc1c9, #4caf50); border-radius: 10px; padding: 1rem;}
+    .negative {background: linear-gradient(90deg, #f05454, #b33939); border-radius: 10px; padding: 1rem;}
+    .neutral {background: linear-gradient(90deg, #607d8b, #455a64); border-radius: 10px; padding: 1rem;}
+    </style>
 """, unsafe_allow_html=True)
 
-# -------------------------------
-# Title and Description
-# -------------------------------
 st.title("🧠 AI Sentiment Essay Generator")
-st.write("Enter a prompt — the AI will detect its sentiment (positive, negative, or neutral) "
-         "and write an expressive paragraph aligned with that emotion 💬.")
+st.caption("Analyze the emotion of your topic and let AI craft an expressive essay aligned with it 🎭")
 
-# -------------------------------
-# Load Models (Cached)
-# -------------------------------
+# ----------------------------------
+# Model Initialization (cached)
+# ----------------------------------
 @st.cache_resource
-def load_sentiment_model():
-    tokenizer = AutoTokenizer.from_pretrained("cardiffnlp/twitter-roberta-base-sentiment")
-    model = AutoModelForSequenceClassification.from_pretrained("cardiffnlp/twitter-roberta-base-sentiment")
-    sentiment_pipeline = pipeline("sentiment-analysis", model=model, tokenizer=tokenizer)
-    return sentiment_pipeline
+def load_models():
+    sentiment_model_name = "cardiffnlp/twitter-roberta-base-sentiment-latest"
+    sentiment_analyzer = pipeline("sentiment-analysis", model=sentiment_model_name)
 
-@st.cache_resource
-def load_generation_model():
-    gen_model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen1.5-0.5B-Chat")
-    gen_tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen1.5-0.5B-Chat")
-    generator = pipeline("text-generation", model=gen_model, tokenizer=gen_tokenizer)
-    return generator
-
-with st.spinner("Loading AI models... ⏳"):
-    sentiment_analyzer = load_sentiment_model()
-    text_generator = load_generation_model()
-
-# -------------------------------
-# Input Section
-# -------------------------------
-user_input = st.text_area("💬 Enter your topic or prompt here:", height=120, placeholder="e.g., Technology and human connection...")
-
-generate_button = st.button("✨ Generate Essay")
-
-# -------------------------------
-# Main Logic
-# -------------------------------
-if generate_button and user_input.strip() != "":
-    with st.spinner("Analyzing sentiment... 🧠"):
-        sentiment_result = sentiment_analyzer(user_input)[0]
-        label = sentiment_result["label"].capitalize()
-        score = round(sentiment_result["score"] * 100, 2)
-
-    # Sentiment-based UI styling
-    if "Positive" in label:
-        color = "#00C851"
-        emoji = "😊"
-        gradient = "linear-gradient(90deg, #00C851, #007E33)"
-    elif "Negative" in label:
-        color = "#ff4444"
-        emoji = "😞"
-        gradient = "linear-gradient(90deg, #ff4444, #CC0000)"
-    else:
-        color = "#33b5e5"
-        emoji = "😐"
-        gradient = "linear-gradient(90deg, #33b5e5, #0099CC)"
-
-    st.markdown(f"""
-        <div style='padding:15px;border-radius:12px;background:{gradient};text-align:center;'>
-        <h3>{emoji} Detected Sentiment: <b>{label}</b> ({score}%)</h3>
-        </div>
-    """, unsafe_allow_html=True)
-
-    # -------------------------------
-    # Generate Sentiment-Aligned Text
-    # -------------------------------
-    prompt = f"Write a {label.lower()} and expressive paragraph about: {user_input}."
-    with st.spinner(f"Generating a {label.lower()} essay... ✍️"):
-        result = text_generator(prompt, max_length=200, do_sample=True, temperature=0.9)
-        generated_text = result[0]["generated_text"]
-
-    # -------------------------------
-    # Display Output
-    # -------------------------------
-    st.subheader("📝 AI-Generated Essay")
-    st.markdown(
-        f"""
-        <div style='padding:20px;background-color:rgba(255,255,255,0.08);border-radius:15px;border-left:6px solid {color};'>
-        <p style='font-size:1.1rem;line-height:1.6em;'>{generated_text}</p>
-        </div>
-        """, unsafe_allow_html=True
-    )
+    text_gen_model_name = "Qwen/Qwen1.5-1.8B"
+    tokenizer = AutoTokenizer.from_pretrained(text_gen_model_name)
+    text_model = AutoModelForCausalLM.from_pretrained(text_gen_model_name)
+    return sentiment_analyzer, tokenizer, text_model
+
+sentiment_analyzer, tokenizer, text_model = load_models()
+
+# ----------------------------------
+# User Input
+# ----------------------------------
+user_prompt = st.text_area("💬 Enter a topic or sentence:", placeholder="e.g., The future of Artificial Intelligence...")
 
-    # -------------------------------
-    # Sidebar Info
-    # -------------------------------
-    with st.sidebar:
-        st.header("🧩 Summary")
-        st.markdown(f"**Prompt:** {user_input}")
-        st.markdown(f"**Sentiment:** {label} {emoji}")
-        st.markdown(f"**Confidence:** {score}%")
-        st.markdown("---")
-        st.info("Tip: Try emotional prompts like *“life after success”* or *“loneliness in a big city.”*")
+generate_btn = st.button("✨ Generate Essay")
 
-else:
-    st.info("💡 Enter your text above and click **Generate Essay** to begin!")
+# ----------------------------------
+# Processing
+# ----------------------------------
+if generate_btn and user_prompt.strip():
+    with st.spinner("Analyzing sentiment..."):
+        sentiment_result = sentiment_analyzer(user_prompt)[0]
 
+        label_map = {
+            "LABEL_0": "Negative",
+            "LABEL_1": "Neutral",
+            "LABEL_2": "Positive"
+        }
+
+        sentiment_label = label_map.get(sentiment_result["label"], sentiment_result["label"])
+        confidence = round(sentiment_result["score"] * 100, 2)
+
+    # Sentiment Display
+    sentiment_emoji = {"Positive": "😊", "Negative": "😠", "Neutral": "😐"}
+    sentiment_class = sentiment_label.lower()
+
+    st.markdown(f"<div class='{sentiment_class}'>", unsafe_allow_html=True)
+    st.markdown(
+        f"<h3 style='text-align:center;'>{sentiment_emoji.get(sentiment_label, '😐')} "
+        f"Detected Sentiment: <b>{sentiment_label}</b> ({confidence}%)</h3>",
+        unsafe_allow_html=True
+    )
+    st.markdown("</div>", unsafe_allow_html=True)
+
+    # ----------------------------------
+    # Generate Essay
+    # ----------------------------------
+    with st.spinner("Generating essay with AI..."):
+        gen_prompt = f"Write a {sentiment_label.lower()} and expressive essay about: {user_prompt}"
+
+        input_ids = tokenizer(gen_prompt, return_tensors="pt").input_ids
+        output_ids = text_model.generate(
+            input_ids,
+            max_new_tokens=250,
+            temperature=0.9,
+            top_p=0.95,
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id
+        )
+        generated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+
+    st.markdown("### 📝 AI-Generated Essay")
+    st.write(generated_text)
+
+    st.markdown("---")
+    st.caption("✨ Powered by RoBERTa + Qwen | Streamlit Frontend")
+
+elif generate_btn:
+    st.warning("⚠️ Please enter a topic or sentence first.")