Upload app.py

app.py

@@ -4,10 +4,7 @@ import torchaudio
 import pandas as pd
 import os
 import torch.nn as nn
-from transformers import (
-    Wav2Vec2ForCTC, Wav2Vec2Processor, AutoModel, AutoTokenizer,
-    WhisperProcessor, WhisperForConditionalGeneration
-)
+from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor, AutoModel, AutoTokenizer
 
 # Import các class mô hình từ file models.py
 from models import MultimodalClassifier, TextClassifier
@@ -24,18 +21,12 @@ LABELS_B = {0: "Đe dọa", 1: "Tức giận", 2: "Tiêu cực thông thường"
 # Đường dẫn (Tương đối với thư mục gốc của Space)
 MODEL_A_PATH = "saved_models/best_model_A.pth"
 MODEL_B_PATH = "saved_models/best_model_B.pth"
-FUZZY_RULES_PATH = "data/datafuzzy29d.csv"
-
-# --- Tải Mô hình STT (ĐÃ THAY ĐỔI SANG WHISPER) ---
-# === SỬA LỖI LẦN 2: Đã cập nhật tên model chính xác ===
-STT_MODEL_ID = "vinai/vinai-whisper-base" 
-print(f"Đang tải mô hình STT Whisper: {STT_MODEL_ID}...")
-# Cần 'language' và 'task' để bộ xử lý biết cách hoạt động
-audio_processor = WhisperProcessor.from_pretrained(STT_MODEL_ID, language="vi", task="transcribe")
-stt_model = WhisperForConditionalGeneration.from_pretrained(STT_MODEL_ID).to(device)
-
-# --- Tải các mô hình nền khác ---
-print("Đang tải mô hình PhoBERT...")
+FUZZY_RULES_PATH = "data/datafuzzy29d.csv" # Đảm bảo tên file này chính xác
+
+# Tải các mô hình nền (từ Hugging Face Hub)
+print("Đang tải các mô hình nền (STT, PhoBERT)...")
+audio_processor = Wav2Vec2Processor.from_pretrained("nguyenvulebinh/wav2vec2-base-vietnamese-250h")
+stt_model = Wav2Vec2ForCTC.from_pretrained("nguyenvulebinh/wav2vec2-base-vietnamese-250h").to(device)
 text_tokenizer = AutoTokenizer.from_pretrained("vinai/phobert-base")
 text_feature_extractor = AutoModel.from_pretrained("vinai/phobert-base").to(device)
 
@@ -59,6 +50,7 @@ try:
     fuzzy_rules_df = pd.read_csv(FUZZY_RULES_PATH, sep=';')
     fuzzy_rules = {}
     for _, row in fuzzy_rules_df.iterrows():
+        # Đảm bảo tên cột khớp với file CSV của bạn
         fuzzy_rules[(row['model_a_label'], row['model_b_label'])] = row['final_label']
     print(f"Đã tải {len(fuzzy_rules)} luật fuzzy.")
 except Exception as e:
@@ -67,47 +59,48 @@ except Exception as e:
 
 print("Tất cả mô hình đã sẵn sàng.")
 
-# --- 2. Định nghĩa Hàm Dự đoán (ĐÃ CẬP NHẬT) ---
+# --- 2. Định nghĩa Hàm Dự đoán ---
+# Hàm này sẽ được Gradio gọi mỗi khi người dùng nhấn "Submit"
 def predict_sentiment(audio_input):
     if audio_input is None:
         return "[Chưa có âm thanh]", "N/A", "N/A", "N/A"
 
     sample_rate, waveform_numpy = audio_input
+    
+    # Đảm bảo waveform là tensor float
     waveform = torch.from_numpy(waveform_numpy).float()
     
+    # Đảm bảo là 1D (mono) hoặc lấy kênh đầu tiên nếu là stereo
     if waveform.ndim > 1:
         waveform = waveform[0]
         
-    # --- Bước 1 & 2 (Gộp): STT và Đặc trưng Audio (Logic của Whisper) ---
+    # Thêm chiều batch (1,)
+    waveform = waveform.unsqueeze(0)
+
+    # --- Bước 1 & 2 (Gộp): STT và Đặc trưng Audio ---
     try:
-        # 1a. Resample (Whisper yêu cầu 16000 Hz)
+        # 1a. Resample
         if sample_rate != 16000:
             resampler = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=16000)
             waveform = resampler(waveform)
 
-        # 1b. Chuẩn bị input audio cho Whisper
-        # Không cần unsqueeze(0) vì processor tự xử lý
-        inputs = audio_processor(waveform, sampling_rate=16000, return_tensors="pt")
-        input_features = inputs.input_features.to(device)
+        # 1b. Chuẩn bị input audio
+        input_values = audio_processor(waveform.squeeze(), return_tensors="pt", sampling_rate=16000).input_values.to(device)
 
-        # 2a. Trích xuất Đặc trưng Audio (cho Model A)
-        # Chúng ta cần chạy encoder của Whisper để lấy hidden states
         with torch.no_grad():
-            encoder_outputs = stt_model.model.encoder(input_features)
-            # Lấy hidden state cuối cùng và tính trung bình
-            audio_feat_A = torch.mean(encoder_outputs.last_hidden_state, dim=1)
+            audio_outputs = stt_model(input_values, output_hidden_states=True)
 
-        # 2b. Trích xuất Văn bản (STT)
-        # Chạy hàm generate() để tạo token ID
-        with torch.no_grad():
-            predicted_ids = stt_model.generate(input_features, language="vi")
-        
-        # Giải mã token ID thành văn bản
-        transcribed_text = audio_processor.batch_decode(predicted_ids, skip_special_tokens=True)[0].lower()
+        # 2a. Trích xuất Văn bản (STT)
+        logits = audio_outputs.logits
+        predicted_ids = torch.argmax(logits, dim=-1)
+        transcribed_text = audio_processor.batch_decode(predicted_ids)[0].lower()
         
         if not transcribed_text:
             transcribed_text = "[Không nhận diện được giọng nói]"
 
+        # 2b. Trích xuất Đặc trưng Audio (cho Model A)
+        audio_feat_A = torch.mean(audio_outputs.hidden_states[-1], dim=1)
+
     except Exception as e:
         return f"[Lỗi xử lý audio: {e}]", "Lỗi Audio", "Lỗi Audio", "Lỗi Audio"
 
@@ -158,8 +151,10 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     
     with gr.Row():
         with gr.Column(scale=2):
+            # === BỔ SUNG TÍNH NĂNG ===
+            # Thêm "microphone" vào sources để cho phép ghi âm
             audio_in = gr.Audio(
-                sources=["upload", "microphone"],
+                sources=["upload", "microphone"],  # Cho phép cả tải lên và ghi âm
                 type="numpy", 
                 label="Tải lên tệp âm thanh hoặc Ghi âm"
             )
@@ -167,6 +162,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         
         with gr.Column(scale=3):
             gr.Markdown("### Kết quả Phân tích")
+            # Các ô output
             text_out = gr.Textbox(label="Văn bản được nhận diện (STT)")
             final_pred_out = gr.Label(label="Kết quả cuối cùng (Nguy cơ)")
             
@@ -174,14 +170,14 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                 pred_A_out = gr.Textbox(label="Dự đoán Model A (Đa phương tiện)")
                 pred_B_out = gr.Textbox(label="Dự đoán Model B (Chỉ văn bản)")
 
+    # Liên kết nút bấm với hàm dự đoán
     submit_btn.click(
         fn=predict_sentiment,
         inputs=audio_in,
         outputs=[text_out, pred_A_out, pred_B_out, final_pred_out]
     )
     
-    gr.Markdown("Lưu ý: Mô hình STT hiện đang sử dụng `vinai/vinai-whisper-base`.")
+    gr.Markdown("Lưu ý: Mô hình STT được tối ưu cho tiếng Việt.")
 
 print("Đang khởi chạy demo...")
-demo.launch()
-
+demo.launch() # Không cần (share=True) khi chạy trên Spaces