datasets.py

# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: Apache-2.0


import bisect
import copy
import os
import logging
import json
from itertools import islice
from collections import Counter, defaultdict
import numpy as np
import random
import networkx as nx
from typing import Dict, List, Tuple, Set
import torch
from transformers import PreTrainedTokenizer

from arguments import DataTrainingArguments
from input_example import InputFeatures, EntityType, RelationType, Entity, Relation, Intent, InputExample, CorefDocument
from base_dataset import BaseDataset
from utils import get_precision_recall_f1
from coreference_metrics import CorefAllMetrics
from input_formats import INPUT_FORMATS
from output_formats import OUTPUT_FORMATS

DATASETS = {}


def register_dataset(dataset_class):
    DATASETS[dataset_class.name] = dataset_class
    return dataset_class


def load_dataset(
        dataset_name: str,
        data_args: DataTrainingArguments,
        tokenizer: PreTrainedTokenizer,
        split: str,
        max_input_length: int,
        max_output_length: int,
        train_subset: float = 1,
        seed: int = None,
        shuffle: bool = True,
        is_eval: bool = False
):
    """
    Load a registered dataset.
    """
    return DATASETS[dataset_name](
        tokenizer=tokenizer,
        max_input_length=max_input_length,
        max_output_length=max_output_length,
        mode=split,
        overwrite_cache=data_args.overwrite_cache,
        train_subset=train_subset,
        seed=seed,
        shuffle=shuffle,
        data_args=data_args,
        is_eval=is_eval,
    )


class JointERDataset(BaseDataset):
    """
    Base class for datasets of joint entity and relation extraction.
    """
    entity_types = None
    relation_types = None

    natural_entity_types = None     # dictionary from entity types given in the dataset to the natural strings to use
    natural_relation_types = None   # dictionary from relation types given in the dataset to the natural strings to use

    default_output_format = 'joint_er'

    def load_cached_data(self, cached_features_file):
        d = torch.load(cached_features_file)
        self.entity_types, self.relation_types, self.examples, self.features = \
            d['entity_types'], d['relation_types'], d['examples'], d['features']

    def save_data(self, cached_features_file):
        torch.save({
            'entity_types': self.entity_types,
            'relation_types': self.relation_types,
            'examples': self.examples,
            'features': self.features,
        }, cached_features_file)

    def load_schema(self):
        """
        Load entity and relation types.

        This is the default implementation which uses the dictionaries natural_entity_types and natural_relation_types.
        """
        if self.natural_entity_types is not None:
            self.entity_types = {short: EntityType(
                short=short,
                natural=natural,
            ) for short, natural in self.natural_entity_types.items()}

        if self.natural_relation_types is not None:
            self.relation_types = {short: RelationType(
                short=short,
                natural=natural,
            ) for short, natural in self.natural_relation_types.items()}

    def load_data_single_split(self, split: str, seed: int = None) -> List[InputExample]:
        """
        Load data for a single split (train, dev, or test).

        This is the default implementation for datasets in the SpERT format
        (see https://github.com/markus-eberts/spert).
        """
        examples = []
        name = self.name if self.data_name is None else self.data_name
        file_path = os.path.join(self.data_dir(), f'{name}_{split}.json')

        with open(file_path, 'r') as f:
            data = json.load(f)
            logging.info(f"Loaded {len(data)} sentences for split {split} of {self.name}")

            for i, x in enumerate(data):
                entities = [
                    Entity(id=j, type=self.entity_types[y['type']], start=y['start'], end=y['end'])
                    for j, y in enumerate(x['entities'])
                ]

                relations = [
                    Relation(
                        type=self.relation_types[y['type']], head=entities[y['head']], tail=entities[y['tail']]
                    )
                    for y in x['relations']
                ]

                tokens = x['tokens']

                example = InputExample(
                    id=f'{split}-{i}',
                    tokens=tokens,
                    entities=entities,
                    relations=relations,
                )

                examples.append(example)

        return examples

    def evaluate_example(self, example: InputExample, output_sentence: str, model=None, tokenizer=None) -> Counter:
        """
        Evaluate an output sentence on a single example of this dataset.
        """
        # extract entities and relations from output sentence
        res = self.output_format.run_inference(
            example,
            output_sentence,
            entity_types=self.entity_types,
            relation_types=self.relation_types,
        )
        predicted_entities, predicted_relations = res[:2]
        if len(res) == 6:
            # the output format provides information about errors
            wrong_reconstruction, label_error, entity_error, format_error = res[2:]
        else:
            # in case the output format does not provide information about errors
            wrong_reconstruction = label_error = entity_error = format_error = False

        predicted_entities_no_type = set([entity[1:] for entity in predicted_entities])

        # load ground truth entities
        gt_entities = set(entity.to_tuple() for entity in example.entities)
        gt_entities_no_type = set([entity[1:] for entity in gt_entities])

        # compute correct entities
        correct_entities = predicted_entities & gt_entities
        correct_entities_no_type = gt_entities_no_type & predicted_entities_no_type

        # load ground truth relations
        gt_relations = set(relation.to_tuple() for relation in example.relations)

        # compute correct relations
        correct_relations = predicted_relations & gt_relations

        assert len(correct_entities) <= len(predicted_entities)
        assert len(correct_entities) <= len(gt_entities)
        assert len(correct_entities_no_type) <= len(predicted_entities_no_type)
        assert len(correct_entities_no_type) <= len(gt_entities_no_type)

        assert len(correct_relations) <= len(predicted_relations)
        assert len(correct_relations) <= len(gt_relations)

        res = Counter({
            'num_sentences': 1,
            'wrong_reconstructions': 1 if wrong_reconstruction else 0,
            'label_error': 1 if label_error else 0,
            'entity_error': 1 if entity_error else 0,
            'format_error': 1 if format_error else 0,
            'gt_entities': len(gt_entities),
            'predicted_entities': len(predicted_entities),
            'correct_entities': len(correct_entities),
            'gt_entities_no_type': len(gt_entities_no_type),
            'predicted_entities_no_type': len(predicted_entities_no_type),
            'correct_entities_no_type': len(correct_entities_no_type),
            'gt_relations': len(gt_relations),
            'predicted_relations': len(predicted_relations),
            'correct_relations': len(correct_relations),
        })

        # add information about each entity/relation type so that we can compute the macro-F1 scores
        if self.entity_types is not None:
            for entity_type in self.entity_types.values():
                predicted = set(entity for entity in predicted_entities if entity[0] == entity_type.natural)
                gt = set(entity for entity in gt_entities if entity[0] == entity_type.natural)
                correct = predicted & gt
                res['predicted_entities', entity_type.natural] = len(predicted)
                res['gt_entities', entity_type.natural] = len(gt)
                res['correct_entities', entity_type.natural] = len(correct)

        if self.relation_types is not None:
            for relation_type in self.relation_types.values():
                predicted = set(relation for relation in predicted_relations if relation[0] == relation_type.natural)
                gt = set(relation for relation in gt_relations if relation[0] == relation_type.natural)
                correct = predicted & gt
                res['predicted_relations', relation_type.natural] = len(predicted)
                res['gt_relations', relation_type.natural] = len(gt)
                res['correct_relations', relation_type.natural] = len(correct)

        return res

    def evaluate_dataset(self, data_args: DataTrainingArguments, model, device, batch_size: int, macro: bool = False) \
            -> Dict[str, float]:
        """
        Evaluate model on this dataset.
        """
        results = Counter()

        for example, output_sentence in self.generate_output_sentences(data_args, model, device, batch_size):
            new_result = self.evaluate_example(
                    example=example,
                    output_sentence=output_sentence,
                    model=model,
                    tokenizer=self.tokenizer,
                )
            results += new_result

        entity_precision, entity_recall, entity_f1 = get_precision_recall_f1(
            num_correct=results['correct_entities'],
            num_predicted=results['predicted_entities'],
            num_gt=results['gt_entities'],
        )

        entity_precision_no_type, entity_recall_no_type, entity_f1_no_type = get_precision_recall_f1(
            num_correct=results['correct_entities_no_type'],
            num_predicted=results['predicted_entities_no_type'],
            num_gt=results['gt_entities_no_type'],
        )

        entity_precision_by_type = []
        entity_recall_by_type = []
        entity_f1_by_type = []

        if macro:
            # compute also entity macro scores
            for entity_type in self.entity_types.values():
                precision, recall, f1 = get_precision_recall_f1(
                    num_correct=results['correct_entities', entity_type.natural],
                    num_predicted=results['predicted_entities', entity_type.natural],
                    num_gt=results['gt_entities', entity_type.natural],
                )
                entity_precision_by_type.append(precision)
                entity_recall_by_type.append(recall)
                entity_f1_by_type.append(f1)

        relation_precision, relation_recall, relation_f1 = get_precision_recall_f1(
            num_correct=results['correct_relations'],
            num_predicted=results['predicted_relations'],
            num_gt=results['gt_relations'],
        )

        res = {
            'wrong_reconstruction': results['wrong_reconstructions'] / results['num_sentences'],
            'label_error': results['label_error'] / results['num_sentences'],
            'entity_error': results['entity_error'] / results['num_sentences'],
            'format_error': results['format_error'] / results['num_sentences'],
            'entity_precision': entity_precision,
            'entity_recall': entity_recall,
            'entity_f1': entity_f1,
            'relation_precision': relation_precision,
            'relation_recall': relation_recall,
            'relation_f1': relation_f1,
            'entity_precision_no_type': entity_precision_no_type,
            'entity_recall_no_type': entity_recall_no_type,
            'entity_f1_no_type': entity_f1_no_type,
        }

        if macro:
            res.update({
                'entity_macro_precision': np.mean(np.array(entity_precision_by_type)),
                'entity_macro_recall': np.mean(np.array(entity_recall_by_type)),
                'entity_macro_f1': np.mean(np.array(entity_f1_by_type)),
            })

        return res


@register_dataset
class Conll04Dataset(JointERDataset):
    """
    CoNLL04 dataset (joint entity and relation extraction).

    Downloaded using https://github.com/markus-eberts/spert/blob/master/scripts/fetch_datasets.sh
    """
    name = 'conll04'

    natural_entity_types = {
        'Loc': 'location',
        'Org': 'organization',
        'Peop': 'person',
        'Other': 'other',
    }

    natural_relation_types = {
        'Work_For': 'works for',
        'Kill': 'kills',
        'OrgBased_In': 'organization based in',
        'Live_In': 'lives in',
        'Located_In': 'located in'
    }


@register_dataset
class ADEDataset(JointERDataset):
    """
    ADE dataset (joint entity and relation extraction).

    Downloaded using https://github.com/markus-eberts/spert/blob/master/scripts/fetch_datasets.sh
    """
    name = 'ade'

    natural_entity_types = {
        'Adverse-Effect': 'disease',
        'Drug': 'drug',
    }

    natural_relation_types = {
        'Adverse-Effect': 'effect',
    }

    def load_data_single_split(self, split: str, seed: int = None) -> List[InputExample]:
        """
        Load data for a single split (train, dev, or test).

        We decide which split to use based on the seed.
        In this way, running episodes 1-10 has the effect of running on all 10 different splits once.
        """
        if seed is None:
            i = 0
        else:
            i = seed % 10

        if split == 'train':
            return super().load_data_single_split(f'split_{i}_train', seed)

        elif split == 'dev':
            return []

        elif split == 'test':
            return super().load_data_single_split(f'split_{i}_test', seed)

    def evaluate_dataset(self, *args, **kwargs):
        """
        Evaluate model on this dataset.

        We include the macro entity scores, since it is standard to report them.
        """
        return super().evaluate_dataset(*args, **kwargs, macro=True)


@register_dataset
class NYTDataset(JointERDataset):
    """
    NYT dataset (joint entity and relation extraction).

    Downloaded from https://github.com/yubowen-ph/JointER/tree/master/dataset/NYT-multi/data
    """
    name = 'nyt'

    natural_entity_types = {
        'PERSON': 'person',
        'LOCATION': 'location',
        'ORGANIZATION': 'organization',
    }

    @staticmethod
    def to_natural_relation_type(relation_type: str) -> str:
        # for example, '/people/person/place_of_birth' -> place of birth
        return relation_type.split('/')[-1].replace('_', ' ')

    def load_schema(self):
        """
        Load entity and relation types.
        """
        # entity types are given explicitly in self.natural_entity_types
        super().load_schema()

        # load relation types from file
        with open(os.path.join(self.data_dir(), f'schemas.json'), 'r') as f:
            types = json.load(f)

            # self.entity_types = {name: EntityType(
            #     natural=self.to_natural_entity_type(name),
            # ) for name in types[2].values()}

            self.relation_types = {name: RelationType(
                natural=self.to_natural_relation_type(name)
            ) for name in types[0].values()}

    def load_data_single_split(self, split: str, seed: int = None) -> List[InputExample]:
        """
        Load data for a single split (train, dev, or test).
        """
        examples = []
        file_path = os.path.join(self.data_dir(), f'{split}.json')

        with open(file_path, 'r') as f:
            data = json.load(f)
            logging.info(f"Loaded {len(data)} sentences for split {split} of {self.name}")

            for i, x in enumerate(data):
                entities = []
                relations = []

                for y in x['spo_details']:
                    entity1_start, entity1_end, entity1_type, relation_type, \
                        entity2_start, entity2_end, entity2_type = y

                    entity1 = Entity(type=self.entity_types[entity1_type], start=entity1_start, end=entity1_end)
                    entity2 = Entity(type=self.entity_types[entity2_type], start=entity2_start, end=entity2_end)

                    try:
                        i1 = entities.index(entity1)
                    except ValueError:
                        # add entity to the list
                        i1 = len(entities)
                        entities.append(entity1)

                    try:
                        i2 = entities.index(entity2)
                    except ValueError:
                        # add entity to the list
                        i2 = len(entities)
                        entities.append(entity2)

                    relation = Relation(
                        type=self.relation_types[relation_type], head=entities[i1], tail=entities[i2],
                    )

                    relations.append(relation)

                tokens = x['tokens']

                example = InputExample(
                    id=f'{split}-{i}',
                    tokens=tokens,
                    entities=entities,
                    relations=relations,
                )

                examples.append(example)

        return examples


@register_dataset
class ACE2005REDataset(JointERDataset):
    """
    ACE2005 dataset (joint entity and relation extraction.

    Processed using https://github.com/luanyi/DyGIE/tree/master/preprocessing
    """
    name = 'ace2005_joint_er'

    natural_entity_types = {
        'PER': 'person',
        'LOC': 'location',
        'ORG': 'organization',
        'VEH': 'vehicle',
        'GPE': 'geographical entity',
        'WEA': 'weapon',
        'FAC': 'facility',
    }

    natural_relation_types = {
        'PHYS': 'located in',
        'ART': 'artifact',
        'ORG-AFF': 'employer',
        'GEN-AFF': 'affiliation',
        'PER-SOC': 'social',
        'PART-WHOLE': 'part of',
    }

    def load_data_single_split(self, split: str, seed: int = None) -> List[InputExample]:
        """
        Load data for a single split (train, dev, or test).
        """
        file_path = os.path.join(self.data_dir(), f'{split}.json')

        examples = []
        num_documents = 0
        num_entities = 0
        num_relations = 0

        with open(file_path, 'r') as f:
            for j, l in enumerate(f):
                document = json.loads(l)
                num_documents += 1
                offset = 0

                for i, tokens in enumerate(document['sentences']):
                    num_entities += len(document['ner'][i])
                    num_relations += len(document['relations'][i])

                    if len(document['ner'][i]) > 0:
                        entities = [
                            Entity(type=self.entity_types[entity_type], start=start-offset, end=end-offset+1)
                            for start, end, entity_type in document['ner'][i]
                        ]

                        relations = []

                        skip = False
                        for start1, end1, start2, end2, relation_type in document['relations'][i]:
                            # find entities
                            if len([e for e in entities if e.start == start1-offset and e.end == end1-offset+1]) > 1 \
                                    or \
                                    len([e for e in entities if e.start == start2-offset and e.end == end2-offset+1]) \
                                    > 1:
                                skip = True
                                break

                            [head] = [e for e in entities if e.start == start1-offset and e.end == end1-offset+1]
                            [tail] = [e for e in entities if e.start == start2-offset and e.end == end2-offset+1]

                            relations.append(
                                Relation(type=self.relation_types[relation_type], head=head, tail=tail)
                            )

                        if not skip:
                            example = InputExample(
                                id=f'{split}-{j}-{i}',
                                tokens=tokens,
                                entities=entities,
                                relations=relations,
                            )
                            examples.append(example)

                    offset += len(tokens)

        logging.info(f'Constructed {len(examples)} examples (from {num_documents} documents) for {self.name} ({split})')
        return examples


class NERDataset(JointERDataset):
    """
    Base class for NER datasets.
    """

    def load_data_single_split(self, split: str, seed: int = None) -> List[InputExample]:
        """
        Load data for a single split (train, dev, or test).
        """
        file_path = os.path.join(self.data_dir(), f'{split}.txt')

        raw_examples = []
        tokens = []
        labels = []
        with open(file_path, 'r') as f:
            for line in f:
                if line.startswith("-DOCSTART-") or line == "" or line == "\n":
                    if tokens:
                        raw_examples.append((tokens, labels))
                        tokens = []
                        labels = []
                else:
                    splits = line.split()
                    tokens.append(splits[0])
                    if len(splits) > 1:
                        label = splits[-1].strip()
                        if label == 'O':
                            label = None
                        labels.append(label)
                    else:
                        labels.append(None)

            if tokens:
                raw_examples.append((tokens, labels))

        logging.info(f"Loaded {len(raw_examples)} sentences for split {split} of {self.name}")

        examples = []
        for i, (tokens, labels) in enumerate(raw_examples):
            assert len(tokens) == len(labels)

            # process labels
            entities = []

            current_entity_start = None
            current_entity_type = None

            for j, label in enumerate(labels + [None]):
                previous_label = labels[j-1] if j > 0 else None
                if (label is None and previous_label is not None) \
                        or (label is not None and previous_label is None) \
                        or (label is not None and previous_label is not None and (
                            label[2:] != previous_label[2:] or label.startswith('B-') or label.startswith('S-')
                        )):
                    if current_entity_start is not None:
                        # close current entity
                        entities.append(Entity(
                            id=len(entities),
                            type=self.entity_types[current_entity_type],
                            start=current_entity_start,
                            end=j,
                        ))

                        current_entity_start = None
                        current_entity_type = None

                    if label is not None:
                        # a new entity begins
                        current_entity_start = j
                        assert any(label.startswith(f'{prefix}-') for prefix in 'BIS')
                        current_entity_type = label[2:]
                        assert current_entity_type in self.entity_types

            example = InputExample(
                id=f'{split}-{i}',
                tokens=tokens,
                entities=entities,
                relations=[],
            )

            examples.append(example)

        return examples

    def evaluate_dataset(self, data_args: DataTrainingArguments, model, device, batch_size: int, macro: bool = False) \
            -> Dict[str, float]:
        """
        Evaluate model on this dataset, and return entity metrics only.
        """
        results = super().evaluate_dataset(data_args, model, device, batch_size, macro=macro)
        return {k: v for k, v in results.items() if k.startswith('entity') and k != 'entity_error'}


@register_dataset
class CoNLL03Dataset(NERDataset):
    """
    CoNLL03 dataset (NER).
    """
    name = 'conll03'

    natural_entity_types = {
        'LOC': 'location',
        'MISC': 'miscellaneous',
        'ORG': 'organization',
        'PER': 'person',
    }


@register_dataset
class OntonotesDataset(NERDataset):
    """
    Ontonotes dataset (NER).
    """
    name = 'ontonotes'

    natural_entity_types = {
        'CARDINAL': 'cardinal',
        'DATE': 'date',
        'EVENT': 'event',
        'FAC': 'facility',
        'GPE': 'country city state',
        'LANGUAGE': 'language',
        'LAW': 'law',
        'LOC': 'location',
        'MONEY': 'monetary',
        'NORP': 'nationality religious political group',
        'ORDINAL': 'ordinal',
        'ORG': 'organization',
        'PERCENT': 'percent',
        'PERSON': 'person',
        'PRODUCT': 'product',
        'QUANTITY': 'quantity',
        'TIME': 'time',
        'WORK_OF_ART': 'work_of_art',
    }


@register_dataset
class SnipsDataset(NERDataset):
    name = 'snips'

    default_output_format = 'joint_icsl'

    natural_entity_types = {
        'entity_name': 'entity name',
        'playlist_owner': 'playlist owner',
        'playlist': 'playlist',
        'music_item': 'music item',
        'artist': 'artist',
        'party_size_description': 'party size description',
        'party_size_number': 'party size number',
        'restaurant_type': 'restaurant type',
        'spatial_relation': 'spatial relation',
        'state': 'state',
        'cuisine': 'cuisine',
        'poi': 'poi',
        'country': 'country',
        'city': 'city',
        'timeRange': 'time range',
        'facility': 'facility',
        'served_dish': 'served dish',
        'condition_description': 'condition description',
        'geographic_poi': 'geographic poi',
        'condition_temperature': 'condition temperature',
        'current_location': 'current location',
        'album': 'album',
        'service': 'service',
        'sort': 'sort',
        'track': 'track',
        'year': 'year',
        'object_name': 'object name',
        'rating_value': 'rating value',
        'best_rating': 'best rating',
        'rating_unit': 'rating unit',
        'object_select': 'object select',
        'object_part_of_series_type': 'object part of series type',
        'movie_name': 'movie name',
        'location_name': 'location name',
        'object_location_type': 'object location type',
        'movie_type': 'movie type'
    }

    natural_intent_types = {
        'AddToPlaylist': 'add to playlist',
        'BookRestaurant': 'book restaurant',
        'GetWeather': 'get weather',
        'PlayMusic': 'play music',
        'RateBook': 'rate book',
        'SearchCreativeWork': 'search creative work',
        'SearchScreeningEvent': 'search screening event'
    }

    def convert_bio_to_entities(self, bio_tag: List[str]) -> Tuple[List[Entity], Entity]:
        entities = []
        current_entity = None
        for ii, el in enumerate(bio_tag):
            if el.startswith('B-'):
                tag_type = el[2:]
                current_entity = Entity(
                    type=EntityType(
                        short=tag_type,
                        natural=self.natural_entity_types[tag_type]
                        if tag_type in self.natural_entity_types else tag_type
                    ),
                    start=ii,
                    end=ii+1,
                )
                entities.append(current_entity)
            elif el.startswith('I-'):
                current_entity.end = ii + 1
        return entities


    def load_data_single_split(self, split: str, seed: int = None) -> List[InputExample]:
        """
        Load data for a single split (train, dev, or test).
        """
        file_path = os.path.join(self.data_dir(), f'{split}.tsv')
        examples = []

        with open(file_path, 'r') as f:
            data = f.readlines()[1:]    # skip header
            for id, example in enumerate(data):
                uid, cid, turn, author, utterance, short_intent, act, slot_labels = example.strip().split('\t')
                tokens = utterance.split()

                slot_labels = slot_labels.split()
                entities = self.convert_bio_to_entities(slot_labels)

                intent = Intent(
                    short=short_intent,
                    natural=self.natural_intent_types[short_intent]
                    if short_intent in self.natural_intent_types.keys() else short_intent
                )

                example = InputExample(
                    id=f'{split}-{id}',
                    tokens=tokens,
                    intent=intent,
                    entities=entities,
                )

                examples.append(example)

        self.entity_types = {
            entity.type.natural: entity.type
            for example in examples for entity in example.entities
        }
        self.intents = {
            example.intent.natural: example.intent
            for example in examples
        }

        return examples


    def evaluate_example(self, example: InputExample, output_sentence: str, model=None, tokenizer=None) -> Counter:
        """
        Evaluate an output sentence on a single example of this dataset.
        """
        # extract entities and relations from output sentence
        res = self.output_format.run_inference(
            example,
            output_sentence,
            entity_types=self.entity_types,
        )
        predicted_intent, predicted_entities, wrong_reconstruction, label_error, format_error = res

        predicted_entities_no_type = set([entity[1:] for entity in predicted_entities])

        # load ground truth entities
        gt_entities = set(entity.to_tuple() for entity in example.entities)
        gt_entities_no_type = set([entity[1:] for entity in gt_entities])

        # compute correct entities
        correct_entities = predicted_entities & gt_entities
        correct_entities_no_type = gt_entities_no_type & predicted_entities_no_type

        # load ground truth intent
        gt_intent = example.intent

        # print(f"Ground truth: {gt_intent} ||| Predicted: {predicted_intent}")

        # compute correct intent
        correct_intent = int(predicted_intent == gt_intent.natural)


        assert len(correct_entities) <= len(predicted_entities)
        assert len(correct_entities) <= len(gt_entities)
        assert len(correct_entities_no_type) <= len(predicted_entities_no_type)
        assert len(correct_entities_no_type) <= len(gt_entities_no_type)

        res = Counter({
            'num_sentences': 1,
            'wrong_reconstructions': 1 if wrong_reconstruction else 0,
            'label_error': 1 if label_error else 0,
            'format_error': 1 if format_error else 0, 
            'predicted_intent': 1 if len(predicted_intent) > 0 else 0,
            'gt_intent': 1,
            'correct_intent': correct_intent,
            'gt_entities': len(gt_entities),
            'predicted_entities': len(predicted_entities),
            'correct_entities': len(correct_entities),
            'gt_entities_no_type': len(gt_entities_no_type),
            'predicted_entities_no_type': len(predicted_entities_no_type),
            'correct_entities_no_type': len(correct_entities_no_type),
        })
        
        if self.intents is not None:
            for intent_type in self.intents.values():
                predicted = int(predicted_intent == intent_type.natural)
                gt = int(gt_intent.natural == intent_type.natural)
                correct = int(predicted_intent == gt_intent.natural)
                res['predicted_intent', intent_type.natural] = predicted
                res['gt_intent', intent_type.natural] = gt
                res['correct_intent', intent_type.natural] = correct

        # add information about each entity/relation type so that we can compute the macro-F1 scores
        if self.entity_types is not None:
            for entity_type in self.entity_types.values():
                predicted = set(entity for entity in predicted_entities if entity[0] == entity_type.natural)
                gt = set(entity for entity in gt_entities if entity[0] == entity_type.natural)
                correct = predicted & gt
                res['predicted_entities', entity_type.natural] = len(predicted)
                res['gt_entities', entity_type.natural] = len(gt)
                res['correct_entities', entity_type.natural] = len(correct)

        return res

    def evaluate_dataset(self, data_args: DataTrainingArguments, model, device, batch_size: int, macro: bool = False) \
            -> Dict[str, float]:
        """
        Evaluate model on this dataset.
        """
        results = Counter()

        for example, output_sentence in self.generate_output_sentences(data_args, model, device, batch_size):
            new_result = self.evaluate_example(
                    example=example,
                    output_sentence=output_sentence,
                    model=model,
                    tokenizer=self.tokenizer,
                )
            results += new_result

        entity_precision, entity_recall, entity_f1 = get_precision_recall_f1(
            num_correct=results['correct_entities'],
            num_predicted=results['predicted_entities'],
            num_gt=results['gt_entities'],
        )

        entity_precision_no_type, entity_recall_no_type, entity_f1_no_type = get_precision_recall_f1(
            num_correct=results['correct_entities_no_type'],
            num_predicted=results['predicted_entities_no_type'],
            num_gt=results['gt_entities_no_type'],
        )

        entity_precision_by_type = []
        entity_recall_by_type = []
        entity_f1_by_type = []

        if macro:
            # compute also entity macro scores
            for entity_type in self.entity_types.values():
                precision, recall, f1 = get_precision_recall_f1(
                    num_correct=results['correct_entities', entity_type.natural],
                    num_predicted=results['predicted_entities', entity_type.natural],
                    num_gt=results['gt_entities', entity_type.natural],
                )
                entity_precision_by_type.append(precision)
                entity_recall_by_type.append(recall)
                entity_f1_by_type.append(f1)

        intent_precision, intent_recall, intent_f1 = get_precision_recall_f1(
            num_correct=results['correct_intent'],
            num_predicted=results['predicted_intent'],
            num_gt=results['gt_intent']
        )

        res = {
            'wrong_reconstruction': results['wrong_reconstructions'] / results['num_sentences'],
            'label_error': results['label_error'] / results['num_sentences'],
            'format_error': results['format_error'] / results['num_sentences'],
            'intent_precision': intent_precision,
            'intent_recall': intent_recall,
            'intent_f1': intent_f1,
            'entity_precision': entity_precision,
            'entity_recall': entity_recall,
            'entity_f1': entity_f1,
            'entity_precision_no_type': entity_precision_no_type,
            'entity_recall_no_type': entity_recall_no_type,
            'entity_f1_no_type': entity_f1_no_type,
        }

        if macro:
            res.update({
                'entity_macro_precision': np.mean(np.array(entity_precision_by_type)),
                'entity_macro_recall': np.mean(np.array(entity_recall_by_type)),
                'entity_macro_f1': np.mean(np.array(entity_f1_by_type)),
            })

        return res


@register_dataset
class ATISDataset(SnipsDataset):
    name = 'atis'

    default_output_format = 'joint_icsl'

    natural_entity_types = {
        'fromloc': 'from', 'toloc': 'to',
        'city_name': 'city', 'state_code': 'state code', 'state_name': 'state name',
        'country_name': 'country name',
        'airport_code': 'airport code', 'airport_name': 'airport name',
        'depart_date': 'depart date', 'arrive_date': 'arrive date',
        'depart_time': 'depart time', 'arrive_time': 'arrive time',
        'return_date': 'return date', 'return_time': 'return time',
        'day_number': 'day number', 'day_name': 'day name', 'days_code': 'days code',
        'month_name': 'month', 'year': 'year',
        'date_relative': 'relative date', 'today_relative': 'relative today',
        'period_of_day': 'period of day', 'time_relative': 'relative time',
        'time': 'time', 'start_time': 'start time', 'end_time': 'end time',
        'cost_relative': 'relative cost',
        'airline_name': 'airline name', 'airline_code': 'airline code',
        'class_type': 'class type',
        'round_trip': 'round trip',
        'fare_basis_code': 'fare basis code',
        'fare_amount': 'fare amount',
        'meal': 'meal', 'meal_code': 'meal code', 'meal_description': 'meal description',
        'flight_mod': 'flight modify', 'mod': 'modify', 'period_mod': 'period modify',
        'stoploc': 'stop location',
        'connect': 'connect',
        'flight_number': 'flight number', 'flight_time': 'flight time', 'flight_stop': 'flight stop',
        'flight_days': 'flight days',
        'aircraft_code': 'aircraft code',
        'or': 'or',
        'restriction_code': 'restriction code',
        'transport_type': 'transport type',
        'economy': 'economy'
    }

    natural_intent_types = {
        'atis_flight': 'flight',
        'atis_airfare': 'airfare',
        'atis_airline': 'airline',
        'atis_aircraft': 'aircraft',
        'atis_flight<DIV>atis_airfare': 'flight and airfare',
        'atis_abbreviation': 'abbreviation',
        'atis_ground_service': 'ground service',
        'atis_meal': 'meal',
        'atis_restriction': 'restriction',
        'atis_quantity': 'quantity',
        'atis_aircraft<DIV>atis_flight<DIV>atis_flight_no': 'aircraft and flight and flight number',
        'atis_airport': 'airport',
        'atis_ground_fare': 'ground fare',
        'atis_airline<DIV>atis_flight_no': 'airline and flight number',
        'atis_flight_time': 'flight time',
        'atis_flight_no': 'flight number',
        'atis_distance': 'distance',
        'atis_city': 'city',
        'atis_capacity': 'capacity',
        'atis_cheapest': 'cheapest',
        'atis_ground_service<DIV>atis_ground_fare': 'ground service and ground fare'
    }

    def convert_bio_to_entities(self, bio_tag: List[str]) -> Tuple[List[Entity], Entity]:
        entities = []
        current_entity = None
        for ii, el in enumerate(bio_tag):
            if el.startswith('B-'):
                tag_type = el[2:]
                if '.' in tag_type:
                    natural = ' '.join([self.natural_entity_types[tag_part]
                            if tag_part in self.natural_entity_types else tag_part
                            for tag_part in tag_type.split('.')])
                else:
                    natural = self.natural_entity_types[tag_type] if tag_type in self.natural_entity_types else tag_type
                current_entity = Entity(
                    type=EntityType(
                        short=tag_type,
                        natural=natural
                    ),
                    start=ii,
                    end=ii+1,
                )
                entities.append(current_entity)
            elif el.startswith('I-'):
                current_entity.end = ii + 1
        return entities


@register_dataset
class ACE2005NERDataset(NERDataset):
    """
    ACE2005 dataset (NER).

    Downloaded from https://github.com/ShannonAI/mrc-for-flat-nested-ner/
    """
    name = 'ace2005_ner'

    natural_entity_types = {
        'PER': 'person',
        'LOC': 'location',
        'ORG': 'organization',
        'VEH': 'vehicle',
        'GPE': 'geographical entity',
        'WEA': 'weapon',
        'FAC': 'facility',
    }

    def load_data_single_split(self, split: str, seed: int = None) -> List[InputExample]:
        """
        Load data for a single split (train, dev, or test).
        """
        file_path = os.path.join(self.data_dir(), f'{split}.ner.json')
        examples = []

        with open(file_path, 'r') as f:
            data = json.load(f)

            for i, x in enumerate(data):
                tokens = x['context'].split()
                entities = []

                for entity_type, l in x['label'].items():
                    for start_end in l:
                        start, end = map(int, start_end.split(';'))
                        end += 1

                        entities.append(Entity(
                            id=len(entities),
                            type=self.entity_types[entity_type],
                            start=start,
                            end=end,
                        ))

                example = InputExample(
                    id=f'{split}-{i}',
                    tokens=tokens,
                    entities=entities,
                    relations=[],
                )

                examples.append(example)

        logging.info(f"Loaded {len(examples)} sentences for split {split} of {self.name}")
        return examples


@register_dataset
class GENIADataset(NERDataset):
    """
    GENIA dataset (NER).
    """
    name = 'genia'

    natural_entity_types = {
        'G#DNA': 'DNA',
        'G#RNA': 'RNA',
        'G#cell_line': 'cell line',
        'G#cell_type': 'cell type',
        'G#protein': 'protein',
    }

    def load_data_single_split(self, split: str, seed: int = None) -> List[InputExample]:
        """
        Load data for a single split (train, dev, or test).
        """
        file_path = os.path.join(self.data_dir(), f'{split}.data')
        examples = []

        with open(file_path, 'r') as f:
            for i, (raw_sentence, _, raw_entities, _) in enumerate(iter(lambda: tuple(islice(f, 4)), ())):
                tokens = raw_sentence.strip().split()
                entities = []

                for raw_entity in raw_entities.strip().split('|'):
                    if len(raw_entity) > 0:
                        span, entity_type = raw_entity.split()
                        start, end = map(int, span.split(','))

                        entities.append(Entity(
                            id=len(entities),
                            type=self.entity_types[entity_type],
                            start=start,
                            end=end,
                        ))

                example = InputExample(
                    id=f'{split}-{i}',
                    tokens=tokens,
                    entities=entities,
                    relations=[],
                )

                examples.append(example)

        logging.info(f"Loaded {len(examples)} sentences for split {split} of {self.name}")
        return examples


@register_dataset
class ACE2005EventTriggerDataset(JointERDataset):
    """
    ACE 2005 dataset (event extraction), trigger extraction component.
    """
    name = 'ace2005event_trigger'
    data_name = 'ace2005event'

    relation_schemas = None

    def load_schema(self):
        types_file_name = os.path.join(self.data_dir(), f'{self.data_name}_types.json')
        with open(types_file_name, 'r') as f:
            types = json.load(f)

            self.entity_types = {name: EntityType(
                short=name,
                natural=x['verbose'],
            ) for name, x in types['entities'].items()}

            self.relation_types = {name: RelationType(
                short=name,
                natural=x['verbose'],
            ) for name, x in types['relations'].items()}

        schema_file_name = os.path.join(self.data_dir(), f'{self.data_name}_schema.json')
        with open(schema_file_name, 'r') as f:
            schema = json.load(f)
            self.relation_schemas = dict()
            for trigger_type, role_types in schema.items():
                trigger_type = self.entity_types[trigger_type].natural
                self.relation_schemas[trigger_type] = \
                    set(self.relation_types[role_type].natural for role_type in role_types)

    def load_data_single_split(self, split: str, seed: int = None) -> List[InputExample]:
        """
        Load data for a single split (train, dev, or test).
        """
        examples = []
        file_path = os.path.join(self.data_dir(), f'{self.data_name}_{split}.json')
        
        with open(file_path, 'r') as f:
            data = json.load(f)
            logging.info(f"Loaded {len(data)} sentences for split {split} of {self.name}")

            for i, x in enumerate(data):
                triggers = [
                    Entity(id=j, type=self.entity_types[y['type']], start=y['start'], end=y['end'])
                    for j, y in enumerate(x['triggers'])
                ]

                tokens = x['tokens']

                example = InputExample(
                    id=f'{split}-{i}',
                    tokens=tokens,
                    entities=triggers,
                    relations=[],
                )

                examples.append(example)

        return examples


@register_dataset
class ACE2005EventArgumentDataset(ACE2005EventTriggerDataset):
    """
    ACE 2005 dataset (event extraction), argument extraction component.
    """
    name = 'ace2005event_argument'
    data_name = 'ace2005event'

    default_input_format = 'ace2005_event_with_trigger'
    default_output_format = 'ace2005_event'

    def load_data_single_split(self, split: str, seed: int = None) -> List[InputExample]:
        """
        Load data for a single split (train, dev, or test).
        """
        examples = []
        name = self.name if self.data_name is None else self.data_name
        file_path = os.path.join(self.data_dir(), f'{name}_{split}.json')

        with open(file_path, 'r') as f:
            data = json.load(f)
            logging.info(f"Loaded {len(data)} sentences for split {split} of {self.name}")

            for i, x in enumerate(data):
                num_triggers = x['triggers']
                for trigger_id in range(min(1, len(num_triggers))):
                    entities = [
                        Entity(id=j, type=self.entity_types[y['type']], start=y['start'], end=y['end'])
                        for j, y in enumerate(x['entities'])
                    ]

                    triggers = [
                        Entity(id=j, type=self.entity_types[y['type']], start=y['start'], end=y['end'])
                        for j, y in enumerate(x['triggers'][trigger_id:trigger_id+1]) if len(x['triggers']) > 0
                    ]
                    assert len(triggers) <= 1, 'no more than 1 trigger'

                    relations = [
                        # here we take the trigger as the tail entity of the relation
                        Relation(
                            type=self.relation_types[y['type']], head=entities[y['head']], tail=triggers[y['tail']]
                        )
                        for y in x['relations'] if y['tail'] == trigger_id and len(num_triggers) > 0
                    ]
                    tokens = x['tokens']

                    example = InputExample(
                        id=f'{split}-{i}',
                        tokens=tokens,
                        entities=entities,
                        triggers=triggers,
                        relations=relations,
                    )

                    examples.append(example)

        return examples

    def evaluate_example(self, example: InputExample, output_sentence: str, model=None, tokenizer=None) -> Counter:
        """
        Evaluate an output sentence on a single example of this dataset.
        """
        # extract natural name of entity and relation types
        predicted_entities, predicted_relations, wrong_reconstruction = \
            self.output_format.run_inference(
                example,
                output_sentence,
                entity_types=self.entity_types,
                relation_types=self.relation_types,
            )

        # filter relation tuples for argument classification
        # since we don't need the entity type to be predicted correct

        def filter_relation_tuple(relation_tuple):
            return relation_tuple[0], relation_tuple[1][1:], relation_tuple[2]

        gt_relations = set(filter_relation_tuple(relation.to_tuple()) for relation in example.relations)
        gt_relations_no_type = set([relation[1:] for relation in gt_relations])

        # load ground truth relations that only have valid relations (exist in relation_schema)
        filtered_predicted_relations = set()
        for relation in predicted_relations:
            if relation[2][0] in self.relation_schemas and relation[0] in self.relation_schemas[relation[2][0]]:
                filtered_predicted_relations.add(filter_relation_tuple(relation))

        predicted_relations = filtered_predicted_relations
        predicted_relations_no_type = set(relation[1:] for relation in predicted_relations)

        # compute correct relations
        correct_relations = predicted_relations & gt_relations
        correct_relations_no_type = predicted_relations_no_type & gt_relations_no_type

        return Counter({
            'num_sentences': 1,
            'wrong_reconstructions': 1 if wrong_reconstruction else 0,
            'gt_relations': len(gt_relations),
            'predicted_relations': len(predicted_relations),
            'correct_relations': len(correct_relations),
            'gt_relations_no_type': len(gt_relations_no_type),
            'predicted_relations_no_type': len(predicted_relations_no_type),
            'correct_relations_no_type': len(correct_relations_no_type),
        })

    def evaluate_dataset(self, data_args: DataTrainingArguments, model, device, batch_size: int, macro: bool = False) \
            -> Dict[str, float]:
        """
        Evaluate model on this dataset.
        """
        results = Counter()

        for example, output_sentence in self.generate_output_sentences(data_args, model, device, batch_size):
            new_result = self.evaluate_example(
                    example=example,
                    output_sentence=output_sentence,
                    tokenizer=self.tokenizer,
                )
            results += new_result

        relation_precision, relation_recall, relation_f1 = get_precision_recall_f1(
            num_correct=results['correct_relations'],
            num_predicted=results['predicted_relations'],
            num_gt=results['gt_relations'],
        )

        relation_precision_no_type, relation_recall_no_type, relation_f1_no_type = get_precision_recall_f1(
            num_correct=results['correct_relations_no_type'],
            num_predicted=results['predicted_relations_no_type'],
            num_gt=results['gt_relations_no_type'],
        )

        res = {
            'relation_precision': relation_precision,
            'relation_recall': relation_recall,
            'relation_f1': relation_f1,
            'relation_precision_no_type': relation_precision_no_type,
            'relation_recall_no_type': relation_recall_no_type,
            'relation_f1_no_type': relation_f1_no_type,
            'num_gt_triggers': results['gt_entities'],
            'num_pred_triggers': results['predicted_entities'],
            'num_gt_relations': results['gt_relations'],
            'num_pred_relations': results['predicted_relations'],
        }

        return res


@register_dataset
class ACE2005EventDataset(ACE2005EventArgumentDataset):
    """
    ACE 2005 dataset (event extraction), for evaluation only.
    """
    name = 'ace2005event'
    task_descriptor = 'ace2005event_trigger'
    default_input_format = 'plain'
    default_output_format = 'joint_er'
    argument_input_format = 'ace2005_event_with_trigger'
    argument_output_format = 'ace2005_event'

    def load_data_single_split(self, split: str, seed: int = None) -> List[InputExample]:
        """
        Load data for a single split (train, dev, or test).
        """
        examples = []
        name = self.name if self.data_name is None else self.data_name
        file_path = os.path.join(self.data_dir(), f'{name}_{split}.json')
        
        with open(file_path, 'r') as f:
            data = json.load(f)
            logging.info(f"Loaded {len(data)} sentences for split {split} of {self.name}")
            
            for i, x in enumerate(data):
                entities = [
                    Entity(id=j, type=self.entity_types[y['type']], start=y['start'], end=y['end'])
                    for j, y in enumerate(x['entities'])
                ]
                
                triggers = [
                    Entity(id=j, type=self.entity_types[y['type']], start=y['start'], end=y['end'])
                    for j, y in enumerate(x['triggers'])
                ]

                relations = [
                    # the trigger is the tail, and the entity is the head
                    Relation(
                        type=self.relation_types[y['type']], head=entities[y['head']], tail=triggers[y['tail']]
                    )
                    for y in x['relations']
                ]
                
                tokens = x['tokens']
                
                example = InputExample(
                    id=f'{split}-{i}',
                    tokens=tokens,
                    entities=entities,
                    triggers=triggers,
                    relations=relations,
                )
                
                examples.append(example)
        
        return examples
    
    def evaluate_argument(self, output_format, example_argument_single_trigger: InputExample, example: InputExample,
                          argument_output_sentence: str) -> Tuple[Set[tuple], Set[tuple], Set[tuple]]:
        """
        Perform argument prediction.
        """
        predicted_entities, predicted_relations, wrong_reconstruction = \
            output_format.run_inference(example_argument_single_trigger,
                                        argument_output_sentence,
                                        entity_types=self.entity_types,
                                        relation_types=self.relation_types)
        
        # filter relation tuples for argument classification
        # since we don't need the entity type to be predicted correct

        def filter_relation_tuple(relation_tuple):
            return relation_tuple[0], relation_tuple[1][1:], relation_tuple[2]

        gt_relations = set(filter_relation_tuple(relation.to_tuple()) for relation in example.relations)

        # load ground truth relations to only have relations that are valid (exist in relation_schema)
        filtered_predicted_relations = set()
        for relation in predicted_relations:
            if relation[2][0] in self.relation_schemas and relation[0] in self.relation_schemas[relation[2][0]]:
                filtered_predicted_relations.add(filter_relation_tuple(relation))
        
        predicted_relations = filtered_predicted_relations
        
        # compute correct relations
        correct_relations = predicted_relations & gt_relations
        
        return predicted_relations, gt_relations, correct_relations

    def evaluate_dataset(self, data_args: DataTrainingArguments, model, device, batch_size: int, macro: bool = False) \
            -> Dict[str, float]:
        """
        Evaluate model on this dataset.
        """
        results = Counter()

        for example, trigger_output_sentence in self.generate_output_sentences(data_args, model, device, batch_size):
            # phase 1: trigger prediction
            trigger_output_format = self.output_format
            predicted_triggers = \
                trigger_output_format.run_inference(
                    example,
                    trigger_output_sentence,
                    entity_types=self.entity_types,
                    relation_types=self.relation_types,
                )[0]
            gt_triggers = set(trigger.to_tuple() for trigger in example.triggers)
            correct_triggers = predicted_triggers & gt_triggers
            predicted_triggers_notype = set()
            gt_triggers_notype = set()
            # trigger tuple format: (type, start, end) -- resetting all types to the same as 'TYPE'
            for trig in predicted_triggers:
                trig_list = list(trig)
                trig_list[0] = 'TYPE'
                predicted_triggers_notype.add(tuple(trig_list))
            for trig in gt_triggers:
                trig_list = list(trig)
                trig_list[0] = 'TYPE'
                gt_triggers_notype.add(tuple(trig_list))
            correct_triggers_notype = predicted_triggers_notype & gt_triggers_notype
            
            # phase 2: argument classification
            all_gt_relations, all_predicted_relations, all_correct_relations = set(), set(), set()
            for trigger in predicted_triggers:
                example_argument_single_trigger = copy.deepcopy(example)
                trigger_type = None
                for trigger_type in self.entity_types:
                    if self.entity_types[trigger_type].natural == trigger[0]: break
                example_argument_single_trigger.triggers = [
                    Entity(type=self.entity_types[trigger_type], start=trigger[1], end=trigger[2])]

                argument_input_format = INPUT_FORMATS[self.argument_input_format]()
                argument_output_format = OUTPUT_FORMATS[self.argument_output_format]()
                example_input = argument_input_format.format_input(example_argument_single_trigger, multitask=True,
                                                                   task_descriptor=ACE2005EventArgumentDataset.name)
                example_input_ids = self.tokenizer.batch_encode_plus(
                    [example_input],
                    max_length=data_args.max_seq_length,
                    return_tensors='pt',
                    padding='max_length',
                    truncation=True
                )
                argument_output = model.generate(
                    example_input_ids['input_ids'].to(device),
                    max_length=data_args.max_output_seq_length_eval,
                    num_beams=data_args.num_beams,
                )[0]  # only one sample
                argument_output_sentence = self.tokenizer.decode(argument_output, skip_special_tokens=True,
                                                                 clean_up_tokenization_spaces=False)

                gt_relations, predicted_relations, correct_relations = \
                    self.evaluate_argument(argument_output_format, example_argument_single_trigger, example,
                                           argument_output_sentence)
                all_gt_relations = all_gt_relations.union(gt_relations)
                all_predicted_relations = all_predicted_relations.union(predicted_relations)
                all_correct_relations = all_correct_relations.union(correct_relations)
                
            all_predicted_relations_notype = set()
            all_gt_relations_notype = set()
            for rel in all_predicted_relations:
                rel_list = list(rel)
                rel_list[0] = 'TYPE'
                all_predicted_relations_notype.add(tuple(rel_list))
            for rel in all_gt_relations:
                rel_list = list(rel)
                rel_list[0] = 'TYPE'
                all_gt_relations_notype.add(tuple(rel_list))

            all_correct_relations_notype = all_predicted_relations_notype & all_gt_relations_notype
            res = Counter({
                'num_sentences': 1,
                'gt_triggers': len(gt_triggers),
                'predicted_triggers': len(predicted_triggers),
                'correct_triggers': len(correct_triggers),
                'correct_triggers_notype': len(correct_triggers_notype),
                'predicted_relations': len(all_predicted_relations),
                'gt_relations': len(all_gt_relations),
                'correct_relations': len(all_correct_relations),
                'correct_relations_notype': len(all_correct_relations_notype)
            })
            
            results += res
        
        trigger_precision, trigger_recall, trigger_f1 = get_precision_recall_f1(
            num_correct=results['correct_triggers'],
            num_predicted=results['predicted_triggers'],
            num_gt=results['gt_triggers'],
        )
        trigger_precision_notype, trigger_recall_notype, trigger_f1_notype = get_precision_recall_f1(
            num_correct=results['correct_triggers_notype'],
            num_predicted=results['predicted_triggers'],
            num_gt=results['gt_triggers'],
        )
        relation_precision, relation_recall, relation_f1 = get_precision_recall_f1(
            num_correct=results['correct_relations'],
            num_predicted=results['predicted_relations'],
            num_gt=results['gt_relations'],
        )
        relation_precision_notype, relation_recall_notype, relation_f1_notype = get_precision_recall_f1(
            num_correct=results['correct_relations_notype'],
            num_predicted=results['predicted_relations'],
            num_gt=results['gt_relations'],
        )
        
        full_results = {
            'relation_precision': relation_precision,
            'relation_recall': relation_recall,
            'relation_f1': relation_f1,
            'relation_precision_notype': relation_precision_notype,
            'relation_recall_notype': relation_recall_notype,
            'relation_f1_notype': relation_f1_notype,
            'trigger_precision': trigger_precision,
            'trigger_recall': trigger_recall,
            'trigger_f1': trigger_f1,
            'trigger_precision_notype': trigger_precision_notype,
            'trigger_recall_notype': trigger_recall_notype,
            'trigger_f1_notype': trigger_f1_notype,
        }
        
        return full_results


@register_dataset
class CoNLL12CorefDataset(BaseDataset):
    """
    CoNLL2012 dataset (coreference resolution).
    """
    name = 'conll12_coref'
    default_output_format = 'coref'

    documents = None    # list of documents

    def load_data_single_split(self, split: str, seed: int = None) -> List[InputExample]:
        """
        Load data for a single split (train, dev, or test).
        """
        file_path = os.path.join(self.data_dir(), f'{split}.json')

        self.documents = {}
        examples = []

        if self.is_eval:
            chunk_size = self.data_args.chunk_size_eval
            chunk_overlap = self.data_args.chunk_overlap_eval
        else:
            chunk_size = self.data_args.chunk_size
            chunk_overlap = self.data_args.chunk_overlap

        with open(file_path, 'r') as f:
            for i, l in enumerate(f):
                raw_document = json.loads(l)
                document_id = f'{split}-{i}'

                tokens_data = raw_document['preprocessing']['segments']['tokens']
                tokens = [x['extent'] for x in tokens_data]
                tokens_start_char = [x['start'] for x in tokens_data]
                tokens_end_char = [x['end'] for x in tokens_data]

                groups = []
                for raw_group in raw_document['annotations']['coreference']['groups']:
                    mentions = []
                    for raw_mention in raw_group['mentions']:
                        # find start and end tokens
                        start = bisect.bisect_left(tokens_start_char, raw_mention['start'])
                        end = bisect.bisect_left(tokens_end_char, raw_mention['end']) + 1
                        mentions.append(Entity(start=start, end=end))

                    groups.append(mentions)

                # create chunks
                chunks = []
                pos = 0
                chunk_id = 0
                while pos < len(tokens):
                    # create a chunk starting at this position
                    chunk_tokens = tokens[pos:pos+chunk_size]

                    chunk_groups = []
                    for group in groups:
                        mentions = [
                            Entity(start=mention.start-pos, end=mention.end-pos, type=mention.type)
                            for mention in group
                            if mention.start >= pos and mention.end <= pos + chunk_size
                        ]
                        if len(mentions) >= 2:
                            chunk_groups.append(mentions)

                    example = InputExample(
                        id=f'{split}-{i}-{chunk_id}',
                        tokens=chunk_tokens,
                        offset=pos,
                        groups=chunk_groups,
                        document_id=document_id,
                        chunk_id=chunk_id,
                    )

                    examples.append(example)
                    chunks.append(example)

                    if pos + chunk_size >= len(tokens):
                        # this chunk arrives until the end, so we can stop
                        break

                    pos += chunk_size - chunk_overlap
                    chunk_id += 1

                self.documents[document_id] = CorefDocument(
                    id=document_id,
                    tokens=tokens,
                    groups=groups,
                    chunks=chunks,
                    chunk_centers=[example.offset + len(example.tokens) // 2 for example in chunks]
                )

        logging.info(f"Loaded {len(self.documents)} documents split in {len(examples)} chunks"
                     f" for split {split} of {self.name}")

        return examples

    @staticmethod
    def get_document_predictions(chunk_data: List[List[tuple]]) -> List[List[Tuple[int, int]]]:
        """
        Aggregate predictions for each chunk into document-level predictions.
        """
        all_edges = set(x for l in chunk_data for x in l)

        graph = nx.Graph()
        graph.add_edges_from(all_edges)

        processed_groups = []
        for component in nx.connected_components(graph):
            processed_group = []
            for start, end in sorted(component, key=lambda x: (x[0], -x[1])):
                # add this entity if it does not overlap with the previous one
                if len(processed_group) == 0 or start >= processed_group[-1][1]:
                    processed_group.append((start, end))

            processed_groups.append(processed_group)

        return [[(start, end) for start, end in group] for group in processed_groups]

    def evaluate_dataset(self, data_args, model, device, batch_size=8, macro=False, by_relation_type=False) \
            -> Dict[str, float]:
        """
        Evaluate model on this dataset.
        """
        documents_to_chunk_data = defaultdict(list)
        predictions = {}

        for example, output_sentence in self.generate_output_sentences(data_args, model, device, batch_size):
            document_id = example.document_id

            data = self.output_format.run_inference(
                example=example,
                output_sentence=output_sentence,
            )

            # add offset to all indices
            offset = example.offset
            data = [tuple(tuple(y + offset for y in x) for x in z) for z in data if z[1] is not None]

            documents_to_chunk_data[document_id].append(data)

            if len(documents_to_chunk_data[document_id]) == len(self.documents[document_id].chunks):
                # process predictions for this document
                predictions[document_id] = self.get_document_predictions(documents_to_chunk_data[document_id])

        predictions_list = []
        labels_list = []
        for document_id, document in self.documents.items():
            predictions_list.append(predictions[document_id])
            labels_list.append([
                [(entity.start, entity.end) for entity in group]
                for group in document.groups
            ])

        metrics = CorefAllMetrics().get_all_metrics(labels_list, predictions_list)
        return {
            f'{metric_name}_{x}': v
            for metric_name, metric_values in metrics['micro'].items()
            for x, v in metric_values.items()
        }


class RelationClassificationDataset(JointERDataset):
    """
    Base class for relation classification datasets, implementing NLL inference.
    """

    def nll_inference(self, example: InputExample, relation_types: List[RelationType], model=None,
                      tokenizer=None) -> RelationType:
        """
        Run inference on a single example of this dataset, searching for the relation which maximizes the likelihood
        of the corresponding output sentence.
        """
        formatted_input = [self.input_format.format_input(example)]
        scores = []

        x = tokenizer.batch_encode_plus(
            formatted_input,
            max_length=self.max_input_length,
            return_tensors='pt',
            padding='max_length',
            truncation=True,
        )

        for relation_type in relation_types:
            # make a copy of the given example, substituting the correct relation with the current one
            candidate_example = copy.deepcopy(example)
            candidate_example.relations[0].type = relation_type

            # construct augmented natural language output
            formatted_output = [self.output_format.format_output(candidate_example)]
            y = tokenizer.batch_encode_plus(
                formatted_output,
                max_length=self.max_output_length,
                return_tensors='pt',
                padding='max_length',
                truncation=True,
            )

            # run through the model to compute the likelihood
            res = model(
                input_ids=x['input_ids'].to(model.device),
                attention_mask=x['attention_mask'].to(model.device),
                labels=y['input_ids'].to(model.device)
            )
            scores.append(res[0].cpu().detach())    # res[0] is the log likelihood

        scores = np.array(scores)
        min_idx = scores.argmin()

        return relation_types[min_idx]


@register_dataset
class FewRelFull(RelationClassificationDataset):
    """
    Full FewRel dataset (relation classification), not episodic.

    Data was downloaded from https://github.com/thunlp/FewRel/tree/master/data
    """
    name = 'FewRelFull'
    data_name = 'FewRel'

    natural_entity_types = {
        'head': 'head',
        'tail': 'tail',
    }   # fake entity types corresponding to head and tail of a relation

    default_input_format = 'rel_input'
    default_output_format = 'rel_output'

    def load_schema(self):
        """
        Load relation types from the pid2name.json file provided with the dataset.
        """
        super().load_schema()   # this is to initialize the fake entity types 'head' and 'tail'

        with open(os.path.join(self.data_dir(), 'pid2name.json'), 'r') as f:
            data = json.load(f)
            self.relation_types = {
                short: RelationType(short=short, natural=description[0])
                for short, description in data.items()
            }

    def load_data_by_relation_type(self, split: str) -> Dict[str, List[InputExample]]:
        """
        Load data for a single split (train or dev) by relation type.

        This is useful for episodic training/evaluation, where we sample N classes at each episode.
        """
        examples_by_type = {}
        file_path = os.path.join(self.data_dir(), f'{self.data_name}_{split}.json')

        with open(file_path, 'r') as f:
            data = json.load(f)
            i = 0

            for type_id in data:
                assert type_id in self.relation_types
                relation_type = self.relation_types[type_id]
                examples = []

                for idx, _data in enumerate(data[type_id]):
                    tokens = _data['tokens']
                    head_entity = _data['h'][2][0]
                    tail_entity = _data['t'][2][0]

                    if len(head_entity) == 1:
                        head_entity = [head_entity[0], head_entity[0]]

                    if len(tail_entity) == 1:
                        tail_entity = [tail_entity[0], tail_entity[0]]

                    head_entity = Entity(id=None, type=self.entity_types['head'],
                                         start=head_entity[0], end=head_entity[1] + 1)

                    tail_entity = Entity(id=None, type=self.entity_types['tail'],
                                         start=tail_entity[0], end=tail_entity[1] + 1)

                    entities = [head_entity, tail_entity]

                    relations = [
                        Relation(
                            type=relation_type, head=head_entity, tail=tail_entity
                        )
                    ]

                    example = InputExample(
                        id=f'{split}-{i}',
                        tokens=tokens,
                        entities=entities,
                        relations=relations,
                    )
                    examples.append(example)
                    i += 1

                examples_by_type[type_id] = examples

        return examples_by_type

    def load_data_single_split(self, split: str, seed: int = None) -> List[InputExample]:
        """
        Load data for a single split (train, dev, or test).
        """
        examples_by_type = self.load_data_by_relation_type(split=split)
        examples = [example for x in examples_by_type.values() for example in x]
        return examples

    def evaluate_dataset(self, data_args: DataTrainingArguments, model, device, batch_size: int, macro: bool = False) \
            -> Dict[str, float]:
        """
        Evaluate model on this dataset, and return relation metrics only.
        """
        results = super().evaluate_dataset(data_args, model, device, batch_size, macro=macro)
        return {k: v for k, v in results.items() if k.startswith('relation')}


@register_dataset
class FewRelEpisodic(FewRelFull):
    """
    Full FewRel dataset (relation classification), episodic.

    Episodic fine-tuning should happen after meta-training on the FewRelFull dataset.
    """
    name = 'FewRelEpisodic'

    default_input_format = 'rel_input'
    default_output_format = 'rel_output'

    target_relation_types = None    # the relation types involved in this episode (there is self.num_ways of them)

    def load_data_single_split(self, split: str, seed: int = None) -> List[InputExample]:
        """
        Load data for a single split (train, dev, or test).
        """
        examples_by_type = self.load_data_by_relation_type(split='dev')   # we use the dev set for episodic experiments

        # set few-shot parameters
        num_ways, num_shots, num_queries = self.data_args.num_ways, self.data_args.num_shots, self.data_args.num_query

        # sample num_ways relation types for this particular episode
        self.target_relation_types = {
            type_id: self.relation_types[type_id]
            for type_id in random.sample(list(examples_by_type.keys()), num_ways)
        }

        logging.info(f'Target relation types for this few-shot episode: '
                     f'{[relation.natural for relation in self.target_relation_types.values()]}')

        support = []
        query = []

        # sample num_shots + num_queries examples from each relation type
        for type_id in self.target_relation_types:
            sampled_examples = random.sample(examples_by_type[type_id], num_shots + num_queries)
            support += sampled_examples[:num_shots]
            query += sampled_examples[num_shots:]

        if split == 'train':
            return support
        else:
            return query

    def evaluate_example(self, example: InputExample, output_sentence: str, model=None, tokenizer=None) -> Counter:
        """
        Evaluate a single example of this dataset, using NLL inference.
        """
        predicted_relation = self.nll_inference(
            example=example,
            relation_types=list(self.target_relation_types.values()),
            model=model,
            tokenizer=tokenizer,
        )

        predicted_relations = {predicted_relation}
        gt_relations = set(relation.type for relation in example.relations)
        correct_relations = gt_relations & predicted_relations

        return Counter({
            'num_sentences': 1,
            'gt_relations': len(gt_relations),
            'predicted_relations': len(predicted_relations),
            'correct_relations': len(correct_relations),
        })


@register_dataset
class TACRED(RelationClassificationDataset):
    name = 'tacred'
    default_input_format = 'rel_input'
    default_output_format = 'rel_output'

    NO_RELATION = 'no relation'

    @staticmethod
    def to_natural(t: str) -> str:
        """
        Convert entity or relation type to a natural text.
        """
        t = t.split(":")
        assert len(t) <= 2, "Unexpected format {}".format(t)
        t = t[1] if len(t) == 2 else t[0]
        t = t.lower()
        t = t.replace("_", " ")
        t = t.replace("/", " ")
        t = t.replace("stateorprovince", "state or province")

        return t

    def load_data_single_split(self, split: str, seed: int = None) -> List[InputExample]:
        """
        Load data for a single split (train, dev, or test).
        """
        file_path = os.path.join(self.data_dir(), f'json/{split}.json')
        examples = []

        # we fill entity/relation types while parsing the data
        self.entity_types = {}
        self.relation_types = {}

        with open(file_path, 'r') as f:
            data = json.load(f)
            logging.info(f"Loaded {len(data)} sentences for split {split} of {self.name}")
            i = 0
            for idx, obj in enumerate(data):
                words = obj['token']
                head_start, head_end, head_type = obj['subj_start'], obj['subj_end'] + 1, obj['subj_type']
                tail_start, tail_end, tail_type = obj['obj_start'], obj['obj_end'] + 1, obj['obj_type']
                relation = obj['relation']

                if head_type not in self.entity_types:
                    self.entity_types[head_type] = EntityType(short=head_type, natural=self.to_natural(head_type))
                if tail_type not in self.entity_types:
                    self.entity_types[tail_type] = EntityType(short=tail_type, natural=self.to_natural(tail_type))

                head_entity = Entity(
                    id=None,
                    type=self.entity_types[head_type],
                    start=head_start,
                    end=head_end
                )
                tail_entity = Entity(
                    id=None,
                    type=self.entity_types[tail_type],
                    start=tail_start,
                    end=tail_end
                )

                entities = [
                    head_entity, tail_entity
                ]

                if relation not in self.relation_types:
                    self.relation_types[relation] = RelationType(short=relation, natural=self.to_natural(relation))

                relations = [
                    Relation(type=self.relation_types[relation], head=head_entity, tail=tail_entity)
                ]

                example = InputExample(
                    id=f'{split}-{i}',
                    tokens=words,
                    entities=entities,
                    relations=relations,
                )
                i += 1
                examples.append(example)

        self.relation_types = {
            relation.type.short: relation.type
            for example in examples for relation in example.relations
        }

        return examples

    def evaluate_example(self, example: InputExample, output_sentence: str, model=None, tokenizer=None,
                         eval_nll=False) -> Counter:
        """
        Evaluate an output sentence on a single example of this dataset.

        We use NLL inference only when the generated output sentence has an unrecognized relation.
        """
        if not eval_nll and not self.eval_nll:
            # evaluate by generating the output sentence
            predicted_entities, predicted_relations = self.output_format.run_inference(
                example,
                output_sentence,
                entity_types=self.entity_types,
                relation_types=self.relation_types,
            )

            predicted_relation_str = next(iter(predicted_relations))[0]

            if predicted_relation_str not in [relation_type.natural for relation_type in self.relation_types.values()]:
                # the output relation is not in the list of possible relations, so we use NLL evaluation instead
                return self.evaluate_example(
                    example=example,
                    output_sentence=output_sentence,
                    model=model,
                    tokenizer=tokenizer,
                    eval_nll=True
                )

        else:
            # NLL evaluation
            predicted_relation_type = self.nll_inference(
                example=example,
                relation_types=list(self.relation_types.values()),
                model=model,
                tokenizer=tokenizer,
            )

            predicted_relation_str = predicted_relation_type.natural

        # load ground truth relation
        gt_relation_str = example.relations[0].type.natural

        if gt_relation_str == self.NO_RELATION and predicted_relation_str == self.NO_RELATION:
            return Counter({
                'num_sentences': 0,
                'gt_relations': 0,
                'predicted_relations': 0,
                'correct_relations': 0,
            })
        elif gt_relation_str == self.NO_RELATION:
            return Counter({
                'num_sentences': 1,
                'gt_relations': 0,
                'predicted_relations': 1,
                'correct_relations': 0,
            })
        elif predicted_relation_str == self.NO_RELATION:
            return Counter({
                'num_sentences': 1,
                'gt_relations': 1,
                'predicted_relations': 0,
                'correct_relations': 0,
            })
        else:
            return Counter({
                'num_sentences': 1,
                'gt_relations': 1,
                'predicted_relations': 1,
                'correct_relations': 1 if predicted_relation_str == gt_relation_str else 0,
            })

    def evaluate_dataset(self, data_args: DataTrainingArguments, model, device, batch_size: int, macro: bool = False) \
            -> Dict[str, float]:
        """
        Evaluate model on this dataset, and return relation metrics only.
        """
        results = super().evaluate_dataset(data_args, model, device, batch_size, macro=macro)
        return {k: v for k, v in results.items() if k.startswith('relation')}


@register_dataset
class CONLL05SRL(NERDataset):
    name = 'CoNLL2005-SRL'
    natural_entity_types = {
        'ARG0': 'first argument',
        'ARG1': 'second argument',
        'V': 'predicate',
    }

    default_input_format = 'srl_input'

    def convert_bio_to_entities(self, bio_tag: List[str]) -> Tuple[List[Entity], Entity]:
        entities = []
        current_entity = None
        predicate = None
        for ii, el in enumerate(bio_tag):
            if el.startswith('B-'):
                tag_type = el[2:]
                current_entity = Entity(
                    type=EntityType(
                        short=tag_type,
                        natural=self.natural_entity_types[tag_type]
                        if tag_type in self.natural_entity_types else tag_type
                    ),
                    start=ii,
                    end=ii+1,
                )
                if tag_type == 'V':
                    predicate = current_entity
                else:
                    entities.append(current_entity)
            elif el.startswith('I-'):
                current_entity.end = ii + 1
        return entities, predicate
  
    def load_data_single_split(self, split: str, seed: int = None) -> List[InputExample]:
        """
        Load data for a single split (train, dev, or test).
        """
        file_path = os.path.join(self.data_dir(), f'conll05.{split}.txt')
        examples = []
        with open(file_path, 'r') as f:
            for i, line in enumerate(f):
                line = line.split('|||')
                assert len(line) == 2
                sentence, tag = line[0].strip(), line[1].strip()
                sentence = sentence.split()[1:]
                tag = tag.split()
                arguments, predicate = self.convert_bio_to_entities(tag)
                example = InputExample(
                    id=f'{split}-{i}',
                    tokens=sentence,
                    entities=arguments,
                    relations=[],
                    sentence_level_entities=[predicate]
                )
                examples.append(example)

        self.entity_types = {
            entity.type.natural: entity.type
            for example in examples for entity in example.entities
        }
        return examples


@register_dataset
class CONLL12SRL(NERDataset):
    name = 'CoNLL2012-SRL'
    natural_entity_types = {
        'ARGM-LOC': 'location'
    }

    default_input_format = 'srl_input'

    @staticmethod
    def get_word_idx(text: str) -> List[int]:
        word_idx = []
        for ii, c in enumerate(text):
            if len(word_idx) == 0:
                word_idx.append(0)
            elif c == ' ':
                word_idx.append(-1)
            elif word_idx[ii - 1] == -1:
                word_idx.append(word_idx[ii - 2] + 1)
            else:
                word_idx.append(word_idx[ii - 1])
        return word_idx
  
    def load_data_single_split(self, split: str, seed: int = None) -> List[InputExample]:
        """
        Load data for a single split (train, dev, or test).
        """
        file_path = os.path.join(self.data_dir(), f'{split}.json')
        i = 0
        examples = []
        with open(file_path, 'r') as f:
            for line in f:
                obj = json.loads(line)
                text = obj['raw_text']
                text_split = text.split()
                char_to_word_idx = self.get_word_idx(text)

                for event in obj['events']:
                    # trigger
                    start_char, end_char = event['trigger']['start'], event['trigger']['end']
                    start_word, end_word = char_to_word_idx[start_char], char_to_word_idx[end_char - 1] + 1
                    predicate = Entity(
                        id=None, type=EntityType(short='V', natural='predicate'),
                        start=start_word, end=end_word
                    )

                    arguments = []
                    for arg in event['arguments']:
                        assert len(arg['values']) == 1
                        values = arg['values'][0]
                        arg_start_char, arg_end_char = values['start'], values['end']
                        arg_start_word, arg_end_word = char_to_word_idx[arg_start_char], \
                            char_to_word_idx[arg_end_char - 1] + 1
                        arg_name = arg['name']
                        argument = Entity(
                            id=None,
                            type=EntityType(short=arg_name, natural=self.natural_entity_types[arg_name]
                                            if arg_name in self.natural_entity_types else arg_name),
                            start=arg_start_word,
                            end=arg_end_word
                        )

                        arguments.append(argument)

                    example = InputExample(
                        id=f'{split}-{i}',
                        tokens=text_split,
                        entities=arguments,
                        relations=[],
                        sentence_level_entities=[predicate]
                    )
                    examples.append(example)
                    i += 1

        self.entity_types = {
            entity.type.natural: entity.type
            for example in examples for entity in example.entities
        }
        return examples


@register_dataset
class CONLL05SRLBrown(CONLL05SRL):
    name = 'CoNLL2005-SRL-Brown'


@register_dataset
class MultiWoz(BaseDataset):
    """
    MultiWoz 2.1 dataset (Dialogue State Tracking).

    To obtain the multi-woz 2.1 dataset in TANL format:
        1. download the pre-processing script from trade-dst
            - git clone https://github.com/jasonwu0731/trade-dst.git
        2. update the dataset_url in file trade-dst/create_data.py line 309 to the url given below.
           This step ensures you are using the MultiWoz 2.1 dataset vs. the 2.0 version.
            - "https://www.repository.cam.ac.uk/bitstream/handle/1810/294507/MULTIWOZ2.1.zip?sequence=1&isAllowed=y"
        3. on lines 320, 321, 322, and 323 change the multi-woz directory to the following:
            - "data/multi-woz/MULTIWOZ2.1/"
        4. run the trade-dst pre-processing script:
            - python create_data.py
        5. copy the script ./trade-dst/utils/fix_label.py to tanl/preprocess_multiwoz
            - cp ./trade-dst/utils/fix_label.py tanl/preprocess_multiwoz
        6. run prepare_multi_woz.py
            - python prepare_multi_woz.py --data-dir ./trade-dst/data
        7. move the saved splits to ./data/multi_woz
            - mv ./trade-dst/data/splits ./data/multi_woz_2.1
    """
    name = 'multi_woz'
    data_name = 'multi_woz_2.1'
    default_input_format = 'plain'
    default_output_format = 'multi_woz'

    def truncate_first_n_tokens(self, examples, max_seq_length, delimiter='▁'):
        output = []
        for x in examples:
            tokens = self.tokenizer.tokenize(x)
            if len(tokens) > max_seq_length:
                x = ''.join(tokens[-1 * max_seq_length + 1:]).replace(delimiter, ' ')
                assert self.tokenizer.tokenize(''.join(tokens).replace(delimiter, ' ')) == tokens
            output.append(x)
        return output

    def load_data_single_split(self, split: str, seed: int = None) -> List[InputExample]:
        """
        Load data for a single split (train, dev, or test).
        """
        examples = []
        file_path = os.path.join(self.data_dir(), f'multi_woz_2.1_{split}_5_domain.json')

        with open(file_path, 'r') as f:
            data = json.load(f)
            num_examples = len(data["examples"])
            logging.info(f"Loaded {num_examples} sentences for split {split} of {self.name}")

            for i, x in enumerate(data["examples"]):
                turn_id = x["turn_id"]
                conv_id = x["ID"]
                dialog = x["dialog_history"]
                tokens = dialog.split(" ")
                belief = x["turn_belief"]
                uid = "{0}-{1}".format(turn_id, conv_id)

                example = InputExample(
                    id=uid,
                    tokens=tokens,
                    belief_state=belief,
                    utterance_tokens=x["turn_uttr"],
                )
                examples.append(example)
        return examples

    def compute_features(self, max_input_length: int, max_output_length: int, multitask: bool = False):
        input_sentences = [self.input_format.format_input(example, multitask=multitask) for example in self.examples]
        output_sentences = [self.output_format.format_output(example) for example in self.examples]

        input_sentences = self.truncate_first_n_tokens(examples=input_sentences,
                                                       max_seq_length=max_input_length)
        output_sentences = self.truncate_first_n_tokens(examples=output_sentences,
                                                        max_seq_length=max_output_length)

        input_tok = self.tokenizer.batch_encode_plus(
            input_sentences,
            max_length=max_input_length,
            return_tensors='pt',
            padding='max_length',
            truncation=True,
        )
        output_tok = self.tokenizer.batch_encode_plus(
            output_sentences,
            max_length=max_output_length,
            return_tensors='pt',
            padding='max_length',
            truncation=True,
        )
        self._warn_max_sequence_length(max_input_length, input_sentences, "input")
        self._warn_max_sequence_length(max_output_length, output_sentences, "output")

        assert input_tok.input_ids.size(0) == output_tok.input_ids.size(0), print(
            f'Size does not match: len(sentences_tok.input_ids)={len(input_tok.input_ids)}, '
            f'len(labels_tok.input_ids)={len(output_tok.input_ids)}'
        )
        features = []
        for sentence_input_ids, att_mask, label_input_ids in zip(input_tok.input_ids, input_tok.attention_mask,
                                                                 output_tok.input_ids):
            features.append(InputFeatures(
                input_ids=sentence_input_ids.tolist(),
                attention_mask=att_mask.tolist(),
                label_ids=label_input_ids.tolist()
            ))

        return features

    def evaluate(self, example: InputExample, output_sentence: str):
        """
        Evaluate a single example.
        """
        gold_belief_set = set(example.belief_state)
        pred_belief_set = \
            self.output_format.run_inference(
                example,
                output_sentence,
            )

        # remove "-"
        pred_belief_set = {elem.replace("-", " ") for elem in pred_belief_set}
        gold_belief_set = {elem.replace("-", " ") for elem in gold_belief_set}
        correct_belief_state = gold_belief_set == pred_belief_set

        return {
            'num_sentences': 1,
            'correct_state': int(correct_belief_state),
            'raw_gold_state': None,
            'raw_pred_state': output_sentence,
            'list_pred_state': list(pred_belief_set),
            'list_gold_state': list(gold_belief_set)
        }

    def evaluate_dataset(self, data_args: DataTrainingArguments, model, device, batch_size: int, macro: bool = False) \
            -> Dict[str, float]:
        """
        Evaluate model on this dataset.
        """

        def compute_accuracy(results_dict):
            num_examples = float(sum(results_dict["num_sentences"]))
            return sum(results_dict["correct_state"]) / num_examples

        results = {
            'num_sentences': [],
            'correct_state': [],
            'raw_gold_state': [],
            'raw_pred_state': [],
            'list_pred_state': [],
            'list_gold_state': []
        }

        for example, output_sentence in self.generate_output_sentences(data_args, model, device, batch_size):
            new_result = self.evaluate(
                example=example,
                output_sentence=output_sentence,
            )

            for k, v in new_result.items():
                results[k].append(v)

        return {
            'joint_accuracy': compute_accuracy(results),
        }