gtyd

app.py

@@ -11,41 +11,39 @@ import fitz
 import json
 from pathlib import Path
 from tqdm import tqdm
-from transformers import pipeline
+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 # Importation pour télécharger et uploader
+from huggingface_hub import hf_hub_download, HfApi
 from datetime import datetime
-import uuid # Pour générer un identifiant unique
+import uuid
 
 # ---------- CONFIG ----------
-MODEL = "mistralai/Mistral-7B-Instruct-v0.3" # Modèle Hugging Face à utiliser
+MODEL = "mistralai/Mistral-7B-Instruct-v0.3"
 
-# Hugging Face Dataset d'où télécharger le fichier ZIP des PDF
-HF_DATASET_REPO_ID_PDFS = "lamekemal/brvm-reports-pdfs" # L'ID de votre dépôt de dataset pour les PDF
-ZIP_FILENAME_IN_DATASET = "brvm_reports.zip" # Le nom du fichier ZIP à l'intérieur de ce dataset
+HF_DATASET_REPO_ID_PDFS = "lamekemal/brvm-reports-pdfs"
+ZIP_FILENAME_IN_DATASET = "brvm_reports.zip"
 
-# Hugging Face Dataset où uploader les résultats JSON
 HF_DATASET_REPO_ID_JSON_OUTPUT = "lamekemal/brvm-reports-json"
 
-# Vous devez ajouter votre jeton Hugging Face (avec rôle "Write" pour l'upload)
-# comme un secret dans les paramètres de votre Hugging Face Space.
-# Le nom du secret DOIT être HF_TOKEN.
-# os.getenv('HF_TOKEN') récupérera automatiquement cette valeur.
 HF_TOKEN = os.getenv('HF_TOKEN')
 
-# Dossiers locaux pour le traitement
-PDF_FOLDER = "brvm_reports"    # Dossier où les PDF seront extraits et traités
-LOCAL_OUT_BASE_FOLDER = "local_json_outputs" # Dossier de base local pour les sorties JSON temporaires
+PDF_FOLDER = "brvm_reports"
+LOCAL_OUT_BASE_FOLDER = "local_json_outputs"
 # ----------------------------
 
+# Prompt optimisé pour le modèle Mistral Instruct
 PROMPT_TEMPLATE = """
-Tu es un expert en finance spécialisé dans la BRVM.
+[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 :
+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].
 
+JSON Schema:
 {{
   "indicateurs": {{
     "brvm_10": {{ "niveau": float, "var_jour_pct": float, "var_annuelle_pct": float }},
@@ -73,27 +71,49 @@ extrait uniquement les données suivantes et retourne un JSON valide au format s
 }}
 
 Contraintes :
-- Ne mets aucun texte hors du JSON
-- Si une donnée est absente, mets null
-- Les nombres utilisent un point comme séparateur décimal
+- Si une donnée est absente, mets null.
+- Les nombres utilisent un point comme séparateur décimal.
 
 Texte du bulletin :
-{texte_pdf}
+{texte_pdf}[/INST]
 """
+# --- Initialisation du pipeline Hugging Face avec quantification ---
+extractor = None # Initialiser à None pour gérer les erreurs de chargement
 
-# --- Initialisation du pipeline Hugging Face ---
 try:
-    print("Chargement du modèle Hugging Face...")
-    # Le pipeline utilisera automatiquement le jeton HF_TOKEN si disponible en tant que secret/variable d'environnement
+    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,
+        model=model,
+        tokenizer=tokenizer,
         device=0 # Utilise le GPU si disponible (0 pour le premier GPU, -1 pour CPU)
     )
-    print("Modèle chargé avec succès.")
+    print("Modèle quantifié chargé avec succès.")
 except Exception as e:
     print(f"Erreur lors du chargement du modèle Hugging Face : {e}")
-    extractor = None
+    # Ne pas initialiser extractor si erreur
 
 # --- Fonctions d'extraction et de traitement ---
 def extract_text(pdf_path):
@@ -109,32 +129,63 @@ def extract_text(pdf_path):
 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)
     
     try:
-        # Ajuster max_new_tokens si les réponses JSON sont tronquées
-        response = extractor(prompt, max_new_tokens=2048) 
+        # 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
+        ) 
         
-        # Le résultat est souvent une liste de dictionnaires, on prend le premier
         output_text = response[0]['generated_text']
 
-        # Il se peut que le modèle répète le prompt, on doit donc extraire la partie JSON
-        json_start = output_text.find('{')
-        json_end = output_text.rfind('}') + 1
+        json_start_marker = "[DEBUT_JSON]"
+        json_end_marker = "[FIN_JSON]"
+
+        # 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
+
+        # 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 = output_text[json_start:json_end]
-            return json.loads(json_text)
+            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": "JSON non trouvé dans la réponse", "raw": output_text}
+            return {"error": "Marqueurs JSON non trouvés", "raw": output_text}
             
     except json.JSONDecodeError:
-        print(f"Erreur JSONDecodeError: {output_text}") # Afficher le texte brut qui a causé l'erreur JSON
-        return {"error": "JSON invalide", "raw": output_text}
+        print(f"Erreur JSONDecodeError: Le texte brut était : {output_text[:500]}...")
+        return {"error": "JSON invalide ou mal formé", "raw": output_text}
     except Exception as e:
-        return {"error": str(e)}
+        return {"error": str(e), "raw": output_text}
 
 def download_and_extract_zip_from_hf_dataset(dataset_repo_id, zip_filename, target_folder):
     """
@@ -144,34 +195,28 @@ def download_and_extract_zip_from_hf_dataset(dataset_repo_id, zip_filename, targ
     print(f"Tentative de téléchargement du fichier '{zip_filename}' "
           f"depuis le dataset Hugging Face : {dataset_repo_id}")
     try:
-        # Télécharger le fichier ZIP depuis le dataset Hugging Face Hub
         local_zip_path = hf_hub_download(
             repo_id=dataset_repo_id,
             filename=zip_filename,
             repo_type="dataset",
-            cache_dir="./hf_cache" # Vous pouvez spécifier un répertoire de cache
+            cache_dir="./hf_cache"
         )
         print(f"Fichier ZIP téléchargé localement : {local_zip_path}")
         
-        # Créer le dossier cible s'il n'existe pas
         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():
-                # Construire le chemin complet du fichier extrait
-                # Éviter les chemins absolus ou les traversées de répertoire pour la sécurité
                 if file_info.filename.startswith('/') or '..' in file_info.filename:
                     print(f"Skipping potentially unsafe path: {file_info.filename}")
                     continue
 
                 extracted_path = Path(target_folder) / file_info.filename
                 
-                # Créer les sous-dossiers si nécessaire
                 if file_info.is_dir():
                     extracted_path.mkdir(parents=True, exist_ok=True)
                 elif file_info.filename.lower().endswith('.pdf'):
-                    # Extraire uniquement les fichiers PDF
                     print(f"Extraction de : {file_info.filename} vers {extracted_path}")
                     try:
                         with extracted_path.open("wb") as outfile:
@@ -189,13 +234,11 @@ def download_and_extract_zip_from_hf_dataset(dataset_repo_id, zip_filename, targ
         print(f"Une erreur est survenue lors du téléchargement ou de l'extraction depuis Hugging Face Dataset : {e}")
         return []
     finally:
-        # Laisser hf_hub_download gérer son propre cache.
         pass
 
 # --- Fonction principale ---
 def main():
     """Fonction principale pour exécuter l'ensemble du processus d'extraction."""
-    # 1. Télécharger et extraire les PDF depuis le Hugging Face Dataset
     downloaded_pdfs = download_and_extract_zip_from_hf_dataset(
         HF_DATASET_REPO_ID_PDFS, ZIP_FILENAME_IN_DATASET, PDF_FOLDER
     )
@@ -204,8 +247,6 @@ def main():
         print("Aucun fichier PDF n'a pu être téléchargé ou extrait. Arrêt du processus.")
         return
 
-    # 2. Préparer le dossier de sortie LOCAL unique pour les JSON
-    # Créer un nom de dossier unique basé sur la date/heure et un UUID
     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)
@@ -223,31 +264,25 @@ def main():
         data = call_huggingface_model(text)
         aggregate.append(data)
 
-        # Sauvegarde par fichier dans le dossier local unique
         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)
 
-    # Sauvegarde du fichier global aggrégé dans le dossier local unique
     with open(local_out_dir / "brvm_aggregate.json", "w", encoding="utf-8") as f:
         json.dump(aggregate, f, ensure_ascii=False, indent=2)
 
     print(f"[OK] Extraction terminée - {len(downloaded_pdfs)} fichiers traités dans {local_out_dir}.")
 
-    # 3. Uploader les résultats vers le Hugging Face Dataset dédié
     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}...")
             
-            # Utilisation de upload_large_folder pour les gros volumes de fichiers
-            api.upload_folder( # Note: upload_folder est recommandé même pour les "large folders" pour les datasets
-                                # C'est le message d'avertissement qui est trompeur.
-                                # L'API gère LFS en interne pour les gros fichiers.
+            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, # Le dossier à créer dans le 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}")
@@ -260,6 +295,5 @@ def main():
         print(f"Les fichiers sont disponibles localement dans le Space à l'adresse : {local_out_dir}")
 
 
-# Point d'entrée du script
 if __name__ == "__main__":
     main()