Update app.py

app.py

@@ -17,6 +17,7 @@ import numpy as np
 import evaluate
 import tempfile
 from huggingface_hub import snapshot_download
+from transformers import pipeline
 
 st.title("📊 Évaluation WER d'un modèle Whisper")
 st.markdown("Ce Space permet d'évaluer la performance WER d'un modèle Whisper sur un dataset audio.")
@@ -26,14 +27,21 @@ st.subheader("1. Choix du modèle")
 model_option = st.radio("Quel modèle veux-tu utiliser ?", (
     "Whisper Large (baseline)",
     "Whisper Large + LoRA (SimpleFrog/whisper_finetuned)",
-    "Whisper Large + LoRA + Post-processing (à venir)"
+    "Whisper Large + LoRA + Post-processing"
 ))
 
 # Section : Lien du dataset
 st.subheader("2. Chargement du dataset Hugging Face")
-dataset_link = st.text_input("Lien du dataset (format: user/dataset_name)", value="SimpleFrog/mon_dataset")
+dataset_link = st.text_input("Lien du dataset (format: user/dataset_name)", value="SimpleFrog/Dataset_Test")
 hf_token = st.text_input("Token Hugging Face (si dataset privé)", type="password")
 
+# Section : Choix du split
+split_option = st.selectbox(
+    "Choix du split à évaluer",
+    options=["Tous", "train", "validation", "test"],
+    index=0  # par défaut "Tous"
+)
+
 # Section : Bouton pour lancer l'évaluation
 start_eval = st.button("🚀 Lancer l'évaluation WER")
 
@@ -43,8 +51,19 @@ if start_eval:
     # 🔹 Télécharger dataset
     with st.spinner("Chargement du dataset..."):
         try:
-            #dataset = load_dataset(dataset_link, data_files="metadata.csv", data_dir=".", split="train", token=hf_token)
-            dataset = load_dataset(dataset_link, split="train", token=hf_token)
+
+            dataset_full = load_dataset(dataset_link, split="train", token=hf_token)
+
+            # 🔹 Filtrage selon la colonne 'split'
+            if split_option != "Tous":
+                dataset = dataset_full.filter(lambda x: x.get("split", "unknown") == split_option)
+            else:
+                dataset = dataset_full
+            
+            if len(dataset) == 0:
+                st.warning(f"Aucun exemple trouvé pour le split sélectionné : '{split_option}'.")
+                st.stop()
+            
         except Exception as e:
             st.error(f"Erreur lors du chargement du dataset : {e}")
             st.stop()
@@ -60,6 +79,21 @@ if start_eval:
         processor = WhisperProcessor.from_pretrained(base_model_name)
         model.eval()
 
+        # Charger le pipeline de Mistral si post-processing demandé
+        if "Post-processing" in model_option:
+            with st.spinner("Chargement du modèle de post-traitement Mistral..."):
+                postproc_pipe = pipeline(
+                    "text2text-generation",
+                    model="mistralai/Mistral-7B-Instruct-v0.2",
+                    device_map="auto",  # ou device=0 si tu veux forcer le GPU
+                    torch_dtype=torch.float16  # optionnel mais plus léger
+                )
+        
+                def postprocess_with_llm(text):
+                    prompt = f"Ce texte est issue d'une translation vocal. L'enregistrement est tiré d'une inspection détaillé de pont et comprend du vocabulaire technique associé. Corriges les éventuelles erreurs de translation : {text}"
+                    result = postproc_pipe(prompt, max_new_tokens=256)[0]["generated_text"]
+                    return result.strip()
+
     # 🔹 Préparer WER metric
     wer_metric = evaluate.load("wer")
 
@@ -71,18 +105,11 @@ if start_eval:
     for example in dataset:
         st.write("Exemple brut :", example)
         try:
-            #audio_path = example["file_name"]  # full path or relative path in AudioFolder
-            #audio_path = os.path.join(repo_local_path, example["file_name"])
+           
             reference = example["text"]
             
             waveform = example["audio"]["array"]
             audio_path = example["audio"]["path"]
-
-            #st.write(example)
-            #st.write("Exemple brut :", dataset[0])
-
-            # Load audio (we assume dataset is structured with 'file_name')
-            #waveform, _ = librosa.load(audio_path, sr=16000)
             
             waveform = np.expand_dims(waveform, axis=0)
             inputs = processor(waveform, sampling_rate=16000, return_tensors="pt")
@@ -91,18 +118,27 @@ if start_eval:
                 pred_ids = model.generate(input_features=inputs.input_features)
             prediction = processor.batch_decode(pred_ids, skip_special_tokens=True)[0]
 
+            # === Post-processing conditionnel ===
+            if "Post-processing" in model_option:
+                postprocessed_prediction = postprocess_with_llm(prediction)
+                final_prediction = postprocessed_prediction
+            else:
+                postprocessed_prediction = "-"
+                final_prediction = prediction
+
             # 🔹 Nettoyage ponctuation pour WER "sans ponctuation"
             def clean(text):
                 return ''.join([c for c in text.lower() if c.isalnum() or c.isspace()]).strip()
 
             ref_clean = clean(reference)
-            pred_clean = clean(prediction)
+            pred_clean = clean(final_prediction)
             wer = wer_metric.compute(predictions=[pred_clean], references=[ref_clean])
 
             results.append({
                 "Fichier": audio_path,
                 "Référence": reference,
-                "Transcription": prediction,
+                "Transcription brute": prediction,
+                "Transcription corrigée": postprocessed_prediction,
                 "WER": round(wer, 4)
             })
 
@@ -110,7 +146,8 @@ if start_eval:
             results.append({
                 "Fichier": example["audio"].get("path", "unknown"),
                 "Référence": "Erreur",
-                "Transcription": f"Erreur: {e}",
+                "Transcription brute": f"Erreur: {e}",
+                "Transcription corrigée": "-",
                 "WER": "-"
             })
 
@@ -122,11 +159,10 @@ if start_eval:
         df.to_csv(tmp_csv.name, index=False)
 
         mean_wer = df[df["WER"] != "-"]["WER"].mean()
+        
         st.markdown(f"### 🎯 WER moyen (sans ponctuation) : `{mean_wer:.3f}`")
 
-        # Bloc placeholder pour post-processing à venir
-        if "Post-processing" in model_option:
-            st.info("🛠️ Le post-processing sera ajouté prochainement ici...")
+
 
 
         # 🔹 Bouton de téléchargement