File size: 2,509 Bytes
6cf191b |
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 |
#############################
# Imports
#############################
# Python modules
import random
from typing import Tuple, Optional, List
# Remote modules
# Local modules
from .kg_base_wrapper import KGBaseHandler
from utils import read_json_file_2_dict
#############################
# Constants
#############################
#############################
# Stuff
#############################
class SwowHandler(KGBaseHandler):
def __init__(self, store_dir='kgs_binding/swow'):
super(SwowHandler, self).__init__()
self.swow: dict = self.load_stored_data(store_dir=store_dir)
def get_relation_types(self) -> List[str]:
return ['related_to']
def load_stored_data(self, filename='swow_knowledge.json', store_dir='kgs_binding/swow'):
self.swow = read_json_file_2_dict(filename, store_dir)
return self.swow
def exists_relation_between(self, concept, other_concept):
connections = self.swow.get(concept)
if not connections:
return False
for connetion in connections:
if connetion == other_concept:
return True
return False
def does_concept_exist(self, concept):
return self.swow.get(concept, None) is not None
def relation_between(self, concept, other_concept) -> Tuple[Optional[str], Optional[str]]:
exists_left_right = self.exists_relation_between(concept, other_concept)
exists_right_left = self.exists_relation_between(other_concept, concept)
relation = None
if exists_left_right or exists_right_left:
relation = 'related_to'
return relation, relation
def get_related_concepts(self, concept) -> Optional[List[str]]:
return self.swow.get(concept, [])
def simple_knowledge_prediction(self, knowledge):
kw = list(knowledge)
idx = random.randint(0, len(knowledge)-1) # 0-1-2
kw[idx] = '<mask>'
textual_knowledge_input = f'{kw[0]} {kw[1]} {kw[2]}'
label = f'{knowledge[0]} {knowledge[1]} {knowledge[2]}'
return f'{textual_knowledge_input},{label}\n', label
def create_mask_knowledge_for_model(self):
with open(f'bart_input/swow_bart.txt', 'w') as f:
for subject, objects in self.swow.items():
for obj in objects:
knowledge = (subject, 'is related to', obj)
w_kw, label = self.simple_knowledge_prediction(knowledge)
f.write(w_kw)
|