wikontic 0.0.3__py3-none-any.whl

wikontic/utils/base_inference_with_db.py

@@ -0,0 +1,329 @@
+from typing import Dict, List, Optional
+import logging
+
+logger = logging.getLogger("BaseInferenceWithDB")
+logger.setLevel(logging.ERROR)
+
+
+class BaseInferenceWithDB:
+    """
+    Base class for inference with database functionality.
+    Contains common methods shared by InferenceWithDB and StructuredInferenceWithDB.
+
+    Note: This is an abstract base class. Child classes must define the following
+    attributes in their __init__ methods:
+        - self.extractor: The extractor instance
+        - self.aligner: The aligner instance
+        - self.triplets_db: The triplets database instance
+    """
+
+    def retrieve_similar_entity_names(
+        self, entity_name: str, k: int, sample_id: Optional[str] = None
+    ) -> List[Dict[str, str]]:
+        """
+        Retrieve similar entity names from the knowledge graph using vector search.
+        Useful to link entities from the question to the knowledge graph.
+        Args:
+            entity_name: The entity name to retrieve similar entity names from.
+            k: The number of similar entity names to retrieve.
+            sample_id: The sample ID of the subgraph to retrieve similar entity names from. If None, perform the search across all samples.
+        Returns:
+            A list of dictionaries with the entity name and entity type.
+        """
+
+        similar_entity_names = self.aligner.retrieve_similar_entity_names(
+            entity_name=entity_name, k=k, sample_id=sample_id
+        )
+        if isinstance(similar_entity_names, dict):
+            similar_entity_names = [e["entity"] for e in similar_entity_names]
+        return similar_entity_names
+
+    def identify_relevant_entities_from_question_with_llm(
+        self, question, sample_id=None, use_entity_types=True
+    ):
+        """
+        Identify relevant entities from question using LLM.
+        Args:
+            question: The question to identify relevant entities from.
+            sample_id: The sample ID of the subgraph to identify relevant entities from. If None, perform the search across all samples.
+        Returns:
+            The relevant entities.
+        """
+
+        entities = self.extractor.extract_entities_from_question(question)
+        identified_entities = []
+        linked_entities = []
+
+        if isinstance(entities, dict):
+            entities = [entities]
+
+        for ent in entities:
+            similar_entities = self.retrieve_similar_entity_names(
+                entity_name=ent, k=10, sample_id=sample_id
+            )
+            logger.log(logging.DEBUG, "Similar entities: %s" % (str(similar_entities)))
+
+            exact_entity_match = [e for e in similar_entities if e == ent]
+            if len(exact_entity_match) > 0:
+                linked_entities.extend(exact_entity_match)
+            else:
+                identified_entities.extend(similar_entities)
+
+        logger.log(
+            logging.DEBUG,
+            "Identified entities from question: %s" % (str(identified_entities)),
+        )
+        logger.log(
+            logging.DEBUG, "Linked entities from question: %s" % (str(linked_entities))
+        )
+        if use_entity_types:
+            linked_identified_entities = self.extractor.identify_relevant_entities(
+                question=question, entity_list=identified_entities
+            )
+        else:
+            linked_identified_entities = (
+                self.extractor.identify_relevant_entities_wo_types(
+                    question=question, entity_list=identified_entities
+                )
+            )
+        linked_entities.extend([e["entity"] for e in linked_identified_entities])
+
+        logger.log(
+            logging.DEBUG,
+            "Linked entities after refinement: %s" % (str(linked_entities)),
+        )
+        return linked_entities
+
+    def get_1_hop_supporting_triplets(
+        self,
+        entities4search: List[str],
+        sample_id=None,
+        use_qualifiers=False,
+        use_filtered_triplets=False,
+    ):
+        """
+        Get the 1-hop supporting triplets for the given entities.
+        Useful to answer the question with the given entities.
+        Can be invoked multiple times for more than 1-hop support.
+        Args:
+            entities4search: The entities to get the 1-hop supporting triplets for.
+            sample_id: The sample ID of the subgraph to get the 1-hop supporting triplets from. If None, perform the search across all samples.
+            use_qualifiers: Whether to use qualifiers.
+            use_filtered_triplets: Whether to use the triplets that violate the ontology constraints along with the valid triplets.
+        Returns:
+            A list of dictionaries with the subject, relation, object, and qualifiers that correspond to the 1-hop supporting triplets for the given entities.
+        """
+        or_conditions = []
+        for ent in entities4search:
+            or_conditions.append({"$and": [{"subject": ent}]})
+            or_conditions.append({"$and": [{"object": ent}]})
+        if sample_id is None:
+            pipeline = [{"$match": {"$or": or_conditions}}]
+        else:
+            pipeline = [{"$match": {"sample_id": sample_id, "$or": or_conditions}}]
+        results = list(
+            self.triplets_db.get_collection(
+                self.aligner.triplets_collection_name
+            ).aggregate(pipeline)
+        )
+
+        if use_filtered_triplets:
+            filtered_results = list(
+                self.triplets_db.get_collection(
+                    self.aligner.ontology_filtered_triplets_collection_name
+                ).aggregate(pipeline)
+            )
+            results.extend(filtered_results)
+
+        if use_qualifiers:
+            supporting_triplets = [
+                {
+                    "subject": item["subject"],
+                    "relation": item["relation"],
+                    "object": item["object"],
+                    "qualifiers": item["qualifiers"],
+                }
+                for item in results
+            ]
+        else:
+            supporting_triplets = [
+                {
+                    "subject": item["subject"],
+                    "relation": item["relation"],
+                    "object": item["object"],
+                }
+                for item in results
+            ]
+        logger.log(
+            logging.DEBUG,
+            "Supporting triplets: %s\n%s" % (str(supporting_triplets), "-" * 100),
+        )
+        return supporting_triplets
+
+    def answer_question_with_llm(
+        self,
+        question,
+        linked_entities,
+        sample_id=None,
+        hop_depth=5,
+        use_filtered_triplets=False,
+        use_qualifiers=False,
+    ):
+        """
+        "Answer a question with relevant entities."
+        Args:
+            question: The question to answer.
+            linked_entities: The linked entities to answer the question.
+            sample_id: The sample ID of the subgraph to answer the question from. If None, perform the search across all samples.
+            use_filtered_triplets: Whether to use filtered triplets.
+            use_qualifiers: Whether to use qualifiers.
+        Returns:
+            The answer to the question.
+        """
+        logger.log(logging.DEBUG, "Linked entities: %s" % (str(linked_entities)))
+
+        entity_set = {e for e in linked_entities}
+        entities4search = list(entity_set)
+        supporting_triplets = []
+
+        for _ in range(hop_depth):
+            new_entities4search = []
+            new_supporting_triplets = self.get_1_hop_supporting_triplets(
+                entities4search, sample_id, use_qualifiers, use_filtered_triplets
+            )
+            supporting_triplets.extend(new_supporting_triplets)
+
+            for doc in supporting_triplets:
+                if doc["subject"] not in entities4search:
+                    new_entities4search.append(doc["subject"])
+                if doc["object"] not in entities4search:
+                    new_entities4search.append(doc["object"])
+                if use_qualifiers:
+                    for q in doc["qualifiers"]:
+                        if q["object"] not in entities4search:
+                            new_entities4search.append(q["object"])
+
+            entities4search = list(set(new_entities4search))
+
+        if use_qualifiers:
+            supporting_triplets = [
+                {
+                    "subject": item["subject"],
+                    "relation": item["relation"],
+                    "object": item["object"],
+                    "qualifiers": item["qualifiers"],
+                }
+                for item in supporting_triplets
+            ]
+        else:
+            supporting_triplets = [
+                {
+                    "subject": item["subject"],
+                    "relation": item["relation"],
+                    "object": item["object"],
+                }
+                for item in supporting_triplets
+            ]
+        logger.log(
+            logging.DEBUG,
+            "Supporting triplets: %s\n%s" % (str(supporting_triplets), "-" * 100),
+        )
+
+        ans = self.extractor.answer_question(
+            question=question, triplets=supporting_triplets
+        )
+        return supporting_triplets, ans
+
+    def answer_with_qa_collapsing(
+        self,
+        question,
+        sample_id=None,
+        max_attempts=5,
+        use_qualifiers=False,
+        use_filtered_triplets=False,
+    ):
+        """
+        "Answer a question with QA collapsing."
+        Args:
+            question: The question to answer.
+            sample_id: The sample ID of the subgraph to answer the question from. If None, perform the search across all samples.
+            max_attempts: The maximum number of attempts to answer the question. Useful to handle complex questions that require multiple hops to answer.
+            use_qualifiers: Whether to use qualifiers.
+            use_filtered_triplets: Whether to use filtered triplets.
+        Returns:
+            The answer to the question.
+        """
+        collapsed_question_answer = ""
+        collapsed_question_sequence = []
+        collapsed_answer_sequence = []
+
+        logger.log(logging.DEBUG, "Question: %s" % (str(question)))
+        collapsed_question = self.extractor.decompose_question(question)
+
+        for i in range(max_attempts):
+            extracted_entities = self.extractor.extract_entities_from_question(
+                collapsed_question
+            )
+            logger.log(
+                logging.DEBUG, "Collapsed question: %s" % (str(collapsed_question))
+            )
+            logger.log(
+                logging.DEBUG, "Extracted entities: %s" % (str(extracted_entities))
+            )
+
+            if len(collapsed_question_answer) > 0:
+                extracted_entities.append(collapsed_question_answer)
+
+            entities4search = []
+            for ent in extracted_entities:
+                similar_entities = self.retrieve_similar_entity_names(
+                    entity_name=ent, k=10, sample_id=sample_id
+                )
+                entities4search.extend([e for e in similar_entities])
+
+            entities4search = list(set(entities4search))
+            logger.log(logging.DEBUG, "Similar entities: %s" % (str(entities4search)))
+
+            supporting_triplets = self.get_1_hop_supporting_triplets(
+                entities4search, sample_id, use_qualifiers, use_filtered_triplets
+            )
+
+            logger.log(
+                logging.DEBUG,
+                "Supporting triplets length: %s" % (str(len(supporting_triplets))),
+            )
+
+            collapsed_question_answer = self.extractor.answer_question(
+                collapsed_question, supporting_triplets
+            )
+            collapsed_question_sequence.append(collapsed_question)
+            collapsed_answer_sequence.append(collapsed_question_answer)
+
+            logger.log(
+                logging.DEBUG, "Collapsed question: %s" % (str(collapsed_question))
+            )
+            logger.log(
+                logging.DEBUG,
+                "Collapsed question answer: %s" % (str(collapsed_question_answer)),
+            )
+
+            is_answered = self.extractor.check_if_question_is_answered(
+                question, collapsed_question_sequence, collapsed_answer_sequence
+            )
+            question_answer_sequence = list(
+                zip(collapsed_question_sequence, collapsed_answer_sequence)
+            )
+
+            if is_answered == "NOT FINAL":
+                collapsed_question = self.extractor.collapse_question(
+                    original_question=question,
+                    question=collapsed_question,
+                    answer=collapsed_question_answer,
+                )
+                continue
+            else:
+                logger.log(logging.DEBUG, "Final answer: %s" % (str(is_answered)))
+                return is_answered
+
+        logger.log(logging.DEBUG, "Final answer: %s" % (str(collapsed_question_answer)))
+        return collapsed_question_answer

wikontic/utils/dynamic_aligner.py

@@ -0,0 +1,281 @@
+from typing import List, Tuple, Set, Dict, Optional
+from transformers import AutoTokenizer, AutoModel
+from dataclasses import dataclass
+from pydantic import BaseModel, ValidationError
+from pymongo import MongoClient, UpdateOne
+from dotenv import load_dotenv, find_dotenv
+import os
+import torch
+
+# os.environ["CUDA_VISIBLE_DEVICES"] = "1"
+_ = load_dotenv(find_dotenv())
+
+
+class EntityAlias(BaseModel):
+    _id: int
+    label: str
+    alias: str
+    sample_id: str
+    alias_text_embedding: List[float]
+
+
+class PropertyAlias(BaseModel):
+    _id: int
+    label: str
+    alias: str
+    sample_id: str
+    alias_text_embedding: List[float]
+
+
+class Aligner:
+    def __init__(self, triplets_db, device="cuda:0"):
+        self.db = triplets_db
+
+        self.entity_aliases_collection_name = "entity_aliases"
+        self.property_aliases_collection_name = "property_aliases"
+
+        self.property_vector_index_name = "property_aliases"
+        self.entities_vector_index_name = "entity_aliases"
+
+        self.initial_triplets_collection_name = "initial_triplets"
+        self.triplets_collection_name = "triplets"
+        self.filtered_triplets_collection_name = "filtered_triplets"
+
+        self.device = torch.device(device)
+        # self.tokenizer = AutoTokenizer.from_pretrained('facebook/contriever', token=os.getenv("HF_KEY"))
+        self.tokenizer = AutoTokenizer.from_pretrained("facebook/contriever")
+        # self.model = AutoModel.from_pretrained('facebook/contriever', token=os.getenv("HF_KEY")).to(self.device)
+        self.model = AutoModel.from_pretrained(
+            "facebook/contriever", use_safetensors=True
+        ).to(self.device)
+
+    def get_embedding(self, text):
+
+        def mean_pooling(token_embeddings, mask):
+            token_embeddings = token_embeddings.masked_fill(
+                ~mask[..., None].bool(), 0.0
+            )
+            sentence_embeddings = (
+                token_embeddings.sum(dim=1) / mask.sum(dim=1)[..., None]
+            )
+            return sentence_embeddings
+
+        if not text or not isinstance(text, str):
+            return None
+
+        inputs = self.tokenizer(
+            [text], padding=True, truncation=True, return_tensors="pt"
+        )
+        outputs = self.model(**inputs.to(self.device))
+        embeddings = mean_pooling(outputs[0], inputs["attention_mask"])
+        return embeddings.detach().cpu().tolist()[0]
+
+    def retrieve_similar_properties(
+        self, target_relation: str, sample_id: str, k: int = 10
+    ) -> List[str]:  # List of property labels
+        """
+        Retrieve and rank properties that match given relation.
+
+        Args:
+            target_relation: The relation to search for
+            k: Number of results to return
+
+        Returns:
+            List of property labels
+        """
+
+        collection = self.db.get_collection(self.property_aliases_collection_name)
+        query_embedding = self.get_embedding(target_relation)
+        if query_embedding is None:
+            return []
+
+        query_k = k * 2
+        max_attempts = 5  #
+        attempt = 0
+        unique_ranked_properties: List[str] = []
+
+        while len(unique_ranked_properties) < k and attempt < max_attempts:
+
+            pipeline = [
+                {
+                    "$vectorSearch": {
+                        "index": self.property_vector_index_name,
+                        "queryVector": query_embedding,
+                        "path": "alias_text_embedding",
+                        "numCandidates": 150,
+                        "limit": query_k if query_k < 150 else 150,
+                        # "filter": {
+                        #                 "sample_id": {"$eq": sample_id},
+                        #             },
+                    }
+                },
+                {
+                    "$project": {
+                        "_id": 0,
+                        "label": 1,
+                        # "alias": 1
+                        # "score": {"$meta": "vectorSearchScore"}
+                    }
+                },
+            ]
+
+            similar_properties = collection.aggregate(pipeline)
+
+            for prop in similar_properties:
+                if prop["label"] not in unique_ranked_properties:
+                    unique_ranked_properties.append(prop["label"])
+                if len(unique_ranked_properties) == k:
+                    break
+
+            query_k *= 2
+            attempt += 1
+
+        return unique_ranked_properties
+
+    def retrieve_similar_entity_names(
+        self, entity_name: str, sample_id: Optional[str] = None, k: int = 10
+    ) -> List[str]:  # List of entity labels
+        """
+        Retrieve and rank entities that match given entity.
+
+        Args:
+            entity_name: The entity to search for
+            k: Number of results to return
+
+        Returns:
+            List of entity labels
+        """
+
+        collection = self.db.get_collection(self.entity_aliases_collection_name)
+        query_embedding = self.get_embedding(entity_name)
+        if query_embedding is None:
+            return []
+
+        query_k = k * 2
+        max_attempts = 5  #
+        attempt = 0
+        unique_ranked_entities: List[str] = []
+
+        while len(unique_ranked_entities) < k and attempt < max_attempts:
+
+            if sample_id is not None:
+                filter = {
+                    "sample_id": {"$eq": sample_id},
+                }
+            else:
+                filter = {}
+
+            pipeline = [
+                {
+                    "$vectorSearch": {
+                        "index": self.entities_vector_index_name,
+                        "queryVector": query_embedding,
+                        "path": "alias_text_embedding",
+                        "numCandidates": 150,
+                        "limit": query_k if query_k < 150 else 150,
+                        "filter": filter,
+                    }
+                },
+                {
+                    "$project": {
+                        "_id": 0,
+                        "label": 1,
+                        # "score": {"$meta": "vectorSearchScore"}
+                    }
+                },
+            ]
+
+            similar_entities = collection.aggregate(pipeline)
+
+            for entity in similar_entities:
+                if entity["label"] not in unique_ranked_entities:
+                    unique_ranked_entities.append(entity["label"])
+                if len(unique_ranked_entities) == k:
+                    break
+
+            query_k *= 2
+            attempt += 1
+
+        return unique_ranked_entities
+
+    def add_entity(self, entity_name, alias, sample_id):
+        collection = self.db.get_collection(self.entity_aliases_collection_name)
+        if not collection.find_one(
+            {"label": entity_name, "alias": alias, "sample_id": sample_id}
+        ):
+
+            collection.insert_one(
+                {
+                    "label": entity_name,
+                    "alias": alias,
+                    "sample_id": sample_id,
+                    "alias_text_embedding": self.get_embedding(alias),
+                }
+            )
+
+    def add_property(self, property_name, alias, sample_id):
+        collection = self.db.get_collection(self.property_aliases_collection_name)
+        if not collection.find_one({"label": property_name, "alias": alias}):
+            collection.insert_one(
+                {
+                    "label": property_name,
+                    "alias": alias,
+                    # "sample_id": sample_id,
+                    "alias_text_embedding": self.get_embedding(alias),
+                }
+            )
+
+    def add_triplets(self, triplets_list, sample_id):
+        collection = self.db.get_collection(self.triplets_collection_name)
+
+        operations = []
+        for triple in triplets_list:
+            triple["sample_id"] = sample_id
+            filter_query = {
+                "subject": triple["subject"],
+                "relation": triple["relation"],
+                "object": triple["object"],
+                "sample_id": triple["sample_id"],
+            }
+            operations.append(
+                UpdateOne(filter_query, {"$setOnInsert": triple}, upsert=True)
+            )
+
+        if operations:
+            collection.bulk_write(operations)
+
+    def add_filtered_triplets(self, triplets_list, sample_id):
+        collection = self.db.get_collection(self.filtered_triplets_collection_name)
+
+        operations = []
+        for triple in triplets_list:
+            triple["sample_id"] = sample_id
+            filter_query = {
+                "subject": triple["subject"],
+                "relation": triple["relation"],
+                "object": triple["object"],
+                "sample_id": triple["sample_id"],
+            }
+            operations.append(
+                UpdateOne(filter_query, {"$setOnInsert": triple}, upsert=True)
+            )
+
+        if operations:
+            collection.bulk_write(operations)
+
+    def add_initial_triplets(self, triplets_list, sample_id):
+        collection = self.db.get_collection(self.initial_triplets_collection_name)
+        operations = []
+        for triple in triplets_list:
+            triple["sample_id"] = sample_id
+            filter_query = {
+                "subject": triple["subject"],
+                "relation": triple["relation"],
+                "object": triple["object"],
+                "sample_id": triple["sample_id"],
+            }
+            operations.append(
+                UpdateOne(filter_query, {"$setOnInsert": triple}, upsert=True)
+            )
+        if operations:
+            collection.bulk_write(operations)