Update app.py

app.py

@@ -1,103 +1,148 @@
 ##########################################
 # Step 0: Import required libraries
 ##########################################
-import streamlit as st  # For building the web application
-from transformers import ( 
-    pipeline, 
-    SpeechT5Processor, 
-    SpeechT5ForTextToSpeech, 
-    SpeechT5HifiGan, 
-    AutoModelForCausalLM, 
-    AutoTokenizer 
-)  # For emotion analysis, text-to-speech, and text generation
-from datasets import load_dataset  # For loading datasets (e.g., speaker embeddings)
-import torch  # For tensor operations
-import soundfile as sf  # For saving audio as .wav files
+import streamlit as st  # Web app framework
+from transformers import (
+    pipeline,
+    SpeechT5Processor,
+    SpeechT5ForTextToSpeech,
+    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
 
 ##########################################
-# Streamlit application title and input
+# Initialize models and resources globally
 ##########################################
-st.title("🚀 Just Comment")  # Application title displayed to users
-st.write("I'm listening to you, my friend～")  # Application description for users
-text = st.text_area("Enter your comment", "")  # Text area for user input of comments
+@st.cache_resource  # Cache resources to reduce reload time
+def load_models():
+    """Load all required models once and cache them"""
+    return {
+        'emotion_classifier': pipeline(
+            "text-classification", 
+            model="Thea231/jhartmann_emotion_finetuning"
+        ),
+        'tts_processor': SpeechT5Processor.from_pretrained("microsoft/speecht5_tts"),
+        'tts_model': SpeechT5ForTextToSpeech.from_pretrained("microsoft/speecht5_tts"),
+        'tts_vocoder': SpeechT5HifiGan.from_pretrained("microsoft/speecht5_hifigan"),
+        'textgen_tokenizer': AutoTokenizer.from_pretrained("Qwen/Qwen1.5-0.5B"),
+        'textgen_model': AutoModelForCausalLM.from_pretrained("Qwen/Qwen1.5-0.5B"),
+        'speaker_embeddings': torch.tensor(
+            load_dataset("Matthijs/cmu-arctic-xvectors", split="validation")[7306]["xvector"]
+        ).unsqueeze(0)
+    }
 
 ##########################################
-# Step 1: Sentiment Analysis Function
+# Streamlit UI Configuration
 ##########################################
-def analyze_dominant_emotion(user_review):
-    """ Analyze the dominant emotion in the user's review using a text classification model. """
-    emotion_classifier = pipeline("text-classification", model="Thea231/jhartmann_emotion_finetuning", return_all_scores=True)  # Load emotion classification model
-    emotion_results = emotion_classifier(user_review)[0]  # Get emotion scores for the review
-    dominant_emotion = max(emotion_results, key=lambda x: x['score'])  # Find the emotion with the highest confidence
-    return dominant_emotion  # Return the dominant emotion (as a dict with label and score)
+def setup_ui():
+    """Configure Streamlit user interface"""
+    st.set_page_config(page_title="🚀 Just Comment", page_icon="💬")
+    st.title("🚀 Just Comment - Smart Response Generator")
+    st.markdown("""
+    <style>
+        .reportview-container {background: #f8f9fa;}
+        .stTextArea textarea {border: 2px solid #dee2e6;}
+    </style>
+    """, unsafe_allow_html=True)
+    return st.text_area("📝 Enter your customer comment:", "", height=150)
 
 ##########################################
-# Step 2: Response Generation Function
+# Enhanced Sentiment Analysis
 ##########################################
-def response_gen(user_review):
-    """ Generate a response based on the sentiment of the user's review. """
-    dominant_emotion = analyze_dominant_emotion(user_review)  # Get dominant emotion for the input
-    emotion_label = dominant_emotion['label'].lower()  # Extract emotion label
-
-    # Define response templates for each emotion
-    emotion_prompts = {
-        "anger": "I appreciate your feedback and apologize for the inconvenience caused by '{review}'. We're committed to resolving this issue promptly and will ensure it doesn't happen again. Thank you for your patience.",
-        "joy": "Thank you for your positive feedback on '{review}'! We're thrilled to hear you had a great experience and hope to serve you again soon.",
-        "disgust": "We regret that your experience with '{review}' did not meet our standards. We will take immediate steps to address this issue and appreciate your understanding.",
-        "fear": "Your safety is our priority. Regarding your concern about '{review}', we ensure that all our products meet strict safety standards. Please feel free to reach out for further assistance.",
-        "neutral": "Thank you for your feedback on '{review}'. We value your input and would love to hear more about your experience to improve our services.",
-        "sadness": "I'm sorry to hear that you were disappointed with '{review}'. We're here to help and would like to offer you a solution tailored to your needs.",
-        "surprise": "We're glad to hear that '{review}' exceeded your expectations! Thank you for sharing your excitement with us."
-    }
-
-    # Format the prompt with the user's review
-    prompt = emotion_prompts.get(emotion_label, "Neutral").format(review=user_review)
-
-    # Load a pre-trained text generation model
-    tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen1.5-0.5B")  # Load tokenizer
-    model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen1.5-0.5B")  # Load model
-    inputs = tokenizer(prompt, return_tensors="pt")  # Tokenize the prompt
-
-    outputs = model.generate(**inputs, max_new_tokens=100)  # Generate a response
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)  # Decode the generated text
-    return response.strip()[:200]  # Return a response trimmed to 200 characters
+def analyze_emotion(text, classifier):
+    """Determine dominant emotion with confidence threshold"""
+    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}
 
 ##########################################
-# Step 3: Text-to-Speech Conversion Function
+# Improved Response Generation
 ##########################################
-def sound_gen(response):
-    """ Convert the generated response to speech and save as a .wav file. """
-    # Load the pre-trained TTS models
-    processor = SpeechT5Processor.from_pretrained("microsoft/speecht5_tts")  # Load processor
-    model = SpeechT5ForTextToSpeech.from_pretrained("microsoft/speecht5_tts")  # Load TTS model
-    vocoder = SpeechT5HifiGan.from_pretrained("microsoft/speecht5_hifigan")  # Load vocoder
+def generate_response(text, models):
+    """Generate context-appropriate response with length control"""
+    emotion = analyze_emotion(text, models['emotion_classifier'])
+    prompt = create_prompt(text, emotion['label'].lower())
+    
+    inputs = models['textgen_tokenizer'](prompt, return_tensors="pt")
+    outputs = models['textgen_model'].generate(
+        **inputs,
+        max_new_tokens=200,  # Increased for longer responses
+        temperature=0.7,     # Balance creativity and focus
+        do_sample=True,
+        top_p=0.9,
+        no_repeat_ngram_size=2
+    )
+    
+    response = models['textgen_tokenizer'].decode(
+        outputs[0][inputs.input_ids.shape[1]:], 
+        skip_special_tokens=True
+    )
+    return postprocess_response(response)
 
-    # Load speaker embeddings (e.g., neutral female voice)
-    embeddings_dataset = load_dataset("Matthijs/cmu-arctic-xvectors", split="validation")  # Load dataset
-    speaker_embeddings = torch.tensor(embeddings_dataset[7306]["xvector"]).unsqueeze(0)  # Get speaker embeddings
+def create_prompt(text, emotion):
+    """Create emotion-specific prompts with structured guidance"""
+    templates = {
+        "anger": (
+            "Complaint: {input}\nRespond by:\n1. Apologizing sincerely\n"
+            "2. Proving solution steps\n3. Offering compensation\nResponse:"
+        ),
+        "joy": (
+            "Positive feedback: {input}\nRespond by:\n1. Thanking customer\n"
+            "2. Highlighting strengths\n3. Suggesting rewards\nResponse:"
+        ),
+        # Add other emotion templates...
+        "neutral": (
+            "Feedback: {input}\nRespond by:\n1. Acknowledging input\n"
+            "2. Providing information\n3. Requesting details\nResponse:"
+        )
+    }
+    return templates.get(emotion, templates['neutral']).format(input=text)
 
-    # Process the input text and generate a spectrogram
-    inputs = processor(text=response, return_tensors="pt")  # Process the text
-    spectrogram = model.generate_speech(inputs["input_ids"], speaker_embeddings)  # Generate spectrogram
+def postprocess_response(text):
+    """Ensure response quality and length"""
+    text = text.split("\n\n")[0].strip()  # Take first complete response
+    if len(text) < 50:  # Minimum length check
+        return "Thank you for your feedback. We'll carefully review your comments and follow up shortly."
+    return text[:300]  # Hard length limit
 
-    # Use the vocoder to generate a waveform
+##########################################
+# Optimized Text-to-Speech
+##########################################
+def generate_speech(text, models):
+    """Convert text to speech with performance optimizations"""
+    inputs = models['tts_processor'](text=text, return_tensors="pt")
+    spectrogram = models['tts_model'].generate_speech(
+        inputs["input_ids"], 
+        models['speaker_embeddings']
+    )
     with torch.no_grad():
-        speech = vocoder(spectrogram)  # Generate speech waveform
-
-    # Save the generated speech as a .wav file
-    sf.write("customer_service_response.wav", speech.numpy(), samplerate=16000)  # Save audio
-    st.audio("customer_service_response.wav")  # Play the audio in Streamlit
+        audio = models['tts_vocoder'](spectrogram)
+    sf.write("response.wav", audio.numpy(), 16000)
+    return "response.wav"
 
 ##########################################
-# Main Function
+# Main Application Logic
 ##########################################
 def main():
-    """ Main function to orchestrate the workflow of sentiment analysis, response generation, and text-to-speech. """
-    if text:  # Check if the user entered a comment
-        response = response_gen(text)  # Generate a response
-        st.write(f"Generated response: {response}")  # Display the generated response
-        sound_gen(response)  # Convert the response to speech and play it
+    """Main execution flow"""
+    models = load_models()  # Load models once
+    user_input = setup_ui()
+    
+    if user_input:
+        with st.spinner("🔍 Analyzing sentiment and generating response..."):
+            response = generate_response(user_input, models)
+            
+        st.subheader("💡 Generated Response:")
+        st.markdown(f"```\n{response}\n```")
+        
+        with st.spinner("🔊 Generating voice response..."):
+            audio_file = generate_speech(response, models)
+            st.audio(audio_file, format="audio/wav")
 
-# Run the main function
 if __name__ == "__main__":
-    main()  # Execute the main function
+    main()