Update app.py

app.py

@@ -1,6 +1,7 @@
+import os
 import torch
 import torchaudio
-import os
+import numpy as np
 import re
 import streamlit as st
 from difflib import SequenceMatcher
@@ -15,47 +16,90 @@ language = "zh"
 pipe = pipeline(
     task="automatic-speech-recognition",
     model=MODEL_NAME,
-    chunk_length_s=60,
-    device=device
+    chunk_length_s=30,
+    device=device,
+    generate_kwargs={
+        "no_repeat_ngram_size": 4,
+        "repetition_penalty": 1.15,
+        "temperature": 0.5,
+        "top_p": 0.97,
+        "top_k": 40,
+        "max_new_tokens": 300,
+        "do_sample": True
+    }
 )
 pipe.model.config.forced_decoder_ids = pipe.tokenizer.get_decoder_prompt_ids(language=language, task="transcribe")
 
 # Load quality rating model
-rating_pipe = pipeline("text-classification", model="tabularisai/multilingual-sentiment-analysis")
+rating_pipe = pipeline("text-classification", model="tabularisai/multilingual-sentiment-analysis", device=device)
+
+def is_similar(a, b, threshold=0.8):
+    return SequenceMatcher(None, a, b).ratio() > threshold
 
-# Sentiment label mapping
-label_map = {"Negative": "Very Poor", "Neutral": "Neutral", "Positive": "Very Good"}
+def remove_repeated_phrases(text):
+    sentences = re.split(r'(?<=[。！？])', text)
+    cleaned_sentences = []
+    for sentence in sentences:
+        if not cleaned_sentences or not is_similar(sentence.strip(), cleaned_sentences[-1].strip()):
+            cleaned_sentences.append(sentence.strip())
+    return " ".join(cleaned_sentences)
 
 def remove_punctuation(text):
     return re.sub(r'[^\w\s]', '', text)
 
 def transcribe_audio(audio_path):
-    transcript = pipe(audio_path)["text"]
-    return remove_punctuation(transcript)
+    waveform, sample_rate = torchaudio.load(audio_path)
+    
+    if waveform.shape[0] > 1:
+        waveform = torch.mean(waveform, dim=0, keepdim=True)
+    
+    waveform = waveform.squeeze(0).numpy()
+    duration = waveform.shape[0] / sample_rate
+    
+    if duration > 60:
+        chunk_size = sample_rate * 55
+        step_size = sample_rate * 50
+        results = []
+        
+        for start in range(0, waveform.shape[0], step_size):
+            chunk = waveform[start:start + chunk_size]
+            if chunk.shape[0] == 0:
+                break
+            transcript = pipe({"sampling_rate": sample_rate, "raw": chunk})["text"]
+            results.append(remove_punctuation(transcript))
+        
+        return remove_punctuation(remove_repeated_phrases(" ".join(results)))
+    
+    return remove_punctuation(remove_repeated_phrases(pipe({"sampling_rate": sample_rate, "raw": waveform})["text"]))
 
 def rate_quality(text):
-    result = rating_pipe(text)[0]
-    return label_map.get(result["label"], "Unknown")
+    chunks = [text[i:i+512] for i in range(0, len(text), 512)]
+    results = rating_pipe(chunks, batch_size=4)
+    
+    label_map = {"Very Negative": "Very Poor", "Negative": "Poor", "Neutral": "Neutral", "Positive": "Good", "Very Positive": "Very Good"}
+    processed_results = [label_map.get(res["label"], "Unknown") for res in results]
+    
+    return max(set(processed_results), key=processed_results.count)
 
 # Streamlit UI
-st.set_page_config(page_title="Cantonese Audio Transcription & Analysis", layout="centered")
-st.title("🗣️ Cantonese Audio Transcriber & Sentiment Analyzer")
-st.markdown("Upload your Cantonese audio file, and we will transcribe and analyze its sentiment.")
+st.title("Audio Transcription and Quality Rating")
+
+uploaded_file = st.file_uploader("Upload an audio file", type=["wav", "mp3", "flac"])
 
-uploaded_file = st.file_uploader("Upload an audio file (WAV, MP3, etc.)", type=["wav", "mp3", "m4a"])
 if uploaded_file is not None:
-    with st.spinner("Processing audio..."):
-        temp_audio_path = "temp_audio.wav"
-        with open(temp_audio_path, "wb") as f:
-            f.write(uploaded_file.getbuffer())
-        transcript = transcribe_audio(temp_audio_path)
-        sentiment = rate_quality(transcript)
-        os.remove(temp_audio_path)
+    st.audio(uploaded_file, format="audio/wav")
+    
+    temp_audio_path = "temp_audio.wav"
+    with open(temp_audio_path, "wb") as f:
+        f.write(uploaded_file.read())
     
-    st.subheader("Transcription")
-    st.text_area("", transcript, height=150)
+    st.write("Processing audio...")
+    transcript = transcribe_audio(temp_audio_path)
+    st.subheader("Transcript")
+    st.write(transcript)
     
-    st.subheader("Sentiment Analysis")
-    st.markdown(f"### 🎭 Sentiment: **{sentiment}**")
+    quality_rating = rate_quality(transcript)
+    st.subheader("Quality Rating")
+    st.write(quality_rating)
     
-    st.success("Processing complete! 🎉")
+    os.remove(temp_audio_path)