Update app.py

app.py

@@ -4,6 +4,7 @@ from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
 import torchaudio
 import os
 import re
+import jieba
 
 # Device setup
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -21,7 +22,7 @@ def transcribe_audio(audio_path):
     if duration > 60:
         results = []
         for start in range(0, int(duration), 50):
-            end = min(start + 60, int(duration))  # Fixed missing parenthesis
+            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)
@@ -48,14 +49,35 @@ def translate(text):
         translations.append(tokenizer.decode(outputs[0], skip_special_tokens=True))
     return " ".join(translations)
 
-# Load sentiment analysis model
+# Load quality rating model
 rating_pipe = pipeline("text-classification", model="Leo0129/CustomModel_dianping-chinese")
 
+def split_text(text, max_length=512):
+    words = list(jieba.cut(text))  
+    chunks, current_chunk = [], ""
+
+    for word in words:
+        if len(current_chunk) + len(word) < max_length:
+            current_chunk += word
+        else:
+            chunks.append(current_chunk)
+            current_chunk = word
+
+    if current_chunk:
+        chunks.append(current_chunk)
+
+    return chunks
+
 def rate_quality(text):
-    result = rating_pipe(text)[0]
-    label = result["label"].split("(")[0].strip().lower()
-    label_map = {"LABEL_0": "Poor", "LABEL_1": "Neutral", "LABEL_2": "Good"}
-    return label_map.get(label, "Unknown")
+    chunks = split_text(text)
+    results = []
+
+    for chunk in chunks:
+        result = rating_pipe(chunk)[0]
+        label_map = {"LABEL_0": "Poor", "LABEL_1": "Neutral", "LABEL_2": "Good"}
+        results.append(label_map.get(result["label"], "Unknown"))
+
+    return max(set(results), key=results.count)  # Return most frequent rating
 
 # Streamlit UI
 st.title("Cantonese Audio Analysis")