<commit_before>import abc
from nalaf.structures.data import Edge
from itertools import product, chain


class EdgeGenerator(object):
    """
    Abstract class for generating edges between two entities. Each edge represents
    a possible relationship between the two entities
    Subclasses that inherit this class should:
    * Be named [Name]EdgeGenerator
    * Implement the abstract method generate
    * Append new items to the list field "edges" of each Part in the dataset
    """

    @abc.abstractmethod
    def generate(self, dataset):
        """
        :type dataset: nalaf.structures.data.Dataset
        """
        return


class SentenceDistanceEdgeGenerator(EdgeGenerator):
    """
    Simple implementation of generating edges between the two entities
    if they are #`distance` sentences away (always same part)
    """

    def __init__(self, entity1_class, entity2_class, relation_type, distance, use_predicted_entities=True):
        # Note: would be nice to implement the word filter too here -- see below

        self.entity1_class = entity1_class
        self.entity2_class = entity2_class
        self.relation_type = relation_type
        self.distance = distance
        self.use_predicted_entities = use_predicted_entities

        self.part_entities = (lambda part: chain(part.annotations, part.predicted_annotations) if self.use_predicted_entities else part.annotations)

    def generate(self, dataset):

        for part in dataset.parts():
            part.edges = []

            for e_1, e_2 in product(
                    (e for e in self.part_entities(part) if e.class_id == self.entity1_class),
                    (e for e in self.part_entities(part) if e.class_id == self.entity2_class)):

                s1_index = part.get_sentence_index_for_annotation(e_1)
                s2_index = part.get_sentence_index_for_annotation(e_2)

                if s2_index < s1_index:
                    s1_index, s2_index = s2_index, s1_index

                if e_2.offset < e_1.offset:
                    e_1, e_2 = e_2, e_1

                pair_distance = s2_index - s1_index
                assert pair_distance >= 0  # Because they must be sorted

                if pair_distance == self.distance:
                    edge = Edge(self.relation_type, e_1, e_2, part, part, s1_index, s2_index)
                    part.edges.append(edge)

<commit_msg>implement same_part all entities edge generator :)<commit_after>import abc
from nalaf.structures.data import Edge
from itertools import product, chain


class EdgeGenerator(object):
    """
    Abstract class for generating edges between two entities. Each edge represents
    a possible relationship between the two entities
    Subclasses that inherit this class should:
    * Be named [Name]EdgeGenerator
    * Implement the abstract method generate
    * Append new items to the list field "edges" of each Part in the dataset
    """

    @abc.abstractmethod
    def generate(self, dataset):
        """
        :type dataset: nalaf.structures.data.Dataset
        """
        return


class SentenceDistanceEdgeGenerator(EdgeGenerator):
    """
    Simple implementation of generating edges between the two entities
    if they are #`distance` sentences away (always same part).

    As a special case, if `distance` is None, the distance is disregarded,
    consequently creating edges for all entities within a part (paragraph or what not).
    """

    def __init__(self, entity1_class, entity2_class, relation_type, distance, use_predicted_entities=True):
        # Note: would be nice to implement the word filter too here -- see below

        self.entity1_class = entity1_class
        self.entity2_class = entity2_class
        self.relation_type = relation_type
        self.distance = distance
        self.use_predicted_entities = use_predicted_entities

        self.part_entities = (lambda part: chain(part.annotations, part.predicted_annotations) if self.use_predicted_entities else part.annotations)

    def generate(self, dataset):

        for part in dataset.parts():
            part.edges = []

            for e_1, e_2 in product(
                    (e for e in self.part_entities(part) if e.class_id == self.entity1_class),
                    (e for e in self.part_entities(part) if e.class_id == self.entity2_class)):

                s1_index = part.get_sentence_index_for_annotation(e_1)
                s2_index = part.get_sentence_index_for_annotation(e_2)

                if s2_index < s1_index:
                    s1_index, s2_index = s2_index, s1_index

                if e_2.offset < e_1.offset:
                    e_1, e_2 = e_2, e_1

                pair_distance = s2_index - s1_index
                assert pair_distance >= 0  # Because they must be sorted

                if pair_distance == self.distance or self.distance is None:
                    edge = Edge(self.relation_type, e_1, e_2, part, part, s1_index, s2_index)
                    part.edges.append(edge)


{'fraction_non_alphanumeric': 0.06941950927588271, 'fraction_numerical': 0.013165769000598444, 'mean_word_length': 2.467496542185339, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '748161', 'n_tokens_mistral': 1409, 'n_tokens_neox': 1364, 'n_words': 427}