| |
| |
|
|
| import json |
| from itertools import chain |
| from transformers import AutoTokenizer |
|
|
| from utility.subtokenize import subtokenize |
|
|
| import os |
| os.environ["TOKENIZERS_PARALLELISM"] = "true" |
|
|
|
|
| def load_dataset(path): |
| data = {} |
| with open(path, encoding="utf8") as f: |
| for sentence in f.readlines(): |
| sentence = json.loads(sentence) |
| data[sentence["id"]] = sentence |
|
|
| if "nodes" not in sentence: |
| sentence["nodes"] = [] |
|
|
| if "edges" not in sentence: |
| sentence["edges"] = [] |
|
|
| for sample in list(data.values()): |
| sample["sentence"] = sample["input"] |
| sample["input"] = sample["sentence"].split(' ') |
| sample["token anchors"], offset = [], 0 |
| for token in sample["input"]: |
| sample["token anchors"].append({"from": offset, "to": offset + len(token)}) |
| offset += len(token) + 1 |
| return data |
|
|
|
|
| def node_generator(data): |
| for d in data.values(): |
| for n in d["nodes"]: |
| yield n, d |
|
|
|
|
| def anchor_ids_from_intervals(data): |
| for node, sentence in node_generator(data): |
| if "anchors" not in node: |
| node["anchors"] = [] |
| node["anchors"] = sorted(node["anchors"], key=lambda a: (a["from"], a["to"])) |
| node["token references"] = set() |
|
|
| for anchor in node["anchors"]: |
| for i, token_anchor in enumerate(sentence["token anchors"]): |
| if token_anchor["to"] <= anchor["from"]: |
| continue |
| if token_anchor["from"] >= anchor["to"]: |
| break |
|
|
| node["token references"].add(i) |
|
|
| node["anchor intervals"] = node["anchors"] |
| node["anchors"] = sorted(list(node["token references"])) |
| del node["token references"] |
|
|
| for sentence in data.values(): |
| sentence["token anchors"] = [[a["from"], a["to"]] for a in sentence["token anchors"]] |
|
|
|
|
| def create_bert_tokens(data, encoder: str): |
| tokenizer = AutoTokenizer.from_pretrained(encoder, use_fast=True) |
|
|
| for sentence in data.values(): |
| sentence["bert input"], sentence["to scatter"] = subtokenize(sentence["input"], tokenizer) |
|
|
|
|
| def create_edges(sentence, label_f=None): |
| N = len(sentence["nodes"]) |
|
|
| sentence["edge presence"] = [N, N, []] |
| sentence["edge labels"] = [N, N, []] |
|
|
| for e in sentence["edges"]: |
| source, target = e["source"], e["target"] |
| label = e["label"] if "label" in e else "none" |
|
|
| if label_f is not None: |
| label = label_f(label) |
|
|
| sentence["edge presence"][-1].append((source, target, 1)) |
| sentence["edge labels"][-1].append((source, target, label)) |
|
|
| edge_counter = len(sentence["edge presence"][-1]) |
| return edge_counter |
|
|