| 40 | |
| 41 | |
| 42 | class Document: |
| 43 | def __init__(self, data, ignore_nonetype=False): |
| 44 | self.id = data["id"] |
| 45 | self.words = data["tokens"] |
| 46 | self.events = [] |
| 47 | self.eid2mentions = {} |
| 48 | if "events" in data: |
| 49 | for e in data["events"]: |
| 50 | self.events += e["mention"] |
| 51 | self.eid2mentions[e["id"]] = [m["id"] for m in e["mention"]] |
| 52 | self.relations = data["temporal_relations"] |
| 53 | else: |
| 54 | self.events = data["event_mentions"] |
| 55 | self.relations = {} |
| 56 | for t in data["TIMEX"]: |
| 57 | self.events.append(t) |
| 58 | self.eid2mentions[t["id"]] = [t["id"]] |
| 59 | |
| 60 | self.sort_events() |
| 61 | self.get_labels(ignore_nonetype) |
| 62 | |
| 63 | def sort_events(self): |
| 64 | self.events = sorted(self.events, key=lambda x: (x["sent_id"], x["offset"][0])) |
| 65 | self.sorted_event_spans = [(event["sent_id"], event["offset"]) for event in self.events] |
| 66 | |
| 67 | def get_labels(self, ignore_none): |
| 68 | pair2rel = {} |
| 69 | for rel in self.relations: |
| 70 | for pair in self.relations[rel]: |
| 71 | for e1 in self.eid2mentions[pair[0]]: |
| 72 | for e2 in self.eid2mentions[pair[1]]: |
| 73 | pair2rel[(e1, e2)] = REL2ID[rel] |
| 74 | if rel in ["SIMULTANEOUS", "BEGINS-ON"]: |
| 75 | pair2rel[(e2, e1)] = REL2ID[rel] |
| 76 | self.labels = [] |
| 77 | for e1 in self.events: |
| 78 | for e2 in self.events: |
| 79 | if e1["id"] == e2["id"]: |
| 80 | continue |
| 81 | if ignore_none: |
| 82 | if abs(e1["sent_id"] - e2["sent_id"]) > 1: |
| 83 | self.labels.append(-100) |
| 84 | else: |
| 85 | self.labels.append(pair2rel.get((e1["id"], e2["id"]), REL2ID["NONE"])) |
| 86 | else: |
| 87 | self.labels.append(pair2rel.get((e1["id"], e2["id"]), REL2ID["NONE"])) |
| 88 | assert len(self.labels) == len(self.events) ** 2 - len(self.events) |
| 89 | |
| 90 | |
| 91 | |