xujinheng666 commited on
Commit
3f91072
·
verified ·
1 Parent(s): 8a8fcb6

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +4 -3
app.py CHANGED
@@ -16,7 +16,7 @@ def load_models():
16
  st.session_state.translation_tokenizer = AutoTokenizer.from_pretrained("botisan-ai/mt5-translate-yue-zh")
17
  st.session_state.translation_model = AutoModelForSeq2SeqLM.from_pretrained("botisan-ai/mt5-translate-yue-zh")
18
 
19
- st.session_state.rating_pipe = pipeline("sentiment-analysis", model="uer/roberta-base-finetuned-dianping-chinese")
20
 
21
  def transcribe_audio(audio_path):
22
  pipe = st.session_state.transcription_pipe
@@ -31,8 +31,9 @@ def translate_text(text):
31
 
32
  def rate_quality(text):
33
  result = st.session_state.rating_pipe(text)[0]
34
- label_map = {"LABEL_0": "Poor", "LABEL_1": "Average", "LABEL_2": "Good"}
35
- return label_map.get(result["label"], "Unknown")
 
36
 
37
  def main():
38
  st.title("Audio Processing & Conversation Quality Rating")
 
16
  st.session_state.translation_tokenizer = AutoTokenizer.from_pretrained("botisan-ai/mt5-translate-yue-zh")
17
  st.session_state.translation_model = AutoModelForSeq2SeqLM.from_pretrained("botisan-ai/mt5-translate-yue-zh")
18
 
19
+ st.session_state.rating_pipe = pipeline("text-classification", model="uer/roberta-base-finetuned-dianping-chinese")
20
 
21
  def transcribe_audio(audio_path):
22
  pipe = st.session_state.transcription_pipe
 
31
 
32
  def rate_quality(text):
33
  result = st.session_state.rating_pipe(text)[0]
34
+ label = result["label"].lower()
35
+ label_map = {"negative": "Poor", "neutral": "Average", "positive": "Good"}
36
+ return label_map.get(label, "Unknown")
37
 
38
  def main():
39
  st.title("Audio Processing & Conversation Quality Rating")