| 27 | return part1, part2 |
| 28 | |
| 29 | class Document: |
| 30 | def __init__(self, data): |
| 31 | self.id = data["id"] |
| 32 | self.words = data["tokens"] |
| 33 | if "events" in data: |
| 34 | self.merged_events = [e["mention"] for e in data["events"]] |
| 35 | else: |
| 36 | self.merged_events = [data["event_mentions"]] |
| 37 | self.populate_event_spans() |
| 38 | def populate_event_spans(self): |
| 39 | events = sorted(flatten(self.merged_events), key=lambda x: (x["sent_id"], x["offset"][0])) |
| 40 | event2id = {e["id"]:idx for idx, e in enumerate(events)} # sorted events to index |
| 41 | self.label_groups = [[event2id[e["id"]] for e in events] for events in self.merged_events] # List[List[int]] each sublist is a group of event index that co-references each other |
| 42 | self.sorted_event_spans = [(event["sent_id"], event["offset"]) for event in events] |
| 43 | |
| 44 | class myDataset(Dataset): |
| 45 | def __init__(self, tokenizer, data_dir, split, max_length=512, sample_rate=None): |