Update app.py

app.py

@@ -2,7 +2,7 @@
 # Step 0: Essential imports
 ##########################################
 import streamlit as st  # Web interface
-from transformers import (  # AI components
+from transformers import (  # AI components: emotion analysis, TTS, and text generation
     pipeline,
     SpeechT5Processor,
     SpeechT5ForTextToSpeech,
@@ -10,14 +10,15 @@ from transformers import (  # AI components
     AutoModelForCausalLM,
     AutoTokenizer
 )
-from datasets import load_dataset  # Voice data
-import torch  # Tensor operations
-import soundfile as sf  # Audio processing
+from datasets import load_dataset  # To load speaker embeddings dataset
+import torch  # For tensor operations
+import soundfile as sf  # For writing audio files
+import sentencepiece  # Required for SpeechT5Processor tokenization
 
 ##########################################
 # Initial configuration (MUST BE FIRST)
 ##########################################
-st.set_page_config(  # Set page config first
+st.set_page_config(  # Configure the web page
     page_title="Just Comment",
     page_icon="💬",
     layout="centered"
@@ -28,10 +29,10 @@ st.set_page_config(  # Set page config first
 ##########################################
 @st.cache_resource(show_spinner=False)
 def _load_components():
-    """Load and cache all models with hardware optimization"""
-    device = "cuda" if torch.cuda.is_available() else "cpu"
+    """Load and cache all models with hardware optimization."""
+    device = "cuda" if torch.cuda.is_available() else "cpu"  # Detect available device
     
-    # Emotion classifier (fast)
+    # Emotion classifier (fast and truncated)
     emotion_pipe = pipeline(
         "text-classification",
         model="Thea231/jhartmann_emotion_finetuning",
@@ -39,7 +40,7 @@ def _load_components():
         truncation=True
     )
     
-    # Text generator (optimized)
+    # Text generator (optimized with FP16 and auto device mapping)
     text_tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen1.5-0.5B")
     text_model = AutoModelForCausalLM.from_pretrained(
         "Qwen/Qwen1.5-0.5B",
@@ -58,7 +59,7 @@ def _load_components():
         torch_dtype=torch.float16
     ).to(device)
     
-    # Preloaded voice profile
+    # Preloaded voice profile for TTS
     speaker_emb = torch.tensor(
         load_dataset("Matthijs/cmu-arctic-xvectors", split="validation")[7306]["xvector"]
     ).unsqueeze(0).to(device)
@@ -78,10 +79,10 @@ def _load_components():
 # User interface components
 ##########################################
 def _show_interface():
-    """Render input interface"""
-    st.title("Just Comment")
-    st.markdown(f"### I'm listening to you, my friend～")
-    return st.text_area(  # Input field
+    """Render the input interface"""
+    st.title("🚀 Just Comment")  # Display the title with a rocket icon
+    st.markdown("### I'm listening to you, my friend～")  # Display the friendly subtitle
+    return st.text_area(  # Return user's comment input
         "📝 Enter your comment:",
         placeholder="Share your thoughts...",
         height=150,
@@ -92,37 +93,39 @@ def _show_interface():
 # Core processing functions
 ##########################################
 def _fast_emotion(text, analyzer):
-    """Rapid emotion detection with input limits"""
-    result = analyzer(text[:256], return_all_scores=True)[0]  # Limit input length
-    emotions = ['sadness', 'joy', 'love', 'anger', 'fear', 'surprise']
+    """Rapid emotion detection with input length limit."""
+    result = analyzer(text[:256], return_all_scores=True)[0]  # Analyze only the first 256 characters for speed
+    valid_emotions = ['sadness', 'joy', 'love', 'anger', 'fear', 'surprise']
+    # Select the emotion from valid ones or default to neutral
     return max(
-        (e for e in result if e['label'].lower() in emotions),
+        (e for e in result if e['label'].lower() in valid_emotions),
         key=lambda x: x['score'],
         default={'label': 'neutral', 'score': 0}
     )
 
 def _build_prompt(text, emotion):
-    """Template-based prompt engineering"""
+    """Template-based prompt engineering in continuous prose (no bullet points)."""
     templates = {
-        "sadness": f"Sadness detected: {{text}}\nRespond with: 1. Empathy 2. Support 3. Solution\nResponse:",
-        "joy": f"Joy detected: {{text}}\nRespond with: 1. Thanks 2. Praise 3. Engagement\nResponse:",
-        "love": f"Love detected: {{text}}\nRespond with: 1. Appreciation 2. Connection 3. Offer\nResponse:",
-        "anger": f"Anger detected: {{text}}\nRespond with: 1. Apology 2. Action 3. Compensation\nResponse:",
-        "fear": f"Fear detected: {{text}}\nRespond with: 1. Reassurance 2. Safety 3. Support\nResponse:",
-        "surprise": f"Surprise detected: {{text}}\nRespond with: 1. Acknowledgement 2. Solution 3. Follow-up\nResponse:",
-        "neutral": f"Feedback: {{text}}\nProfessional response:\n1. Acknowledgement\n2. Assistance\n3. Next steps\nResponse:"
+        "sadness": "I sensed sadness in your comment: {text}. We are truly sorry and are here to support you.",
+        "joy": "Your comment radiates joy: {text}. Thank you for your bright feedback; we look forward to serving you even better.",
+        "love": "Your message exudes love: {text}. We appreciate your heartfelt words and cherish our connection with you.",
+        "anger": "I understand your comment reflects anger: {text}. Please accept our sincere apologies as we work to resolve your concerns.",
+        "fear": "It seems you feel fear in your comment: {text}. We want to reassure you that your safety and satisfaction are our priority.",
+        "surprise": "Your comment conveys surprise: {text}. We are delighted by your experience and will strive to exceed your expectations.",
+        "neutral": "Thank you for your comment: {text}. We remain committed to providing you with outstanding service."
     }
-    return templates[emotion.lower()].format(text=text[:200])  # Input truncation
+    # Build and return a continuous prompt with the user comment truncated to 200 characters
+    return templates.get(emotion.lower(), templates["neutral"]).format(text=text[:200])
 
 def _generate_response(text, models):
-    """Optimized text generation pipeline"""
-    # Emotion detection
-    emotion = _fast_emotion(text, models["emotion"])
-    
-    # Prompt construction
-    prompt = _build_prompt(text, emotion["label"])
-    
-    # Generate text
+    """Optimized text generation pipeline using the detected emotion."""
+    # Detect the dominant emotion quickly
+    detected = _fast_emotion(text, models["emotion"])
+    # Build prompt based on detected emotion (continuous sentences)
+    prompt = _build_prompt(text, detected["label"])
+    print(f"Generated prompt: {prompt}")  # Print prompt using f-string for debugging
+
+    # Generate text using the Qwen model
     inputs = models["text_tokenizer"](
         prompt,
         return_tensors="pt",
@@ -130,65 +133,63 @@ def _generate_response(text, models):
         truncation=True
     ).to(models["device"])
     
+    # Generate the response ensuring balanced length (approximately 50-200 tokens)
     output = models["text_model"].generate(
         inputs.input_ids,
-        max_new_tokens=120,  # Balanced length
+        max_new_tokens=120,  # Upper bound tokens for answer
+        min_length=50,       # Lower bound to ensure completeness
         temperature=0.7,
         top_p=0.9,
         do_sample=True,
         pad_token_id=models["text_tokenizer"].eos_token_id
     )
     
-    # Process output
+    input_len = inputs.input_ids.shape[1]  # Determine the length of the prompt tokens
     full_text = models["text_tokenizer"].decode(output[0], skip_special_tokens=True)
+    # Extract only the generated portion after "Response:" if present
     response = full_text.split("Response:")[-1].strip()
-    
-    # Ensure completeness
-    if "." in response:
-        response = response.rsplit(".", 1)[0] + "."
-    return response[:200] or "Thank you for your feedback. We'll respond shortly."
+    print(f"Generated response: {response}")  # Debug print using f-string
+    # Return response ensuring it is within 50-200 words (approximation by character length here)
+    return response[:200]  # Truncate to 200 characters as an approximation
 
 def _text_to_speech(text, models):
-    """High-speed audio synthesis"""
+    """Efficiently synthesize speech for the given text."""
     inputs = models["tts_processor"](
-        text=text[:150],  # Limit text length
+        text=text[:150],  # Limit text length for TTS to 150 characters
         return_tensors="pt"
     ).to(models["device"])
     
-    with torch.inference_mode():  # Accelerated inference
+    with torch.inference_mode():  # Fast, no-grad inference
         spectrogram = models["tts_model"].generate_speech(
             inputs["input_ids"],
             models["speaker_emb"]
         )
         audio = models["tts_vocoder"](spectrogram)
     
-    sf.write("output.wav", audio.cpu().numpy(), 16000)
+    sf.write("output.wav", audio.cpu().numpy(), 16000)  # Save generated audio as .wav at 16kHz
     return "output.wav"
 
 ##########################################
 # Main application flow
 ##########################################
 def main():
-    """Primary execution controller"""
-    # Load components
-    components = _load_components()
-    
-    # Show interface
-    user_input = _show_interface()
+    """Primary execution controller."""
+    components = _load_components()  # Load all models and components
+    user_input = _show_interface()  # Render input interface and capture user comment
     
-    if user_input:
-        # Text generation
-        with st.spinner("🔍 Analyzing..."):
-            response = _generate_response(user_input, components)
-        
-        # Display result
-        st.subheader(f"📄 Response")
-        st.markdown(f"```\n{response}\n```")  # f-string formatted
-        
-        # Audio generation
-        with st.spinner("🔊 Synthesizing..."):
-            audio_path = _text_to_speech(response, components)
-            st.audio(audio_path, format="audio/wav")
+    if user_input:  # If a comment is provided
+        with st.spinner("🔍 Generating response..."):
+            generated_response = _generate_response(user_input, components)  # Generate response based on emotion and text
+        st.subheader("📄 Response")
+        st.markdown(
+            f"<p style='color:#3498DB; font-size:20px;'>{generated_response}</p>",
+            unsafe_allow_html=True
+        )  # Display the generated response in styled format
+        with st.spinner("🔊 Synthesizing audio..."):
+            audio_file = _text_to_speech(generated_response, components)  # Convert response to speech
+            st.audio(audio_file, format="audio/wav", start_time=0)  # Embed auto-playing audio player
+        print(f"Final generated response: {generated_response}")  # Debug output using f-string
 
+# Run the main function when the script is executed
 if __name__ == "__main__":
-    main()
+    main()  # Call the main function