Update app.py

app.py

@@ -3,69 +3,80 @@ import torch
 from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
 import torchaudio
 import os
+import re
 
-def load_models():
-    st.session_state.transcription_pipe = pipeline(
-        task="automatic-speech-recognition",
-        model="alvanlii/whisper-small-cantonese",
-        chunk_length_s=60,
-        device="cuda" if torch.cuda.is_available() else "cpu"
-    )
-    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("botisan-ai/mt5-translate-yue-zh")
-    st.session_state.translation_model = AutoModelForSeq2SeqLM.from_pretrained("botisan-ai/mt5-translate-yue-zh")
+# Device setup
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-    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")
+# Load Whisper model for transcription
+MODEL_NAME = "alvanlii/whisper-small-cantonese"
+language = "zh"
+pipe = pipeline(task="automatic-speech-recognition", model=MODEL_NAME, chunk_length_s=60, device=device)
+pipe.model.config.forced_decoder_ids = pipe.tokenizer.get_decoder_prompt_ids(language=language, task="transcribe")
 
 def transcribe_audio(audio_path):
-    pipe = st.session_state.transcription_pipe
+    waveform, sample_rate = torchaudio.load(audio_path)
+    duration = waveform.shape[1] / sample_rate
+    
+    if duration > 60:
+        results = []
+        for start in range(0, int(duration), 50):
+            end = min(start + 60, int(duration)
+            chunk = waveform[:, start * sample_rate:end * sample_rate]
+            temp_filename = f"temp_chunk_{start}.wav"
+            torchaudio.save(temp_filename, chunk, sample_rate)
+            result = pipe(temp_filename)["text"]
+            results.append(result)
+            os.remove(temp_filename)
+        return " ".join(results)
+    
     return pipe(audio_path)["text"]
 
-def translate_text(text):
-    tokenizer = st.session_state.translation_tokenizer
-    model = st.session_state.translation_model
-    inputs = tokenizer(text, return_tensors="pt")
-    outputs = model.generate(inputs["input_ids"], max_length=1000, num_beams=5)
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+# Load translation model
+tokenizer = AutoTokenizer.from_pretrained("botisan-ai/mt5-translate-yue-zh")
+model = AutoModelForSeq2SeqLM.from_pretrained("botisan-ai/mt5-translate-yue-zh").to(device)
+
+def split_sentences(text):
+    return [s for s in re.split(r'(?<=[。！？])', text) if s]
+
+def translate(text):
+    sentences = split_sentences(text)
+    translations = []
+    for sentence in sentences:
+        inputs = tokenizer(sentence, return_tensors="pt").to(device)
+        outputs = model.generate(inputs["input_ids"], max_length=1000, num_beams=5)
+        translations.append(tokenizer.decode(outputs[0], skip_special_tokens=True))
+    return " ".join(translations)
+
+# Load sentiment analysis model
+rating_pipe = pipeline("sentiment-analysis", model="uer/roberta-base-finetuned-dianping-chinese")
 
 def rate_quality(text):
-    result = st.session_state.rating_pipe(text)[0]
+    result = rating_pipe(text)[0]
     label = result["label"].split("(")[0].strip().lower()
-    # label = result["label"]
     label_map = {"negative": "Poor", "neutral": "Average", "positive": "Good"}
     return label_map.get(label, "Unknown")
 
-def main():
-    st.title("Audio Processing & Conversation Quality Rating")
+# Streamlit UI
+st.title("Cantonese Audio Analysis")
+st.write("Upload a Cantonese audio file to transcribe, translate, and rate the conversation quality.")
+
+uploaded_file = st.file_uploader("Upload Audio File", type=["wav", "mp3", "flac"])
+
+if uploaded_file is not None:
+    st.audio(uploaded_file, format="audio/wav")
+    temp_audio_path = "uploaded_audio.wav"
+    with open(temp_audio_path, "wb") as f:
+        f.write(uploaded_file.getbuffer())
     
-    if "transcription_pipe" not in st.session_state:
-        with st.spinner("Loading models..."):
-            load_models()
+    st.write("### Processing...")
+    transcript = transcribe_audio(temp_audio_path)
+    st.write("**Transcript:**", transcript)
     
-    uploaded_file = st.file_uploader("Upload an audio file", type=["wav", "mp3", "m4a"])
+    translated_text = translate(transcript)
+    st.write("**Translation:**", translated_text)
     
-    if uploaded_file is not None:
-        with st.spinner("Processing audio..."):
-            file_path = "temp_audio.wav"
-            with open(file_path, "wb") as f:
-                f.write(uploaded_file.read())
-            
-            transcript = transcribe_audio(file_path)
-            translation = translate_text(transcript)
-            rating = rate_quality(translation)
-            
-            os.remove(file_path)
-        
-        st.subheader("Transcription")
-        st.write(transcript)
-        
-        st.subheader("Translation (Chinese)")
-        st.write(translation)
-        
-        st.subheader("Conversation Quality Rating")
-        st.write(rating)
-
-if __name__ == "__main__":
-    main()
+    quality_rating = rate_quality(translated_text)
+    st.write("**Quality Rating:**", quality_rating)
+    
+    os.remove(temp_audio_path)