gemini

app.py

@@ -1,49 +1,57 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
 """
-Extraction BRVM via Hugging Face Hub sur Space
-Avec téléchargement de PDF depuis un Hugging Face Dataset
-et upload des résultats JSON vers un autre Hugging Face Dataset
-Auteur : Gemini
+Extraction de données de bulletins BRVM via Hugging Face.
+
+Ce script optimisé :
+1. Télécharge un fichier ZIP de rapports PDF depuis un Hugging Face Dataset.
+2. Extrait le texte de chaque PDF.
+3. Utilise un modèle Mistral quantifié pour extraire les données structurées en JSON.
+4. Gère les erreurs de manière robuste (parsing, extraction).
+5. Uploade les résultats JSON vers un autre Hugging Face Dataset.
+
+Auteur: Gemini (avec optimisations)
 """
 
-import fitz
+import fitz  # PyMuPDF
 import json
 from pathlib import Path
 from tqdm import tqdm
 import torch
 from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
-import requests
-import zipfile
-import io
 import os
 from huggingface_hub import hf_hub_download, HfApi
 from datetime import datetime
 import uuid
+import zipfile
 
-# ---------- CONFIG ----------
-MODEL = "mistralai/Mistral-7B-Instruct-v0.3"
+# ---------- CONFIGURATION CENTRALE ----------
+# --- Modèle et Quantification ---
+MODEL_ID = "mistralai/Mistral-7B-Instruct-v0.3"
 
-HF_DATASET_REPO_ID_PDFS = "lamekemal/brvm-reports-pdfs"
+# --- Dépôts Hugging Face ---
+HF_DATASET_PDFS_REPO_ID = "lamekemal/brvm-reports-pdfs"
+HF_DATASET_JSON_REPO_ID = "lamekemal/brvm-reports-json"
 ZIP_FILENAME_IN_DATASET = "brvm_reports.zip"
 
-HF_DATASET_REPO_ID_JSON_OUTPUT = "lamekemal/brvm-reports-json"
-
-HF_TOKEN = os.getenv('HF_TOKEN')
+# --- Token et Dossiers Locaux ---
+HF_TOKEN = os.getenv('HF_TOKEN') # Le token est récupéré des variables d'environnement (secrets)
+LOCAL_PDF_FOLDER = Path("brvm_reports_extracted")
+LOCAL_JSON_OUTPUT_BASE_FOLDER = Path("brvm_json_outputs")
+LOCAL_CACHE_DIR = Path("./hf_cache")
 
-PDF_FOLDER = "brvm_reports"
-LOCAL_OUT_BASE_FOLDER = "local_json_outputs"
-# ----------------------------
+# --- Paramètres du Pipeline ---
+# Assez grand pour contenir le JSON complet, même pour des rapports denses.
+MAX_NEW_TOKENS = 8192
 
-# Prompt optimisé pour le modèle Mistral Instruct
+# --- Prompt Optimisé ---
+# Plus direct, plus concis, et sans marqueurs custom.
+# Le modèle est instruit de ne retourner QUE le JSON, ce qui simplifie le parsing.
 PROMPT_TEMPLATE = """
-[INST] Tu es un expert en finance spécialisé dans la BRVM.
-À partir du texte ci-dessous issu d’un bulletin officiel de la cote BRVM,
-EXTRAIT UNIQUEMENT les données suivantes et RETOURNE UN JSON VALIDE au format strict.
-Ta réponse NE DOIT CONTENIR AUCUN AUTRE TEXTE QUE LE JSON.
-Commence ta réponse par [DEBUT_JSON] et termine la par [FIN_JSON].
+[INST]Tu es un expert en analyse de données financières de la BRVM. Extrais les informations du texte suivant et retourne-les sous la forme d'un objet JSON unique et valide. Ta réponse doit commencer par `{` et se terminer par `}`. N'inclus aucun texte, explication ou formatage en dehors de l'objet JSON.
 
-JSON Schema:
+**JSON Schema attendu :**
+```json
 {{
   "indicateurs": {{
     "brvm_10": {{ "niveau": float, "var_jour_pct": float, "var_annuelle_pct": float }},
@@ -69,230 +77,240 @@ JSON Schema:
     {{ "code": string, "emetteur": string, "coupon_pct": float, "echeance": string, "nominal": float, "cours": float }}
   ]
 }}
+```
 
-Contraintes :
-- Si une donnée est absente, mets null.
-- Les nombres utilisent un point comme séparateur décimal.
+**Contraintes :**
+- Si une donnée est manquante, utilise la valeur `null`.
+- Utilise un point `.` comme séparateur décimal.
 
-Texte du bulletin :
+**Texte du bulletin à analyser :**
 {texte_pdf}[/INST]
 """
-# --- Initialisation du pipeline Hugging Face avec quantification ---
-extractor = None # Initialiser à None pour gérer les erreurs de chargement
 
-try:
-    print("Chargement du modèle Hugging Face avec quantification 4-bit...")
-    
-    # Configuration de la quantification 4-bit
-    bnb_config = BitsAndBytesConfig(
-        load_in_4bit=True,
-        bnb_4bit_quant_type="nf4", # NormalFloat4 quantization
-        bnb_4bit_compute_dtype=torch.bfloat16, # Type de données pour les calculs
-        bnb_4bit_use_double_quant=True, # Double quantification pour plus de précision
-    )
-
-    # Charger le tokenizer et le modèle avec la configuration de quantification
-    tokenizer = AutoTokenizer.from_pretrained(MODEL)
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL,
-        quantization_config=bnb_config,
-        device_map="auto", # Permet de charger le modèle sur le(s) GPU disponible(s)
-        torch_dtype=torch.bfloat16 # Spécifier le dtype pour le chargement initial
-    )
-    
-    # Créer le pipeline avec le modèle et le tokenizer chargés
-    if tokenizer.pad_token is None:
-        tokenizer.pad_token = tokenizer.eos_token # Ou un autre token approprié si nécessaire
-
-    extractor = pipeline(
-        "text-generation",
-        model=model,
-        tokenizer=tokenizer,
-        device=0 # Utilise le GPU si disponible (0 pour le premier GPU, -1 pour CPU)
-    )
-    print("Modèle quantifié chargé avec succès.")
-except Exception as e:
-    print(f"Erreur lors du chargement du modèle Hugging Face : {e}")
-    # Ne pas initialiser extractor si erreur
-
-# --- Fonctions d'extraction et de traitement ---
-def extract_text(pdf_path):
-    """Extraction du texte brut à partir d'un PDF avec PyMuPDF."""
-    try:
-        doc = fitz.open(pdf_path)
-        text = "\n".join(page.get_text() for page in doc)
-        return text
-    except Exception as e:
-        print(f"[Erreur PDF] {pdf_path} : {e}")
-        return ""
+# ---------- FONCTIONS ----------
 
-def call_huggingface_model(text):
-    """Appel du modèle Hugging Face pour extraire le JSON à partir du texte."""
-    if extractor is None:
-        print("[Erreur] Le pipeline d'extraction n'a pas été initialisé. Impossible d'appeler le modèle.")
-        return {"error": "Modèle non chargé."}
-    
-    prompt = PROMPT_TEMPLATE.format(texte_pdf=text)
-    
+def initialize_model_pipeline():
+    """
+    Charge le modèle et le tokenizer avec une quantification 4-bit et crée le pipeline.
+    Retourne le pipeline ou None en cas d'erreur.
+    """
     try:
-        # Augmenter considérablement max_new_tokens pour éviter la troncature
-        # Augmenté de 4096 à 8192
-        response = extractor(
-            prompt, 
-            max_new_tokens=8192, 
-            do_sample=False,
-            num_return_sequences=1
-        ) 
-        
-        output_text = response[0]['generated_text']
+        print(f"Chargement du modèle '{MODEL_ID}' avec quantification 4-bit...")
+
+        # Configuration de la quantification pour réduire l'empreinte mémoire
+        bnb_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_quant_type="nf4",       # Type de quantification (NormalFloat4) - bon équilibre performance/précision
+            bnb_4bit_compute_dtype=torch.bfloat16, # Type pour les calculs, bfloat16 est rapide sur les GPU récents
+            bnb_4bit_use_double_quant=True, # Améliore la précision avec une surcoût mémoire minime
+        )
 
-        json_start_marker = "[DEBUT_JSON]"
-        json_end_marker = "[FIN_JSON]"
+        tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+        model = AutoModelForCausalLM.from_pretrained(
+            MODEL_ID,
+            quantization_config=bnb_config,
+            device_map="auto", # Répartit automatiquement le modèle sur les ressources disponibles (GPU/CPU)
+            torch_dtype=torch.bfloat16,
+            trust_remote_code=True # Nécessaire pour certains modèles
+        )
 
-        # La réponse du modèle Mistral peut inclure le prompt, donc on cherche le début du JSON APRÈS la fin du prompt ou le début du marqueur
-        # et la fin du JSON AVANT la fin du marqueur.
-        
-        # Trouver la fin du prompt pour commencer la recherche du JSON
-        prompt_end_index = output_text.rfind("[/INST]") # C'est le marqueur de fin d'instruction
+        # Assurer que le pad_token est défini pour éviter les avertissements
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
 
-        # Si le marqueur [/INST] est trouvé, commencer la recherche du JSON après.
-        if prompt_end_index != -1:
-            search_start_index = prompt_end_index + len("[/INST]")
-            text_to_parse = output_text[search_start_index:]
-        else:
-            # Si le marqueur [/INST] n'est pas trouvé (moins probable avec un modèle instruct),
-            # on prend toute la chaîne et on se base sur les marqueurs [DEBUT_JSON]/[FIN_JSON]
-            text_to_parse = output_text
-        
-        json_start = text_to_parse.find(json_start_marker)
-        json_end = text_to_parse.rfind(json_end_marker)
-        
-        if json_start != -1 and json_end != -1:
-            json_text_with_markers = text_to_parse[json_start + len(json_start_marker):json_end]
-            
-            actual_json_start = json_text_with_markers.find('{')
-            actual_json_end = json_text_with_markers.rfind('}') + 1
-
-            if actual_json_start != -1 and actual_json_end != -1:
-                json_content = json_text_with_markers[actual_json_start:actual_json_end]
-                return json.loads(json_content)
-            else:
-                return {"error": "JSON non trouvé entre les accolades après les marqueurs", "raw": output_text}
-        else:
-            return {"error": "Marqueurs JSON non trouvés", "raw": output_text}
-            
-    except json.JSONDecodeError:
-        print(f"Erreur JSONDecodeError: Le texte brut était : {output_text[:500]}...")
-        return {"error": "JSON invalide ou mal formé", "raw": output_text}
+        extractor_pipeline = pipeline(
+            "text-generation",
+            model=model,
+            tokenizer=tokenizer,
+        )
+        print("✅ Modèle et pipeline chargés avec succès.")
+        return extractor_pipeline
     except Exception as e:
-        return {"error": str(e), "raw": output_text}
+        print(f"❌ Erreur critique lors du chargement du modèle : {e}")
+        return None
 
-def download_and_extract_zip_from_hf_dataset(dataset_repo_id, zip_filename, target_folder):
+def download_and_extract_pdfs(repo_id, zip_filename, target_folder, cache_folder):
     """
-    Télécharge un fichier ZIP depuis un Hugging Face Dataset et extrait les PDF
-    dans un dossier cible. Crée le dossier cible s'il n'existe pas.
+    Télécharge un ZIP depuis un dataset HF et en extrait les fichiers PDF.
     """
-    print(f"Tentative de téléchargement du fichier '{zip_filename}' "
-          f"depuis le dataset Hugging Face : {dataset_repo_id}")
+    print(f"Téléchargement de '{zip_filename}' depuis le dataset '{repo_id}'...")
     try:
         local_zip_path = hf_hub_download(
-            repo_id=dataset_repo_id,
+            repo_id=repo_id,
             filename=zip_filename,
             repo_type="dataset",
-            cache_dir="./hf_cache"
+            cache_dir=cache_folder
         )
-        print(f"Fichier ZIP téléchargé localement : {local_zip_path}")
-        
-        Path(target_folder).mkdir(parents=True, exist_ok=True)
+        print(f"Fichier ZIP téléchargé dans : {local_zip_path}")
+
+        target_folder.mkdir(parents=True, exist_ok=True)
         
         extracted_files = []
         with zipfile.ZipFile(local_zip_path, 'r') as z:
-            for file_info in z.infolist():
-                if file_info.filename.startswith('/') or '..' in file_info.filename:
-                    print(f"Skipping potentially unsafe path: {file_info.filename}")
+            for member in z.infolist():
+                # Sécurité : Ignorer les chemins absolus ou contenant '..'
+                if member.is_dir() or member.filename.startswith('/') or '..' in member.filename:
                     continue
-
-                extracted_path = Path(target_folder) / file_info.filename
                 
-                if file_info.is_dir():
-                    extracted_path.mkdir(parents=True, exist_ok=True)
-                elif file_info.filename.lower().endswith('.pdf'):
-                    print(f"Extraction de : {file_info.filename} vers {extracted_path}")
-                    try:
-                        with extracted_path.open("wb") as outfile:
-                            outfile.write(z.read(file_info.filename))
-                        extracted_files.append(extracted_path)
-                    except Exception as e:
-                        print(f"Erreur lors de l'extraction de {file_info.filename}: {e}")
-                else:
-                    print(f"Ignoré (non PDF) : {file_info.filename}")
-
-        print(f"Extraction terminée. {len(extracted_files)} fichiers PDF extraits.")
-        return extracted_files
+                if member.filename.lower().endswith('.pdf'):
+                    target_path = target_folder / Path(member.filename).name
+                    with z.open(member) as source, open(target_path, "wb") as target:
+                        target.write(source.read())
+                    extracted_files.append(target_path)
 
+        print(f"✅ {len(extracted_files)} fichiers PDF extraits dans '{target_folder}'.")
+        return extracted_files
     except Exception as e:
-        print(f"Une erreur est survenue lors du téléchargement ou de l'extraction depuis Hugging Face Dataset : {e}")
+        print(f"❌ Erreur lors du téléchargement ou de l'extraction : {e}")
         return []
-    finally:
-        pass
 
-# --- Fonction principale ---
-def main():
-    """Fonction principale pour exécuter l'ensemble du processus d'extraction."""
-    downloaded_pdfs = download_and_extract_zip_from_hf_dataset(
-        HF_DATASET_REPO_ID_PDFS, ZIP_FILENAME_IN_DATASET, PDF_FOLDER
-    )
-    
-    if not downloaded_pdfs:
-        print("Aucun fichier PDF n'a pu être téléchargé ou extrait. Arrêt du processus.")
-        return
+def extract_text_from_pdf(pdf_path):
+    """
+    Extrait le contenu textuel brut d'un fichier PDF.
+    """
+    try:
+        with fitz.open(pdf_path) as doc:
+            return "\n".join(page.get_text() for page in doc)
+    except Exception as e:
+        print(f"⚠️  Avertissement : Impossible de lire le PDF '{pdf_path.name}'. Erreur : {e}")
+        return ""
 
-    unique_folder_name = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + "_" + str(uuid.uuid4())[:8]
-    local_out_dir = Path(LOCAL_OUT_BASE_FOLDER) / unique_folder_name
-    local_out_dir.mkdir(parents=True, exist_ok=True)
+def parse_json_from_model_output(raw_output):
+    """
+    Extrait une chaîne JSON de la sortie brute du modèle.
+    Stratégie robuste : trouve la première '{' et la dernière '}'.
+    """
+    try:
+        # La réponse du modèle inclut le prompt, on ne garde que le texte généré
+        # Le marqueur [/INST] sépare le prompt de la réponse.
+        generated_text = raw_output.split("[/INST]")[-1]
 
-    aggregate = []
-    
-    print(f"Début du traitement de {len(downloaded_pdfs)} fichiers PDF extraits...")
+        start_index = generated_text.find('{')
+        end_index = generated_text.rfind('}')
 
-    for pdf_file in tqdm(downloaded_pdfs, desc="Traitement PDFs"):
-        text = extract_text(pdf_file)
-        if not text.strip():
-            print(f"Le fichier {pdf_file.name} est vide ou l'extraction de texte a échoué, on le passe.")
-            continue
+        if start_index != -1 and end_index != -1 and end_index > start_index:
+            json_str = generated_text[start_index : end_index + 1]
+            return json.loads(json_str)
+        else:
+            raise ValueError("Accolades JSON non trouvées dans la sortie.")
             
-        data = call_huggingface_model(text)
-        aggregate.append(data)
+    except json.JSONDecodeError as e:
+        print(f"Erreur de décodage JSON : {e}")
+        # En cas d'erreur, on retourne un objet d'erreur avec la sortie brute pour le débogage.
+        return {"error": "JSONDecodeError", "details": str(e), "raw_output": raw_output}
+    except Exception as e:
+        print(f"Erreur inattendue lors du parsing : {e}")
+        return {"error": "ParsingFailed", "details": str(e), "raw_output": raw_output}
+
 
-        out_path = local_out_dir / f"{pdf_file.stem}.json"
-        with open(out_path, "w", encoding="utf-8") as f:
-            json.dump(data, f, ensure_ascii=False, indent=2)
+def process_single_pdf(pdf_path, pipeline):
+    """
+    Traite un seul fichier PDF : extraction de texte, appel du modèle et parsing du JSON.
+    """
+    print(f"--- Traitement de : {pdf_path.name} ---")
+    
+    # 1. Extraire le texte
+    text = extract_text_from_pdf(pdf_path)
+    if not text.strip():
+        return {"error": "PDF vide ou illisible", "source_file": pdf_path.name}
+
+    # 2. Préparer le prompt
+    prompt = PROMPT_TEMPLATE.format(texte_pdf=text[:30000]) # Tronquer pour être sûr de ne pas dépasser la limite de contexte
+
+    # 3. Appeler le modèle
+    try:
+        response = pipeline(
+            prompt,
+            max_new_tokens=MAX_NEW_TOKENS,
+            do_sample=False, # Pour des résultats déterministes
+            return_full_text=False, # Ne retourne que le texte généré
+            pad_token_id=pipeline.tokenizer.eos_token_id # Évite un avertissement
+        )
+        raw_output = response[0]['generated_text']
+        
+        # 4. Parser le JSON
+        data = parse_json_from_model_output(f"[INST]{prompt}[/INST]{raw_output}") # Reconstituer pour le parser
+        data['source_file'] = pdf_path.name # Ajouter la source pour la traçabilité
+        return data
 
-    with open(local_out_dir / "brvm_aggregate.json", "w", encoding="utf-8") as f:
-        json.dump(aggregate, f, ensure_ascii=False, indent=2)
+    except Exception as e:
+        print(f"❌ Erreur lors de l'appel du pipeline pour '{pdf_path.name}': {e}")
+        return {"error": "PipelineExecutionError", "details": str(e), "source_file": pdf_path.name}
 
-    print(f"[OK] Extraction terminée - {len(downloaded_pdfs)} fichiers traités dans {local_out_dir}.")
 
-    if HF_TOKEN:
-        try:
-            api = HfApi(token=HF_TOKEN)
-            print(f"Début de l'upload des résultats vers {HF_DATASET_REPO_ID_JSON_OUTPUT}/{unique_folder_name}...")
-            
-            api.upload_folder(
-                folder_path=str(local_out_dir),
-                repo_id=HF_DATASET_REPO_ID_JSON_OUTPUT,
-                repo_type="dataset",
-                path_in_repo=unique_folder_name,
-                commit_message=f"Extraction BRVM du {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}"
-            )
-            print(f"[OK] Upload terminé vers https://huggingface.co/datasets/{HF_DATASET_REPO_ID_JSON_OUTPUT}/tree/main/{unique_folder_name}")
-        except Exception as e:
-            print(f"[ERREUR] Échec de l'upload des résultats vers Hugging Face Dataset : {e}")
-            print("Veuillez vérifier que le HF_TOKEN est correctement configuré avec les permissions d'écriture.")
-            print(f"Et que le dépôt '{HF_DATASET_REPO_ID_JSON_OUTPUT}' existe et que vous y avez accès.")
-    else:
-        print("[AVERTISSEMENT] HF_TOKEN non configuré. Les résultats ne seront PAS uploadés vers le Dataset Hugging Face.")
-        print(f"Les fichiers sont disponibles localement dans le Space à l'adresse : {local_out_dir}")
+def upload_results_to_hf(local_folder, repo_id, hf_token):
+    """
+    Uploade le contenu d'un dossier vers un dataset sur le Hub Hugging Face.
+    """
+    if not hf_token:
+        print("⚠️  Avertissement : HF_TOKEN non configuré. L'upload est ignoré.")
+        print(f"Les résultats sont disponibles localement dans : {local_folder}")
+        return
+
+    try:
+        api = HfApi(token=hf_token)
+        repo_url = api.create_repo(repo_id, repo_type="dataset", exist_ok=True).repo_url
+        
+        commit_message = f"Rapport d'extraction BRVM du {datetime.now().strftime('%Y-%m-%d %H:%M')}"
+        
+        print(f"Début de l'upload de '{local_folder}' vers '{repo_id}'...")
+        api.upload_folder(
+            folder_path=str(local_folder),
+            repo_id=repo_id,
+            repo_type="dataset",
+            commit_message=commit_message
+        )
+        print(f"✅ Upload terminé avec succès ! Consultez les résultats sur : {repo_url}")
+    except Exception as e:
+        print(f"❌ Erreur lors de l'upload vers Hugging Face : {e}")
+        print("Veuillez vérifier votre HF_TOKEN et les permissions d'écriture sur le dépôt.")
+
+
+def main():
+    """
+    Fonction principale orchestrant le processus complet.
+    """
+    # Initialisation du modèle en premier pour échouer rapidement si nécessaire
+    extractor_pipeline = initialize_model_pipeline()
+    if not extractor_pipeline:
+        return # Arrêt si le modèle ne peut pas être chargé
+
+    # Téléchargement et extraction des PDFs
+    pdf_files = download_and_extract_pdfs(
+        HF_DATASET_PDFS_REPO_ID, 
+        ZIP_FILENAME_IN_DATASET, 
+        LOCAL_PDF_FOLDER, 
+        LOCAL_CACHE_DIR
+    )
+    if not pdf_files:
+        print("Aucun PDF à traiter. Arrêt du script.")
+        return
+
+    # Création d'un dossier de sortie unique pour cette exécution
+    run_id = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + "_" + uuid.uuid4().hex[:8]
+    local_output_dir = LOCAL_JSON_OUTPUT_BASE_FOLDER / run_id
+    local_output_dir.mkdir(parents=True, exist_ok=True)
+    print(f"Les résultats JSON seront sauvegardés dans : {local_output_dir}")
+
+    all_results = []
+    for pdf_path in tqdm(pdf_files, desc="Traitement des PDFs"):
+        result = process_single_pdf(pdf_path, extractor_pipeline)
+        all_results.append(result)
+
+        # Sauvegarder le résultat individuel
+        output_json_path = local_output_dir / f"{pdf_path.stem}.json"
+        with open(output_json_path, "w", encoding="utf-8") as f:
+            json.dump(result, f, ensure_ascii=False, indent=2)
+
+    # Sauvegarder le fichier agrégé
+    aggregate_file_path = local_output_dir / "_aggregate_results.json"
+    with open(aggregate_file_path, "w", encoding="utf-8") as f:
+        json.dump(all_results, f, ensure_ascii=False, indent=2)
+    
+    print(f"Traitement terminé. {len(pdf_files)} fichiers traités.")
+
+    # Upload des résultats
+    upload_results_to_hf(local_output_dir, HF_DATASET_JSON_REPO_ID, HF_TOKEN)
 
 
 if __name__ == "__main__":