Update core/extractor.py

core/extractor.py

@@ -15,8 +15,8 @@ class Entity(BaseModel):
 class Relationship(BaseModel):
     source: str = Field(alias="from", description="ID de l'entité source.")
     target: str = Field(alias="to", description="ID de l'entité cible.")
-    type: str = Field(description="Action ou lien (ex: USES, CORRELATED_WITH, AUTHORED_BY).")
-    description: str = Field(description="Explication du lien.")
+    type: str = Field(description="Action ou lien sémantique (ex: USES, CORRELATED_WITH, AUTHORED_BY).")
+    description: str = Field(description="Explication du lien ou Détails.")
 
 class KnowledgeGraph(BaseModel):
     entities: List[Entity]
@@ -35,30 +35,89 @@ class ExtractorEngine:
         self.model = st.session_state.llm_model
         self.json_schema = json.dumps(KnowledgeGraph.model_json_schema(), indent=2)
 
-    def extract_graph(self, text: str):
-        # Votre version du prompt renforcée
-        system_prompt = """Tu es un système d'extraction de graphe de connaissance hautement fiable. 
-        Ton objectif est d'extraire toutes les entités et relations pertinentes du texte fourni. 
-        Réponds TOUJOURS uniquement en JSON. Le JSON DOIT respecter le schéma spécifié ci-dessous 
-        SANS AUCUNE EXPLICATION SUPPLÉMENTAIRE."""
+        
+    def extract_long_text(self, text: str, temperature: float, chunk_size: int = 3000):
+        # Découpage du texte
+        chunks = [text[i:i + chunk_size] for i in range(0, len(text), chunk_size)]
+        final_graph = {"entities": [], "relationships": []}
+        entity_map = {} # Pour fusionner les noms identiques
 
-        user_prompt = f"Schéma:\n{self.json_schema}\n\nTexte:\n{text[:4000]}\n\nRéponse JSON:"
+        for chunk in chunks:
+            raw_res = self._run_inference(chunk, temperature)
+            if raw_res:
+                # Fusion des entités par nom
+                for ent in raw_res.get("entities", []):
+                    name_key = ent["name"].lower().strip()
+                    if name_key not in entity_map:
+                        new_id = f"E{len(entity_map) + 1}"
+                        entity_map[name_key] = new_id
+                        ent["id"] = new_id
+                        final_graph["entities"].append(ent)
+                
+                # Ré-attribution des IDs dans les relations
+                for rel in raw_res.get("relationships", []):
+                    # (Note: simplifiée ici, nécessite que l'IA respecte les noms dans le chunk)
+                    final_graph["relationships"].append(rel)
+        
+        return final_graph
 
-        try:
-            inputs = self.tokenizer.apply_chat_template(
-                [{"role": "system", "content": system_prompt}, {"role": "user", "content": user_prompt}],
-                tokenize=True, add_generation_prompt=True, return_tensors="pt"
-            ).to("cpu")
 
-            with torch.no_grad():
-                outputs = self.model.generate(inputs, max_new_tokens=1500, temperature=0.1)
+        
+    def _run_inference(self, text: str, temperature: float):
+            """
+            Exécute l'inférence sur un segment de texte en utilisant le prompt renforcé 
+            et la température réglable depuis l'interface.
+            """
+            # Utilisation de votre prompt système renforcé pour une fiabilité maximale
+            system_prompt = """Tu es un système d'extraction de graphe de connaissance hautement fiable. 
+            Ton objectif est d'extraire toutes les entités et relations pertinentes du texte fourni. 
+            Réponds TOUJOURS uniquement en JSON. Le JSON DOIT respecter le schéma spécifié ci-dessous 
+            SANS AUCUNE EXPLICATION SUPPLÉMENTAIRE."""
+    
+            # Construction du prompt utilisateur avec le segment de texte
+            user_prompt = f"Schéma JSON STRICT à respecter:\n{self.json_schema}\n\nTexte source:\n<<<{text}>>>\n\nRéponse JSON:"
             
-            result = self.tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True)
-            return json.loads(self._clean(result))
-        except Exception as e:
-            st.error(f"Erreur IA : {e}")
-            return None
-
-    def _clean(self, t):
-        return t.strip().replace("```json", "").replace("```", "")
+            try:
+                # Application du template de chat spécifique à Qwen
+                inputs = self.tokenizer.apply_chat_template(
+                    [
+                        {"role": "system", "content": system_prompt}, 
+                        {"role": "user", "content": user_prompt}
+                    ],
+                    tokenize=True, 
+                    add_generation_prompt=True, 
+                    return_tensors="pt"
+                ).to("cpu")
+    
+                # Génération avec les paramètres optimisés
+                with torch.no_grad():
+                    outputs = self.model.generate(
+                        inputs, 
+                        max_new_tokens=1500,     # Augmenté pour ne pas couper les gros JSON
+                        temperature=temperature, # Dynamique via le slider Streamlit
+                        do_sample=True if temperature > 0.1 else False, 
+                        pad_token_id=self.tokenizer.eos_token_id
+                    )
+                
+                # Décodage de la réponse
+                generated_text = self.tokenizer.decode(
+                    outputs[0][inputs.shape[1]:], 
+                    skip_special_tokens=True
+                )
+                
+                # Nettoyage et conversion en dictionnaire Python
+                clean_content = self._clean(generated_text)
+                return json.loads(clean_content)
+                
+            except Exception as e:
+                # En cas d'erreur de parsing ou de génération sur ce segment
+                st.warning(f"Avertissement sur un segment : {e}")
+                return None
+    
+        def _clean(self, t):
+            """Nettoyage rigoureux des balises Markdown et espaces superflus."""
+            t = t.strip()
+            if t.startswith("```"):
+                t = t.replace("```json", "").replace("```", "")
+            return t.strip()