Update app.py

app.py

@@ -1,7 +1,7 @@
 ##########################################
 # Step 0: Import required libraries
 ##########################################
-import streamlit as st
+import streamlit as st  # For web interface
 from transformers import (
     pipeline,
     SpeechT5Processor,
@@ -9,14 +9,14 @@ from transformers import (
     SpeechT5HifiGan,
     AutoModelForCausalLM,
     AutoTokenizer
-)
-from datasets import load_dataset
-import torch
-import soundfile as sf
-import sentencepiece
+)  # AI model components
+from datasets import load_dataset  # For voice embeddings
+import torch  # Tensor computations
+import soundfile as sf  # Audio file handling
+import re  # Regular expressions for text processing
 
 ##########################################
-# Initial configuration (MUST be first)
+# Initial configuration
 ##########################################
 st.set_page_config(
     page_title="Just Comment",
@@ -29,18 +29,32 @@ st.set_page_config(
 # Global model loading with caching
 ##########################################
 @st.cache_resource(show_spinner=False)
-def load_models():
-    """Load and cache all ML models"""
+def _load_models():
+    """Load and cache all ML models with optimized settings"""
     return {
+        # Emotion classification pipeline
         'emotion': pipeline(
-            "text-classification", 
-            model="Thea231/jhartmann_emotion_finetuning"
+            "text-classification",
+            model="Thea231/jhartmann_emotion_finetuning",
+            truncation=True  # Enable text truncation for long inputs
+        ),
+        
+        # Text generation components
+        'textgen_tokenizer': AutoTokenizer.from_pretrained(
+            "Qwen/Qwen1.5-0.5B",
+            use_fast=True  # Enable fast tokenization
+        ),
+        'textgen_model': AutoModelForCausalLM.from_pretrained(
+            "Qwen/Qwen1.5-0.5B",
+            torch_dtype=torch.float16  # Use half-precision for faster inference
         ),
-        'textgen_tokenizer': AutoTokenizer.from_pretrained("Qwen/Qwen1.5-0.5B"),
-        'textgen_model': AutoModelForCausalLM.from_pretrained("Qwen/Qwen1.5-0.5B"),
+        
+        # Text-to-speech 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"),
+        
+        # Preloaded speaker embeddings
         'speaker_embeddings': torch.tensor(
             load_dataset("Matthijs/cmu-arctic-xvectors", split="validation")[7306]["xvector"]
         ).unsqueeze(0)
@@ -49,96 +63,121 @@ def load_models():
 ##########################################
 # UI Components
 ##########################################
-def render_interface():
-    """Create user interface"""
-    st.title("Just Comment")
+def _display_interface():
+    """Render user interface elements"""
+    st.title("🚀 Just Comment")
     st.markdown("### I'm listening to you, my friend～")
     
     return st.text_area(
         "📝 Enter your comment:",
-        placeholder="Share your thoughts...",
+        placeholder="Type your message here...",
         height=150,
         key="user_input"
     )
 
 ##########################################
-# Core Logic Components
+# Core Processing Functions
 ##########################################
-def analyze_emotion(text, classifier):
-    """Determine emotion with quick analysis"""
+def _analyze_emotion(text, classifier):
+    """Identify dominant emotion with confidence threshold"""
     results = classifier(text, return_all_scores=True)[0]
-    valid_emotions = ['sadness', 'joy', 'love', 'anger', 'fear', 'surprise']
+    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):
-    """Complete prompt templates for all 6 emotions"""
+def _generate_prompt(text, emotion):
+    """Create structured prompts for all emotion types"""
     prompt_templates = {
         "sadness": (
-            "Customer expressed sadness: {input}\n"
-            "Respond with:\n1. Empathetic acknowledgment\n"
-            "2. Supportive statement\n3. Concrete help offer\nResponse:"
+            "Sadness detected: {input}\n"
+            "Required response structure:\n"
+            "1. Empathetic acknowledgment\n2. Support offer\n3. Solution proposal\n"
+            "Response:"
         ),
         "joy": (
-            "Positive feedback: {input}\n"
-            "Respond with:\n1. Enthusiastic thanks\n"
-            "2. Specific compliment\n3. Future engagement\nResponse:"
+            "Joy detected: {input}\n"
+            "Required response structure:\n"
+            "1. Enthusiastic thanks\n2. Positive reinforcement\n3. Future engagement\n"
+            "Response:"
         ),
         "love": (
-            "Customer showed affection: {input}\n"
-            "Respond with:\n1. Warm appreciation\n"
-            "2. Community building\n3. Exclusive offer\nResponse:"
+            "Affection detected: {input}\n"
+            "Required response structure:\n"
+            "1. Warm appreciation\n2. Community focus\n3. Exclusive benefit\n"
+            "Response:"
         ),
         "anger": (
-            "Angry complaint: {input}\n"
-            "Respond with:\n1. Sincere apology\n"
-            "2. Solution steps\n3. Compensation\nResponse:"
+            "Anger detected: {input}\n"
+            "Required response structure:\n"
+            "1. Sincere apology\n2. Action steps\n3. Compensation\n"
+            "Response:"
         ),
         "fear": (
-            "Customer expressed concerns: {input}\n"
-            "Respond with:\n1. Reassurance\n"
-            "2. Safety measures\n3. Support channels\nResponse:"
+            "Concern detected: {input}\n"
+            "Required response structure:\n"
+            "1. Reassurance\n2. Safety measures\n3. Support options\n"
+            "Response:"
         ),
         "surprise": (
-            "Unexpected feedback: {input}\n"
-            "Respond with:\n1. Acknowledge uniqueness\n"
-            "2. Creative solution\n3. Follow-up plan\nResponse:"
+            "Surprise detected: {input}\n"
+            "Required response structure:\n"
+            "1. Acknowledge uniqueness\n2. Creative solution\n3. Follow-up\n"
+            "Response:"
         )
     }
     return prompt_templates.get(emotion.lower(), "").format(input=text)
 
-def process_response(output_text):
-    """Optimized response processing"""
-    output_text = output_text.split("Response:")[-1].strip()
-    return output_text[:200]  # Strict length control
+def _process_response(raw_text):
+    """Clean and format generated response"""
+    # Extract text after last "Response:" marker
+    processed = raw_text.split("Response:")[-1].strip()
+    
+    # Remove incomplete sentences
+    if '.' in processed:
+        processed = processed.rsplit('.', 1)[0] + '.'
+    
+    # Ensure length between 50-200 characters
+    return processed[:200].strip() if len(processed) > 50 else "Thank you for your feedback. We value your input and will respond shortly."
 
-def generate_text_response(user_input, models):
-    """Efficient text generation"""
-    emotion = analyze_emotion(user_input, models['emotion'])
-    prompt = generate_prompt(user_input, emotion['label'])
+def _generate_text_response(input_text, models):
+    """Generate optimized text response with timing controls"""
+    # Emotion analysis
+    emotion = _analyze_emotion(input_text, models['emotion'])
+    
+    # Prompt engineering
+    prompt = _generate_prompt(input_text, emotion['label'])
     
-    inputs = models['textgen_tokenizer'](prompt, return_tensors="pt")
+    # Text generation with optimized parameters
+    inputs = models['textgen_tokenizer'](prompt, return_tensors="pt").to('cpu')
     outputs = models['textgen_model'].generate(
         inputs.input_ids,
-        max_new_tokens=150,  # Reduced for speed
+        max_new_tokens=100,  # Strict token limit
         temperature=0.7,
+        top_p=0.9,
         do_sample=True,
-        top_p=0.9
+        pad_token_id=models['textgen_tokenizer'].eos_token_id
     )
     
-    return process_response(
+    return _process_response(
         models['textgen_tokenizer'].decode(outputs[0], skip_special_tokens=True)
     )
 
-def generate_audio_response(text, models):
-    """Optimized TTS conversion"""
+def _generate_audio_response(text, models):
+    """Convert text to speech with performance optimizations"""
+    # Process text input
     inputs = models['tts_processor'](text=text, return_tensors="pt")
+    
+    # Generate spectrogram
     spectrogram = models['tts_model'].generate_speech(
         inputs["input_ids"],
         models['speaker_embeddings']
     )
-    with torch.no_grad():
+    
+    # Generate waveform with optimizations
+    with torch.no_grad():  # Disable gradient calculation
         waveform = models['tts_vocoder'](spectrogram)
+    
+    # Save audio file
     sf.write("response.wav", waveform.numpy(), samplerate=16000)
     return "response.wav"
 
@@ -146,21 +185,25 @@ def generate_audio_response(text, models):
 # Main Application Flow
 ##########################################
 def main():
-    ml_models = load_models()
-    user_input = render_interface()
+    """Primary execution flow"""
+    # Load models once
+    ml_models = _load_models()
+    
+    # Display interface
+    user_input = _display_interface()
     
     if user_input:
-        # Text Generation
-        with st.spinner("🔍 Analyzing emotions..."):
-            text_response = generate_text_response(user_input, ml_models)
+        # Text generation stage
+        with st.spinner("🔍 Analyzing emotions and generating response..."):
+            text_response = _generate_text_response(user_input, ml_models)
         
-        # Display Results
+        # Display results
         st.subheader("📄 Generated Response")
-        st.success(text_response)
+        st.markdown(f"```\n{text_response}\n```")  # f-string formatted output
         
-        # Audio Generation
-        with st.spinner("🔊 Generating voice..."):
-            audio_file = generate_audio_response(text_response, ml_models)
+        # Audio generation stage
+        with st.spinner("🔊 Converting to speech..."):
+            audio_file = _generate_audio_response(text_response, ml_models)
             st.audio(audio_file, format="audio/wav")
 
 if __name__ == "__main__":