Update app.py

app.py

@@ -1,201 +1,128 @@
 ##########################################
-# Step 0: 导入必需的库
+# Step 0: Import required libraries
 ##########################################
-import streamlit as st
-from transformers import pipeline, SpeechT5Processor, SpeechT5ForTextToSpeech, SpeechT5HifiGan, AutoModelForCausalLM, AutoTokenizer, pipeline
-from datasets import load_dataset
-from IPython.display import Audio, display
-import torch
-import soundfile as sf
-
-from google.colab import drive
-
-from huggingface_hub import login
-
-# Streamlit application title
-st.title("Comment reply for you")
-st.write("automative reply")
-
-# Text input for user to enter the comment
-text = st.text_area("Enter your comment", "")
+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
 
 ##########################################
-# Step 1：情感分析 - 分析用户评论的情感倾向
+# Streamlit application title and input
 ##########################################
+st.title("Comment Reply for You")  # Application title
+st.write("Generate automatic replies for user comments")  # Application description
+text = st.text_area("Enter your comment", "")  # Text input for user to enter comments
 
-
+##########################################
+# Step 1: Sentiment Analysis Function
+##########################################
 def analyze_dominant_emotion(user_review):
-    # Initialize the emotion classifier pipeline
+    """
+    Analyze the dominant emotion in the user's review using a text classification model.
+    """
     emotion_classifier = pipeline(
-         "text-classification",
-         model="Thea231/jhartmann_emotion_finetuning",
+        "text-classification", 
+        model="Thea231/jhartmann_emotion_finetuning", 
         return_all_scores=True
-     )
-
-    # Get emotion predictions for the input review
-    emotion_results = emotion_classifier(user_review)[0]  # Get first result (single input case)
-
-    # Extract the emotion with highest confidence score
-    dominant_emotion = max(emotion_results, key=lambda x: x['score'])
-
+    )  # Load pre-trained 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
 
-# 提取置信度最高的情感标签（可选）
-# dominant_emotion = analyze_dominant_emotion(user_review)
-# print("\n主导情感：", dominant_emotion['label'], f"(置信度: {dominant_emotion['score']:.2f})")
-
-
-
 ##########################################
-# Step 2：回复生成 - 根据情感生成自动回复
+# Step 2: Response Generation Function
 ##########################################
-
-
-def prompt_gen(user_review):
-  dominant_emotion = analyze_dominant_emotion(user_review)
-  emotion_strategies = {
-      "anger": {
-          "prompt": (
-              "Customer complaint: '{review}'\n\n"
-              "As a customer service representative, craft a professional response that:\n"
-              "- Begins with sincere apology and acknowledgment\n"
-              "- Clearly explains solution process with concrete steps\n"
-              "- Offers appropriate compensation/redemption\n"
-              "- Keeps tone humble and solution-focused (3-4 sentences)\n\n"
-              "Response:"
-          )
-      },
-      "disgust": {
-          "prompt": (
-              "Customer quality concern: '{review}'\n\n"
-              "As a customer service representative, craft a response that:\n"
-              "- Immediately acknowledges the product issue\n"
-              "- Explains quality control measures being taken\n"
-              "- Provides clear return/replacement instructions\n"
-              "- Offers goodwill gesture (3-4 sentences)\n\n"
-              "Response:"
-          )
-      },
-      "fear": {
-          "prompt": (
-              "Customer safety concern: '{review}'\n\n"
-              "As a customer service representative, craft a reassuring response that:\n"
-              "- Directly addresses the safety worries\n"
-              "- References relevant certifications/standards\n"
-              "- Offers dedicated support contact\n"
-              "- Provides satisfaction guarantee (3-4 sentences)\n\n"
-              "Response:"
-          )
-      },
-      "joy": {
-          "prompt": (
-              "Customer review: '{review}'\n\n"
-              "As a customer service representative, craft a concise response that:\n"
-              "- Specifically acknowledges both positive and constructive feedback\n"
-              "- Briefly mentions loyalty/referral programs\n"
-              "- Ends with shopping invitation (3-4 sentences)\n\n"
-              "Response:"
-          )
-      },
-      "neutral": {
-          "prompt": (
-              "Customer feedback: '{review}'\n\n"
-              "As a customer service representative, craft a balanced response that:\n"
-              "- Provides additional relevant product information\n"
-              "- Highlights key service features\n"
-              "- Politely requests more detailed feedback\n"
-              "- Maintains professional tone (3-4 sentences)\n\n"
-              "Response:"
-          )
-      },
-      "sadness": {
-          "prompt": (
-              "Customer disappointment: '{review}'\n\n"
-              "As a customer service representative, craft an empathetic response that:\n"
-              "- Shows genuine understanding of the issue\n"
-              "- Proposes personalized recovery solution\n"
-              "- Offers extended support options\n"
-              "- Maintains positive outlook (3-4 sentences)\n\n"
-              "Response:"
-          )
-      },
-      "surprise": {
-          "prompt": (
-              "Customer enthusiastic feedback: '{review}'\n\n"
-              "As a customer service representative, craft a response that:\n"
-              "- Matches customer's positive energy appropriately\n"
-              "- Highlights unexpected product benefits\n"
-              "- Invites to user community/events\n"
-              "- Maintains brand voice (3-4 sentences)\n\n"
-              "Response:"
-          )
-      }
-  }
-  # 生成回复Prompt
-  template = emotion_strategies[dominant_emotion['label'].lower()]["prompt"]
-  prompt = template.format(review=user_review)
-  print(prompt)
-  return prompt
-
-
-
 def response_gen(user_review):
-  prompt = prompt_gen(user_review)
-  # 加载Llama-3作为text generation模型
-  tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B")
-  model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B")
-  inputs = tokenizer(prompt, return_tensors="pt")
-  outputs = model.generate(**inputs, max_new_tokens=100)
-
-  input_length = inputs.input_ids.shape[1]
-  response = tokenizer.decode(outputs[0][input_length:], skip_special_tokens=True)
-  # print(response)
-  return response
-
-
+    """
+    Generate a response based on the sentiment of the user's review.
+    """
+    # Use Llama-based model to create a response based on a generated prompt
+    dominant_emotion = analyze_dominant_emotion(user_review)  # Get the dominant emotion
+    emotion_label = dominant_emotion['label'].lower()  # Extract emotion label
+    
+    # Define response templates for each emotion
+    emotion_prompts = {
+        "anger": (
+            "Customer complaint: '{review}'\n\n"
+            "As a customer service representative, write a response that:\n"
+            "- Sincerely apologizes for the issue\n"
+            "- Explains how the issue will be resolved\n"
+            "- Offers compensation where appropriate\n\n"
+            "Response:"
+        ),
+        "joy": (
+            "Customer review: '{review}'\n\n"
+            "As a customer service representative, write a positive response that:\n"
+            "- Thanks the customer for their feedback\n"
+            "- Acknowledges both positive and constructive comments\n"
+            "- Invites them to explore loyalty programs\n\n"
+            "Response:"
+        ),
+        # Add other emotions as needed...
+    }
+    
+    # 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 (replace 'meta-llama/Llama-3.2-1B' with an available model)
+    tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B")
+    model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B")
+    inputs = tokenizer(prompt, return_tensors="pt")  # Tokenize the prompt
+    outputs = model.generate(**inputs, max_new_tokens=100)  # Generate a response
+    
+    input_length = inputs.input_ids.shape[1]  # Length of the input text
+    response = tokenizer.decode(outputs[0][input_length:], skip_special_tokens=True)  # Decode the generated text
+    return response
 
 ##########################################
-# Step 3：语音生成 - 根据回复合成语音
+# Step 3: Text-to-Speech Conversion Function
 ##########################################
 def sound_gen(response):
-  # 加载模型和处理器
-  #processor = SpeechT5Processor.from_pretrained("microsoft/speecht5_tts")
-  #speech_model = SpeechT5ForTextToSpeech.from_pretrained("microsoft/speecht5_tts")
-  #vocoder = SpeechT5HifiGan.from_pretrained("microsoft/speecht5_hifigan")
-
-  processor = SpeechT5Processor.from_pretrained("microsoft/speecht5_tts")
-  model = SpeechT5ForTextToSpeech.from_pretrained("microsoft/speecht5_tts")
-  vocoder = SpeechT5HifiGan.from_pretrained("microsoft/speecht5_hifigan")
-
-  # 创建默认的说话人嵌入
-  embeddings_dataset = load_dataset("Matthijs/cmu-arctic-xvectors", split="validation")
-  speaker_embeddings = torch.tensor(embeddings_dataset[7306]["xvector"]).unsqueeze(0)  # 女性中性语音
-
-  # 文本预处理和语音合成
-  inputs = processor(text=response, return_tensors="pt")
-  spectrogram = model.generate_speech(inputs["input_ids"], speaker_embeddings)
-
-  # 使用声码器生成波形音频
-  with torch.no_grad():
-      speech = vocoder(spectrogram)
-
-  # 保存为WAV文件(16kHz采样率)
-  sf.write("customer_service_response.wav", speech.numpy(), samplerate=16000)
-
-  print("语音生成完成，已保存为 customer_service_response.wav")
-  return
-
-
-    st.write("customer_service_response.wav", speech.numpy(), samplerate=16000)
-    st.text("I wanna tell you that")
-    st.audio("customer_service_response.wav")
+    """
+    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")
+    model = SpeechT5ForTextToSpeech.from_pretrained("microsoft/speecht5_tts")
+    vocoder = SpeechT5HifiGan.from_pretrained("microsoft/speecht5_hifigan")
+    
+    # Load speaker embeddings (e.g., neutral female voice)
+    embeddings_dataset = load_dataset("Matthijs/cmu-arctic-xvectors", split="validation")
+    speaker_embeddings = torch.tensor(embeddings_dataset[7306]["xvector"]).unsqueeze(0)
+    
+    # Process the input text and generate a spectrogram
+    inputs = processor(text=response, return_tensors="pt")
+    spectrogram = model.generate_speech(inputs["input_ids"], speaker_embeddings)
+    
+    # Use the vocoder to generate a waveform
+    with torch.no_grad():
+        speech = vocoder(spectrogram)
+    
+    # Save the generated speech as a .wav file
+    sf.write("customer_service_response.wav", speech.numpy(), samplerate=16000)
+    st.audio("customer_service_response.wav")  # Play the audio in Streamlit
 
+##########################################
+# Main Function
+##########################################
 def main():
-  user_review = "I love the fast delivery, but the product quality could be better."
-  response = response_gen(user_review)
-  print(response)
-  sound_gen(response)
-  return
-
+    """
+    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
+
+# Run the main function
 if __name__ == "__main__":
-    main()
+    main()