CTINexus/LP/Linker.py at main · 6armyflag6/CTINexus · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
import json
import os
from omegaconf import DictConfig, OmegaConf
from LLMLinker import LLMLinker
from collections import defaultdict

class Linker:

    def __init__(self, config: DictConfig, inFile):

        self.config = config
        self.inFile = inFile
        infile_path = os.path.join(self.config.inSet, self.inFile)
        self.graph = {}

        with open(infile_path, 'r') as f:
            self.js = json.load(f)
            self.aligned_triplets = self.js["EA"]["aligned_triplets"]

        for triplet in self.aligned_triplets:
            subject_entity_id = triplet["subject"]["entity_id"]
            object_entity_id = triplet["object"]["entity_id"]

            if subject_entity_id not in self.graph:
                self.graph[subject_entity_id] = []
            if object_entity_id not in self.graph:
                self.graph[object_entity_id] = []

            self.graph[subject_entity_id].append(object_entity_id)
            self.graph[object_entity_id].append(subject_entity_id)

        self.subgraphs = self.find_disconnected_subgraphs()
        self.main_nodes = []

        for i, subgraph in enumerate(self.subgraphs):
            main_node_entity_id = self.get_main_node(subgraph)
            main_node = self.get_node(main_node_entity_id)
            print(f"subgraph {i}: main node: {main_node['entity_text']}")
            self.main_nodes.append(main_node)

        self.topic_node = self.get_topic_node(self.subgraphs)
        self.main_nodes = [node for node in self.main_nodes if node["entity_id"] != self.topic_node["entity_id"]]
        self.js["LP"] = LLMLinker(self).link()
        self.js["LP"]["topic_node"] = self.topic_node
        self.js["LP"]["main_nodes"] = self.main_nodes
        self.js["LP"]["subgraphs"] = [list(subgraph) for subgraph in self.subgraphs]
        self.js["LP"]["subgraph_num"] = len(self.subgraphs)
        outfolder = self.config.outSet
        os.makedirs(outfolder, exist_ok=True)
        outfile_path = os.path.join(outfolder, self.inFile)

        with open(outfile_path, 'w') as f:
            json.dump(self.js, f, indent=4)


    def find_disconnected_subgraphs(self):

        self.visited = set()
        subgraphs = []

        for start_node in self.graph.keys():

            if start_node not in self.visited:
                current_subgraph = set()
                self.dfs_collect(start_node, current_subgraph)
                subgraphs.append(current_subgraph)

        return subgraphs


    def dfs_collect(self, node, current_subgraph):

        if node in self.visited:
            return

        self.visited.add(node)
        current_subgraph.add(node)

        for neighbour in self.graph[node]:
            self.dfs_collect(neighbour, current_subgraph)


    def get_main_node(self, subgraph):
        outdegrees = defaultdict(int)
        self.directed_graph = {}

        for triplet in self.aligned_triplets:
            subject_entity_id = triplet["subject"]["entity_id"]
            object_entity_id = triplet["object"]["entity_id"]

            if subject_entity_id not in self.directed_graph:
                self.directed_graph[subject_entity_id] = []

            self.directed_graph[subject_entity_id].append(object_entity_id)
            outdegrees[subject_entity_id] += 1
            outdegrees[object_entity_id] += 1

        max_outdegree = 0
        main_node = None

        for node in subgraph:

            if outdegrees[node] > max_outdegree:
                max_outdegree = outdegrees[node]
                main_node = node

        return main_node


    def get_node(self, entity_id):

        for triplet in self.aligned_triplets:

            for key, node in triplet.items():

                if key in ["subject", "object"]:

                    if node["entity_id"] == entity_id:
                        return node


    def get_topic_node(self, subgraphs):

        max_node_num = 0

        for subgraph in subgraphs:

            if len(subgraph) > max_node_num:
                max_node_num = len(subgraph)
                main_subgraph = subgraph

        return self.get_node(self.get_main_node(main_subgraph))