Update app.py

app.py

@@ -1,145 +1,164 @@
-import os  # <-- N'oublie pas d'importer os pour os.getenv !
+import os
 import gradio as gr
-import requests
 import base64
 import io
+import re
 from PIL import Image
-import fitz  # PyMuPDF pour gérer les PDFs
+import fitz  # PyMuPDF
+
+# --- CONFIGURATION ---
+# Récupère la clé HF depuis les secrets du Space (sécurisé)
+HF_TOKEN = os.getenv("HF_TOKEN")
+if not HF_TOKEN:
+    raise ValueError(
+        "❌ Token HF manquant. Ajoutez HF_TOKEN dans les 'Repository Secrets' de votre Space."
+    )
 
-# --- CONFIGURATION DE L'API MISTRAL ---
-# Récupère la clé depuis les secrets (sécurisé)
-MISTRAL_API_KEY = os.getenv("MISTRAL_API_KEY")
-if not MISTRAL_API_KEY:
-    raise ValueError("❌ Clé API Mistral manquante. Ajoutez-la dans les 'Repository Secrets' de votre Space Hugging Face.")
-API_URL = "https://api.mistral.ai/v1/chat/completions"
+MODEL_ID  = "google/gemma-4-E4B-it"
+API_URL   = f"https://api-inference.huggingface.co/models/{MODEL_ID}/v1/chat/completions"
 
 # --- FONCTIONS AUXILIAIRES ---
-def encode_image(image):
-    """Convertit une image PIL en base64 pour l'envoi à l'API."""
-    buffer = io.BytesIO()
-    image.save(buffer, format="JPEG")  # Convertit en JPEG
-    return base64.b64encode(buffer.getvalue()).decode()  # Encode en base64
 
-def pdf_to_image(pdf_path):
-    """
-    Convertit la première page d'un PDF en image PIL.
-    Args:
-        pdf_path (str): Chemin vers le fichier PDF.
-    Returns:
-        PIL.Image: Image de la première page du PDF.
-    """
-    doc = fitz.open(pdf_path)  # Ouvre le PDF
-    first_page = doc.load_page(0)  # Charge la première page (index 0)
-    pix = first_page.get_pixmap()  # Récupère les pixels de la page
-    img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)  # Crée une image PIL
+def encode_image(image: Image.Image) -> str:
+    """Convertit une image PIL en base64 JPEG."""
+    buffer = io.BytesIO()
+    image.save(buffer, format="JPEG", quality=85)
+    return base64.b64encode(buffer.getvalue()).decode()
+
+
+def pdf_to_image(pdf_path: str) -> Image.Image:
+    """Convertit la première page d'un PDF en image PIL."""
+    doc = fitz.open(pdf_path)
+    page = doc.load_page(0)
+    mat  = fitz.Matrix(150 / 72, 150 / 72)   # 150 DPI — lisible et léger
+    pix  = page.get_pixmap(matrix=mat)
+    img  = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
+    doc.close()
     return img
 
+
+def sanitize_filename(name: str) -> str:
+    """Nettoie le nom retourné par le modèle pour en faire un nom de fichier valide."""
+    name = name.strip().strip('"').strip("'")
+    name = re.sub(r"[^\w\s\-]", "", name)
+    name = re.sub(r"\s+", "_", name)
+    return name[:80] if name else "document_sans_nom"
+
+
 # --- FONCTION PRINCIPALE ---
-def generate_filename(file, custom_prompt):
+
+def generate_filename(file, custom_prompt: str):
     """
-    Génère un nom de fichier structuré à partir d'une image ou d'un PDF (première page uniquement).
+    Génère un nom de fichier à partir d'une image ou d'un PDF (première page).
     Args:
-        file (gr.File): Fichier uploadé (image ou PDF).
-        custom_prompt (str): Format personnalisé pour le nom de fichier (optionnel).
+        file     : Fichier uploadé (image ou PDF).
+        custom_prompt : Prompt/format personnalisé pour le nom.
     Returns:
-        str: Nom de fichier généré ou message d'erreur.
+        tuple: (aperçu PIL Image, nom généré str)
     """
     # 1. Vérification du fichier
     if file is None:
-        return "❌ Veuillez uploader un fichier (image ou PDF)."
-
-    # 2. Vérification de la clé API
-    if not MISTRAL_API_KEY or MISTRAL_API_KEY == "VOTRE_CLE_API_MISTRAL":
-        return "❌ Clé API manquante. Obtenez une clé gratuite sur [console.mistral.ai](https://console.mistral.ai)"
+        return None, "❌ Veuillez uploader un fichier (image ou PDF)."
 
     try:
-        # 3. Traitement du fichier (image ou PDF)
-        if file.name.lower().endswith('.pdf'):
-            # Cas PDF : extraire la première page et la convertir en image
+        # 2. Traitement selon le type de fichier
+        if file.name.lower().endswith(".pdf"):
             image = pdf_to_image(file.name)
         else:
-            # Cas image : ouvrir directement
-            image = Image.open(file.name)
+            image = Image.open(file.name).convert("RGB")
 
-        # 4. Encodage de l'image en base64
+        # 3. Encodage base64
         image_base64 = encode_image(image)
 
-        # 5. Préparation du prompt pour Mistral
+        # 4. Prompt
         if custom_prompt and custom_prompt.strip():
-            prompt = custom_prompt.strip()  # Utilise le format personnalisé
+            prompt = custom_prompt.strip()
         else:
-            prompt = """Analyse ce document. Extrait le nom des parties et la date.
-Génère un nom de fichier au format: Affaire_[X]_contre_[Y]_[JJ_MM_AAAA]
-Exemple: Affaire_Dupont_contre_Martin_12_12_2026
-Réponds uniquement par le nom de fichier, rien d'autre."""
+            prompt = (
+                "Analyse ce document. Génère un nom de fichier court et descriptif "
+                "en snake_case (3 à 6 mots max).\n"
+                "Réponds UNIQUEMENT par le nom de fichier, sans extension, sans explication."
+            )
 
-        # 6. Appel à l'API Mistral
+        # 5. Appel à l'Inference API HF (même structure que ton ancien script Mistral)
+        import requests
         response = requests.post(
             API_URL,
             headers={
-                "Authorization": f"Bearer {MISTRAL_API_KEY}",
-                "Content-Type": "application/json"
+                "Authorization": f"Bearer {HF_TOKEN}",
+                "Content-Type": "application/json",
             },
             json={
-                "model": "pixtral-12b-2409",
-                "messages": [{
-                    "role": "user",
-                    "content": [
-                        {"type": "text", "text": prompt},
-                        {"type": "image_url", "image_url": f"data:image/jpeg;base64,{image_base64}"}
-                    ]
-                }],
-                "max_tokens": 100,
-                "temperature": 0.1
+                "model": MODEL_ID,
+                "messages": [
+                    {
+                        "role": "user",
+                        "content": [
+                            {"type": "text", "text": prompt},
+                            {
+                                "type": "image_url",
+                                "image_url": {
+                                    "url": f"data:image/jpeg;base64,{image_base64}"
+                                },
+                            },
+                        ],
+                    }
+                ],
+                "max_tokens": 60,
+                "temperature": 0.1,
             },
-            timeout=60
+            timeout=60,
         )
 
-        # 7. Traitement de la réponse
+        # 6. Traitement de la réponse
         if response.status_code == 200:
-            result = response.json()
-            return result["choices"][0]["message"]["content"].strip()
+            raw   = response.json()["choices"][0]["message"]["content"].strip()
+            clean = sanitize_filename(raw)
+            return image, f"{clean}.pdf"
         else:
-            return f"❌ Erreur API Mistral: Code {response.status_code} - {response.text}"
+            return image, f"❌ Erreur API HF : {response.status_code} — {response.text}"
 
     except Exception as e:
-        return f"❌ Erreur: {str(e)}"
+        return None, f"❌ Erreur : {str(e)}"
+
 
 # --- INTERFACE GRADIO ---
-with gr.Blocks(title="Nommage de documents juridiques") as demo:
+
+with gr.Blocks(title="PDF Auto-Namer") as demo:
     gr.Markdown("""
-    # ⚖️ Générateur de noms pour documents juridiques
-    **Basé sur Mistral Pixtral** - Gratuit
-    Upload un **fichier (image ou PDF)** → la première page sera analysée pour générer un nom structuré.
+    # 📄 PDF Auto-Namer
+    **Basé sur Gemma 4 E4B (multimodal)** via Hugging Face Inference API
+    Upload un **fichier image ou PDF** → la première page est analysée → un nom est généré.
     """)
 
     with gr.Row():
         with gr.Column():
-            # Champ pour uploader un fichier (image ou PDF)
-            file = gr.File(
+            file_input = gr.File(
                 label="Document (Image ou PDF)",
                 file_types=[".jpg", ".jpeg", ".png", ".pdf"],
-                height=400
+                height=200,
             )
-            # Champ pour le format personnalisé
-            custom_prompt = gr.Textbox(
+            prompt_input = gr.Textbox(
                 label="Format personnalisé (optionnel)",
-                placeholder='Ex: "Format: DOS_[X]_VS_[Y]_[DDMMYYYY]"',
-                lines=2
+                placeholder='Ex: "Format: Contrat_[TYPE]_[CLIENT]_[JJMMAAAA]"',
+                lines=3,
             )
-            # Bouton pour lancer l'analyse
-            btn = gr.Button("Générer le nom", variant="primary")
+            btn = gr.Button("🚀 Générer le nom", variant="primary")
 
         with gr.Column():
-            # Zone pour afficher le résultat
-            output = gr.Textbox(label="Nom généré", lines=5)
+            preview = gr.Image(label="👁️ Aperçu — Page 1", type="pil")
+            output  = gr.Textbox(label="📝 Nom de fichier généré", lines=2)
 
-    # Lien entre le bouton et la fonction
     btn.click(
         fn=generate_filename,
-        inputs=[file, custom_prompt],
-        outputs=output
+        inputs=[file_input, prompt_input],
+        outputs=[preview, output],
+    )
+
+    gr.Markdown(
+        "---\n"
+        "ℹ️ *Modèle : `google/gemma-4-E4B-it` · "
+        "La clé `HF_TOKEN` doit être définie dans les Secrets du Space.*"
     )
 
-# --- LANCEMENT DE L'APPLICATION ---
 demo.launch(show_error=True)