File size: 2,746 Bytes
5bd8e7f |
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 |
import random
import string
def _random_id():
return "n" + ''.join(random.choice(string.ascii_uppercase + string.digits) for _ in range(5)).lower()
def _construct_knowledge_graph_spec_node(extrapolated_entity: str):
return {
"id": _random_id(),
"fields": {
"name": extrapolated_entity,
"type": "entity"
}
}
def _construct_knowledge_graph_spec_link(source: str, target: str, extrapolated_relationship: str):
return {
"id": _random_id(),
"source": source,
"target": target,
"fields": {
"type": extrapolated_relationship
}
}
def convert_to_knowledge_graph_spec(model_results):
nodes = []
links = []
node_name_to_id_map = {}
link_set = set()
for srl_results in model_results:
for srl_result in srl_results:
subject = None
verb = None
object = None
for tuple in srl_result:
if tuple[1] == "ARG0":
subject = tuple
if tuple[1] == "PRED":
verb = tuple
if tuple[1] == "ARG1":
object = tuple
if subject and verb and object:
source_node = _construct_knowledge_graph_spec_node(subject[0])
target_node = _construct_knowledge_graph_spec_node(object[0])
source_node_id = source_node["id"]
source_node_name = source_node["fields"]["name"]
target_node_id = target_node["id"]
target_node_name = target_node["fields"]["name"]
if source_node_name not in node_name_to_id_map.keys():
node_name_to_id_map[source_node_name] = source_node_id
nodes.append(source_node)
if target_node_name not in node_name_to_id_map.keys():
node_name_to_id_map[target_node_name] = target_node_id
nodes.append(target_node)
link: str = source_node_name + target_node_name + verb[0]
if link not in link_set:
links.append(
_construct_knowledge_graph_spec_link(
node_name_to_id_map[source_node_name],
node_name_to_id_map[target_node_name],
verb[0]
)
)
link_set.add(link)
subject = None
verb = None
object = None
return {
"nodes": nodes,
"links": links
}
|