Update app.py

app.py

@@ -1,6 +1,6 @@
 # -*- coding: utf-8 -*-
-"""RobotsMali_ASR_Demo.ipynb - Script Final pour Démo Fluide et Stable
-Version corrigée du SyntaxError.
+"""RobotsMali_ASR_Demo.ipynb - Script FINAL
+Traitement complet de l'audio sans découpage, avec barre de progression Gradio et post-correction.
 """
 import gradio as gr
 import time
@@ -11,6 +11,7 @@ import numpy as np
 
 # --- IMPORTS NEMO ---
 import nemo.collections.asr as nemo_asr
+import nemo.collections.nlp as nemo_nlp
 # --------------------
 
 # ----------------------------------------------------------------------
@@ -20,77 +21,79 @@ ROBOTSMALI_MODELS = [
     "RobotsMali/soloba-ctc-0.6b-v0",  
     "RobotsMali/soloni-114m-tdt-ctc-v1",
     "RobotsMali/soloni-114m-tdt-ctc-V0",
-    "RobotsMali/stt-bm-quartznet5x5-V0",
+    "RobotsMali/stt-bm-quartznet5x5-V0", # Modèles souvent en erreur (selon les logs), mais inclus.
     "RobotsMali/stt-bm-quartznet5x5-v1",
     "RobotsMali/soloba-ctc-0.6b-v1"
 ]
 
-CHUNK_DURATION_SEC = 25  # Durée par segment (secondes) pour économiser la RAM
-SR_TARGET = 16000        # Taux d'échantillonnage cible pour NeMo ASR (16kHz)
+SR_TARGET = 16000 # Taux d'échantillonnage cible pour NeMo ASR (16kHz)
 
-# Cache pour stocker les modèles NeMo chargés.
+# Modèle de post-traitement pour restaurer la ponctuation et la casse
+PUNCT_MODEL_NAME = "nemo/nlp/punctuation_and_capitalization" 
+
+# Caches
 asr_pipelines = {}
+punct_pipeline = None
 
 # ----------------------------------------------------------------------
-# 1. FONCTIONS DE GESTION DES MODÈLES (CHARGEMENT + WARM-UP)
+# 1. FONCTIONS DE GESTION DES MODÈLES (CHARGEMENT & CACHE)
 # ----------------------------------------------------------------------
 def load_pipeline(model_name):
-    """
-    Charge le modèle NeMo, le met en cache et effectue un warm-up.
-    """
+    """Charge un modèle ASR NeMo, le met en cache et effectue un warm-up."""
     if model_name not in asr_pipelines:
         print(f"-> Tentative de chargement du modèle NeMo: {model_name}...")
         temp_warmup_file = "dummy_warmup.wav"
 
         try:
-            # 🚀 CHARGEMENT NEMO
             model_instance = nemo_asr.models.ASRModel.from_pretrained(model_name=model_name)
             model_instance.eval() 
-
             asr_pipelines[model_name] = model_instance
             print(f"-> Modèle NeMo {model_name} chargé avec succès.")
 
-            # ----------------------------------------------------
-            # WARM-UP (Inférence à blanc)
-            # ----------------------------------------------------
+            # WARM-UP
             print(f"   [Warmup] Exécution d'une inférence à blanc...")
-            
             dummy_audio = np.random.randn(SR_TARGET).astype(np.float32) 
             sf.write(temp_warmup_file, dummy_audio, SR_TARGET)
-
             model_instance.transcribe([temp_warmup_file], batch_size=1)
-            
             print(f"   [Warmup] Terminé.")
 
         except Exception as e:
-            if model_name in asr_pipelines:
-                 del asr_pipelines[model_name]
+            if model_name in asr_pipelines: del asr_pipelines[model_name]
             print(f"!!! Erreur de chargement NeMo pour {model_name}: {e}")
             raise RuntimeError(f"Impossible de charger le modèle {model_name}. Détail: {e}")
             
         finally:
-            if os.path.exists(temp_warmup_file):
-                os.remove(temp_warmup_file)
+            if os.path.exists(temp_warmup_file): os.remove(temp_warmup_file)
             
     return asr_pipelines.get(model_name)
 
+def load_punct_model():
+    """Charge le modèle de ponctuation/casse et le met en cache."""
+    global punct_pipeline
+    if punct_pipeline is None:
+        print(f"-> Tentative de chargement du modèle de ponctuation: {PUNCT_MODEL_NAME}...")
+        try:
+            punct_pipeline = nemo_nlp.models.PunctuationCapitalizationModel.from_pretrained(model_name=PUNCT_MODEL_NAME)
+            punct_pipeline.eval()
+            print("-> Modèle de ponctuation chargé avec succès.")
+        except Exception as e:
+            print(f"!!! AVERTISSEMENT: Échec du chargement du modèle de ponctuation {PUNCT_MODEL_NAME}. La sortie restera brute. Détail: {e}")
+    return punct_pipeline
+
 # ----------------------------------------------------------------------
-# 2. FONCTION PRINCIPALE D'INFÉRENCE AVEC STREAMING ET DÉCOUPAGE
+# 2. FONCTION PRINCIPALE D'INFÉRENCE (TRAITEMENT COMPLET AVEC PROGRESSION)
 # ----------------------------------------------------------------------
 def transcribe_audio(model_name: str, audio_path: str):
     """
-    Effectue la transcription ASR avec découpage (chunking) et streaming d'état.
+    Effectue la transcription ASR de l'audio complet avec une barre de progression simulée.
     """
     if audio_path is None:
         yield "⚠️ Veuillez d'abord télécharger ou enregistrer un fichier audio."
         return
-    if not ROBOTSMALI_MODELS:
-        yield "Liste de modèles ASR indisponible."
-        return
     
     start_time = time.time()
     model_short_name = model_name.split('/')[-1]
-    temp_chunk_paths = []
+    temp_full_path = f"temp_nemo_input_{os.path.basename(audio_path)}.wav"
 
     try:
         # ----------------------------------------------------------------
@@ -99,109 +102,109 @@ def transcribe_audio(model_name: str, audio_path: str):
         yield f"**[1/4] CHARGEMENT AUDIO...** Préparation du fichier original (Mono @ 16kHz). ⚙️"
         
         full_audio_data, sr = librosa.load(audio_path, sr=SR_TARGET, mono=True)
-        
         total_duration = len(full_audio_data) / SR_TARGET
-        samples_per_chunk = int(CHUNK_DURATION_SEC * SR_TARGET)
+        
+        # Correction de la forme audio (squeeze) pour éviter l'erreur de "Output shape mismatch"
+        segment_data = full_audio_data.squeeze() 
+        sf.write(temp_full_path, segment_data, SR_TARGET)
         
         # ----------------------------------------------------------------
-        # ÉTAPE 2 : CHARGEMENT/VÉRIFICATION DU MODÈLE ET DÉCOUPAGE
+        # ÉTAPE 2 : CHARGEMENT/VÉRIFICATION DU MODÈLE
         # ----------------------------------------------------------------
-        yield f"**[2/4] PRÉ-CALCUL...** Chargement du modèle et découpage ({total_duration:.1f}s en segments de {CHUNK_DURATION_SEC}s). 🧠"
+        yield f"**[2/4] PRÉ-CALCUL...** Chargement du modèle. Durée de l'audio : {total_duration:.1f}s. 🧠"
         
         asr_model = load_pipeline(model_name)
         
-        # Logique de DÉCOUPAGE
-        audio_segments = []
-        for i in range(0, len(full_audio_data), samples_per_chunk):
-            audio_segments.append(full_audio_data[i:i + samples_per_chunk])
-            
-        num_chunks = len(audio_segments)
-        full_transcription_text = ""
-        
         # ----------------------------------------------------------------
-        # ÉTAPE 3 : TRANSCRIPTION PAR SEGMENT
+        # ÉTAPE 3 : TRANSCRIPTION COMPLÈTE (AVEC BARRE DE PROGRESSION)
         # ----------------------------------------------------------------
+        yield f"**[3/4] TRANSCRIPTION EN COURS...** Démarrage de l'inférence. ⏳" 
+
+        # --- BARRE DE PROGRESSION SIMULÉE ---
+        # Affiche une progression visuelle pendant l'attente de l'inférence GPU
+        for progress_percent in range(0, 91, 10):
+            time.sleep(0.3) 
+            # Utilise gr.Progress pour une barre stylée en haut de l'interface
+            yield gr.Progress(progress_percent, total=100, desc=f"Progression ASR ({progress_percent}%)")
+
+        yield f"**[3/4] FINALISATION...** Inférence en cours sur le GPU. 🚀" 
+        # ---------------------------------------------
         
-        for idx, segment_data in enumerate(audio_segments):
-            
-            yield f"**[3/4] TRANSCRIPTION EN COURS...** Analyse du segment {idx + 1}/{num_chunks}. ⏳" 
-            
-            # --- CORRECTION DE LA FORME AUDIO (squeeze) ---
-            segment_data = segment_data.squeeze() 
-            
-            # Écriture du chunk temporaire
-            chunk_path = f"{os.path.splitext(os.path.basename(audio_path))[0]}_chunk_{idx}.wav"
-            sf.write(chunk_path, segment_data, SR_TARGET)
-            temp_chunk_paths.append(chunk_path)
+        # 🚀 INFÉRENCE NEMO
+        transcriptions = asr_model.transcribe([temp_full_path], batch_size=1)
+        
+        # --- GESTION DE L'OBJET HYPOTHESIS ---
+        transcription_text_final = ""
+        if transcriptions and transcriptions[0]:
+            hyp_object = transcriptions[0] 
             
-            # 🚀 INFÉRENCE NEMO
-            transcriptions = asr_model.transcribe([chunk_path], batch_size=1)
+            # Gère les différents formats de sortie de NeMo
+            if hasattr(hyp_object, 'text'):
+                transcription_text_final = hyp_object.text.strip()
+            elif isinstance(hyp_object, str):
+                transcription_text_final = hyp_object.strip()
+            elif isinstance(hyp_object, list) and hasattr(hyp_object[0], 'text'):
+                transcription_text_final = hyp_object[0].text.strip()
             
-            # --- GESTION DE L'OBJET HYPOTHESIS ---
-            segment_text = ""
-            if transcriptions and transcriptions[0]:
-                hyp_object = transcriptions[0] 
-                
-                if hasattr(hyp_object, 'text'):
-                    segment_text = hyp_object.text.strip()
-                elif isinstance(hyp_object, str):
-                    segment_text = hyp_object.strip()
-                elif isinstance(hyp_object, list) and hasattr(hyp_object[0], 'text'):
-                    segment_text = hyp_object[0].text.strip()
-                
-            if not segment_text:
-                segment_text = "[Transcription vide]"
-                
-            full_transcription_text += segment_text + "\n\n" 
+        if not transcription_text_final:
+            transcription_text_final = "[Transcription vide ou échec ASR]"
             
+        
         # ----------------------------------------------------
-        # ÉTAPE 4 : RÉSULTAT FINAL
+        # ÉTAPE 4 : POST-TRAITEMENT ET AFFICHAGE FINAL
         # ----------------------------------------------------
         end_time = time.time()
         duration = end_time - start_time
-        
-        transcription_text_final = full_transcription_text.strip()
-        
+        processed_text = transcription_text_final
+
+        # --- POST-TRAITEMENT (PONCTUATION & CASSE) ---
+        punct_model = load_punct_model()
+        if punct_model and transcription_text_final != "[Transcription vide ou échec ASR]":
+            yield f"**[4/4] POST-TRAITEMENT...** Correction de la ponctuation et de la casse pour la lisibilité. ✨"
+            yield gr.Progress(100, total=100, desc="Progression ASR (100%)") # Termine la barre
+
+            try:
+                corrected_list = punct_model.add_punctuation_capitalization([transcription_text_final])
+                if corrected_list:
+                    processed_text = corrected_list[0].strip()
+            except Exception as pc_error:
+                print(f"!!! Échec du post-traitement de ponctuation : {pc_error}")
+                yield "⚠️ Échec de la correction de ponctuation. Affichage du texte brut."
+
         # 1. EN-TÊTE D'INFORMATION
         output = f"**Modèle Utilisé :** `{model_short_name}` (NeMo)\n"
         output += f"**Durée de l'Audio :** {total_duration:.1f} secondes\n"
         output += f"**Temps de Traitement Total :** {duration:.2f} secondes\n"
-        output += f"**DÉCOUPAGE :** {CHUNK_DURATION_SEC} secondes ({num_chunks} segments)\n"
         output += f"***\n"
         
         # 2. PRÉSENTATION LYRICS PROPRE
         output += "**RÉSULTAT DE LA TRANSCRIPTION (Lyrics) :**\n"
-        # Préparation du texte pour le Markdown (Remplacement avant le yield)
-        formatted_lyrics = transcription_text_final.replace('\n\n', '\n>>> ')
-        output += f">>> {formatted_lyrics}"
+        # Formatage du texte pour l'affichage Markdown
+        formatted_lyrics = processed_text.replace('\n', ' ').strip().replace('. ', '.\n\n>>> ').replace('? ', '?\n\n>>> ')
+        if not formatted_lyrics.startswith('>>> '):
+            formatted_lyrics = '>>> ' + formatted_lyrics
+
+        output += formatted_lyrics 
         
         # 3. NOTE FINALE
-        output += "\n\n*Note : Audio converti en **Mono @ 16kHz** pour la transcription.*"
-        
+        output += "\n\n*Traitement complet de l'audio sans découpage (chunking).* "
+
         yield output
     
     except RuntimeError as e:
         yield f"❌ Erreur critique lors du chargement : {str(e)}"
     
     except Exception as e:
-        # --- CORRECTION DE SYNTAXE APPLIQUÉE ICI ---
-        # Affiche le texte partiel en cas d'erreur
-        if 'full_transcription_text' in locals() and full_transcription_text:
-             partial_text = full_transcription_text.strip().replace('\n\n', '\n>>> ')
-             yield f"❌ Erreur lors de la transcription, le traitement s'est arrêté. Texte partiel:\n>>> {partial_text}"
-             
-        yield f"❌ Erreur générale : {e}"
+        yield f"❌ Erreur générale lors de la transcription complète : {e}"
         
     finally:
         # Nettoyage
-        for chunk_path in temp_chunk_paths:
-            if os.path.exists(chunk_path):
-                os.remove(chunk_path)
-        print(f"-> {len(temp_chunk_paths)} fichiers temporaires de segments supprimés.")
+        if os.path.exists(temp_full_path):
+            os.remove(temp_full_path)
 
 
 # ----------------------------------------------------------------------
-# 4. PRÉ-CHARGEMENT ET INTERFACE GRADIO
+# 3. PRÉ-CHARGEMENT ET INTERFACE GRADIO
 # ----------------------------------------------------------------------
 
 INITIAL_DESCRIPTION = "Sélectionnez un modèle ASR de RobotsMali, puis enregistrez ou téléchargez un fichier audio pour obtenir la transcription."
@@ -213,6 +216,7 @@ if ROBOTSMALI_MODELS:
         default_model_short_name = default_model.split('/')[-1]
         INITIAL_DESCRIPTION = (
             f"✅ Le modèle par défaut `{default_model_short_name}` (NeMo) a été **préchargé et réchauffé** avec succès. "
+            f"**Attention :** Le traitement se fait sur l'audio complet. Les longs fichiers peuvent planter la RAM. "
             f"Téléchargez ou enregistrez votre audio pour transcrire."
         )
     except RuntimeError as e:
@@ -243,9 +247,9 @@ interface = gr.Interface(
     fn=transcribe_audio,
     inputs=[model_dropdown, audio_input],
     outputs=text_output,
-    title="🤖 RobotsMali ASR Multi-Modèles (Démo NeMo Fluide)",
+    title="🤖 RobotsMali ASR Multi-Modèles (Traitement Complet)",
     description=INITIAL_DESCRIPTION, 
     allow_flagging="never")
 
 print("Lancement de l'interface Gradio...")
-interface.launch(share=True)
+interface.launch(share=True)