Update app.py

app.py

@@ -1,16 +1,17 @@
-import streamlit as st
+import gradio as gr
 import json
 import asyncio
 import edge_tts
+import os
 from huggingface_hub import InferenceClient
 
-# Configuration
+# --- CONFIGURATION ---
 EXTRACTOR_MODEL = "meta-llama/Meta-Llama-3-8B-Instruct"
 PERSONALITY_MODEL = "HuggingFaceH4/zephyr-7b-beta"
 
-# Synthetic Data: 30 Messages reflecting personality traits (Big Five / MBTI context)
-DEFAULT_CHAT_LOG = """
-1. User: I feel completely drained after large social gatherings. I need alone time to recharge.
+# Default Data: 30 Personality-focused messages
+DEFAULT_LOGS = """
+1. User: I feel completely drained after large social gatherings.
 2. User: I enjoy abstract theories more than practical details.
 3. User: My desk is always messy, but I know where everything is.
 4. User: I often worry that I said the wrong thing in conversations.
@@ -42,138 +43,131 @@ DEFAULT_CHAT_LOG = """
 30. User: I want to understand my true purpose in life.
 """
 
-st.set_page_config(page_title="Personality Analysis Engine", layout="wide")
-
-# Authentication Sidebar
-with st.sidebar:
-    st.header("Configuration")
-    hf_token = st.text_input("Hugging Face Token", type="password")
+# --- LOGIC: MEMORY EXTRACTION ---
+def extract_memory(chat_logs, hf_token):
     if not hf_token:
-        st.warning("Please enter your Hugging Face Token to proceed.")
-        st.stop()
-
-client = InferenceClient(token=hf_token)
-
-# Module A: Memory Extraction
-def extract_memory(chat_logs):
-    # System prompt designed to extract psychological profile
-    system_prompt = """
-    You are a Psychological Analyst. Analyze the chat logs and extract a structured user profile.
+        return json.dumps({"error": "Missing HF Token"}, indent=2)
     
-    Required JSON Keys:
-    1. "traits": Big Five or MBTI indicators (e.g., Introversion, High Openness, Agreeableness).
-    2. "values": What the user prioritizes (e.g., Logic, Harmony, Efficiency, Autonomy).
-    3. "struggles": Recurring points of friction (e.g., Social anxiety, Procrastination, Conflict avoidance).
+    client = InferenceClient(token=hf_token)
     
-    Output strictly valid JSON only.
+    system_prompt = """
+    Analyze the chat logs and extract a psychological profile.
+    JSON KEYS: "traits" (Big Five/MBTI), "values" (Logic, Harmony, etc), "struggles" (Anxiety, Procrastination).
+    OUTPUT: Valid JSON only.
     """
     
     prompt = f"<|system|>\n{system_prompt}</s>\n<|user|>\n{chat_logs}</s>\n<|assistant|>"
     
     try:
-        with st.spinner("Analyzing psychological patterns..."):
-            response = client.text_generation(
-                model=EXTRACTOR_MODEL,
-                prompt=prompt,
-                max_new_tokens=600,
-                temperature=0.1
-            )
-            # Parse JSON from response
-            text = response.strip()
-            start = text.find("{")
-            end = text.rfind("}") + 1
-            if start != -1 and end != -1:
-                return json.loads(text[start:end])
-            else:
-                return {"error": "Failed to parse JSON output"}
+        response = client.text_generation(
+            model=EXTRACTOR_MODEL,
+            prompt=prompt,
+            max_new_tokens=600,
+            temperature=0.1
+        )
+        # Parse JSON
+        text = response.strip()
+        start = text.find("{")
+        end = text.rfind("}") + 1
+        return json.dumps(json.loads(text[start:end]), indent=2)
     except Exception as e:
-        return {"error": str(e)}
+        return json.dumps({"error": str(e)}, indent=2)
 
-# Module B: Personality Engine (Text Generation)
-def generate_response(message, memory, persona):
-    # Define distinct psychological personas
-    personas = {
-        "Analytical Psychologist": "Role: Jungian Analyst. Tone: Objective, deep, intellectual. Focus: Unconscious patterns, archetypes, and logic.",
-        "Empathetic Counselor": "Role: Humanistic Therapist. Tone: Warm, validating, gentle. Focus: Emotional safety, acceptance, and feelings.",
-        "Pragmatic Coach": "Role: Behavioral Coach. Tone: Direct, action-oriented, blunt. Focus: Solutions, efficiency, and breaking habits."
+# --- LOGIC: PERSONALITY RESPONSE ---
+async def generate_response_and_audio(message, memory_json, persona, hf_token):
+    if not hf_token:
+        return "Please enter HF Token.", None
+        
+    client = InferenceClient(token=hf_token)
+    
+    try:
+        memory = json.loads(memory_json)
+    except:
+        memory = {}
+
+    # 1. Define Persona Prompts
+    prompts = {
+        "Calm Mentor": "Role: Mentor. Tone: Calm, wise, patient. Focus: Growth, long-term perspective. Voice: Deep & Slow.",
+        "Witty Friend": "Role: Friend. Tone: Witty, casual, sarcastic, fun. Focus: Relatability, humor. Voice: Fast & Energetic.",
+        "Therapist-style": "Role: Therapist. Tone: Empathetic, soft, validating. Focus: Emotions, safety. Voice: Soft & Warm."
     }
     
+    # 2. Context Injection
     context = f"""
-    USER PSYCH PROFILE:
+    USER PROFILE:
     - Traits: {memory.get('traits', 'Unknown')}
     - Values: {memory.get('values', 'Unknown')}
     - Struggles: {memory.get('struggles', 'Unknown')}
     """
     
     messages = [
-        {"role": "system", "content": f"{personas[persona]}\n\n{context}"},
+        {"role": "system", "content": f"{prompts[persona]}\n\n{context}"},
         {"role": "user", "content": message}
     ]
     
-    try:
-        res = client.chat_completion(
-            model=PERSONALITY_MODEL,
-            messages=messages,
-            max_tokens=250,
-            temperature=0.7
-        )
-        return res.choices[0].message.content
-    except Exception as e:
-        return f"Error generating text: {e}"
-
-# Module C: Audio Engine (Edge-TTS)
-async def generate_audio(text, persona):
-    # Map personas to suitable voice types
+    # 3. Generate Text
+    res = client.chat_completion(
+        model=PERSONALITY_MODEL,
+        messages=messages,
+        max_tokens=250,
+        temperature=0.7
+    )
+    text_reply = res.choices[0].message.content
+    
+    # 4. Generate Audio (Edge-TTS)
     voice_map = {
-        "Analytical Psychologist": "en-US-ChristopherNeural", # Deep, serious male voice
-        "Empathetic Counselor": "en-US-AvaNeural",            # Soft, soothing female voice
-        "Pragmatic Coach": "en-US-EricNeural"                 # Energetic, firm male voice
+        "Calm Mentor": "en-US-ChristopherNeural",
+        "Witty Friend": "en-US-EricNeural",
+        "Therapist-style": "en-US-AvaNeural"
     }
     
-    voice = voice_map.get(persona, "en-US-AriaNeural")
     output_file = "response.mp3"
-    
-    communicate = edge_tts.Communicate(text, voice)
+    communicate = edge_tts.Communicate(text_reply, voice_map.get(persona, "en-US-AriaNeural"))
     await communicate.save(output_file)
-    return output_file
-
-# Main UI Layout
-st.title("Human Personality Engine")
+    
+    return text_reply, output_file
 
-col1, col2 = st.columns([1, 1])
+# Wrapper for Gradio (handling async)
+def process_interaction(message, memory_json, persona, hf_token):
+    return asyncio.run(generate_response_and_audio(message, memory_json, persona, hf_token))
 
-# Left Column: Data Analysis
-with col1:
-    st.subheader("1. Profile Extraction")
-    with st.expander("View Raw Chat Logs"):
-        st.text(DEFAULT_CHAT_LOG)
-        
-    if st.button("Analyze User Profile"):
-        memory_data = extract_memory(DEFAULT_CHAT_LOG)
-        st.session_state['memory'] = memory_data
-        st.success("Analysis Complete")
+# --- GRADIO UI ---
+with gr.Blocks(title="Personality Engine") as demo:
+    gr.Markdown("# Personality Engine: Tones")
     
-    if 'memory' in st.session_state:
-        st.json(st.session_state['memory'])
-
-# Right Column: Interaction
-with col2:
-    st.subheader("2. Agent Interaction")
+    with gr.Row():
+        token_input = gr.Textbox(label="Hugging Face Token", type="password", placeholder="hf_...")
     
-    if 'memory' not in st.session_state:
-        st.info("Please analyze the user profile first.")
-    else:
-        user_input = st.text_input("Input Message:", "I feel overwhelmed by all the details.")
-        persona = st.selectbox("Select Agent Persona:", ["Analytical Psychologist", "Empathetic Counselor", "Pragmatic Coach"])
-        
-        if st.button("Generate Response"):
-            # Generate Text
-            with st.spinner("Generating text response..."):
-                reply = generate_response(user_input, st.session_state['memory'], persona)
-            st.markdown(f"**{persona}:** {reply}")
+    with gr.Row():
+        # Left Column: Memory
+        with gr.Column(scale=1):
+            gr.Markdown("### 1. Memory Extraction")
+            logs_input = gr.Textbox(label="Chat Logs (Input)", value=DEFAULT_LOGS, lines=10)
+            extract_btn = gr.Button("Analyze Profile")
+            memory_output = gr.Code(label="Extracted Memory (JSON)", language="json")
             
-            # Generate Audio
-            with st.spinner("Synthesizing speech..."):
-                asyncio.run(generate_audio(reply, persona))
+            extract_btn.click(extract_memory, inputs=[logs_input, token_input], outputs=memory_output)
+
+        # Right Column: Interaction
+        with gr.Column(scale=1):
+            gr.Markdown("### 2. Personality Interaction")
+            user_msg = gr.Textbox(label="Your Message", value="I'm feeling overwhelmed by the future.")
+            persona_select = gr.Radio(
+                ["Calm Mentor", "Witty Friend", "Therapist-style"], 
+                label="Select Tone", 
+                value="Calm Mentor"
+            )
+            generate_btn = gr.Button("Generate Response")
             
-            st.audio("response.mp3")
+            with gr.Group():
+                text_output = gr.Markdown(label="Response Text")
+                audio_output = gr.Audio(label="Voice Response")
+            
+            generate_btn.click(
+                process_interaction, 
+                inputs=[user_msg, memory_output, persona_select, token_input], 
+                outputs=[text_output, audio_output]
+            )
+
+if __name__ == "__main__":
+    demo.queue().launch()