Update app.py

app.py

@@ -18,6 +18,7 @@ import evaluate
 import tempfile
 from huggingface_hub import snapshot_download
 from transformers import pipeline
+import openai
 
 
 st.title("📊 Évaluation WER d'un modèle Whisper")
@@ -28,7 +29,8 @@ 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"
+    "Whisper Large + LoRA + Post-processing Mistral 7B",
+    "Whisper Large + LoRA + Post-processing GPT-4o"
 ))
 
 # Section : Lien du dataset
@@ -36,6 +38,8 @@ st.subheader("2. Chargement du dataset Hugging Face")
 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")
 
+openai_api_key = st.text_input("Clé API OpenAI (pour GPT-4o)", type="password")
+
 if hf_token:
     from huggingface_hub import login
     login(hf_token)
@@ -97,7 +101,7 @@ if start_eval:
         model.eval()
 
         # Charger le pipeline de Mistral si post-processing demandé
-        if "Post-processing" in model_option:
+        if "Post-processing Mistral" in model_option:
             with st.spinner("Chargement du modèle de post-traitement Mistral..."):
                 postproc_pipe = pipeline(
                     "text2text-generation",
@@ -111,6 +115,21 @@ if start_eval:
                     prompt = f"Tu es CorrecteurAI, une AI française qui permet de corriger les erreurs de saisie vocal. La translation d'un enregistrement audio tiré d'une inspection détaillé de pont t'es envoyé et tu renvoies le texte identique mais avec les éventuelles corrections si des erreurs sont détectés. Le texte peut comprendre du vocabulaire technique associé aux ouvrages d'art. Renvoies uniquement le texte corrigé en français et sans autre commentaire. Voici le texte : {text}"
                     result = postproc_pipe(prompt, max_new_tokens=256)[0]["generated_text"]
                     return result.strip()
+
+    
+    #fonction process GPT4o
+    def postprocess_with_gpt4o(text, api_key):
+        openai.api_key = api_key
+        response = openai.ChatCompletion.create(
+            model="gpt-4o",
+            messages=[
+                {"role": "system", "content": "Tu es un assistant qui corrige et ponctue des transcriptions vocales françaises sans changer le sens du texte. Répond uniquement avec le texte corrigé."},
+                {"role": "user", "content": f"Corrige ce texte : {text}"}
+            ],
+            temperature=0.3,
+            max_tokens=512
+        )
+        return response.choices[0].message["content"].strip()
     
 
     # 🔹 Préparer WER metric
@@ -138,11 +157,23 @@ if start_eval:
             prediction = processor.batch_decode(pred_ids, skip_special_tokens=True)[0]
 
             # === Post-processing conditionnel ===
-            if "Post-processing" in model_option:
+            if "Post-processing Mistral" in model_option:
                 st.write("⏳ Post-processing avec Mistral...")
                 postprocessed_prediction = postprocess_with_llm(prediction)
-                st.write("✅ Post-processing terminé.")
+                st.write("✅ Terminé.")
                 final_prediction = postprocessed_prediction
+            
+            elif "Post-processing GPT-4o" in model_option:
+                if not openai_api_key:
+                    st.error("Clé API OpenAI requise pour GPT-4o.")
+                    st.stop()
+                st.write("🤖 Post-processing avec GPT-4o...")
+                try:
+                    postprocessed_prediction = postprocess_with_gpt4o(prediction, openai_api_key)
+                except Exception as e:
+                    postprocessed_prediction = f"[Erreur GPT-4o: {e}]"
+                final_prediction = postprocessed_prediction
+            
             else:
                 postprocessed_prediction = "-"
                 final_prediction = prediction