wikontic 0.0.3__py3-none-any.whl

wikontic/utils/structured_inference_with_db.py

@@ -0,0 +1,561 @@
+from unidecode import unidecode
+import re
+import warnings
+from typing import Dict, List, Tuple
+from langchain.tools import tool
+import logging
+
+from .base_inference_with_db import BaseInferenceWithDB
+
+warnings.filterwarnings("ignore")
+logger = logging.getLogger("StructuredInferenceWithDB")
+logger.setLevel(logging.ERROR)
+
+
+class StructuredInferenceWithDB(BaseInferenceWithDB):
+    def __init__(self, extractor, aligner, triplets_db):
+        self.extractor = extractor
+        self.aligner = aligner
+        self.triplets_db = triplets_db
+
+        self.extract_triplets_with_ontology_filtering_tool = tool(
+            self.extract_triplets_with_ontology_filtering
+        )
+        self.extract_triplets_with_ontology_filtering_and_add_to_db_tool = tool(
+            self.extract_triplets_with_ontology_filtering_and_add_to_db
+        )
+        self.retrieve_similar_entity_names_tool = tool(
+            self.retrieve_similar_entity_names
+        )
+        self.identify_relevant_entities_from_question_tool = tool(
+            self.identify_relevant_entities_from_question_with_llm
+        )
+        self.get_1_hop_supporting_triplets_tool = tool(
+            self.get_1_hop_supporting_triplets
+        )
+        # 1st step extraction without database
+
+    def _refine_entity_types(self, text, triplet):
+        """
+        Refine entity types using LLM.
+        """
+        candidate_subj_type_ids, candidate_obj_type_ids = (
+            self.aligner.retrieve_similar_entity_types(triplet=triplet)
+        )
+
+        candidate_entity_type_id_2_label = self.aligner.retrieve_entity_type_labels(
+            candidate_subj_type_ids + candidate_obj_type_ids
+        )
+
+        candidate_entity_type_label_2_id = {
+            entity_label: entity_id
+            for entity_id, entity_label in candidate_entity_type_id_2_label.items()
+        }
+
+        candidate_subject_types = [
+            candidate_entity_type_id_2_label[t] for t in candidate_subj_type_ids
+        ]
+        candidate_object_types = [
+            candidate_entity_type_id_2_label[t] for t in candidate_obj_type_ids
+        ]
+
+        # no need to refine if the triplet's types are in the candidate types
+        if (
+            triplet["subject_type"] in candidate_subject_types
+            and triplet["object_type"] in candidate_object_types
+        ):
+            refined_subject_type, refined_object_type = (
+                triplet["subject_type"],
+                triplet["object_type"],
+            )
+            refined_subject_type_id = candidate_entity_type_label_2_id[
+                triplet["subject_type"]
+            ]
+            refined_object_type_id = candidate_entity_type_label_2_id[
+                triplet["object_type"]
+            ]
+
+        else:
+            # if the triplet's subject type is in the candidate types,
+            # then only refine the subject type
+            if triplet["subject_type"] in candidate_subject_types:
+                candidate_subject_types = [triplet["subject_type"]]
+            # if the triplet's object type is in the candidate types,
+            # then only refine the object type
+            if triplet["object_type"] in candidate_object_types:
+                candidate_object_types = [triplet["object_type"]]
+
+            self.extractor.reset_error_state()
+            refined_entity_types = self.extractor.refine_entity_types(
+                text=text,
+                triplet=triplet,
+                candidate_subject_types=candidate_subject_types,
+                candidate_object_types=candidate_object_types,
+            )
+            refined_subject_type, refined_object_type = (
+                refined_entity_types["subject_type"],
+                refined_entity_types["object_type"],
+            )
+
+            refined_subject_type_id = (
+                candidate_entity_type_label_2_id[refined_subject_type]
+                if refined_subject_type in candidate_subject_types
+                else None
+            )
+
+            refined_object_type_id = (
+                candidate_entity_type_label_2_id[refined_object_type]
+                if refined_object_type in candidate_object_types
+                else None
+            )
+
+        return (
+            refined_subject_type,
+            refined_subject_type_id,
+            refined_object_type,
+            refined_object_type_id,
+        )
+
+    def _get_candidate_entity_properties(
+        self, triplet: Dict[str, str], subj_type_ids: List[str], obj_type_ids: List[str]
+    ) -> Tuple[List[Tuple[str, str]], Dict[str, dict]]:
+        """
+        Retrieve candidate properties and their labels/constraints.
+        """
+        # Get the list of tuples (<property_id>, <property_direction>)
+        properties: List[Tuple[str, str]] = (
+            self.aligner.retrieve_properties_for_entity_type(
+                target_relation=triplet["relation"],
+                object_types=obj_type_ids,
+                subject_types=subj_type_ids,
+                k=10,
+            )
+        )
+        # Get dict {<prop_id>:
+        #           {"label": <prop_label>,
+        #           "valid_subject_type_ids": <valid_subject_type_ids>,
+        #           "valid_object_type_ids": <valid_object_type_ids>}}
+        prop_2_label_and_constraint = (
+            self.aligner.retrieve_properties_labels_and_constraints(
+                property_id_list=[p[0] for p in properties]
+            )
+        )
+        return properties, prop_2_label_and_constraint
+
+    def _refine_relation(
+        self, text, triplet, refined_subject_type_id, refined_object_type_id
+    ):
+        """
+        Refine relation using LLM.
+        """
+        # if refined subject and object types are in the candidate types,
+        # then refine the relation
+        if refined_subject_type_id and refined_object_type_id:
+            relation_direction_candidate_pairs, prop_2_label_and_constraint = (
+                self._get_candidate_entity_properties(
+                    triplet=triplet,
+                    subj_type_ids=[refined_subject_type_id],
+                    obj_type_ids=[refined_object_type_id],
+                )
+            )
+            candidate_relations = [
+                prop_2_label_and_constraint[p[0]]["label"]
+                for p in relation_direction_candidate_pairs
+            ]
+            # no need to refine
+            # if the triplet's relation is in the candidate relations
+            if triplet["relation"] in candidate_relations:
+                refined_relation = triplet["relation"]
+            else:
+                self.extractor.reset_error_state()
+                refined_relation = self.extractor.refine_relation(
+                    text=text, triplet=triplet, candidate_relations=candidate_relations
+                )["relation"]
+        # if refined subject and object types are not in the candidate types,
+        # leave relation as it is
+        else:
+            refined_relation = triplet["relation"]
+            candidate_relations = []
+
+        # if refined relation is in the candidate relations,
+        # then identify the relation direction
+        if refined_relation in candidate_relations:
+            refined_relation_id_candidates = [
+                p_id
+                for p_id in prop_2_label_and_constraint
+                if prop_2_label_and_constraint[p_id]["label"] == refined_relation
+            ]
+            refined_relation_id = refined_relation_id_candidates[0]
+            refined_relation_directions = [
+                p[1]
+                for p in relation_direction_candidate_pairs
+                if p[0] == refined_relation_id
+            ]
+            refined_relation_direction = (
+                "direct" if "direct" in refined_relation_directions else "inverse"
+            )
+
+            prop_subject_type_ids = [
+                prop_2_label_and_constraint[prop]["valid_subject_type_ids"]
+                for prop in prop_2_label_and_constraint
+                if prop_2_label_and_constraint[prop]["label"] == refined_relation
+            ][0]
+            prop_object_type_ids = [
+                prop_2_label_and_constraint[prop]["valid_object_type_ids"]
+                for prop in prop_2_label_and_constraint
+                if prop_2_label_and_constraint[prop]["label"] == refined_relation
+            ][0]
+
+        else:
+            refined_relation_direction = "direct"
+            refined_relation_id = None
+            prop_subject_type_ids = []
+            prop_object_type_ids = []
+
+        return (
+            refined_relation,
+            refined_relation_id,
+            refined_relation_direction,
+            prop_subject_type_ids,
+            prop_object_type_ids,
+        )
+
+    def _validate_backbone(
+        self,
+        refined_subject_type: str,
+        refined_object_type: str,
+        refined_relation: str,
+        refined_object_type_id: str,
+        refined_subject_type_id: str,
+        refined_relation_id: str,
+        valid_subject_type_ids: List[str],
+        valid_object_type_ids: List[str],
+    ):
+        """
+        Check if the selected backbone_triplet's types and relation are in the valid sets.
+        """
+
+        exception_msg = ""
+        if not refined_relation_id:
+            exception_msg += "Refined relation not in candidate relations\n"
+        if not refined_subject_type_id:
+            exception_msg += "Refined subject type not in candidate subject types\n"
+        if not refined_object_type_id:
+            exception_msg += "Refined object type not in candidate object types\n"
+
+        if exception_msg != "":
+            return False, exception_msg
+
+        else:
+
+            subject_type_hierarchy = self.aligner.retrieve_entity_type_hierarchy(
+                refined_subject_type
+            )
+            object_type_hierarchy = self.aligner.retrieve_entity_type_hierarchy(
+                refined_object_type
+            )
+
+            if valid_subject_type_ids == ["ANY"]:
+                valid_subject_type_ids = subject_type_hierarchy
+            if valid_object_type_ids == ["ANY"]:
+                valid_object_type_ids = object_type_hierarchy
+
+            if any(
+                [t in subject_type_hierarchy for t in valid_subject_type_ids]
+            ) and any([t in object_type_hierarchy for t in valid_object_type_ids]):
+                return True, exception_msg
+            else:
+                exception_msg += "Triplet backbone violates property constraints\n"
+                return False, exception_msg
+
+    def _refine_entity_name(self, text, triplet, sample_id, is_object=False):
+        """
+        Refine entity names using type constraints.
+        """
+        self.extractor.reset_error_state()
+        if is_object:
+            entity = unidecode(triplet["object"])
+            entity_type = triplet["object_type"]
+            entity_hierarchy = self.aligner.retrieve_entity_type_hierarchy(entity_type)
+        else:
+            entity = unidecode(triplet["subject"])
+            entity_type = triplet["subject_type"]
+            entity_hierarchy = []
+
+        # do not change time or quantity entities (of objects!)
+        if any([t in ["Q186408", "Q309314"] for t in entity_hierarchy]):
+            updated_entity = entity
+        else:
+            # if not time or quantity entities -> retrieve similar entities by type and name similarity
+            similar_entities = self.aligner.retrieve_entity_by_type(
+                entity_name=entity, entity_type=entity_type, sample_id=sample_id
+            )
+            # if there are similar entities -> refine entity name
+            if len(similar_entities) > 0:
+                # if exact match found -> return the exact match
+                if entity in similar_entities:
+                    updated_entity = similar_entities[entity]
+                else:
+                    # if not exact match -> refine entity name
+                    updated_entity = self.extractor.refine_entity(
+                        text=text,
+                        triplet=triplet,
+                        candidates=list(similar_entities.values()),
+                        is_object=is_object,
+                    )
+                    # unidecode the updated entity
+                    updated_entity = unidecode(updated_entity)
+                    # if the updated entity is None (meaning that LLM didn't find any similar entities)
+                    # -> return the original entity
+                    if re.sub(r"[^\w\s]", "", updated_entity) == "None":
+                        updated_entity = entity
+            else:
+                # if no similar entities -> return the original entity
+                updated_entity = entity
+
+        self.aligner.add_entity(
+            entity_name=updated_entity,
+            alias=entity,
+            entity_type=entity_type,
+            sample_id=sample_id,
+        )
+
+        return updated_entity
+
+    def extract_triplets_with_ontology_filtering(
+        self, text, sample_id=None, source_text_id=None
+    ):
+        """
+        Extract and refine knowledge graph triplets from text using LLM.
+
+        Args:
+            text (str): Input text to extract triplets from
+            sample_id (str): Sample ID - used to distinguish graphs resulted
+                in different launches/by different users
+            source_text_id (str): Optional, - used to distinguish texts from
+                different sources (for example, different paragraphs of the same text)
+        Returns:
+            tuple:
+                (initial_triplets, final_triplets, filtered_triplets, ontology_filtered_triplets)
+        """
+        self.extractor.reset_tokens()
+        self.extractor.reset_messages()
+        self.extractor.reset_error_state()
+
+        extracted_triplets = self.extractor.extract_triplets_from_text(text)
+
+        initial_triplets = []
+        for triplet in extracted_triplets["triplets"]:
+            triplet["prompt_token_num"], triplet["completion_token_num"] = (
+                self.extractor.calculate_used_tokens()
+            )
+            triplet["source_text_id"] = source_text_id
+            triplet["sample_id"] = sample_id
+            initial_triplets.append(triplet.copy())
+
+        final_triplets = []
+        filtered_triplets = []
+        ontology_filtered_triplets = []
+
+        for triplet in extracted_triplets["triplets"]:
+            self.extractor.reset_tokens()
+            try:
+                logger.log(logging.DEBUG, "Triplet: %s\n%s" % (str(triplet), "-" * 100))
+
+                # _____________ Refine entity types  __________
+
+                (
+                    refined_subject_type,
+                    refined_subject_type_id,
+                    refined_object_type,
+                    refined_object_type_id,
+                ) = self._refine_entity_types(text=text, triplet=triplet)
+
+                # ________________ Refine relation ________________
+                (
+                    refined_relation,
+                    refined_relation_id,
+                    refined_relation_direction,
+                    prop_subject_type_ids,
+                    prop_object_type_ids,
+                ) = self._refine_relation(
+                    text=text,
+                    triplet=triplet,
+                    refined_subject_type_id=refined_subject_type_id,
+                    refined_object_type_id=refined_object_type_id,
+                )
+
+                if refined_relation_direction == "inverse":
+                    refined_subject_type_id, refined_object_type_id = (
+                        refined_object_type_id,
+                        refined_subject_type_id,
+                    )
+
+                # __________ Refine entity names ___________
+                backbone_triplet = {
+                    "subject": (
+                        triplet["subject"]
+                        if refined_relation_direction == "direct"
+                        else triplet["object"]
+                    ),
+                    "relation": refined_relation,
+                    "object": (
+                        triplet["object"]
+                        if refined_relation_direction == "direct"
+                        else triplet["subject"]
+                    ),
+                    "subject_type": (
+                        refined_subject_type
+                        if refined_relation_direction == "direct"
+                        else refined_object_type
+                    ),
+                    "object_type": (
+                        refined_object_type
+                        if refined_relation_direction == "direct"
+                        else refined_subject_type
+                    ),
+                }
+
+                backbone_triplet["qualifiers"] = triplet["qualifiers"]
+                if refined_subject_type_id:
+                    backbone_triplet["subject"] = self._refine_entity_name(
+                        text, backbone_triplet, sample_id, is_object=False
+                    )
+
+                if refined_object_type_id:
+                    backbone_triplet["object"] = self._refine_entity_name(
+                        text, backbone_triplet, sample_id, is_object=True
+                    )
+
+                logger.log(
+                    logging.DEBUG,
+                    "Original subject name: %s\n%s"
+                    % (str(backbone_triplet["subject"]), "-" * 100),
+                )
+                logger.log(
+                    logging.DEBUG,
+                    "Original object name: %s\n%s"
+                    % (str(backbone_triplet["object"]), "-" * 100),
+                )
+                logger.log(
+                    logging.DEBUG,
+                    "Refined subject name: %s\n%s"
+                    % (str(backbone_triplet["subject"]), "-" * 100),
+                )
+                logger.log(
+                    logging.DEBUG,
+                    "Refined object name: %s\n%s"
+                    % (str(backbone_triplet["object"]), "-" * 100),
+                )
+
+                (
+                    backbone_triplet["prompt_token_num"],
+                    backbone_triplet["completion_token_num"],
+                ) = self.extractor.calculate_used_tokens()
+                backbone_triplet["source_text_id"] = source_text_id
+                backbone_triplet["sample_id"] = sample_id
+
+                # ___________________________ Validate backbone triplet ___________________________
+                backbone_triplet_valid, backbone_triplet_exception_msg = (
+                    self._validate_backbone(
+                        backbone_triplet["subject_type"],
+                        backbone_triplet["object_type"],
+                        backbone_triplet["relation"],
+                        refined_object_type_id,
+                        refined_subject_type_id,
+                        refined_relation_id,
+                        prop_subject_type_ids,
+                        prop_object_type_ids,
+                    )
+                )
+
+                if backbone_triplet_valid:
+                    final_triplets.append(backbone_triplet.copy())
+                    logger.log(
+                        logging.DEBUG,
+                        "Final triplet: %s\n%s" % (str(backbone_triplet), "-" * 100),
+                    )
+                else:
+                    logger.log(
+                        logging.ERROR,
+                        "Final triplet is ontology filtered: %s\n%s"
+                        % (str(backbone_triplet), "-" * 100),
+                    )
+                    logger.log(
+                        logging.ERROR,
+                        "Exception: %s" % (str(backbone_triplet_exception_msg)),
+                    )
+                    logger.log(
+                        logging.ERROR, "Refined relation: %s" % (str(refined_relation))
+                    )
+                    logger.log(
+                        logging.ERROR,
+                        "Refined subject type: %s" % (str(refined_subject_type)),
+                    )
+                    logger.log(
+                        logging.ERROR,
+                        "Refined object type: %s" % (str(refined_object_type)),
+                    )
+
+                    backbone_triplet["exception_text"] = backbone_triplet_exception_msg
+                    ontology_filtered_triplets.append(backbone_triplet.copy())
+
+            except Exception as e:
+                backbone_triplet = triplet.copy()
+                (
+                    backbone_triplet["prompt_token_num"],
+                    backbone_triplet["completion_token_num"],
+                ) = self.extractor.calculate_used_tokens()
+                backbone_triplet["source_text_id"] = source_text_id
+                backbone_triplet["sample_id"] = sample_id
+                backbone_triplet["exception_text"] = str(e)
+                filtered_triplets.append(backbone_triplet.copy())
+                logger.log(
+                    logging.INFO,
+                    "Filtered triplet: %s\n%s" % (str(backbone_triplet), "-" * 100),
+                )
+                logger.log(logging.INFO, "Exception: %s" % (str(e)))
+
+        return (
+            initial_triplets,
+            final_triplets,
+            filtered_triplets,
+            ontology_filtered_triplets,
+        )
+
+    def extract_triplets_with_ontology_filtering_and_add_to_db(
+        self, text, sample_id=None, source_text_id=None
+    ):
+        """
+        Extract and refine knowledge graph triplets from text using LLM, then add them to the database.
+        Args:
+            text (str): Input text to extract triplets from
+            sample_id (str): Sample ID - used to distinguish graphs resulted in different launches/by different users
+            source_text_id (str): Optional, - used to distinguish text from different sources (e.g., different paragraphs of the same text)
+        Returns:
+            tuple: (initial_triplets, final_triplets, filtered_triplets, ontology_filtered_triplets)
+        """
+        (
+            initial_triplets,
+            final_triplets,
+            filtered_triplets,
+            ontology_filtered_triplets,
+        ) = self.extract_triplets_with_ontology_filtering(
+            text, sample_id=sample_id, source_text_id=source_text_id
+        )
+        if len(initial_triplets) > 0:
+            self.aligner.add_initial_triplets(initial_triplets, sample_id=sample_id)
+        if len(final_triplets) > 0:
+            self.aligner.add_triplets(final_triplets, sample_id=sample_id)
+        if len(filtered_triplets) > 0:
+            self.aligner.add_filtered_triplets(filtered_triplets, sample_id=sample_id)
+        if len(ontology_filtered_triplets) > 0:
+            self.aligner.add_ontology_filtered_triplets(
+                ontology_filtered_triplets, sample_id=sample_id
+            )
+        return (
+            initial_triplets,
+            final_triplets,
+            filtered_triplets,
+            ontology_filtered_triplets,
+        )

wikontic-0.0.3.dist-info/METADATA

@@ -0,0 +1,111 @@
+Metadata-Version: 2.4
+Name: wikontic
+Version: 0.0.3
+Summary: Extract a knowledge graph with LLM from texts and perform QA over the resulted KG
+Author-email: Alla Chepurova <chepurova.data@gmail.com>
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/screemix/Wikontic
+Project-URL: Issues, https://github.com/screemix/Wikontic/issues
+Classifier: Programming Language :: Python :: 3
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: streamlit
+Requires-Dist: numpy
+Requires-Dist: pyvis
+Requires-Dist: python-dotenv
+Requires-Dist: pymongo
+Requires-Dist: openai
+Requires-Dist: tenacity
+Requires-Dist: pathlib
+Requires-Dist: typing
+Requires-Dist: unidecode
+Requires-Dist: torch>=2.4.0
+Requires-Dist: transformers
+Requires-Dist: dataclasses
+Requires-Dist: pydantic
+Requires-Dist: accelerate
+Requires-Dist: langchain
+Dynamic: license-file
+
+![Wikontic logo](/media/wikontic.png)
+
+# Wikontic
+
+**Build ontology-aware, Wikidata-aligned knowledge graphs from raw text using LLMs**
+
+---
+
+## 🚀 Overview
+
+Knowledge Graphs (KGs) provide structured, verifiable representations of knowledge, enabling fact grounding and empowering large language models (LLMs) with up-to-date, real-world information. However, creating high-quality KGs from open-domain text is challenging due to issues like redundancy, inconsistency, and lack of alignment with formal ontologies.
+
+**Wikontic** is a multi-stage pipeline for constructing ontology-aligned KGs from unstructured text using LLMs and Wikidata. It extracts candidate triples from raw text, then refines them through ontology-based typing, schema validation, and entity deduplication—resulting in compact, semantically coherent graphs.
+
+---
+
+## 📁 Repository Structure
+
+- `preprocessing/constraint-preprocessing.ipynb`  
+  Jupyter notebook for collecting constraint rules from Wikidata.
+
+- `utils/`  
+  Utilities for LLM-based triple extraction and alignment with Wikidata ontology rules.
+
+
+- `utils/openai_utils.py`  
+  `LLMTripletExtractor` class for LLM-based triple extraction.
+
+
+### To use ontology:
+
+- `utils/ontology_mappings/`  
+  JSON files containing ontology mappings from Wikidata.
+
+- `utils/structured_inference_with_db.py`  
+  - `StructuredInferenceWithDB` class: triple extraction and qa functions
+
+- `utils/structured_aligner.py`
+  -  `Aligner` class: ontology alignment and entity name refinement
+
+
+### Not to use ontology:
+- `utils/inference_with_db.py`
+  - `InferenceWithDB` class: triple extraction and qa functions
+
+- `utils/dynamic_aligner.py`
+  -  `Aligner` class: entity and relation name refinement
+
+### Evaluation:
+- `inference_and_eval`
+	- Scripts for building KGs for MuSiQue and HotPot datasets and evaluation of QA performance
+- `analysis`
+  - Notebooks with downstream analysis of the resulted KG
+
+### Use Wikontic as a service:
+
+- `pages/` and `Wikontic.py`  
+  Code for the web service for knowledge graph extraction and visualization.
+
+- `Dockerfile`  
+  For building a containerized web service.
+
+
+---
+
+## 🏁 Getting Started
+
+1. **Set up the ontology and KG databases:**
+   ```
+   ./setup_db.sh
+   ```
+
+2. **Launch the web service:**
+   ```
+   streamlit run Wikontic.py
+   ```
+
+---
+
+Enjoy building knowledge graphs with Wikontic!