Update app.py

app.py

@@ -15,8 +15,9 @@ def load_models():
     
     st.session_state.translation_tokenizer = AutoTokenizer.from_pretrained("botisan-ai/mt5-translate-yue-zh")
     st.session_state.translation_model = AutoModelForSeq2SeqLM.from_pretrained("botisan-ai/mt5-translate-yue-zh")
-    
-    st.session_state.rating_pipe = pipeline("text-classification", model="jackietung/bert-base-chinese-finetuned-sentiment")
+
+    st.session_state.rating_pipe = pipeline("sentiment-analysis", model="uer/roberta-base-finetuned-dianping-chinese")
+    # st.session_state.rating_pipe = pipeline("text-classification", model="jackietung/bert-base-chinese-finetuned-sentiment")
 
 def transcribe_audio(audio_path):
     pipe = st.session_state.transcription_pipe
@@ -31,7 +32,8 @@ def translate_text(text):
 
 def rate_quality(text):
     result = st.session_state.rating_pipe(text)[0]
-    label = result["label"]
+    label = result["label"].split("(")[0].strip().lower()
+    # label = result["label"]
     label_map = {"負面": "Poor", "中性": "Average", "正面": "Good"}
     return label_map.get(label, "Unknown")