Update app.py

app.py

@@ -1,7 +1,7 @@
 ##########################################
 # Step 0: Import required libraries
 ##########################################
-import streamlit as st  # Web app framework
+import streamlit as st
 from transformers import (
     pipeline,
     SpeechT5Processor,
@@ -9,17 +9,17 @@ from transformers import (
     SpeechT5HifiGan,
     AutoModelForCausalLM,
     AutoTokenizer
-)  # NLP and TTS models
-from datasets import load_dataset  # Speaker embeddings
-import torch  # Tensor operations
-import soundfile as sf  # Audio file handling
-import sentencepiece  # Tokenization dependency
+)
+from datasets import load_dataset
+import torch
+import soundfile as sf
+import sentencepiece
 
 ##########################################
 # Initial configuration (MUST be first)
 ##########################################
 st.set_page_config(
-    page_title="🚀 Just Comment",
+    page_title="Just Comment",
     page_icon="💬",
     layout="centered",
     initial_sidebar_state="collapsed"
@@ -32,22 +32,15 @@ st.set_page_config(
 def load_models():
     """Load and cache all ML models"""
     return {
-        # Emotion classifier
         'emotion': pipeline(
             "text-classification", 
             model="Thea231/jhartmann_emotion_finetuning"
         ),
-        
-        # Text generation models
         'textgen_tokenizer': AutoTokenizer.from_pretrained("Qwen/Qwen1.5-0.5B"),
         'textgen_model': AutoModelForCausalLM.from_pretrained("Qwen/Qwen1.5-0.5B"),
-        
-        # TTS components
         'tts_processor': SpeechT5Processor.from_pretrained("microsoft/speecht5_tts"),
         'tts_model': SpeechT5ForTextToSpeech.from_pretrained("microsoft/speecht5_tts"),
         'tts_vocoder': SpeechT5HifiGan.from_pretrained("microsoft/speecht5_hifigan"),
-        
-        # Speaker embeddings
         'speaker_embeddings': torch.tensor(
             load_dataset("Matthijs/cmu-arctic-xvectors", split="validation")[7306]["xvector"]
         ).unsqueeze(0)
@@ -57,13 +50,13 @@ def load_models():
 # UI Components
 ##########################################
 def render_interface():
-    """Create user interface elements"""
-    st.title("🚀 I'm listening to you, my friend～")
-    st.caption("Analyzes feedback and tell tailored responses")
+    """Create user interface"""
+    st.title("Just Comment")
+    st.markdown("### I'm listening to you, my friend～")
     
     return st.text_area(
-        "📝 Paste customer feedback here:",
-        placeholder="The product arrived damaged...",
+        "📝 Enter your comment:",
+        placeholder="Share your thoughts...",
         height=150,
         key="user_input"
     )
@@ -72,88 +65,80 @@ def render_interface():
 # Core Logic Components
 ##########################################
 def analyze_emotion(text, classifier):
-    """Determine dominant emotion with confidence threshold"""
+    """Determine emotion with quick analysis"""
     results = classifier(text, return_all_scores=True)[0]
-    top_emotion = max(results, key=lambda x: x['score'])
-    return top_emotion if top_emotion['score'] > 0.6 else {'label': 'neutral', 'score': 1.0}
+    valid_emotions = ['sadness', 'joy', 'love', 'anger', 'fear', 'surprise']
+    filtered = [e for e in results if e['label'].lower() in valid_emotions]
+    return max(filtered, key=lambda x: x['score'])
 
 def generate_prompt(text, emotion):
-    """Create structured prompts for different emotions"""
+    """Complete prompt templates for all 6 emotions"""
     prompt_templates = {
-        "anger": (
-            "Customer complaint: {input}\n"
-            "Respond with:\n"
-            "1. Apology\n2. Solution steps\n3. Compensation offer\n"
-            "Response:"
+        "sadness": (
+            "Customer expressed sadness: {input}\n"
+            "Respond with:\n1. Empathetic acknowledgment\n"
+            "2. Supportive statement\n3. Concrete help offer\nResponse:"
         ),
         "joy": (
             "Positive feedback: {input}\n"
-            "Respond with:\n"
-            "1. Appreciation\n2. Highlight strengths\n3. Loyalty benefits\n"
-            "Response:"
+            "Respond with:\n1. Enthusiastic thanks\n"
+            "2. Specific compliment\n3. Future engagement\nResponse:"
+        ),
+        "love": (
+            "Customer showed affection: {input}\n"
+            "Respond with:\n1. Warm appreciation\n"
+            "2. Community building\n3. Exclusive offer\nResponse:"
+        ),
+        "anger": (
+            "Angry complaint: {input}\n"
+            "Respond with:\n1. Sincere apology\n"
+            "2. Solution steps\n3. Compensation\nResponse:"
+        ),
+        "fear": (
+            "Customer expressed concerns: {input}\n"
+            "Respond with:\n1. Reassurance\n"
+            "2. Safety measures\n3. Support channels\nResponse:"
         ),
-        "neutral": (
-            "Customer comment: {input}\n"
-            "Respond with:\n"
-            "1. Acknowledge feedback\n2. Offer assistance\n3. Next steps\n"
-            "Response:"
+        "surprise": (
+            "Unexpected feedback: {input}\n"
+            "Respond with:\n1. Acknowledge uniqueness\n"
+            "2. Creative solution\n3. Follow-up plan\nResponse:"
         )
     }
-    return prompt_templates.get(emotion.lower(), prompt_templates['neutral']).format(input=text)
+    return prompt_templates.get(emotion.lower(), "").format(input=text)
 
 def process_response(output_text):
-    """Ensure response quality and proper formatting"""
-    # Remove incomplete sentences
-    if '.' in output_text:
-        output_text = output_text.rsplit('.', 1)[0] + '.'
-    
-    # Length constraints
-    output_text = output_text[:300].strip()  # Hard limit at 300 characters
-    
-    # Fallback for short responses
-    if len(output_text) < 50:
-        return "Thank you for your feedback. We'll review this and contact you shortly."
-    
-    return output_text
+    """Optimized response processing"""
+    output_text = output_text.split("Response:")[-1].strip()
+    return output_text[:200]  # Strict length control
 
 def generate_text_response(user_input, models):
-    """Generate and validate text response"""
-    # Emotion analysis
+    """Efficient text generation"""
     emotion = analyze_emotion(user_input, models['emotion'])
-    
-    # Prompt engineering
     prompt = generate_prompt(user_input, emotion['label'])
     
-    # Text generation
     inputs = models['textgen_tokenizer'](prompt, return_tensors="pt")
     outputs = models['textgen_model'].generate(
         inputs.input_ids,
-        max_new_tokens=100,
+        max_new_tokens=150,  # Reduced for speed
         temperature=0.7,
         do_sample=True,
         top_p=0.9
     )
     
-    # Decode and process
-    full_response = models['textgen_tokenizer'].decode(outputs[0], skip_special_tokens=True)
-    return process_response(full_response.split("Response:")[-1].strip())
+    return process_response(
+        models['textgen_tokenizer'].decode(outputs[0], skip_special_tokens=True)
+    )
 
 def generate_audio_response(text, models):
-    """Convert text to speech"""
-    # Process text input
+    """Optimized TTS conversion"""
     inputs = models['tts_processor'](text=text, return_tensors="pt")
-    
-    # Generate spectrogram
     spectrogram = models['tts_model'].generate_speech(
         inputs["input_ids"],
         models['speaker_embeddings']
     )
-    
-    # Generate waveform
     with torch.no_grad():
         waveform = models['tts_vocoder'](spectrogram)
-    
-    # Save and return audio
     sf.write("response.wav", waveform.numpy(), samplerate=16000)
     return "response.wav"
 
@@ -161,25 +146,20 @@ def generate_audio_response(text, models):
 # Main Application Flow
 ##########################################
 def main():
-    # Load models once
     ml_models = load_models()
-    
-    # Render UI
     user_input = render_interface()
     
-    # Process input
     if user_input:
-        # Text generation
-        with st.status("🔍 Analyzing feedback...", expanded=True) as status:
+        # Text Generation
+        with st.spinner("🔍 Analyzing emotions..."):
             text_response = generate_text_response(user_input, ml_models)
-            status.update(label="✅ Analysis Complete", state="complete")
         
-        # Display text response
-        st.subheader("📝 Generated Response")
-        st.markdown(f"```\n{text_response}\n```")
+        # Display Results
+        st.subheader("📄 Generated Response")
+        st.success(text_response)
         
-        # Audio generation
-        with st.spinner("🔊 Generating voice response..."):
+        # Audio Generation
+        with st.spinner("🔊 Generating voice..."):
             audio_file = generate_audio_response(text_response, ml_models)
             st.audio(audio_file, format="audio/wav")