Update core/extractor.py

core/extractor.py

@@ -5,6 +5,7 @@ from typing import List
 from pydantic import BaseModel, Field
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from gliner import GLiNER # N'oubliez pas d'ajouter 'gliner' dans requirements.txt
+from core.extractor import ExtractorEngine
 
 # --- SCHÉMAS DE DONNÉES ---
 class Entity(BaseModel):
@@ -75,36 +76,46 @@ class ExtractorEngine:
         return final_graph
 
     def _get_dynamic_labels(self, text: str):
-        # Prompt universel qui ne change JAMAIS, quel que soit le document
-        prompt = (
-            "En tant qu'expert en analyse de données, identifie TOUS les types d'entités "
-            "nécessaires pour reconstruire ce document sans perte d'information. "
-            "Inclus les acteurs, les objets, les actions, les chiffres clés, les dates et les lieux. "
-            "Réponds uniquement par une liste de mots séparés par des virgules."
-        )
+        """
+        Analyse le texte intégral pour générer des catégories d'extraction 
+        exhaustives et uniques.
+        """
+        # Prompt pour une analyse totale et sans perte
+        prompt = f"""Tu es un analyste expert en extraction de connaissances. 
+    Analyse l'intégralité du texte ci-dessous et liste tous les types d'entités (catégories) 
+    nécessaires pour reconstruire ce document sous forme de graphe sans perte de précision.
+    
+    Cherche : Acteurs, Méthodologies, Chiffres clés, Unités de mesure, Dates, Lieux, 
+    Variables, Fichiers sources, et Conditions contractuelles.
+    
+    TEXTE COMPLET :
+    {text}
+    
+    Réponds uniquement par une liste de mots simples séparés par des virgules :"""
+    
         inputs = self.tokenizer(prompt, return_tensors="pt").to("cpu")
-            with torch.no_grad():
-                # On laisse un peu plus de tokens pour une liste riche
-                outputs = self.model.generate(**inputs, max_new_tokens=150)
-            
-            res = self.tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True)
-            
-            # --- LOGIQUE DE NETTOYAGE ET DÉDUPLICATION ---
-            raw_labels = res.split(",")
-            clean_labels = []
-            seen = set()
+        with torch.no_grad():
+            # On laisse un peu plus de tokens pour une liste riche
+            outputs = self.model.generate(**inputs, max_new_tokens=150)
         
-            for l in raw_labels:
-                # Nettoyage : retrait des espaces, mise en minuscule pour comparer
-                label = l.strip().replace(".", "").replace("\n", "")
-                if len(label) > 2:
-                    # On normalise (singulier et minuscule) pour éviter les doublons
-                    norm_label = label.lower().rstrip('s') 
-                    if norm_label not in seen:
-                        seen.add(norm_label)
-                        clean_labels.append(label.capitalize()) # On garde un joli format (ex: "Montant")
+        res = self.tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True)
         
-            return clean_labels
+        # --- LOGIQUE DE NETTOYAGE ET DÉDUPLICATION ---
+        raw_labels = res.split(",")
+        clean_labels = []
+        seen = set()
+    
+        for l in raw_labels:
+            # Nettoyage : retrait des espaces, mise en minuscule pour comparer
+            label = l.strip().replace(".", "").replace("\n", "")
+            if len(label) > 2:
+                # On normalise (singulier et minuscule) pour éviter les doublons
+                norm_label = label.lower().rstrip('s') 
+                if norm_label not in seen:
+                    seen.add(norm_label)
+                    clean_labels.append(label.capitalize()) # On garde un joli format (ex: "Montant")
+    
+        return clean_labels
 
     def _run_inference_with_entities(self, text: str, gliner_ents: list, temperature: float):
         """Phase de liaison : le LLM crée le graphe JSON final."""