Update app.py

app.py

@@ -1,110 +1,155 @@
 # -*- coding: utf-8 -*-
-import os, tempfile, glob, gc
+import os
+import tempfile
+import glob
+import shutil
 import torch
 import gradio as gr
 from huggingface_hub import snapshot_download
 from nemo.collections import asr as nemo_asr
 
+# Configuration du matériel
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 SEGMENT_DURATION = 10.0
 
-# ⚠️ CPU Basic → modèle léger uniquement
+# Dictionnaire des modèles
 MODELS = {
-    "Soloni V3 (Recommandé CPU)": ("RobotsMali/soloni-114m-tdt-ctc-v3", "rnnt")
+    "Soloni V3 (Recommandé CPU)": "RobotsMali/soloni-114m-tdt-ctc-v3"
 }
 
+# Cache pour éviter de recharger le modèle à chaque clic
 _cache = {}
 
 # ==========================
-# MODEL LOADER (identique logique vidéo)
+# CHARGEMENT DU MODÈLE
 # ==========================
-def get_model(name):
-    if name in _cache:
-        return _cache[name]
+def get_model(model_name):
+    if model_name in _cache:
+        return _cache[model_name]
 
-    repo, _ = MODELS[name]
-    folder = snapshot_download(repo, local_dir_use_symlinks=False)
+    repo_id = MODELS[model_name]
+    print(f"⏳ Téléchargement du modèle depuis {repo_id}...")
+    
+    # Téléchargement depuis Hugging Face
+    folder = snapshot_download(repo_id, local_dir_use_symlinks=False)
 
-    nemo_file = next(
-        os.path.join(folder, f)
-        for f in os.listdir(folder)
-        if f.endswith(".nemo")
-    )
+    # Recherche automatique du fichier .nemo dans le dossier téléchargé
+    try:
+        nemo_file = next(
+            os.path.join(folder, f)
+            for f in os.listdir(folder)
+            if f.endswith(".nemo")
+        )
+    except StopIteration:
+        raise FileNotFoundError("Aucun fichier .nemo trouvé dans le dépôt.")
 
+    print(f"📦 Restauration du modèle NeMo sur {DEVICE}...")
+    
+    # Utilisation de ASRModel pour l'auto-détection (CTC/RNNT)
     model = nemo_asr.models.ASRModel.restore_from(
         nemo_file,
         map_location=torch.device(DEVICE)
     )
 
     model.eval()
-
-    _cache[name] = model
+    _cache[model_name] = model
     return model
 
 # ==========================
-# PIPELINE (même logique que vidéo)
+# PIPELINE DE TRANSCRIPTION
 # ==========================
-def pipeline(audio_file, model_name):
-    tmp_dir = tempfile.mkdtemp()
+def pipeline(audio_path, model_name):
+    if not audio_path:
+        yield "❌ Aucun fichier audio fourni", ""
+        return
 
+    # Création d'un dossier temporaire unique pour les segments
+    tmp_dir = tempfile.mkdtemp()
+    
     try:
-        if not audio_file:
-            yield "❌ Aucun fichier audio", None
-            return
-
-        yield "⏳ Segmentation...", None
+        yield "⏳ Préparation de l'audio (FFmpeg)...", ""
 
-        # Découpe en segments
-        os.system(
-            f"ffmpeg -y -i '{audio_file}' -f segment -segment_time {SEGMENT_DURATION} "
-            f"-ac 1 -ar 16000 -c copy {tmp_dir}/seg_%03d.wav"
+        # Normalisation et segmentation de l'audio
+        # -ac 1 (mono), -ar 16000 (16kHz requis par NeMo)
+        command = (
+            f"ffmpeg -y -i '{audio_path}' -f segment -segment_time {SEGMENT_DURATION} "
+            f"-ac 1 -ar 16000 {tmp_dir}/seg_%03d.wav > /dev/null 2>&1"
         )
+        os.system(command)
 
-        segments = sorted(glob.glob(f"{tmp_dir}/seg_*.wav"))
+        # Liste et tri des segments générés
+        segments = sorted(glob.glob(os.path.join(tmp_dir, "seg_*.wav")))
+        
+        # Filtrer les fichiers trop petits (silences ou erreurs FFmpeg)
         segments = [s for s in segments if os.path.getsize(s) > 1000]
 
-        yield f"🎙️ Transcription {len(segments)} segments...", None
+        if not segments:
+            yield "❌ Erreur lors de la segmentation audio", ""
+            return
+
+        yield f"🎙️ Transcription de {len(segments)} segments en cours...", ""
 
+        # Récupération du modèle
         model = get_model(model_name)
 
+        # Inférence
         with torch.inference_mode():
             results = model.transcribe(
                 segments,
                 batch_size=4,
-                num_workers=0
+                num_workers=0 # Important sur CPU pour éviter les fuites de mémoire
             )
 
-        text = "\n".join(results)
+        # Reconstruction du texte
+        # Note: selon le modèle, results peut être une liste de chaînes ou un objet complexe
+        if isinstance(results, tuple):
+            text_results = results[0]
+        else:
+            text_results = results
+
+        final_text = " ".join(text_results)
 
-        yield "✅ Transcription terminée", text
+        yield "✅ Transcription terminée", final_text
 
     except Exception as e:
-        yield f"❌ Erreur : {str(e)}", None
+        yield f"❌ Erreur système : {str(e)}", ""
 
     finally:
+        # Nettoyage automatique du dossier temporaire
         if os.path.exists(tmp_dir):
-            import shutil
             shutil.rmtree(tmp_dir)
 
 # ==========================
-# UI
+# INTERFACE GRADIO (UI)
 # ==========================
-with gr.Blocks() as demo:
-    gr.Markdown("# 🤖 RobotsMali ASR – CPU Edition")
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("""
+    # 🤖 RobotsMali ASR – Edition CPU
+    Déposez un fichier audio pour obtenir une transcription automatique.
+    """)
 
     with gr.Row():
         with gr.Column():
-            audio_input = gr.Audio(type="filepath", label="Audio")
+            audio_input = gr.Audio(type="filepath", label="Fichier Audio (WAV, MP3, etc.)")
             model_input = gr.Dropdown(
                 choices=list(MODELS.keys()),
-                value="Soloni V3 (Recommandé CPU)"
+                value="Soloni V3 (Recommandé CPU)",
+                label="Sélection du Modèle"
             )
-            btn = gr.Button("🚀 Transcrire")
+            btn = gr.Button("🚀 Lancer la Transcription", variant="primary")
 
         with gr.Column():
-            status = gr.Markdown("Prêt.")
-            output = gr.Textbox(lines=15)
+            status = gr.Markdown("### Statut : Prêt")
+            output = gr.Textbox(label="Résultat de la transcription", lines=12, show_copy_button=True)
+
+    # Interaction
+    btn.click(
+        fn=pipeline,
+        inputs=[audio_input, model_input],
+        outputs=[status, output]
+    )
 
-    btn.click(pipeline, [audio_input, model_input], [status, output])
+    gr.Markdown("--- \n *Optimisé pour les environnements à ressources limitées.*")
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()