| import os |
| import zipfile |
| from configuration import get_aida_yago_tsv_file_path, get_resources_dir |
|
|
| TRAIN_START_LINE = "-DOCSTART- (1 EU)" |
| TESTA_START_LINE = "-DOCSTART- (947testa CRICKET)" |
| TESTB_START_LINE = "-DOCSTART- (1163testb SOCCER)" |
|
|
| CANONICAL_REDIRECTS = None |
|
|
|
|
| class AnnotationRecord: |
| def __init__(self, line): |
| """ |
| Lines with tabs are tokens the are part of a mention: |
| - column 1 is the token |
| - column 2 is either B (beginning of a mention) or I (continuation of a mention) |
| - column 3 is the full mention used to find entity candidates |
| - column 4 is the corresponding YAGO2 entity (in YAGO encoding, i.e. unicode characters are backslash encoded and spaces are replaced by underscores, see also the tools on the YAGO2 website), OR --NME--, denoting that there is no matching entity in YAGO2 for this particular mention, or that we are missing the connection between the mention string and the YAGO2 entity. |
| - column 5 is the corresponding Wikipedia URL of the entity (added for convenience when evaluating against a Wikipedia based method) |
| - column 6 is the corresponding Wikipedia ID of the entity (added for convenience when evaluating against a Wikipedia based method - the ID refers to the dump used for annotation, 2010-08-17) |
| - column 7 is the corresponding Freebase mid, if there is one (thanks to Massimiliano Ciaramita from Google Zürich for creating the mapping and making it available to us) |
| """ |
| data_columns = line.split('\t') |
| self.token = None |
| self.begin_inside_tag = None |
| self.full_mention = None |
| self.yago_entity = None |
| self.wikipedia_url = None |
| self.wikipedia_id = None |
| self.freebase_mid = None |
| self.candidates = None |
| if data_columns: |
| self.token = data_columns[0] |
| if len(data_columns) > 1: |
| self.begin_inside_tag = data_columns[1] |
| if len(data_columns) > 2: |
| self.full_mention = data_columns[2] |
| if len(data_columns) > 3: |
| self.yago_entity = data_columns[3] |
| if len(data_columns) > 4: |
| self.wikipedia_url = data_columns[4] |
| if len(data_columns) > 5: |
| self.wikipedia_id = data_columns[5] |
| if len(data_columns) > 6: |
| self.freebase_mid = data_columns[6] |
|
|
| def set_candidates(self, candidate_record): |
| self.candidates = candidate_record |
| self.candidates.non_considered_word_count -= 1 |
|
|
| def __str__(self): |
| res = "" |
| t = [self.token, self.begin_inside_tag, self.full_mention, self.yago_entity, self.wikipedia_url, |
| self.wikipedia_id, self.freebase_mid] |
| for ind, e in enumerate(t): |
| if not e: |
| continue |
| if ind < len(t) - 1: |
| res += e + "|" |
| else: |
| res += e |
| if res[-1] == "|": |
| res = res[:-1] |
| return res |
|
|
|
|
| class Document: |
| def __init__(self, document_id): |
| self.document_id = document_id |
| self.annotations = [] |
| self.current_annotation = [] |
|
|
| def add_annotation(self, line, candidates): |
| if not line: |
| self.flush_current_annotation() |
| else: |
| ar = AnnotationRecord(line) |
| for c in candidates: |
| if c.non_considered_word_count < 1: |
| continue |
| if c.orig_text == ar.full_mention: |
| ar.set_candidates(c) |
| break |
| self.current_annotation.append(ar) |
|
|
| def flush_current_annotation(self): |
| self.annotations.append(self.current_annotation) |
| self.current_annotation = [] |
|
|
|
|
| class Candidate: |
| def __init__(self, candidate_line): |
| self.id = "" |
| self.in_count = 0 |
| self.out_count = 0 |
| self.links = 0 |
| self.url = "" |
| self.name = "" |
| self.normal_name = "" |
| self.normal_wiki_title = "" |
| self.predicted_type = "" |
| for item in candidate_line.split('\t'): |
| if item == 'CANDIDATE' or not item.strip(): |
| continue |
| elif item.startswith('id:'): |
| self.id = item[3:] |
| elif item.startswith('inCount:'): |
| self.in_count = int(item[8:]) |
| elif item.startswith('outCount:'): |
| self.out_count = int(item[9:]) |
| elif item.startswith('links:'): |
| self.links = item[6:] |
| elif item.startswith('url:'): |
| self.url = item[4:] |
| elif item.startswith('name:'): |
| self.name = item[5:] |
| elif item.startswith('normalName:'): |
| self.normal_name = item[11:] |
| elif item.startswith('normalWikiTitle:'): |
| self.normal_wiki_title = item[16:] |
| elif item.startswith('predictedType:'): |
| self.predicted_type = item[14:] |
| else: |
| raise ValueError(f"Undefined PPRforNED CANDIDATE column: {item}") |
|
|
| def __str__(self): |
| return f"id: {self.id}\twiki_page: {self.url.replace('http://en.wikipedia.org/wiki/', '')}" |
|
|
|
|
| class CandidateRecord: |
| def __init__(self, entity_header): |
| self.candidates = [] |
| self.text = "" |
| self.normal_name = "" |
| self.predicted_type = "" |
| self.q = False |
| self.qid = "" |
| self.docid = -1 |
| self.orig_text = "" |
| self.non_considered_word_count = 0 |
| self.url = "" |
| for item in entity_header.split('\t'): |
| if item == 'ENTITY': |
| continue |
| elif item.startswith('text:'): |
| self.text = item[5:] |
| elif item.startswith('normalName:'): |
| self.normal_name = item[11:] |
| elif item.startswith('predictedType:'): |
| self.predicted_type = item[14:] |
| elif item.startswith('q:'): |
| self.q = bool(item[2:]) |
| elif item.startswith('qid:'): |
| self.qid = item[4:] |
| elif item.startswith('docId:'): |
| self.docid = int(item[6:]) - 1 |
| elif item.startswith('origText:'): |
| self.orig_text = item[9:] |
| self.non_considered_word_count = len(self.orig_text.split()) |
| elif item.startswith('url:'): |
| self.url = item[4:] |
| else: |
| raise ValueError(f"Undefined PPRforNED column: {item}") |
|
|
| def add_candidate(self, candidate_line): |
| self.candidates.append(Candidate(candidate_line)) |
|
|
| def __str__(self): |
| cnds = '\n\t'.join([str(x) for x in self.candidates]) |
| return f"doc_id: {self.docid}\toriginal_text: {self.orig_text}\tcandidates:\n\t{cnds}" |
|
|
|
|
| def get_candidates(ppr_for_ned_candidates_zip, last_document_id): |
| candidates_string = ppr_for_ned_candidates_zip.read(str(last_document_id + 1)).decode("utf-8").split("\n") |
| candidates = [] |
| for c_line in candidates_string: |
| if not c_line.strip(): |
| continue |
| if c_line.startswith("ENTITY"): |
| candidates.append(CandidateRecord(c_line)) |
| elif c_line.startswith("CANDIDATE"): |
| assert len(candidates) |
| candidates[-1].add_candidate(c_line) |
| else: |
| raise ValueError("This must be unreachable!") |
| return candidates |
|
|
| class AIDADataset: |
| def __init__(self): |
| super(AIDADataset, self).__init__() |
| self.dataset = None |
| self.data_path = str(get_aida_yago_tsv_file_path().absolute()) |
| assert os.path.exists(self.data_path), f"The passed dataset address: {self.data_path} does not exist" |
| self.load_dataset() |
|
|
| def load_dataset(self): |
| ppr_for_ned_candidates_zip = zipfile.ZipFile(get_resources_dir() / "data" / "PPRforNED.zip", "r") |
| annotations = [[], [], []] |
| current_document = None |
| current_document_candidates = None |
| data_split_id = -1 |
| last_document_id = 0 |
| with open(self.data_path, "r", encoding="utf-8") as data_file: |
| for ind, line in enumerate(data_file): |
| line = line.strip() |
| if line.startswith("-DOCSTART-"): |
| if current_document: |
| annotations[data_split_id].append(current_document) |
| last_document_id += 1 |
| if line == TRAIN_START_LINE or line == TESTA_START_LINE or line == TESTB_START_LINE: |
| data_split_id += 1 |
| current_document = Document(last_document_id) |
| current_document_candidates = get_candidates(ppr_for_ned_candidates_zip, last_document_id) |
| else: |
| current_document.add_annotation(line, current_document_candidates) |
| if current_document: |
| annotations[data_split_id].append(current_document) |
| self.dataset = {"train": annotations[0], "testa": annotations[1], "testb": annotations[2]} |
| ppr_for_ned_candidates_zip.close() |
|
|
