Update app.py

app.py

@@ -237,26 +237,18 @@ plt.close()
         print("🟠 Raw output du modèle :")
         print(outputs)
 
-        # Nettoyage
         generated_code = codegen_tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
 
-        # Si le modèle recopie le prompt, on garde seulement avant la répétition
-        if "### Génère uniquement du code Python fonctionnel" in generated_code:
-            generated_code = generated_code.split("### Génère uniquement du code Python fonctionnel")[0].strip()
-
-        # Protection si jamais il reste un morceau incomplet
-        if generated_code.strip().endswith("import") or generated_code.strip().endswith("sns."):
-            generated_code = generated_code.rsplit("\n", 1)[0]
+        # Nettoyage
+        generated_code = re.sub(r"(import matplotlib.pyplot as plt\nimport seaborn as sns\n)+", "import matplotlib.pyplot as plt\nimport seaborn as sns\n", generated_code)
+        generated_code = generated_code.split("###")[0].strip()  # <- Ajout pour enlever tout ce qui suit s'il y a des '###'
 
-        # Nettoyage doublon imports matplotlib/seaborn
-        generated_code = re.sub(r"(import matplotlib\.pyplot as plt\nimport seaborn as sns\n)+", "import matplotlib.pyplot as plt\nimport seaborn as sns\n", generated_code)
+        if generated_code.strip().endswith("sns.") or len(generated_code.strip()) < 20:
+            return JSONResponse(content={"error": "Erreur : Code généré invalide ou incomplet."}, status_code=500)
 
         print("🔵 Code généré propre :")
         print(generated_code)
 
-        if not generated_code.strip():
-            return JSONResponse(content={"error": "Erreur : Code généré vide."}, status_code=500)
-
         try:
             compile(generated_code, "<string>", "exec")
         except SyntaxError as e: