Update app.py

app.py

@@ -1,57 +1,88 @@
-# app.py
 import gradio as gr
-import spacy
 import torch
 import os
-from transformers import pipeline
-from pyannote.audio import Pipeline as DiarizationPipeline
-
-# Chargement du modèle Whisper via transformers
-asr_pipeline = pipeline("automatic-speech-recognition", model="openai/whisper-small", chunk_length_s=30)
-
-# Chargement du modèle spaCy pour NER
-nlp = spacy.load("fr_core_news_md")
-
-# Diarisation avec PyAnnote (si HF_TOKEN dispo dans l'environnement)
-hf_token = os.getenv("HF_TOKEN")
-if hf_token:
-    diar_pipeline = DiarizationPipeline.from_pretrained("pyannote/speaker-diarization-3.1", use_auth_token=hf_token)
-else:
-    diar_pipeline = None
-
-def process_audio(file):
-    # Transcription avec Whisper via transformers
-    result = asr_pipeline(file)
-    transcription = result["text"]
-
-    # Diarisation
-    if diar_pipeline:
-        diar_result = diar_pipeline(file)
-        diar_str = "\n".join([
-            f"{turn.start:.1f}s - {turn.end:.1f}s : {speaker}"
-            for turn, _, speaker in diar_result.itertracks(yield_label=True)
-        ])
-    else:
-        diar_str = "Diarisation non disponible (ajoutez votre HF_TOKEN dans les secrets)"
-
-    # NER
-    doc = nlp(transcription)
-    entities = [(ent.text, ent.label_) for ent in doc.ents]
-    ent_str = "\n".join([f"{text} ({label})" for text, label in entities]) if entities else "Aucune entité détectée"
-
-    return transcription, diar_str, ent_str
-
-demo = gr.Interface(
+from whisperx import load_model, load_align_model, align
+from resemblyzer import preprocess_wav, VoiceEncoder
+from sklearn.cluster import AgglomerativeClustering
+import librosa
+import numpy as np
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+compute_type = "float16" if device == "cuda" else "int8"
+
+whisper_model = load_model("medium", device=device, compute_type=compute_type)
+align_model, metadata = load_align_model(language_code="fr", device=device)
+voice_encoder = VoiceEncoder()
+
+def get_speaker_segments(audio_path, window_size=1.0, step_size=0.5, num_speakers=2):
+    wav, sr = librosa.load(audio_path, sr=16000, mono=True)
+    wav = librosa.util.normalize(wav)
+    duration = librosa.get_duration(y=wav, sr=sr)
+
+    segments = []
+    embeddings = []
+
+    for start in np.arange(0, duration - window_size, step_size):
+        end = start + window_size
+        clip = wav[int(start * sr):int(end * sr)]
+        if len(clip) == 0:
+            continue
+        try:
+            embed = voice_encoder.embed_utterance(clip)
+            embeddings.append(embed)
+            segments.append((start, end))
+        except Exception as e:
+            print(f"⚠️ Skipped segment {start}-{end}s: {e}")
+
+    if len(embeddings) < 2:
+        print("⚠️ Pas assez de segments pour la diarisation. Diarisation annulée.")
+        return [{"start": 0, "end": duration, "speaker": "speaker_00"}]
+
+    clustering = AgglomerativeClustering(n_clusters=num_speakers)
+    labels = clustering.fit_predict(embeddings)
+
+    speaker_segments = []
+    for (start, end), label in zip(segments, labels):
+        speaker_segments.append({"start": start, "end": end, "speaker": f"speaker_{label:02d}"})
+
+    return speaker_segments
+
+def process_audio(audio_file):
+    tmp_path = audio_file
+
+    # Step 1: Transcription
+    result = whisper_model.transcribe(tmp_path, language="fr", word_timestamps=False, verbose=False)
+
+    # Step 2: Diarisation via resemblyzer
+    speaker_segments = get_speaker_segments(tmp_path)
+
+    # Step 3: Alignement mot à mot
+    result_aligned = align(result["segments"], align_model, metadata, tmp_path, return_char_alignments=False)
+
+    # Attribution speaker
+    for segment in result_aligned["segments"]:
+        segment_start = segment["start"]
+        speaker_found = next((sp["speaker"] for sp in speaker_segments if sp["start"] <= segment_start <= sp["end"]), "speaker_??")
+        segment["speaker"] = speaker_found
+
+    # Format final
+    final_output = ""
+    for seg in result_aligned["segments"]:
+        speaker = seg["speaker"]
+        start = f"{seg['start']:.2f}s"
+        end = f"{seg['end']:.2f}s"
+        text = seg['text'].strip()
+        final_output += f"[{start} - {end}] {speaker}: {text}\n"
+
+    return final_output
+
+iface = gr.Interface(
     fn=process_audio,
-    inputs=gr.Audio(type="filepath", label="Audio (.mp3/.wav)"),
-    outputs=[
-        gr.Textbox(label="📝 Transcription Whisper"),
-        gr.Textbox(label="🗣️ Diarisation (PyAnnote)"),
-        gr.Textbox(label="🧠 Entités Nommées (spaCy)")
-    ],
-    title="🔎 Pipeline Audio Intelligent",
-    description="Transcription, Diarisation, et Extraction d'Entités Nommées sur un fichier audio français."
+    inputs=gr.Audio(type="filepath", label="Audio (.wav, .mp3...)"),
+    outputs=gr.Textbox(label="Transcription + Diarisation + Alignement"),
+    title="🎙️ Transcription enrichie avec WhisperX + Resemblyzer",
+    description="Transcription française, diarisation légère (sans token), alignement mot à mot."
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    iface.launch()