Update app.py

app.py

@@ -13,12 +13,10 @@ def load_models():
     )
     st.session_state.transcription_pipe.model.config.forced_decoder_ids = st.session_state.transcription_pipe.tokenizer.get_decoder_prompt_ids(language="zh", task="transcribe")
     
-    st.session_state.translation_tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-zh-en")
-    st.session_state.translation_model = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-zh-en")
+    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.summary_pipe = pipeline("summarization", model="Falconsai/text_summarization")
-    
-    st.session_state.rating_pipe = pipeline("text-classification", model="distilbert/distilbert-base-uncased-finetuned-sst-2-english")
+    st.session_state.rating_pipe = pipeline("text-classification", model="uer/roberta-base-finetuned-dianping-chinese")
 
 def transcribe_audio(audio_path):
     pipe = st.session_state.transcription_pipe
@@ -31,9 +29,6 @@ def translate_text(text):
     outputs = model.generate(inputs["input_ids"], max_length=1000, num_beams=5)
     return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-def summarize_text(text):
-    return st.session_state.summary_pipe(text)[0]['summary_text']
-
 def rate_quality(text):
     result = st.session_state.rating_pipe(text)[0]
     label_map = {"LABEL_0": "Poor", "LABEL_1": "Average", "LABEL_2": "Good"}
@@ -56,7 +51,6 @@ def main():
             
             transcript = transcribe_audio(file_path)
             translation = translate_text(transcript)
-            summary = summarize_text(translation)
             rating = rate_quality(translation)
             
             os.remove(file_path)
@@ -64,12 +58,9 @@ def main():
         st.subheader("Transcription")
         st.write(transcript)
         
-        st.subheader("Translation (English)")
+        st.subheader("Translation (Chinese)")
         st.write(translation)
         
-        st.subheader("Summary")
-        st.write(summary)
-        
         st.subheader("Conversation Quality Rating")
         st.write(rating)