Update app.py

app.py

@@ -1,15 +1,13 @@
 # app.py
 import streamlit as st
 from groq import Groq
-from transformers import pipeline
-import re
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
+import torch
 import random
 from datetime import datetime
 from reportlab.lib.pagesizes import letter
 from reportlab.pdfgen import canvas
 import io
-import torch
-from transformers import AutoModelForSequenceClassification, AutoTokenizer
 
 # Initialize components
 try:
@@ -20,7 +18,10 @@ except KeyError:
 
 # Load personality model
 try:
-    personality_model = AutoModelForSequenceClassification.from_pretrained("KevSun/Personality_LM", ignore_mismatched_sizes=True)
+    personality_model = AutoModelForSequenceClassification.from_pretrained(
+        "KevSun/Personality_LM", 
+        ignore_mismatched_sizes=True
+    )
     personality_tokenizer = AutoTokenizer.from_pretrained("KevSun/Personality_LM")
 except Exception as e:
     st.error(f"Model loading error: {str(e)}")
@@ -52,6 +53,11 @@ st.markdown("""
   gap: 10px;
   margin: 20px 0;
 }
+.response-box {
+  border-left: 3px solid #4CAF50;
+  padding: 10px;
+  margin: 10px 0;
+}
 </style>
 """, unsafe_allow_html=True)
 
@@ -67,9 +73,20 @@ QUESTION_BANK = [
     {"text": "How do you recharge after a stressful day? 🧘", "type": "serious", "trait": "neuroticism"}
 ]
 
+def initialize_session():
+    """Initialize all session state variables"""
+    if 'questions' not in st.session_state:
+        st.session_state.questions = random.sample(QUESTION_BANK, 5)
+        for i in range(5):
+            st.session_state[f'q{i}'] = ""
+        st.session_state.responses = []
+    if 'current_q' not in st.session_state:
+        st.session_state.current_q = 0
+    if 'show_post' not in st.session_state:
+        st.session_state.show_post = False
+
 def get_dynamic_questions():
     """Generate random mix of funny/serious questions"""
-    random.shuffle(QUESTION_BANK)
     return random.sample(QUESTION_BANK, 5)
 
 def analyze_personality(text):
@@ -84,14 +101,14 @@ def analyze_personality(text):
 def generate_social_post(platform, traits):
     """Generate platform-specific social post"""
     prompts = {
-        "instagram": "Create an Instagram post with 3 emojis and 2 hashtags about:",
-        "linkedin": "Create a professional LinkedIn post about personal growth with 1 emoji:",
-        "facebook": "Create a friendly Facebook post with 2 emojis:",
-        "whatsapp": "Create a casual WhatsApp status with 2 emojis:"
+        "instagram": "Create an Instagram post with 3 emojis and 2 hashtags about personal growth:",
+        "linkedin": "Create a professional LinkedIn post about self-improvement with 1 emoji:",
+        "facebook": "Create a friendly Facebook post about personality insights with 2 emojis:",
+        "whatsapp": "Create a casual WhatsApp status about self-discovery with 2 emojis:"
     }
     prompt = f"""{prompts[platform]}
-    My personality strengths: {traits}
-    Make it {['playful and visual', 'professional', 'friendly', 'casual'][list(prompts.keys()).index(platform)]}"""
+    Based on these personality traits: {traits}
+    Make it {['visual and trendy', 'professional', 'friendly', 'casual'][list(prompts.keys()).index(platform)]}"""
     
     response = groq_client.chat.completions.create(
         model="mixtral-8x7b-32768",
@@ -100,13 +117,8 @@ def generate_social_post(platform, traits):
     )
     return response.choices[0].message.content
 
-# Session state management
-if 'questions' not in st.session_state:
-    st.session_state.questions = get_dynamic_questions()
-if 'current_q' not in st.session_state:
-    st.session_state.current_q = 0
-if 'show_post' not in st.session_state:
-    st.session_state.show_post = False
+# Initialize session state
+initialize_session()
 
 # Main UI
 st.title("🧠 PsychBot Pro")
@@ -126,11 +138,12 @@ if st.session_state.current_q < len(st.session_state.questions):
         user_input = st.text_input("Your response:", key=f"q{st.session_state.current_q}")
         
         if st.button("Next ➡️"):
+            st.session_state.responses.append(user_input)
             st.session_state.current_q += 1
             st.rerun()
 else:
     # Generate personality report
-    combined_text = " ".join([st.session_state[f"q{i}"] for i in range(len(st.session_state.questions))])
+    combined_text = " ".join(st.session_state.responses)
     traits = analyze_personality(combined_text)
     
     st.balloons()
@@ -138,34 +151,42 @@ else:
     
     # Personality visualization
     cols = st.columns(5)
-    traits = {k: v for k, v in sorted(traits.items(), key=lambda item: item[1], reverse=True)}
-    for i, (trait, score) in enumerate(traits.items()):
+    sorted_traits = sorted(traits.items(), key=lambda x: x[1], reverse=True)
+    for i, (trait, score) in enumerate(sorted_traits):
         cols[i].metric(label=trait.upper(), value=f"{score*100:.0f}%")
     
     # Social post generation
+    st.markdown("---")
     if st.button("📱 Generate Social Media Post"):
         st.session_state.show_post = True
         
     if st.session_state.show_post:
         platforms = ["instagram", "linkedin", "facebook", "whatsapp"]
-        selected = st.radio("Choose platform:", platforms, horizontal=True)
+        selected = st.radio("Choose platform:", platforms, format_func=lambda x: x.capitalize(), horizontal=True)
         
         if post := generate_social_post(selected, traits):
             st.markdown(f"""
             <div class="social-post">
                 <h4>🎨 {selected.capitalize()} Post Draft</h4>
-                <p>{post}</p>
-                <button onclick="navigator.clipboard.writeText(`{post}`)">📋 Copy</button>
+                <div class="response-box">{post}</div>
+                <button onclick="navigator.clipboard.writeText(`{post}`)" style="margin-top:10px;">📋 Copy Text</button>
             </div>
             """, unsafe_allow_html=True)
 
+    # Restart conversation
+    if st.button("🔄 Start New Analysis"):
+        for key in list(st.session_state.keys()):
+            del st.session_state[key]
+        st.rerun()
+
 # Sidebar
 with st.sidebar:
     st.markdown("## 🌈 Features")
     st.markdown("""
-    - 🎭 Dynamic personality assessment
+    - 🎭 Personality assessment
     - 🤖 AI-generated social posts
     - 📊 Visual trait analysis
-    - 💬 Mix of fun/serious questions
+    - 💬 Dynamic questions
+    - 📥 PDF report download
     """)
-    st.image("https://i.imgur.com/7Q4X4yN.png", width=200)
+