Guideline following Large Language Model for Information Extraction
from . import tasks


__all__ = ["tasks"]

src/__init__.py
hitz-zentroa/GoLLIE
from dataclasses import dataclass, field
from typing import List, Optional


@dataclass
class ModelArguments:
    """
    Arguments pertaining to which model/config/tokenizer we are going to fine-tune, or train from scratch.
    """

    model_name_or_path: Optional[str] = field(
        default=None,
        metadata={"help": "The local path or huggingface hub name of the model and tokenizer to use."},
    )

    torch_dtype: Optional[str] = field(
        default=None,
        metadata={
            "help": (
                "Override the default `torch.dtype` and load the model under this"
                " dtype. If `auto` is passed, the dtype will be automatically derived"
                " from the model's weights. We will override this if we use quantization."
            ),
            "choices": ["auto", "bfloat16", "float16", "float32"],
        },
    )

    use_lora: bool = field(
        default=False,
        metadata={
            "help": (
                "Whether to use LoRA. If True, the model will be trained with LoRA: https://arxiv.org/abs/2106.09685"
            )
        },
    )

    quantization: Optional[int] = field(
        default=None,
        metadata={
            "help": (
                "Whether to use '4' or '8' bit quantization. Requires bitsandbytes library:"
                " https://github.com/TimDettmers/bitsandbytes. This parameter is only used for training."
            )
        },
    )

    quantization_inference: Optional[int] = field(
        default=None,
        metadata={
            "help": (
                "Whether to use '4' or '8' bit quantization. Requires bitsandbytes library:"
                " https://github.com/TimDettmers/bitsandbytes. This parameter is only used for inference."
            )
        },
    )

    lora_weights_name_or_path: Optional[str] = field(
        default=None,
        metadata={
            "help": (
                "If the model has been trained with LoRA, "
                "path or huggingface hub name or local path to the pretrained weights."
            )
        },
    )

    lora_r: Optional[int] = field(
        default=8,
        metadata={"help": "Lora attention dimension."},
    )

    lora_alpha: Optional[float] = field(
        default=16,
        metadata={"help": "The alpha parameter for Lora scaling."},
    )
    lora_dropout: Optional[float] = field(
        default=0.05,
        metadata={"help": "The dropout probability for Lora layers."},
    )

    lora_target_modules: Optional[List[str]] = field(
        default_factory=list,
        metadata={
            "help": (
                "The target modules to which LoRA will be applied. If not specified, We"
                " will use the default modules for the model in huggingface PEFT library."
            )
        },
    )

    force_auto_device_map: bool = field(
        default=False,
        metadata={
            "help": (
                "Whether to force the use of the auto device map. If set to True, the model will be split across "
                "GPUs and CPU to fit the model in memory. If set to False, a full copy of the model will be loaded "
                "into each GPU. Defaults to False."
            )
        },
    )

    use_better_transformer: bool = field(
        default=False,
        metadata={
            "help": (
                "Whether to transform the model using Better Transformer library:"
                "https://huggingface.co/docs/optimum/bettertransformer/overview. Requires optimum"
                "'https://huggingface.co/docs/optimum/installation'. Defaults to False. This flag is "
                "only supported for inference, we will override it for training as BetterTransformer does not"
                "support custom attention masks during training."
            )
        },
    )

    trust_remote_code: bool = field(
        default=False,
        metadata={"help": "Trust the remote code from HuggingFace model hub. Defaults to False Defaults to False."},
    )

    use_flash_attention: bool = field(
        default=True,
        metadata={
            "help": (
                "Whether to use FlashAttention. If True, the model will be trained with FlashAttention."
                "Flash attention must be installed, see: https://github.com/Dao-AILab/flash-attention "
                "for more details."
            )
        },
    )

    max_memory_MB: Optional[int] = field(
        default=None,
        metadata={
            "help": (
                "Free memory per gpu in MB. Used to compute the device map when force_auto_device_map is set to True."
                "Defaults to None."
            )
        },
    )

    merge_lora_after_training: bool = field(
        default=False,
        metadata={
            "help": (
                "Whether to merge LoRA layers after training. If True, the model will be trained with LoRA and then"
                " the LoRA layers will be merged into the model. We will save the merged model in"
                " {output_dir}/merged_model. Defaults to False."
            )
        },
    )

    merge_lora_before_eval: bool = field(
        default=False,
        metadata={
            "help": (
                "Whether to merge LoRA layers before eval. If True, the model and LoRA layers will be loaded, merged"
                " and saved into {model_name_or_path}/merged_model. Then we will load the merged model using the"
                " provided quantization and torch.dtype values. Note: This flag is only useful is you want to use 4"
                " bit or 8 bit quantization for inference. bf16, fp16 and fp32 models are automatically merged when"
                " loaded for inference, there is no need to save the merged model and reload it. Defaults to False."
            )
        },
    )

    keep_merged_model_after_eval: bool = field(
        default=False,
        metadata={
            "help": (
                "If 'merge_lora_before_eval' is set to True, whether to keep the merged model after eval, or delete it"
                " to save disk space. Merged model will be saved in {model_name_or_path}/merged_model and will use"
                " the same disk space as the original model. While storing only the LoRA layers only takes a few MBs."
                " Defaults to False."
            )
        },
    )


@dataclass
class DataTrainingArguments:
    """
    Arguments pertaining to what data we are going to input our model for training and eval.
    """

    dataset_dir: str = field(
        default="~/CoLLIE/data/processed",
        metadata={"help": "The tasks to train on. Can be a list of tasks or a single task."},
    )

    train_tasks: List[str] = field(
        default=None,
        metadata={"help": "The tasks to train on. Can be a list of tasks or a single task."},
    )

    validation_tasks: List[str] = field(
        default=None,
        metadata={"help": "The tasks to train on. Can be a list of tasks or a single task."},
    )

    test_tasks: List[str] = field(
        default=None,
        metadata={"help": "The tasks to test on. Can be a list of tasks or a single task."},
    )

    max_seq_length: int = field(
        default=512,
        metadata={
            "help": (
                "The maximum total input sequence length after tokenization. Sequences"
                " longer than this will be truncated, sequences shorter will be padded."
            )
        },
    )

    ignore_pad_token_for_loss: bool = field(
        default=True,
        metadata={
            "help": "Whether to ignore the tokens corresponding to padded labels in the loss computation or not."
        },
    )

    use_dev_inference: bool = field(
        default=False,
        metadata={"help": "Use the development set for inference instead of the test set."},
    )

    evaluate_all_checkpoints: bool = field(
        default=False,
        metadata={"help": "Evaluate all checkpoints in the model directory."},
    )

    prompt_loss_weight: float = field(
        default=0.05,
        metadata={
            "help": (
                "The weight of the prompt tokens in the loss. If set to '0.05' the prompt tokens will have a total"
                " weight of 5% in the loss while the result tokens will have a total weight of 95%. Only used for"
                " computing the loss in the training data. Defaults to `0.05`."
            )
        },
    )

    max_examples_per_task_train: int = field(
        default=None,
        metadata={
            "help": (
                "The maximum number of examples to use per task for training. "
                "If set to None, all examples will be"
                " used."
            )
        },
    )

    max_examples_per_task_val: int = field(
        default=None,
        metadata={
            "help": (
                "The maximum number of examples to use per task for development. "
                "If set to None, all examples will be used."
            )
        },
    )

    max_examples_per_task_test: int = field(
        default=None,
        metadata={
            "help": (
                "The maximum number of examples to use per task for inference. "
                "If set to None, all examples will be used."
            )
        },
    )

src/config.py
hitz-zentroa/GoLLIE
import glob
import json
import logging
import math
import os
import random
from dataclasses import dataclass
from functools import partial
from itertools import chain
from multiprocessing import Pool
from typing import Any, Dict, Iterator, List, Optional, Sized, Union

import numpy as np
from rich.progress import Progress, SpinnerColumn, TimeElapsedColumn
from torch.utils.data import Dataset

from transformers import BatchEncoding, PreTrainedTokenizerBase
from transformers.utils import PaddingStrategy


def batch(iterable: Sized, n=1) -> Iterator:
    """
    Yield successive n-sized chunks from iterable.

    Args:
        iterable (`Sized`):
            The iterable to split.
        n (`int`, optional):
            The size of the chunks. Defaults to `1`.

    Yields:
        `Iterator`:
            An iterator with the chunks.
    """
    l: int = len(iterable)
    p: int = math.ceil(l / n)
    for ndx in range(0, l, p):
        yield iterable[ndx : min(ndx + p, l)]


def prepare_data(
    example: str,
    tokenizer: PreTrainedTokenizerBase,
    is_encoder_decoder: bool = False,
    max_length: int = 2048,
    inference: bool = False,
    prompt_loss_weight: float = 0.05,
) -> BatchEncoding:
    """
    Prepare data for training or inference.

    Args:
        example (`str`):
            The example to prepare.
        tokenizer (`PreTrainedTokenizerBase`):
            The tokenizer to use.
        is_encoder_decoder (`bool`, optional):
            Whether the model is an encoder-decoder model. Defaults to `False`.
        max_length (`int`, optional):
            The maximum length of the input. Defaults to `2048`.
        inference (`bool`, optional):
            Whether to prepare the data for inference. During inference labels
            are not included in model inputs. Defaults to `False`.
        prompt_loss_weight (`float`, optional):
            The weight of the prompt tokens in the loss. If set to '0.05' the prompt tokens will have a total weight
            of 5% in the loss while the result tokens will have a total weight of 95%. Defaults to `0.05`.

    Returns:
        `BatchEncoding`: `BatchEncoding` with the prepared data.
    """

    if is_encoder_decoder:
        prompt, result = example.split("result =")
        prompt = prompt + "result ="
        prompt = prompt.strip()
        result = result.strip()

        model_inputs = tokenizer(
            text=prompt,
            max_length=max_length,
            truncation=True,
            padding=False,
            return_tensors=None,
            add_special_tokens=True,
        )
        if not inference:
            model_inputs["labels"] = tokenizer(
                text_target=result,
                max_length=max_length,
                truncation=True,
                padding=False,
                return_tensors=None,
                add_special_tokens=True,
            )["input_ids"]

            model_inputs["loss_weight_mask"] = np.ones(len(model_inputs["labels"]), dtype=np.float32)

    else:
        if inference:
            prompt = example.split("result =")[0] + "result ="
            model_inputs = tokenizer(
                text=prompt,
                max_length=max_length,
                truncation=True,
                padding=False,
                return_tensors=None,
                add_special_tokens=True,
            )

            # Remove the last token if it is an eos token
            if model_inputs["input_ids"][-1] == tokenizer.eos_token_id:
                model_inputs["input_ids"] = model_inputs["input_ids"][:-1]
                model_inputs["attention_mask"] = model_inputs["attention_mask"][:-1]

        else:
            model_inputs = tokenizer(
                text=example,
                max_length=max_length,
                truncation=True,
                padding=False,
                return_tensors=None,
                add_special_tokens=True,
            )

            # Make sure the `eos_token_id` is added at the end
            # This bug is reported at https://github.com/huggingface/transformers/issues/22794
            if model_inputs["input_ids"][-1] != tokenizer.eos_token_id:
                model_inputs["input_ids"].append(tokenizer.eos_token_id)
                model_inputs["attention_mask"].append(1)

            model_inputs["labels"] = model_inputs["input_ids"].copy()

            # Find the prompt length
            prompt = example.split("result =")[0] + "result ="
            prompt = tokenizer(
                text=prompt,
                max_length=max_length,
                truncation=True,
                padding=False,
                return_tensors=None,
                add_special_tokens=True,
            )["input_ids"]

            # Remove the last token if it is an eos token
            if prompt[-1] == tokenizer.eos_token_id:
                prompt = prompt[:-1]

            if len(prompt) > len(model_inputs["labels"]):
                raise ValueError(
                    f"Prompt is longer than the input, something went wrong. Prompt: {prompt}, input:"
                    f" {model_inputs['input_ids']}"
                )

            # Create the weight mask
            loss_weight_mask = np.ones(len(model_inputs["labels"]), dtype=np.float32)

            # The sum of the loss of the prompt tokens should be equal to 'prompt_loss_weight' percent of the total loss
            len_prompt = len(prompt)
            len_result = len(model_inputs["labels"]) - len_prompt
            prompt_token_weight = len_result * prompt_loss_weight  # 'prompt_loss_weight' percent of the total loss
            try:
                prompt_token_weight = prompt_token_weight * (
                    len_result / (len_result * (1 - prompt_loss_weight))
                )  # Scale so result tokens can have 1.0 weight
                prompt_token_weight = prompt_token_weight / len_prompt  # Divide by the number of prompt tokens
            except ZeroDivisionError:
                logging.warning(
                    "Found division by zero in prompt token weight calculation. You might have an empty prompt, empty"
                    f" result, or both. Example with error: {example}. Setting prompt token weight to 0.0."
                )
                prompt_token_weight = 0.0

            for i in range(len(prompt)):
                loss_weight_mask[i] = prompt_token_weight

            model_inputs["loss_weight_mask"] = loss_weight_mask

    if "token_type_ids" in model_inputs:
        # LLaMa tokenizer adds token type ids, but we don't need them
        model_inputs.pop("token_type_ids")

    return model_inputs


def batch_tokenization(
    tokenizer: PreTrainedTokenizerBase,
    dataset_name: str,
    is_encoder_decoder: bool,
    max_length: int,
    inference: bool,
    prompt_loss_weight: float,
    examples: List[str],
    process_no: int,
) -> List[BatchEncoding]:
    """
    Batch tokenization function.

    Args:
        tokenizer (`PreTrainedTokenizerBase`):
            The tokenizer to use.
        dataset_name (`str`):
            The name of the dataset.
        is_encoder_decoder (`bool`):
            Whether the model is an encoder-decoder model.
        max_length (`int`):
            The maximum length of the input.
        inference (`bool`):
            Whether to prepare the data for inference. If model
            `is_encoder_decoder=False`, inputs ids will be truncated to don't include the
            results section of the example. Labels will still include the full correct
            example. If model `is_encoder_decoder=True`, this parameter is ignored.
        prompt_loss_weight (`float`):
            The weight of the prompt tokens in the loss. If set to '0.05' the prompt tokens will have a total weight
            of 5% in the loss while the result tokens will have a total weight of 95%. Defaults to `0.05`.
        examples (`List[str]`):
            The examples to tokenize.
        process_no (`int`):
            The process number.

    Returns:
        `List[BatchEncoding]`:
            List of BatchEncoding with the prepared data.
    """
    tokenized_examples: List[BatchEncoding] = []
    if process_no == 0:
        with Progress(
            SpinnerColumn(),
            *Progress.get_default_columns(),
            TimeElapsedColumn(),
        ) as progress:
            task = progress.add_task(f"Tokenizing {dataset_name}", total=len(examples))

            for example in examples:
                tokenized_examples.append(
                    prepare_data(
                        example=example,
                        tokenizer=tokenizer,
                        is_encoder_decoder=is_encoder_decoder,
                        max_length=max_length,
                        inference=inference,
                        prompt_loss_weight=prompt_loss_weight,
                    )
                )
                progress.update(task, advance=1)
    else:
        tokenized_examples = [
            prepare_data(
                example=example,
                tokenizer=tokenizer,
                is_encoder_decoder=is_encoder_decoder,
                max_length=max_length,
                inference=inference,
                prompt_loss_weight=prompt_loss_weight,
            )
            for example in examples
        ]

    return tokenized_examples


class CollieDataset(Dataset):
    """
    Dataset for Collie.

    Args:
        tokenizer (`PreTrainedTokenizerBase`):
            The tokenizer to use.
        dataset_path (`str`):
            The path to the jsonl file containing the dataset.
        is_encoder_decoder (`bool`, optional):
            Whether the model is an encoder-decoder model. Defaults to `False`.
        max_length (`int`, optional):
            The maximum length of the input. Defaults to `2048`.
        inference (`bool`, optional):
            Whether to prepare the data for inference. If model
            `is_encoder_decoder=False`, inputs ids will be truncated to don't include
            the results section of the example. Labels will still include the full
            correct example. If model `is_encoder_decoder=True`, this parameter is
            ignored. Defaults to `False`.
        prompt_loss_weight (`float`, optional):
            The weight of the prompt tokens in the loss. If set to '0.05' the prompt tokens will have a total weight
            of 5% in the loss while the result tokens will have a total weight of 95%. Defaults to `0.05`.
        num_workers (`int`, optional):
            The number of workers to use for tokenization. Defaults to
            `min(os.cpu_count(), 16)`.
        max_examples (`Optional[int]`, optional):
            The maximum number of examples to load. Defaults to `None`. If `None` all
            examples will be loaded. If `max_examples` is smaller is set we will randomly
            sample `max_examples` examples from the dataset.
    """

    def __init__(
        self,
        tokenizer: PreTrainedTokenizerBase,
        dataset_path: str,
        is_encoder_decoder: bool = False,
        max_length: int = 2048,
        inference: bool = False,
        prompt_loss_weight: float = 0.0,
        num_workers: int = min(os.cpu_count(), 16),
        max_examples: Optional[int] = None,
    ):
        self.is_encoder_decoder = is_encoder_decoder
        self.max_length = max_length
        self.inference = inference
        self.max_examples = max_examples

        assert (
            prompt_loss_weight >= 0.0 and prompt_loss_weight < 1.0
        ), f"Prompt loss weight must be in [0, 1). Found {prompt_loss_weight}."

        self.prompt_loss_weight = prompt_loss_weight

        try:
            self.dataset_name, self.task_name, self.split, extension = os.path.basename(dataset_path).split(".")
        except ValueError:
            raise ValueError(
                f"Something is wrong with the dataset path {dataset_path}. Please check it and ensure "
                "it follows the format `dataset_name.task_name.split.jsonl`"
            )

        # Find pre-computed epoch datasets for training
        self.dataset_dict: Dict[int, List[BatchEncoding]] = {}
        self.dataset_keys: List[int] = []
        self.current_dataset_key: int = 0

        if self.split == "train":
            epoch_datasets = glob.glob(
                os.path.join(
                    os.path.dirname(dataset_path), f"{self.dataset_name}.{self.task_name}.{self.split}.*.jsonl"
                )
            )
            if epoch_datasets:
                epoch_datasets.sort(key=lambda x: int(x.split(".")[-2]))
                logging.info(f"Found {len(epoch_datasets)} pre-computed epoch datasets.")
                for dataset in epoch_datasets:
                    try:
                        _, _, _, epoch, _ = os.path.basename(dataset).split(".")
                    except ValueError:
                        logging.warning(f"Error loading pre-computed epoch {dataset} . Skipping...")
                        continue

                    self.dataset_dict[int(epoch)] = self.compute_tokenized_examples(
                        dataset_path=dataset,
                        num_workers=num_workers,
                        tokenizer=tokenizer,
                    )

                    self.dataset_keys.append(int(epoch))

                # Truncate datasers to ensure all datasets have the same length
                # min_length = min([len(x) for x in self.dataset_dict.values()])
                # for key in self.dataset_dict.keys():
                #    if len(self.dataset_dict[key]) > min_length:
                #        logging.warning(f"Truncating dataset {key} from {len(self.dataset_dict[key])} to {min_length}")
                #    self.dataset_dict[key] = self.dataset_dict[key][:min_length]

                # Oversample datasets to ensure all datasets have the same length
                max_length = max([len(x) for x in self.dataset_dict.values()])
                for key in self.dataset_dict.keys():
                    if len(self.dataset_dict[key]) < max_length:
                        logging.warning(
                            f"Oversampling dataset {key} from {len(self.dataset_dict[key])} to {max_length}"
                        )
                        num_samples = max_length - len(self.dataset_dict[key])
                        random_samples = random.choices(self.dataset_dict[key], k=num_samples)
                        self.dataset_dict[key].extend(random_samples)

                self.current_dataset_key = self.dataset_keys[0]

        if len(self.dataset_dict) == 0:
            self.dataset_dict[0] = self.compute_tokenized_examples(
                dataset_path=dataset_path,
                num_workers=num_workers,
                tokenizer=tokenizer,
            )
            self.dataset_keys.append(0)
            self.current_dataset_key = self.dataset_keys[0]

        logging.info(f"Loaded {[len(x) for x in self.dataset_dict.values()]} examples from {dataset_path}")

    def compute_tokenized_examples(
        self,
        dataset_path,
        num_workers,
        tokenizer,
    ) -> List[BatchEncoding]:
        """
        Compute the tokenized examples.

        Args:
            dataset_path (`str`):
                The path to the jsonl file containing the dataset.
            num_workers (`int`):
                The number of workers to use for tokenization.
            tokenizer (`PreTrainedTokenizerBase`):
                The tokenizer to use.

        Returns:
            `List[BatchEncoding]`:
                List of BatchEncoding with the prepared data.

        """

        with open(dataset_path, "r", encoding="utf8") as f:
            examples = f.readlines()

        examples = [json.loads(example.strip())["text"] for example in examples]

        if self.max_examples is not None and self.max_examples < len(examples):
            examples = random.sample(examples, self.max_examples)

        if num_workers <= 1:
            return batch_tokenization(
                tokenizer=tokenizer,
                dataset_name=".".join([self.dataset_name, self.task_name, self.split]),
                is_encoder_decoder=self.is_encoder_decoder,
                max_length=self.max_length,
                inference=self.inference,
                prompt_loss_weight=self.prompt_loss_weight,
                examples=examples,
                process_no=0,
            )
        else:
            tokenizer_fn = partial(
                batch_tokenization,
                tokenizer,
                ".".join([self.dataset_name, self.task_name, self.split]),
                self.is_encoder_decoder,
                self.max_length,
                self.inference,
                self.prompt_loss_weight,
            )
            with Pool(num_workers) as p:
                tokenized_examples = p.starmap(
                    tokenizer_fn,
                    zip(batch(examples, num_workers), range(num_workers)),
                )

            return list(chain.from_iterable(tokenized_examples))

    def __len__(self) -> int:
        return len(self.dataset_dict[self.current_dataset_key])

    def __getitem__(self, idx) -> List[BatchEncoding]:
        return self.dataset_dict[self.current_dataset_key][idx].copy()

    def rotate_split(self):
        """
        Rotate the current dataset to the next one.
        """
        self.current_dataset_key = self.dataset_keys[
            (self.dataset_keys.index(self.current_dataset_key) + 1) % len(self.dataset_keys)
        ]

        if len(self.dataset_dict) > 1:
            logging.info(
                f' Dataset {".".join([self.dataset_name, self.task_name, self.split])} rotated to split'
                f" {self.current_dataset_key}"
            )


@dataclass
class DataCollatorForCoLLIE:
    """
    Adapted from transformers.DataCollatorForSeq2Seq to handle CoLLIE data.

    Data collator that will dynamically pad the inputs received, as well as the labels.

    Args:
        tokenizer ([`PreTrainedTokenizer`] or [`PreTrainedTokenizerFast`]):
            The tokenizer used for encoding the data.
        model ([`PreTrainedModel`]):
            The model that is being trained. If set and has the *prepare_decoder_input_ids_from_labels*, use it to
            prepare the *decoder_input_ids*

            This is useful when using *label_smoothing* to avoid calculating loss twice.
        padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `True`):
            Select a strategy to pad the returned sequences (according to the model's padding side and padding index)
            among:

            - `True` or `'longest'` (default): Pad to the longest sequence in the batch (or no padding if only a single
              sequence is provided).
            - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum
              acceptable input length for the model if that argument is not provided.
            - `False` or `'do_not_pad'`: No padding (i.e., can output a batch with sequences of different lengths).
        max_length (`int`, *optional*):
            Maximum length of the returned list and optionally padding length (see above).
        pad_to_multiple_of (`int`, *optional*):
            If set will pad the sequence to a multiple of the provided value.

            This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability >=
            7.5 (Volta).
        label_pad_token_id (`int`, *optional*, defaults to -100):
            The id to use when padding the labels (-100 will be automatically ignored by PyTorch loss functions).
        return_tensors (`str`):
            The type of Tensor to return. Allowable values are "np", "pt" and "tf".
    """

    tokenizer: PreTrainedTokenizerBase
    model: Optional[Any] = None
    padding: Union[bool, str, PaddingStrategy] = True
    max_length: Optional[int] = None
    pad_to_multiple_of: Optional[int] = None
    label_pad_token_id: int = -100
    return_tensors: str = "pt"

    def __call__(self, features, return_tensors=None):
        if return_tensors is None:
            return_tensors = self.return_tensors
        labels = [feature["labels"] for feature in features] if "labels" in features[0].keys() else None
        loss_weight_mask = (
            [feature["loss_weight_mask"] for feature in features] if "loss_weight_mask" in features[0].keys() else None
        )
        # We have to pad the labels before calling `tokenizer.pad` as this method won't pad them and needs them of the
        # same length to return tensors.
        if labels is not None:
            max_label_length = max(len(l) for l in labels)
            if self.pad_to_multiple_of is not None:
                max_label_length = (
                    (max_label_length + self.pad_to_multiple_of - 1)
                    // self.pad_to_multiple_of
                    * self.pad_to_multiple_of
                )

            padding_side = self.tokenizer.padding_side
            for feature in features:
                remainder = [self.label_pad_token_id] * (max_label_length - len(feature["labels"]))
                if isinstance(feature["labels"], list):
                    feature["labels"] = (
                        feature["labels"] + remainder if padding_side == "right" else remainder + feature["labels"]
                    )
                elif padding_side == "right":
                    feature["labels"] = np.concatenate([feature["labels"], remainder]).astype(np.int64)
                else:
                    feature["labels"] = np.concatenate([remainder, feature["labels"]]).astype(np.int64)

        if loss_weight_mask is not None:
            max_loss_weight_mask_length = max(len(l) for l in loss_weight_mask)
            if self.pad_to_multiple_of is not None:
                max_loss_weight_mask_length = (
                    (max_loss_weight_mask_length + self.pad_to_multiple_of - 1)
                    // self.pad_to_multiple_of
                    * self.pad_to_multiple_of
                )

            padding_side = self.tokenizer.padding_side
            for feature in features:
                remainder = [0.0 if self.label_pad_token_id == -100 else 1.0] * (
                    max_loss_weight_mask_length - len(feature["loss_weight_mask"])
                )
                if isinstance(feature["loss_weight_mask"], list):
                    feature["loss_weight_mask"] = (
                        feature["loss_weight_mask"] + remainder
                        if padding_side == "right"
                        else remainder + feature["loss_weight_mask"]
                    )
                elif padding_side == "right":
                    feature["loss_weight_mask"] = np.concatenate([feature["loss_weight_mask"], remainder]).astype(
                        np.float32
                    )
                else:
                    feature["loss_weight_mask"] = np.concatenate([remainder, feature["loss_weight_mask"]]).astype(
                        np.float32
                    )

        features = self.tokenizer.pad(
            features,
            padding=self.padding,
            max_length=self.max_length,
            pad_to_multiple_of=self.pad_to_multiple_of,
            return_tensors=return_tensors,
        )

        # prepare decoder_input_ids
        if (
            labels is not None
            and self.model is not None
            and hasattr(self.model, "prepare_decoder_input_ids_from_labels")
        ):
            decoder_input_ids = self.model.prepare_decoder_input_ids_from_labels(labels=features["labels"])
            features["decoder_input_ids"] = decoder_input_ids

        return features

src/dataset/dataset.py
hitz-zentroa/GoLLIE
import glob
import importlib
import json
import logging
import os
import sys
from typing import Any, Dict, List, Type

from tqdm import tqdm

from src.config import DataTrainingArguments, ModelArguments
from src.tasks import task_id_to_prompts
from src.tasks.utils_typing import AnnotationList
from transformers import HfArgumentParser, Seq2SeqTrainingArguments


class ResultLogger:
    """
    A class to log the results of a task.

    Args:
        task_name: The name of the task.
    """

    def __init__(self, task_name: str) -> None:
        self.task_name = task_name
        self.sentences: List[str] = []
        self.golds: List[AnnotationList] = []
        self.predictions: List[AnnotationList] = []
        self.impossible_to_parse: int = 0
        self.hallucinated_predictions: int = 0
        self.valid_predictions: int = 0
        self.total_predictions: int = 0
        self.gold_predictions: int = 0

    def add_sentence(self, sentence: str, gold_labels: AnnotationList, pred_labels: AnnotationList) -> None:
        """
        Add a sentence to the logger.

        Args:
            sentence: The sentence to label.
            gold_labels: The 'AnnotationList' gold labels.
            pred_labels: The 'AnnotationList' predicted labels.
        """
        filtered_pred_labels = pred_labels.filter_hallucinations(sentence)
        self.impossible_to_parse += 1 if filtered_pred_labels.parse_error else 0
        self.hallucinated_predictions += filtered_pred_labels.hallucinated_no
        self.valid_predictions += len(filtered_pred_labels)
        self.total_predictions += len(pred_labels)
        self.gold_predictions += len(gold_labels)

        self.sentences.append(sentence)
        self.golds.append(gold_labels)
        self.predictions.append(filtered_pred_labels)

    def compute_metrics(self, scorer: Any) -> Dict[str, Dict[str, Dict[str, float]]]:
        """
        Compute the metrics for the task.

        Args:
            scorer: The scorer to use.

        Returns:
            A dictionary containing the metrics.
        """
        if len(self.sentences) == 0:
            raise ValueError(f"No sentences were added to the {self.task_name} logger, we cannot compute metrics.")

        scores: Dict[str, Dict[str, Dict[str, float]]] = scorer(reference=self.golds, predictions=self.predictions)
        results: Dict[str, Dict[str, Dict[str, float]]] = {
            "predictions_stats": {
                "impossible_to_parse": {
                    "total": self.impossible_to_parse,
                    "percentage%": round((self.impossible_to_parse / len(self.sentences)) * 100, 4),
                },
                "hallucinated_predictions": {
                    "total": self.hallucinated_predictions,
                    "percentage%": (
                        0.0
                        if self.total_predictions == 0
                        else round((self.hallucinated_predictions / self.total_predictions) * 100, 4)
                    ),
                },
                "total": {"predictions": self.valid_predictions, "gold": self.gold_predictions},
            },
        }
        scores.update(results)
        return scores

    def print_predictions(self, output_path: str):
        """
        Print the predictions to a file in json format. In the following format:
        {
            "sentence": {
                "golds": "gold labels",
                "predictions": "predicted labels"
            }
        }

        Args:
            output_path: The path to the output file.
        """
        if len(self.sentences) == 0:
            raise ValueError(f"No sentences were added to the {self.task_name} logger, we cannot print predictions.")
        with open(output_path, "w", encoding="utf8") as f:
            for sentence, prediction, gold in zip(self.sentences, self.predictions, self.golds):
                example = {sentence: {"golds": gold.to_string(), "predictions": prediction.to_string()}}
                json.dump(example, ensure_ascii=False, indent=4, fp=f)
                f.write("\n")


def import_prompts(task_module: str) -> None:
    """Imports everything from a module.

    Args:
        task_module (str): the module to import everything from.
    """
    # get a handle on the module
    mdl = importlib.import_module(task_module)

    # # is there an __all__?  if so respect it
    # if "__all__" in mdl.__dict__:
    #     names = mdl.__dict__["__all__"]
    # else:
    #     # otherwise we import all names that don't begin with _
    #     names = [x for x in mdl.__dict__ if not x.startswith("_")]
    names = {x: y for x, y in mdl.__dict__.items() if not x.startswith("_")}

    # now drag them in
    globals().update(dict(names.items()))


def get_class(class_path: str) -> Type:
    components = class_path.split(".")
    mod = __import__(components[0])
    for comp in components[1:]:
        mod = getattr(mod, comp)

    return mod


def fix_prompt_outputs(text: str) -> str:
    """Fixes some known structural bugs.

    Args:
        text (str): Bugged output.

    Returns:
        str: Corrected output.
    """

    text = text.replace(")\n ", "),\n ")
    return text


def remove_hallucinations(unlabelled_sentence: str, predictions: List[Any]) -> List[Any]:
    """Removes predictions that are not in the unlabelled sentence.
    Args:
        unlabelled_sentence (str): The unlabelled sentence.
        predictions (List[Any]): The list of predictions.
    Returns:
        List[Any]: The list of predictions that are in the unlabelled sentence.
    """
    accepted_list: List[Any] = []
    for prediction in predictions:
        if prediction.exists_in(unlabelled_sentence):
            accepted_list.append(prediction)
    return accepted_list


def evaluate(
    model_args: ModelArguments,
    data_args: DataTrainingArguments,
    training_args: Seq2SeqTrainingArguments,
    checkpoint_path: str = None,
) -> Dict[str, Dict[str, Dict[str, Dict[str, float]]]]:
    """This function evaluates the output of a model.

    Args:
        model_args (ModelArguments): Model arguments. See `ModelArguments` docs.
        data_args (DataTrainingArguments): Data arguments. See `DataTrainingArguments` docs.

    Returns:
        Dict[str, Dict[str, Dict[str, Dict[str, float]]]]: A dictionary containing the scores for each task
        present in the dataset.
    """
    if checkpoint_path is None:
        output_dir = training_args.output_dir
    else:
        output_dir = checkpoint_path

    predictions_dir = os.path.join(output_dir, "predictions")

    gold_data_dir = data_args.dataset_dir
    all_scores = {}

    scores_file_name = os.path.join(output_dir, "task_scores.json")
    scores_file_name_summary = os.path.join(output_dir, "task_scores_summary.json")
    with tqdm(total=len(data_args.test_tasks), desc="Evaluating") as pbar:
        for task in data_args.test_tasks:
            pbar.set_description(f"Evaluating {task}")
            gold_path = os.path.join(
                gold_data_dir,
                f"{task}.{'test' if not data_args.use_dev_inference else 'dev'}.jsonl",
            )
            pred_path = os.path.join(predictions_dir, task) + ".predictions.jsonl"

            if not os.path.exists(gold_path):
                raise FileNotFoundError(f"File not found: '{gold_path}'")

            if not os.path.exists(pred_path):
                raise FileNotFoundError(f"File not found: '{pred_path}'")

            task_logger = ResultLogger(task)
            task_module = None
            scorer = None

            with open(gold_path, "rt", encoding="utf8") as gold_f, open(pred_path, "rt", encoding="utf8") as pred_f:
                for gold_line, pred_line in zip(gold_f, pred_f):
                    gold_line = json.loads(gold_line)
                    pred_line = json.loads(pred_line)

                    if not task_module:
                        task_module = task_id_to_prompts(gold_line["task_id"])
                        import_prompts(task_module)

                    if not scorer:
                        scorer = get_class(gold_line["scorer_cls"])()

                    gold_labels: AnnotationList = AnnotationList.from_output(
                        str(gold_line["labels"]), task_module=task_module
                    )

                    pred_labels = pred_line["model_prediction"].strip().split("result = ")[-1]
                    pred_labels: AnnotationList = AnnotationList.from_output(str(pred_labels), task_module=task_module)

                    task_logger.add_sentence(
                        sentence=gold_line["unlabelled_sentence"], gold_labels=gold_labels, pred_labels=pred_labels
                    )

            task_metrics = task_logger.compute_metrics(scorer)
            all_scores[task] = task_metrics
            # rich.print(f"{task} scores: {task_metrics}")
            task_logger.print_predictions(output_path=os.path.join(predictions_dir, f"{task}.eval_file.json"))
            pbar.update(1)

    with open(scores_file_name, "wt", encoding="utf8") as f:
        json.dump(all_scores, f, indent=4, ensure_ascii=False)

    for task, score in all_scores.items():
        for metric, value in score.items():
            if "class_scores" in value:
                value.pop("class_scores")

    with open(scores_file_name_summary, "wt", encoding="utf8") as f:
        json.dump(all_scores, f, indent=4, ensure_ascii=False)

    logging.info(
        f"Scores saved in: {scores_file_name}. A summary without class scores saved in: {scores_file_name_summary}"
    )

    return all_scores


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)

    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, Seq2SeqTrainingArguments))

    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
        # If we pass only one argument to the script, and it's the path to a json file,
        # let's parse it to get our arguments.
        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))

    elif len(sys.argv) == 2 and sys.argv[1].endswith(".yaml"):
        # If we pass only one argument to the script, and it's the path to a yaml file,
        # let's parse it to get our arguments.
        model_args, data_args, training_args = parser.parse_yaml_file(yaml_file=os.path.abspath(sys.argv[1]))
    else:
        model_args, data_args, training_args = parser.parse_args_into_dataclasses()

    if data_args.test_tasks is not None:
        if not data_args.evaluate_all_checkpoints:
            evaluate(
                model_args,
                data_args,
                training_args,
            )
        else:
            checkpoints = [
                c
                for c in glob.glob(
                    os.path.join(training_args.output_dir, "checkpoint-*"),
                )
                if os.path.isdir(c)
            ]

            checkpoints = sorted(checkpoints, key=lambda x: int(x.split("-")[-1]))

            logging.info(
                f"Found {len(checkpoints)} checkpoints in {training_args.output_dir}:"
                f" {checkpoints} . We will evaluate each of them."
            )

            for checkpoint in checkpoints:
                logging.info(f"Evaluating {checkpoint}")
                evaluate(
                    model_args,
                    data_args,
                    training_args,
                    checkpoint_path=checkpoint,
                )

src/evaluate.py
hitz-zentroa/GoLLIE
import copy
import json
import logging
import multiprocessing as mp
import os
from argparse import ArgumentParser
from functools import partial
from itertools import cycle
from typing import Type

import datasets
from tqdm import tqdm


def get_class(class_path: str) -> Type:
    components = class_path.split(".")
    mod = __import__(components[0])
    for comp in components[1:]:
        mod = getattr(mod, comp)

    return mod


def multicpu_generator(args, tqdm_position, config):
    dataloader_cls = get_class(config["dataloader_cls"])
    sampler_cls = get_class(config["sampler_cls"])
    seeds = config.get("seed", 0)
    if isinstance(seeds, int):
        seeds = [seeds]
    label_noise = config.get("label_noise_prob", 0.0)
    if isinstance(label_noise, float):
        label_noise = cycle([label_noise])

    if "train_file" in config:
        dataloader = dataloader_cls(config["train_file"], **config)
        for ie_task in config["tasks"]:
            for seed, noise_prob in zip(seeds, label_noise):
                config["seed"] = seed
                config["label_noise_prob"] = noise_prob
                # Avoid multiple values for keyword argument
                _kwargs = {**config, **config["task_configuration"][ie_task]}
                sampler = sampler_cls(
                    dataloader,
                    task=ie_task,
                    split="train",
                    **_kwargs,
                )

                output_name = f"{config['dataset_name'].lower()}.{ie_task.lower()}.train.{seed}.jsonl"

                if os.path.exists(os.path.join(args.output_dir, output_name)) and not args.overwrite_output_dir:
                    logging.warning(f"Skipping {output_name} because it already exists.")
                    continue

                with open(os.path.join(args.output_dir, output_name), "w") as _file, tqdm(
                    total=len(dataloader),
                    desc=f"{config['dataset_name']}-{ie_task}-train-{seed}",
                    position=tqdm_position,
                ) as progress:
                    ids = []
                    for elem in sampler:
                        _file.write(f"{json.dumps(elem, ensure_ascii=False)}\n")
                        if ids != elem["ids"]:
                            ids = elem["ids"]
                            progress.update(len(ids))

                logging.info(f"Data saved to {os.path.abspath(os.path.join(args.output_dir, output_name))}")

    if "dev_file" in config:
        config["seed"] = 0
        dataloader = dataloader_cls(config["dev_file"], **config)
        for task in config["tasks"]:
            _kwargs = {**config, **config["task_configuration"][task]}
            sampler = sampler_cls(
                dataloader,
                task=task,
                split="dev",
                **_kwargs,
            )

            output_name = f"{config['dataset_name'].lower()}.{task.lower()}.dev.jsonl"

            if os.path.exists(os.path.join(args.output_dir, output_name)) and not args.overwrite_output_dir:
                logging.warning(f"Skipping {output_name} because it already exists.")
                continue

            with open(os.path.join(args.output_dir, output_name), "w") as _file, tqdm(
                total=len(dataloader),
                desc=f"{config['dataset_name']}-{task}-dev",
                position=tqdm_position,
            ) as progress:
                ids = []
                for elem in sampler:
                    _file.write(f"{json.dumps(elem, ensure_ascii=False)}\n")
                    if ids != elem["ids"]:
                        ids = elem["ids"]
                        progress.update(len(ids))

            logging.info(f"Data saved to {os.path.abspath(os.path.join(args.output_dir, output_name))}")

    if "test_file" in config:
        config["seed"] = 0
        dataloader = dataloader_cls(config["test_file"], **config)
        for task in config["tasks"]:
            _kwargs = {**config, **config["task_configuration"][task]}
            sampler = sampler_cls(
                dataloader,
                task=task,
                split="test",
                **_kwargs,
            )

            output_name = f"{config['dataset_name'].lower()}.{task.lower()}.test.jsonl"

            if os.path.exists(os.path.join(args.output_dir, output_name)) and not args.overwrite_output_dir:
                logging.warning(f"Skipping {output_name} because it already exists.")
                continue

            with open(os.path.join(args.output_dir, output_name), "w") as _file, tqdm(
                total=len(dataloader),
                desc=f"{config['dataset_name']}-{task}-test",
                position=tqdm_position,
            ) as progress:
                ids = []
                for elem in sampler:
                    _file.write(f"{json.dumps(elem, ensure_ascii=False)}\n")
                    if ids != elem["ids"]:
                        ids = elem["ids"]
                        progress.update(len(ids))

            logging.info(f"Data saved to {os.path.abspath(os.path.join(args.output_dir, output_name))}")


def main(args):
    os.makedirs(args.output_dir, exist_ok=True)

    config_files = args.configs
    # We generate a new config for each train split and task, so we also parallelize over each split and task
    configs = []
    splits = ["train_file", "dev_file", "test_file"]
    data_info = {}
    for config_file in config_files:
        with open(config_file, "rt") as f:
            config = json.load(f)

        # Remove guidelines if baseline
        config["remove_guidelines"] = args.baseline
        config["include_examples_prob"] = float(args.include_examples)
        if args.remove_masking:
            config["label_noise_prob"] = 0.0

        if args.remove_dropout:
            for task in config["tasks"]:
                config["task_configuration"][task]["guideline_dropout"] = 0.0

        # We generate a new config for each train split and task
        tasks = config["tasks"]
        for split in splits:
            for task in tasks:
                new_config = copy.deepcopy(config)
                if split in new_config:
                    for other_split in splits:
                        if other_split != split:
                            if other_split in new_config:
                                new_config.pop(other_split)
                    new_config["tasks"] = [task]

                    if config["dataset_name"] not in data_info:
                        data_info[config["dataset_name"]] = {}

                    if task not in data_info[config["dataset_name"]]:
                        data_info[config["dataset_name"]][task] = {
                            "train_file": False,
                            "dev_file": False,
                            "test_file": False,
                        }
                    data_info[config["dataset_name"]][task][split] = True
                    configs.append(new_config)

    generator_fn = partial(
        multicpu_generator,
        args,
    )

    logging.warning(f"We will generate the following data: {json.dumps(data_info, indent=4)})")

    with mp.Pool(processes=min(os.cpu_count(), len(configs))) as pool:
        pool.starmap(generator_fn, enumerate(configs))

    logging.warning(f"Data saved to {os.path.abspath(args.output_dir)}")


if __name__ == "__main__":
    logging.basicConfig(level=logging.WARNING)
    datasets.logging.set_verbosity_error()
    datasets.logging.disable_progress_bar()
    parser = ArgumentParser("generate_data", description="Generate Code formatted data.")

    parser.add_argument(
        "-c",
        "--configs",
        nargs="+",
        dest="configs",
        type=str,
        help="The list of configuration files.",
    )
    parser.add_argument(
        "-o",
        "--output",
        type=str,
        dest="output_dir",
        default="data/processed",
        help="Output directory where files will be saved.",
    )
    parser.add_argument(
        "--overwrite_output_dir",
        action="store_true",
        help="Whether to overwrite the output dir.",
    )
    parser.add_argument(
        "--baseline",
        action="store_true",
        default=False,
        help="Whether to generate baseline data.",
    )
    parser.add_argument(
        "--include_examples",
        action="store_true",
        default=False,
        help="Whether to include examples in the data.",
    )
    parser.add_argument(
        "--remove_dropout",
        action="store_true",
        default=False,
        help="Remove guideline dropout for the ablation analysis.",
    )
    parser.add_argument(
        "--remove_masking",
        action="store_true",
        default=False,
        help="Remove guideline masking for the ablation analysis.",
    )

    args = parser.parse_args()
    main(args)

src/generate_data.py
hitz-zentroa/GoLLIE
import json
import logging
import os
from typing import List, Optional, Tuple, Union

import torch

from transformers import (
    AutoConfig,
    AutoModelForCausalLM,
    AutoModelForSeq2SeqLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    PreTrainedModel,
    PreTrainedTokenizerBase,
)
from transformers.models.auto.modeling_auto import (
    MODEL_FOR_CAUSAL_LM_MAPPING_NAMES,
    MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES,
)
from transformers.utils import is_ipex_available

from .model_utils import find_all_linear_names, get_trainable_parameters


def get_device_map(
    force_auto_device_map: bool, max_memory_MB: int = None, use_better_transformer: bool = False
) -> (str, Union[int, List[int]]):
    """
    Get the device map to use for loading the model

    Args:
        force_auto_device_map (`bool`):
            Whether to force the use of the auto device map. If set to True, the model will be split across
            GPUs and CPU to fit the model in memory. If set to False, a full copy of the model will be loaded
            into each GPU.
        max_memory_MB (`int`):
            Free memory per gpu in MB. Used to compute the device map when force_auto_device_map is set to True.
        use_better_transformer (`bool`, optional):
            Whether to transform the model using Better Transformer library:
            https://huggingface.co/docs/optimum/bettertransformer/overview. Requires optimum
            'https://huggingface.co/docs/optimum/installation'. Defaults to False.

    Returns:
        `str`:
            The device map to use for loading the model
    """
    if force_auto_device_map:
        if os.environ.get("LOCAL_RANK") is not None:
            # raise ValueError(
            #    "Found DDP environment and force_auto_device_map is set to True, this configuration "
            #    "is not supported. If you want to use DPP, set force_auto_device_map to False, so "
            #    "a copy of the model is loaded in each GPU. If you want the split the model across "
            #    "GPUs (force_auto_device_map=True), do not use DDP (launch your script with "
            #    "pyton -m src/run.py config.json). If you are not in a DDP environment but you see "
            #    "this error, you might have manually set the environment variable 'LOCAL_WORLD_SIZE' to a "
            #    "number different than 1, please, remove this environment variable or set it to 1"
            # )
            if torch.cuda.is_available():
                n_gpus = torch.cuda.device_count()
            elif is_ipex_available() and torch.xpu.is_available():
                n_gpus = torch.xpu.device_count()
            else:
                logging.warning("You are in a DDP environment but no GPU is available, this may cause errors later on")
                n_gpus = 0

            max_memory = {i: max_memory_MB for i in range(n_gpus)}
            local_rank = int(os.environ.get("LOCAL_RANK", "0"))
            device_map = {"": local_rank}
            max_memory = {"": max_memory[local_rank]} if max_memory_MB is not None else None

        else:
            logging.warning(
                "Using auto device map, we will split the model across GPUs and CPU to fit the model in memory."
            )
            device_map = "auto"
            max_memory = max_memory_MB
    else:
        max_memory = None
        word_size = int(os.environ.get("LOCAL_WORLD_SIZE", 1))
        if word_size > 1:
            logging.warning(
                "Found DDP environment and force_auto_device_map is set to False, we will load a copy of the model "
                "on each GPU."
            )
            device_map = None  # {"": int(os.environ.get("LOCAL_RANK", 0))}

        else:
            if not use_better_transformer:
                device_map = None
            else:
                logging.warning("Setting device map to 'auto' to use Better Transformers library.")
                device_map = "auto"

    logging.info(f"We will load the model using the following device map: {device_map} and max_memory: {max_memory}")

    return device_map, max_memory


def merge_lora_model(
    weights_path: str,
    lora_weights_name_or_path: str,
    output_path: str,
    torch_dtype: Optional[str] = None,
):
    """
    Given a model path and the path to the LoRA weights, merge the LoRA weights into the model and save the merged model
    weights_path (`str`):
            The path to your local model weights and tokenizer. You can also provide a
            huggingface hub model name.
    lora_weights_name_or_path (`str`):
        If the model has been trained with LoRA, path or huggingface hub name to the
        pretrained weights. Defaults to `None`.
    output_path (`str`):
        The path to the output directory where the merged model will be saved.
    torch_dtype (`Optional[str]`, optional):
        The torch dtype to use for the model. If set to `"auto"`, the dtype will be
        automatically derived.
    """

    logging.info(f"We will merge the LoRA weights from {lora_weights_name_or_path} into the model {weights_path}")
    model, tokenizer = load_model(
        inference=True,
        model_weights_name_or_path=weights_path,
        use_lora=True,
        lora_weights_name_or_path=lora_weights_name_or_path,
        torch_dtype=torch_dtype,
    )

    model.config.save_pretrained(output_path)
    model.save_pretrained(output_path)
    tokenizer.save_pretrained(output_path)

    logging.info(f"Model merged and saved in {output_path}")


def load_model(
    inference: bool,
    model_weights_name_or_path: str,
    quantization: Optional[int] = None,
    use_lora: bool = False,
    lora_weights_name_or_path: Optional[str] = None,
    lora_target_modules: Optional[List[str]] = None,
    lora_r: Optional[int] = 8,
    lora_alpha: Optional[int] = 16,
    lora_dropout: Optional[float] = 0.05,
    torch_dtype: Optional[str] = "bfloat16",
    force_auto_device_map: bool = False,
    use_gradient_checkpointing: bool = False,
    trust_remote_code: bool = False,
    use_flash_attention: bool = True,
    use_better_transformer: bool = False,
    fsdp_training: bool = False,
    max_memory_MB: Optional[int] = None,
) -> Tuple[PreTrainedModel, PreTrainedTokenizerBase]:
    """
    Load any Decoder model for training.

    Args:
        inference (`bool`):
            Whether to load the model for inference or training. If set to `True`, the model will be loaded
            in evaluation mode. In this case, if use_lora is set to `True`, you must provide the path to the
            LoRA weights. Defaults to `False`.
        model_weights_name_or_path (`str`):
            The path to your local model weights and tokenizer or huggingface model name.
        quantization (`int`, optional):
            '4' or '8' for 4 bits or 8 bits quantization or None for 16/32bits training. Defaults to `None`.

            Requires bitsandbytes library: https://github.com/TimDettmers/bitsandbytes
        use_lora (`bool`, optional):
            Whether to use LORA. Defaults to False.

            See https://arxiv.org/pdf/2106.09685.pdf for more details.

            Requires huggingface PEFT library: https://github.com/huggingface/peft
        lora_weights_name_or_path (`Optional[str]`, optional):
            The name or path to the pre-trained LORA model weights. You can also provide
            a huggingface hub model name to load the weights from there. If not provided,
            the weights will be initialized randomly, this requires training the model.
            Defaults to `None`.
        lora_target_modules (`Optional[List[str]]`, optional):
            The list of modules to apply LORA to. If not provided, we will use PEFT
            default modules. Defaults to `None`.
        lora_r (`Optional[int]`, optional):
            Lora attention dimension. Defaults to `8`.
        lora_alpha (`Optional[int]`, optional):
            The alpha parameter for Lora scaling. Defaults to `16`.
        lora_dropout (`Optional[float]`, optional):
            The dropout probability for Lora layers. Defaults to 0.05.
        torch_dtype (`Optional[str]`, optional):
            Override the default `torch.dtype` and load the model under this dtype. If
            `auto` is passed, the dtype will be automatically derived from the model's
            weights. Defaults to `None`.
        force_auto_device_map (`bool`, optional):
            Whether to force the use of the auto device map. If set to True, the model will be split across
            GPUs and CPU to fit the model in memory. If set to False, a full copy of the model will be loaded
            into each GPU. Defaults to False.
        use_gradient_checkpointing (`bool`, optiona):
            Whether to use gradient checkpointing for training
        trust_remote_code (`bool`, optional):
            Trust the remote code from HuggingFace model hub. Defaults to False.
        use_flash_attention (`bool`, optional):
            Whether to use Flash Attention. Defaults to True. Flash attention must be installed, see:
            'https://github.com/Dao-AILab/flash-attention' for more details.
            GoLLIE models have been trained with Flash Attention, setting this flag to False can cause
            the model to perform worse due to numerical differences in the attention weights.
        use_better_transformer (`bool`, optional):
            Whether to transform the model using Better Transformer library:
            https://huggingface.co/docs/optimum/bettertransformer/overview. Requires optimum
            'https://huggingface.co/docs/optimum/installation'. Only supported for inference!
            Defaults to False.
        fsdp_training: (`bool`, optional):
            Whether Fully Sharded Data Parallelism is enabled for training. Defaults to False.
            Used to prevent casting layers to fp32 if the model is already in fp16, which causes
            an error: ValueError: Must flatten tensors with uniform dtype but got torch.float16 and torch.float32
        max_memory_MB (`int`):
            Free memory per gpu in MB. Used to compute the device map when force_auto_device_map is set to True.
    Raises:
        `ValueError`:
            is raised when `int8_quantization=True` but `use_lora=False`.

    Returns:
        `Tuple[PreTrainedModel, PreTrainedTokenizerBase]`:
            The loaded model and tokenizer.
    """

    if not use_flash_attention and inference:
        logging.warning(
            "\n\n==========================================================\n\n"
            "You are not using Flash Attention. The released\n"
            "GoLLIE models have been trained with Flash Attention,\n"
            "setting this flag to False can cause the model to perform\n"
            "worse due to numerical differences in the attention weights."
            "=====================================================\n\n"
        )

    # Sanity checks

    if isinstance(quantization, str):
        quantization = int(quantization)
    assert (quantization is None) or (
        quantization in [4, 8]
    ), f"Quantization must be 4 or 8, or None for FP32/FP16 training. You passed: {quantization}"

    if not inference and quantization is not None and not use_lora:
        raise ValueError(
            "'Quantization' == 4/8 is only supported with LoRA. If you want "
            "to train a 4/8bits quantified model, you must set `use_lora=True`. If you want to "
            "use a 4/8 bits optimizer, set `quantization=None` and choose a 4/8 bit optimizer using 'optim' "
            "argument (e.g 'adamw_bnb_8bit', 'lion_8bit', 'paged_adamw_8bit', ...)."
        )

    if inference and use_lora and lora_weights_name_or_path is None:
        raise ValueError("You must provide the path to the LoRA weights when loading the model for inference.")

    if use_better_transformer and not inference:
        logging.warning(
            "Better Transformer is only supported for inference. Better Transformers does not support "
            "attention mask for training, therefore it is not compatible with CoLLIE training. See "
            "https://huggingface.co/docs/optimum/bettertransformer/overview for more details. We will "
            "set use_better_transformer=False."
        )
        use_better_transformer = False

    if use_better_transformer and use_flash_attention:
        raise ValueError(
            "You cannot use both Flash Attention and Better Transformer flags. Flash Attention is already part of"
            " Better Transformers, so you can just set use_better_transformer=True to use Flash Attention. The Flash"
            " Attention flag is intended for patching HuggingFace models."
        )

    if lora_weights_name_or_path is not None and not use_lora:
        logging.warning("You provided a path to LoRA weights but use_lora is set to False. We will set use_lora=True.")
        use_lora = True

    logging.info(f"Loading model model from {model_weights_name_or_path}")

    # Get the device map config

    device_map, max_memory = get_device_map(
        force_auto_device_map=force_auto_device_map,
        max_memory_MB=max_memory_MB,
        use_better_transformer=use_better_transformer,
    )

    # Load the model config

    if use_lora:
        config = AutoConfig.from_pretrained(
            model_weights_name_or_path,
            trust_remote_code=trust_remote_code,
            pretraining_tp=1,  # Fix mat1 and mat2 shapes cannot be multiplied  error with LLaMA-2
            # See https://github.com/huggingface/transformers/pull/24906
        )
    else:
        config = AutoConfig.from_pretrained(
            model_weights_name_or_path,
            trust_remote_code=trust_remote_code,
        )

    # Load the model tokenizer

    tokenizer: PreTrainedTokenizerBase = AutoTokenizer.from_pretrained(
        model_weights_name_or_path,
        add_eos_token=True,
        trust_remote_code=trust_remote_code,
    )

    if tokenizer.pad_token_id is None:
        if "<|padding|>" in tokenizer.get_vocab():
            # StabilityLM specific fix
            tokenizer.add_special_tokens({"pad_token": "<|padding|>"})
        elif tokenizer.unk_token is not None:
            logging.warning("Tokenizer does not have a pad token, we will use the unk token as pad token.")
            tokenizer.pad_token_id = tokenizer.unk_token_id
        else:
            logging.warning("Tokenizer does not have a pad token. We will use the eos token as pad token.")
            tokenizer.pad_token_id = tokenizer.eos_token_id

    # Load the model weights

    #  Get the quantization config
    quant_args = {}
    torch_dtype = torch_dtype if torch_dtype in ["auto", None] else getattr(torch, torch_dtype)

    if quantization is not None:
        quant_args = {"load_in_4bit": True} if quantization == 4 else {"load_in_8bit": True}
        if quantization == 4:
            bnb_config = BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_use_double_quant=True,
                bnb_4bit_quant_type="nf4",
                bnb_4bit_compute_dtype=torch.bfloat16 if torch_dtype in ["auto", None] else torch_dtype,
            )

        else:
            bnb_config = BitsAndBytesConfig(
                load_in_8bit=True,
            )
        logging.info(f"Bits and Bytes config: {json.dumps(bnb_config.to_dict(),indent=4,ensure_ascii=False)}")
    else:
        logging.info(f"Loading model with dtype: {torch_dtype}")
        bnb_config = None

    #  Get the correct load function for each model_type
    if config.model_type in MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES:
        logging.warning(
            f"Model {model_weights_name_or_path} is a encoder-decoder model. We will load it as a Seq2SeqLM model."
        )

        if config.model_type == "t5" or config.model_type == "mt5":
            logging.warning(
                f"PLEASE READ!!! Model {model_weights_name_or_path} is a T5 model."
                " T5/mT5/Flan-T5/UL2/Flan-UL2 released by google lack the token"
                " representation for new line tokens and multiple spaces, they are"
                " currently NOT supported in CoLLIE. We will attempt to load the model"
                " anyway but you might encounter unexpected behavior or errors."
            )

        load_fn = AutoModelForSeq2SeqLM
        model_type = "seq2seq"

    elif config.model_type in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES:
        logging.warning(
            f"Model {model_weights_name_or_path} is an decoder-only model. We will load it as a CausalLM model."
        )

        if config.model_type == "llama" and use_flash_attention:
            from src.model.patch_models.modeling_flash_llama import LlamaForCausalLM as LlamaForCausalLMFlash

            logging.warning("Using Flash Attention for LLaMA model.")
            load_fn = LlamaForCausalLMFlash
            use_flash_attention = False  # Do not path the model twice

        else:
            load_fn = AutoModelForCausalLM

        tokenizer.padding_side = "left"
        model_type = "causal"

    else:
        raise ValueError(
            f"Model {model_weights_name_or_path} of type {config.model_type} is not supported by CoLLIE."
            "Supported models are:\n"
            f"Seq2SeqLM: {MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES}\n"
            f"CausalLM: {MODEL_FOR_CAUSAL_LM_MAPPING_NAMES}\n"
        )

    #  Load the model weights
    model: PreTrainedModel = load_fn.from_pretrained(
        pretrained_model_name_or_path=model_weights_name_or_path,
        device_map=device_map,
        max_memory=max_memory,
        quantization_config=bnb_config,
        torch_dtype=torch_dtype,
        config=config,
        trust_remote_code=trust_remote_code,
        **quant_args,
    )

    # Path the model to use flash attention using OpenAssistant patching function
    if use_flash_attention:
        from src.model.patch_models.patching import patch_model

        logging.info("Patching model to use flash attention")
        patch_model(model, resid_pdrop=None, flash_attention=True)

    logging.info(f"Model dtype: {model.dtype}")
    logging.info("Total model memory footprint: " + str(model.get_memory_footprint() / 1e6) + " MB")

    # Prepare the model for k-bit training and enable gradient checkpointing
    if quantization is not None and not inference:
        # Custom prepare_model_for_kbit_training fuction that does not convert weights to float32
        # Our fuction is faster and more memory efficient, altough it may introduce some numerical instability
        # we have not observed any issues in our experiments. If you encounter any issues, you can comment
        # the following line and uncomment the next one.
        from .model_utils import prepare_model_for_kbit_training

        # from peft import prepare_model_for_kbit_training

        model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=use_gradient_checkpointing)
    else:
        if use_gradient_checkpointing and not inference:
            model.gradient_checkpointing_enable()

    # Load LoRA weights
    if use_lora:
        from peft import LoraConfig, PeftModel, TaskType, get_peft_model

        if not inference:
            model.enable_input_require_grads()  #  Enables the gradients for the input embeddings

        if lora_weights_name_or_path is None:
            logging.info("No pretrained LORA weights provided, we will initialize the weights randomly.")

            if lora_target_modules is None or (lora_target_modules is not None and len(lora_target_modules) == 0):
                logging.warning(
                    "No target modules provided,  will use the default modules for the"
                    " model in huggingface PEFT library. "
                )
                lora_target_modules = None

            if lora_target_modules == ["all"]:
                logging.warning(
                    "You provided 'all' as target modules, we will use all the model to which LoRA can be applied."
                )
                lora_target_modules = find_all_linear_names(model, quantization=quantization)

            lora_config = LoraConfig(
                r=lora_r,
                lora_alpha=lora_alpha,
                lora_dropout=lora_dropout,
                bias="none",
                task_type=TaskType.CAUSAL_LM if model_type == "causal" else TaskType.SEQ_2_SEQ_LM,
                target_modules=lora_target_modules,
            )

            model = get_peft_model(model, lora_config)

        else:
            logging.info(f"Loading pretrained LORA weights from {lora_weights_name_or_path}")

            model = PeftModel.from_pretrained(model, lora_weights_name_or_path)

        logging.info(f"\nLoRA config:\n{model.peft_config}\n")

    """
    Convert the model layers to the correct dtype
    If LoRA and bf16 is used, we convert the LoRA layers to bf16 for faster training
    """

    if use_lora:
        from peft.tuners.lora import LoraLayer

        for name, module in model.named_modules():
            if isinstance(module, LoraLayer):
                if torch_dtype == torch.bfloat16:
                    logging.debug(f"Converting LoRA layer {name} to {torch_dtype}")
                    module = module.to(torch.bfloat16)

    if not fsdp_training:
        for name, module in model.named_modules():
            if "norm" in name:
                logging.debug(f"Converting layer {name} to {torch.float32}")
                module = module.to(torch.float32)
            if "lm_head" in name or "embed_tokens" in name:
                if hasattr(module, "weight"):
                    if torch_dtype == torch.bfloat16 and module.weight.dtype == torch.float32:
                        logging.debug(f"Converting layer {name} to {torch_dtype}")
                        module = module.to(torch.bfloat16)

    if inference:
        if use_lora:
            if quantization is None:
                # If we are not using quantization, we merge the LoRA layers into the model for faster inference.
                # This is not possible if we are using 4/8 bit quantization.
                logging.info("Merging LoRA layers into the model for faster inference.")
                model = model.merge_and_unload()
            else:
                logging.info(
                    "Quantization is enabled, we will not merge LoRA layers into the model. Inference will be slower."
                )
    else:
        trainable_params, total_params, trainable_percentage = get_trainable_parameters(model)
        logging.info(
            f"---> Trainable params: {trainable_params} || all params: {total_params} ||"
            f" trainable%: {round(trainable_percentage,6)}\n"
        )

    return model, tokenizer

src/model/load_model.py
hitz-zentroa/GoLLIE
from typing import Optional, Tuple

import torch

from transformers import PreTrainedModel


def get_trainable_parameters(model: PreTrainedModel) -> Tuple[int, int, float]:
    """
    Prints the number of trainable parameters in the model.

    Args:
        model (`PreTrainedModel`):
            The model to print the number of trainable parameters for.

    Returns:
        `Tuple[int, int, float]`:
            The number of trainable parameters, the total number of parameters and the
            percentage of trainable parameters.
    """
    trainable_params = 0
    all_param = 0
    for _, param in model.named_parameters():
        num_params = param.numel()
        # if using DS Zero 3 and the weights are initialized empty
        if num_params == 0 and hasattr(param, "ds_numel"):
            num_params = param.ds_numel

        all_param += num_params
        if param.requires_grad:
            trainable_params += num_params

    return trainable_params, all_param, 100 * trainable_params / all_param


def prepare_model_for_kbit_training(
    model, use_gradient_checkpointing=True, layer_norm_names=["layer_norm"], output_embedding_layer_name="lm_head"
):
    r"""
    This method wraps the entire protocol for preparing a model before running a training. This includes:
        1- Cast the layernorm in fp32 2- making output embedding layer require grads 3- Add the upcasting of the lm
        head to fp32

    Args:
        model, (`transformers.PreTrainedModel`):
            The loaded model from `transformers`
    """
    loaded_in_kbit = getattr(model, "is_loaded_in_8bit", False) or getattr(model, "is_loaded_in_4bit", False)

    for name, param in model.named_parameters():
        # freeze base model's layers
        param.requires_grad = False

        if loaded_in_kbit:
            # cast layer norm in fp32 for stability for 8bit models
            if param.ndim == 1 and any(layer_norm_name in name for layer_norm_name in layer_norm_names):
                param.data = param.data.to(torch.float32)

    if loaded_in_kbit and use_gradient_checkpointing:
        # For backward compatibility
        if hasattr(model, "enable_input_require_grads"):
            model.enable_input_require_grads()
        else:

            def make_inputs_require_grad(module, input, output):
                output.requires_grad_(True)

            model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)

        # enable gradient checkpointing for memory efficiency
        model.gradient_checkpointing_enable()

    if hasattr(model, output_embedding_layer_name):
        output_embedding_layer = getattr(model, output_embedding_layer_name)
        input_dtype = output_embedding_layer.weight.dtype

        class CastOutputToFloat(torch.nn.Sequential):
            r"""
            Manually cast to the expected dtype of the lm_head as sometimes there is a final layer norm that is casted
            in fp32

            """

            def forward(self, x):
                return super().forward(x.to(input_dtype)).to(torch.float32)

        setattr(model, output_embedding_layer_name, CastOutputToFloat(output_embedding_layer))

    return model


def find_all_linear_names(model, quantization: Optional[int] = None):
    if quantization is None:
        cls = torch.nn.Linear
    elif quantization == 4:
        from bitsandbytes.nn import Linear4bit

        cls = Linear4bit
    elif quantization == 8:
        from bitsandbytes.nn import Linear8bitLt

        cls = Linear8bitLt
    else:
        raise ValueError(f"Unknown quantization type: {quantization}")

    lora_module_names = set()
    for name, module in model.named_modules():
        if isinstance(module, cls):
            names = name.split(".")
            lora_module_names.add(names[0] if len(names) == 1 else names[-1])

    if "lm_head" in lora_module_names:  # needed for 16-bit
        lora_module_names.remove("lm_head")
    return list(lora_module_names)

src/model/model_utils.py
hitz-zentroa/GoLLIE
# coding=utf-8
# From https://huggingface.co/togethercomputer/LLaMA-2-7B-32K/blob/main/modeling_flash_llama.py
# With seqlen fix from Alex Birch: https://huggingface.co/togethercomputer/LLaMA-2-7B-32K/discussions/17
# With dtype Fix by Oscar Sainz
# With Beam Search Fix by Iker García-Ferrero

# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
#
# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
# and OPT implementations in this library. It has been modified from its
# original forms to accommodate minor architectural differences compared
# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
""" PyTorch LLaMA model."""
from typing import List, Optional, Tuple, Union

import torch
import torch.nn.functional as F
import torch.utils.checkpoint
from torch import nn
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss

from transformers.activations import ACT2FN
from transformers.modeling_outputs import (
    BaseModelOutputWithPast,
    CausalLMOutputWithPast,
    SequenceClassifierOutputWithPast,
)
from transformers.modeling_utils import PreTrainedModel
from transformers.models.llama.configuration_llama import LlamaConfig
from transformers.utils import (
    add_start_docstrings,
    add_start_docstrings_to_model_forward,
    logging,
    replace_return_docstrings,
)


try:
    from flash_attn.bert_padding import pad_input, unpad_input
    from flash_attn.flash_attn_interface import (
        flash_attn_kvpacked_func,
        flash_attn_varlen_kvpacked_func,
    )

    flash_attn_v2_installed = True
    print(">>>> Flash Attention installed")
except ImportError:
    flash_attn_v2_installed = False
    raise ImportError("Please install Flash Attention: `pip install flash-attn --no-build-isolation`")

try:
    from flash_attn.layers.rotary import apply_rotary_emb_func

    flash_rope_installed = True
    print(">>>> Flash RoPE installed")
except ImportError:
    flash_rope_installed = False
    raise ImportError(
        "Please install RoPE kernels: `pip install"
        " git+https://github.com/HazyResearch/flash-attention.git#subdirectory=csrc/rotary`"
    )

logger = logging.get_logger(__name__)

_CONFIG_FOR_DOC = "LlamaConfig"


# @torch.jit.script
def rmsnorm_func(hidden_states, weight, variance_epsilon):
    input_dtype = hidden_states.dtype
    hidden_states = hidden_states.to(torch.float32)
    variance = hidden_states.pow(2).mean(-1, keepdim=True)
    hidden_states = hidden_states * torch.rsqrt(variance + variance_epsilon)
    return (weight * hidden_states).to(input_dtype)


class LlamaRMSNorm(nn.Module):
    def __init__(self, hidden_size, eps=1e-6):
        """
        LlamaRMSNorm is equivalent to T5LayerNorm
        """
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.register_buffer(
            "variance_epsilon",
            torch.tensor(eps),
            persistent=False,
        )

    def forward(self, hidden_states):
        return rmsnorm_func(hidden_states, self.weight, self.variance_epsilon)


class FlashRotaryEmbedding(torch.nn.Module):
    """
    The rotary position embeddings from RoFormer_ (Su et. al).
    A crucial insight from the method is that the query and keys are
    transformed by rotation matrices which depend on the relative positions.
    Other implementations are available in the Rotary Transformer repo_ and in
    GPT-NeoX_, GPT-NeoX was an inspiration
    .. _RoFormer: https://arxiv.org/abs/2104.09864
    .. _repo: https://github.com/ZhuiyiTechnology/roformer
    .. _GPT-NeoX: https://github.com/EleutherAI/gpt-neox
    If scale_base is not None, this implements XPos (Sun et al., https://arxiv.org/abs/2212.10554).
    A recommended value for scale_base is 512: https://github.com/HazyResearch/flash-attention/issues/96
    Reference: https://github.com/sunyt32/torchscale/blob/main/torchscale/component/xpos_relative_position.py
    """

    def __init__(
        self,
        dim: int,
        base=10000.0,
        interleaved=False,
        scale_base=None,
        scaling_factor=1.0,
        pos_idx_in_fp32=True,
        device=None,
    ):
        """
        interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead
            of 1st half and 2nd half (GPT-NeoX style).
        pos_idx_in_fp32: if True, the position indices [0.0, ..., seqlen - 1] are in fp32,
            otherwise they might be in lower precision.
            This option was added because previously (before 2023-07-02), when we construct
            the position indices, we use the dtype of self.inv_freq. In most cases this would
            be fp32, but if the model is trained in pure bf16 (not mixed precision), then
            self.inv_freq would be bf16, and the position indices are also in bf16.
            Because of the limited precision of bf16 (e.g. 1995.0 is rounded to 2000.0), the
            embeddings for some positions will coincide.
            To maintain compatibility with models previously trained in pure bf16,
            we add this option.
        scaling_factor: RotaryEmbedding extended with linear scaling.
        """
        super().__init__()
        self.dim = dim
        self.base = float(base)
        self.pos_idx_in_fp32 = pos_idx_in_fp32
        # Generate and save the inverse frequency buffer (non trainable)
        inv_freq = self._compute_inv_freq(device)
        self.register_buffer("inv_freq", inv_freq, persistent=False)
        self.interleaved = interleaved
        self.scale_base = scale_base
        self.scaling_factor = scaling_factor
        scale = (
            (torch.arange(0, dim, 2, device=device, dtype=torch.float32) + 0.4 * dim) / (1.4 * dim)
            if scale_base is not None
            else None
        )
        self.register_buffer("scale", scale)

        self._seq_len_cached = 0
        self._cos_cached = None
        self._sin_cached = None
        self._cos_k_cached = None
        self._sin_k_cached = None

    def _compute_inv_freq(self, device=None):
        return 1 / (self.base ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim))

    def _update_cos_sin_cache(self, seqlen, device=None, dtype=None):
        # Reset the tables if the sequence length has changed,
        # if we're on a new device (possibly due to tracing for instance),
        # or if we're switching from inference mode to training
        if (
            seqlen > self._seq_len_cached
            or self._cos_cached.device != device
            or self._cos_cached.dtype != dtype
            or (self.training and self._cos_cached.is_inference())
        ):
            self._seq_len_cached = seqlen
            # We want fp32 here, not self.inv_freq.dtype, since the model could be loaded in bf16
            # And the output of arange can be quite large, so bf16 would lose a lot of precision.
            # However, for compatibility reason, we add an option to use the dtype of self.inv_freq.
            if self.pos_idx_in_fp32:
                t = torch.arange(seqlen, device=device, dtype=torch.float32)
                t /= self.scaling_factor
                # We want fp32 here as well since inv_freq will be multiplied with t, and the output
                # will be large. Having it in bf16 will lose a lot of precision and cause the
                # cos & sin output to change significantly.
                # We want to recompute self.inv_freq if it was not loaded in fp32
                if self.inv_freq.dtype != torch.float32:
                    inv_freq = self.inv_freq.to(torch.float32)
                else:
                    inv_freq = self.inv_freq
            else:
                t = torch.arange(seqlen, device=device, dtype=self.inv_freq.dtype)
                t /= self.scaling_factor
                inv_freq = self.inv_freq
            # Don't do einsum, it converts fp32 to fp16 under AMP
            # freqs = torch.einsum("i,j->ij", t, self.inv_freq)
            freqs = torch.outer(t, inv_freq)
            if self.scale is None:
                self._cos_cached = torch.cos(freqs).to(dtype)
                self._sin_cached = torch.sin(freqs).to(dtype)
            else:
                power = (
                    torch.arange(seqlen, dtype=self.scale.dtype, device=self.scale.device) - seqlen // 2
                ) / self.scale_base
                scale = self.scale.to(device=power.device) ** power.unsqueeze(-1)
                # We want the multiplication by scale to happen in fp32
                self._cos_cached = (torch.cos(freqs) * scale).to(dtype)
                self._sin_cached = (torch.sin(freqs) * scale).to(dtype)
                self._cos_k_cached = (torch.cos(freqs) / scale).to(dtype)
                self._sin_k_cached = (torch.sin(freqs) / scale).to(dtype)

    def forward(self, q: torch.Tensor, k: torch.Tensor, seqlen_offset: int = 0) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        q: (batch, seqlen, nheads, headdim)
        k: (batch, seqlen, nheads, headdim)
        seqlen_offset: can be used in generation where the qkv being passed in is only the last
        token in the batch.
        """
        self._update_cos_sin_cache(q.shape[1] + seqlen_offset, device=q.device, dtype=q.dtype)
        if self.scale is None:
            return apply_rotary_emb_func(
                q,
                self._cos_cached[seqlen_offset:],
                self._sin_cached[seqlen_offset:],
                self.interleaved,
                True,  # inplace=True
            ), apply_rotary_emb_func(
                k,
                self._cos_cached[seqlen_offset:],
                self._sin_cached[seqlen_offset:],
                self.interleaved,
                True,  # inplace=True
            )
        else:
            assert False


class LlamaMLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.hidden_size = config.hidden_size
        self.intermediate_size = config.intermediate_size
        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
        self.act_fn = ACT2FN[config.hidden_act]

    def forward(self, x):
        if self.config.pretraining_tp > 1:
            slice = self.intermediate_size // self.config.pretraining_tp
            gate_proj_slices = self.gate_proj.weight.split(slice, dim=0)
            up_proj_slices = self.up_proj.weight.split(slice, dim=0)
            down_proj_slices = self.down_proj.weight.split(slice, dim=1)

            gate_proj = torch.cat(
                [F.linear(x, gate_proj_slices[i]) for i in range(self.config.pretraining_tp)], dim=-1
            )
            up_proj = torch.cat([F.linear(x, up_proj_slices[i]) for i in range(self.config.pretraining_tp)], dim=-1)

            intermediate_states = (self.act_fn(gate_proj) * up_proj).split(slice, dim=2)
            down_proj = [
                F.linear(intermediate_states[i], down_proj_slices[i]) for i in range(self.config.pretraining_tp)
            ]
            down_proj = sum(down_proj)
        else:
            down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))

        return down_proj


@torch.jit.script
def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
    """
    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
    """
    batch, slen, _, num_key_value_heads, head_dim = hidden_states.shape
    if n_rep == 1:
        return hidden_states
    hidden_states = hidden_states[:, :, :, :, None, :].expand(batch, slen, 2, num_key_value_heads, n_rep, head_dim)
    return hidden_states.reshape(batch, slen, 2, num_key_value_heads * n_rep, head_dim)


class LlamaAttention(nn.Module):
    """Multi-headed attention from 'Attention Is All You Need' paper"""

    def __init__(self, config: LlamaConfig):
        super().__init__()
        self.config = config
        self.hidden_size = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.head_dim = self.hidden_size // self.num_heads
        self.num_key_value_heads = config.num_key_value_heads
        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
        self.max_position_embeddings = config.max_position_embeddings

        if (self.head_dim * self.num_heads) != self.hidden_size:
            raise ValueError(
                f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
                f" and `num_heads`: {self.num_heads})."
            )
        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)

        self.register_buffer(
            "norm_factor",
            torch.sqrt(torch.tensor(self.head_dim, dtype=torch.float32)).to(torch.get_default_dtype()),
            persistent=False,
        )

        if self.config.rope_scaling is None:
            scaling_factor = 1
        else:
            scaling_type = self.config.rope_scaling["type"]
            scaling_factor = self.config.rope_scaling["factor"]
            assert scaling_type == "linear"

        self.rotary_emb = FlashRotaryEmbedding(
            self.head_dim,
            base=10000,
            interleaved=False,
            scaling_factor=scaling_factor,
        )

    def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
        return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_value: Optional[Tuple[torch.Tensor]] = None,
        output_attentions: bool = False,
        use_cache: bool = False,
        is_padded_inputs: Optional[bool] = False,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
        bsz, q_len, h_size = hidden_states.size()

        has_layer_past = past_key_value is not None

        if has_layer_past:
            past_kv = past_key_value[0]
            past_len = past_key_value[1]
        else:
            past_len = 0

        if self.config.pretraining_tp > 1:
            key_value_slicing = (self.num_key_value_heads * self.head_dim) // self.config.pretraining_tp
            query_slices = self.q_proj.weight.split(
                (self.num_heads * self.head_dim) // self.config.pretraining_tp, dim=0
            )
            key_slices = self.k_proj.weight.split(key_value_slicing, dim=0)
            value_slices = self.v_proj.weight.split(key_value_slicing, dim=0)

            q = [F.linear(hidden_states, query_slices[i]) for i in range(self.config.pretraining_tp)]
            q = torch.cat(q, dim=-1)

            k = [F.linear(hidden_states, key_slices[i]) for i in range(self.config.pretraining_tp)]
            k = torch.cat(k, dim=-1)

            v = [F.linear(hidden_states, value_slices[i]) for i in range(self.config.pretraining_tp)]
            v = torch.cat(v, dim=-1)

        else:
            q = self.q_proj(hidden_states)
            k = self.k_proj(hidden_states)
            v = self.v_proj(hidden_states)

        q = q.view(bsz, q_len, self.num_heads, self.head_dim)
        k = k.view(bsz, q_len, self.num_key_value_heads, self.head_dim)
        v = v.view(bsz, q_len, self.num_key_value_heads, self.head_dim)

        q, k = self.rotary_emb(q, k, past_len)

        kv = torch.stack([k, v], 2)
        kv = repeat_kv(kv, self.num_key_value_groups)

        # Make sure both are same dtype
        if q.dtype != kv.dtype:
            kv = kv.to(q.dtype)

        # Cache QKV values
        if has_layer_past:
            new_len = past_len + q.size(1)
            if new_len > past_kv.size(1):
                past_kv = torch.cat(
                    [past_kv, torch.empty(bsz, 256, 2, kv.size(3), kv.size(4), dtype=kv.dtype, device=kv.device)], 1
                )
            past_kv[:, past_len:new_len] = kv
            kv = past_kv[:, :new_len]
        else:
            past_kv = kv

        past_key_value = (past_kv, past_len + q.size(1)) if use_cache else None

        if is_padded_inputs:
            # varlen, ignore padding tokens, efficient for large batch with many paddings

            assert attention_mask is not None

            unpadded_kv, indices_k, cu_seqlens_k, max_seqlen_k = unpad_input(kv, attention_mask)
            unpadded_q, indices_q, cu_seqlens_q, max_seqlen_q = unpad_input(q, attention_mask[:, -q.size(1) :])

            # Make sure both are same dtype
            if unpadded_q.dtype != unpadded_kv.dtype:
                unpadded_kv = unpadded_kv.to(unpadded_q.dtype)

            attn_outputs = flash_attn_varlen_kvpacked_func(
                unpadded_q,
                unpadded_kv,
                cu_seqlens_q,
                cu_seqlens_k,
                max_seqlen_q,
                max_seqlen_k,
                dropout_p=0.0,
                softmax_scale=1.0 / self.norm_factor,
                causal=(not has_layer_past),
                return_attn_probs=output_attentions,
            )

            attn_output = attn_outputs[0] if output_attentions else attn_outputs
            attn_output = pad_input(attn_output, indices_q, bsz, q_len).reshape(bsz, q_len, h_size)
            attn_weights = attn_outputs[2] if output_attentions else None

        else:
            # Make sure both are same dtype
            if q.dtype != kv.dtype:
                kv = kv.to(q.dtype)

            # no padding tokens, more efficient
            attn_outputs = flash_attn_kvpacked_func(
                q,
                kv,
                dropout_p=0.0,
                softmax_scale=1.0 / self.norm_factor,
                causal=(not has_layer_past),
                return_attn_probs=output_attentions,
            )

            attn_output = attn_outputs[0] if output_attentions else attn_outputs
            attn_output = attn_output.reshape(bsz, q_len, h_size)
            attn_weights = attn_outputs[2] if output_attentions else None

        if self.config.pretraining_tp > 1:
            attn_output = attn_output.split(self.hidden_size // self.config.pretraining_tp, dim=2)
            o_proj_slices = self.o_proj.weight.split(self.hidden_size // self.config.pretraining_tp, dim=1)
            attn_output = sum([F.linear(attn_output[i], o_proj_slices[i]) for i in range(self.config.pretraining_tp)])
        else:
            attn_output = self.o_proj(attn_output)

        if not output_attentions:
            attn_weights = None

        return attn_output, attn_weights, past_key_value


class LlamaDecoderLayer(nn.Module):
    def __init__(self, config: LlamaConfig):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.self_attn = LlamaAttention(config=config)
        self.mlp = LlamaMLP(config)
        self.input_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_value: Optional[Tuple[torch.Tensor]] = None,
        is_padded_inputs: Optional[bool] = False,
        output_attentions: Optional[bool] = False,
        use_cache: Optional[bool] = False,
    ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
        """
        Args:
            hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
            attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
                `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
            output_attentions (`bool`, *optional*):
                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
                returned tensors for more detail.
            use_cache (`bool`, *optional*):
                If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
                (see `past_key_values`).
            past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
        """

        residual = hidden_states

        hidden_states = self.input_layernorm(hidden_states)

        # Self Attention
        hidden_states, self_attn_weights, present_key_value = self.self_attn(
            hidden_states=hidden_states,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_value=past_key_value,
            output_attentions=output_attentions,
            use_cache=use_cache,
            is_padded_inputs=is_padded_inputs,
        )
        hidden_states = residual + hidden_states

        # Fully Connected
        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        hidden_states = self.mlp(hidden_states)
        hidden_states = residual + hidden_states

        outputs = (hidden_states,)

        if output_attentions:
            outputs += (self_attn_weights,)

        if use_cache:
            outputs += (present_key_value,)

        return outputs


LLAMA_START_DOCSTRING = r"""
    This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
    library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
    etc.)
    This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
    Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
    and behavior.
    Parameters:
        config ([`LlamaConfig`]):
            Model configuration class with all the parameters of the model. Initializing with a config file does not
            load the weights associated with the model, only the configuration. Check out the
            [`~PreTrainedModel.from_pretrained`] method to load the model weights.
"""


@add_start_docstrings(
    "The bare LLaMA Model outputting raw hidden-states without any specific head on top.",
    LLAMA_START_DOCSTRING,
)
class LlamaPreTrainedModel(PreTrainedModel):
    config_class = LlamaConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = True
    _no_split_modules = ["LlamaDecoderLayer"]
    _skip_keys_device_placement = "past_key_values"

    def _init_weights(self, module):
        std = self.config.initializer_range
        if isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=std)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=std)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()

    def _set_gradient_checkpointing(self, module, value=False):
        if isinstance(module, LlamaModel):
            module.gradient_checkpointing = value


LLAMA_INPUTS_DOCSTRING = r"""
    Args:
        input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
            Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
            it.
            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
            [`PreTrainedTokenizer.__call__`] for details.
            [What are input IDs?](../glossary#input-ids)
        attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
            Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
            - 1 for tokens that are **not masked**,
            - 0 for tokens that are **masked**.
            [What are attention masks?](../glossary#attention-mask)
            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
            [`PreTrainedTokenizer.__call__`] for details.
            If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see
            `past_key_values`).
            If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]
            and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more
            information on the default strategy.
            - 1 indicates the head is **not masked**,
            - 0 indicates the head is **masked**.
        position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
            Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
            config.n_positions - 1]`.
            [What are position IDs?](../glossary#position-ids)
        past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
            Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape
            `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of shape
            `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`.
            Contains pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
            blocks) that can be used (see `past_key_values` input) to speed up sequential decoding.
            If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
            don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
            `decoder_input_ids` of shape `(batch_size, sequence_length)`.
        inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
            Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
            is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
            model's internal embedding lookup matrix.
        use_cache (`bool`, *optional*):
            If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
            `past_key_values`).
        output_attentions (`bool`, *optional*):
            Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
            tensors for more detail.
        output_hidden_states (`bool`, *optional*):
            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
            more detail.
        return_dict (`bool`, *optional*):
            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
"""


@add_start_docstrings(
    "The bare LLaMA Model outputting raw hidden-states without any specific head on top.",
    LLAMA_START_DOCSTRING,
)
class LlamaModel(LlamaPreTrainedModel):
    """
    Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`LlamaDecoderLayer`]
    Args:
        config: LlamaConfig
    """

    def __init__(self, config: LlamaConfig):
        super().__init__(config)
        self.padding_idx = config.pad_token_id
        self.vocab_size = config.vocab_size

        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
        self.layers = nn.ModuleList([LlamaDecoderLayer(config) for _ in range(config.num_hidden_layers)])
        self.norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)

        self.gradient_checkpointing = False
        # Initialize weights and apply final processing
        self.post_init()

    def get_input_embeddings(self):
        return self.embed_tokens

    def set_input_embeddings(self, value):
        self.embed_tokens = value

    @add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING)
    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[List[torch.FloatTensor]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        is_padded_inputs: Optional[bool] = False,
    ) -> Union[Tuple, BaseModelOutputWithPast]:
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        use_cache = use_cache if use_cache is not None else self.config.use_cache

        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        # retrieve input_ids and inputs_embeds
        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
        elif input_ids is not None:
            batch_size, seq_length = input_ids.shape
        elif inputs_embeds is not None:
            batch_size, seq_length, _ = inputs_embeds.shape
        else:
            raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")

        seq_length_with_past = seq_length
        past_key_values_length = 0

        if past_key_values is not None:
            past_key_values_length = past_key_values[0][0].shape[2]
            seq_length_with_past = seq_length_with_past + past_key_values_length

        position_ids = None

        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids)

        hidden_states = inputs_embeds

        if self.gradient_checkpointing and self.training:
            if use_cache:
                logger.warning_once(
                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
                )
                use_cache = False

        # decoder layers
        all_hidden_states = () if output_hidden_states else None
        all_self_attns = () if output_attentions else None
        next_decoder_cache = () if use_cache else None

        for idx, decoder_layer in enumerate(self.layers):
            if output_hidden_states:
                all_hidden_states += (hidden_states,)

            past_key_value = past_key_values[idx] if past_key_values is not None else None

            if self.gradient_checkpointing and self.training:

                def create_custom_forward(module):
                    def custom_forward(*inputs):
                        # None for past_key_value
                        return module(*inputs, output_attentions, None)

                    return custom_forward

                layer_outputs = torch.utils.checkpoint.checkpoint(
                    create_custom_forward(decoder_layer),
                    hidden_states,
                    attention_mask,
                    position_ids,
                    None,
                    is_padded_inputs,
                )
            else:
                layer_outputs = decoder_layer(
                    hidden_states,
                    attention_mask=attention_mask,
                    position_ids=position_ids,
                    past_key_value=past_key_value,
                    output_attentions=output_attentions,
                    use_cache=use_cache,
                    is_padded_inputs=is_padded_inputs,
                )

            hidden_states = layer_outputs[0]

            if use_cache:
                next_decoder_cache += (layer_outputs[2 if output_attentions else 1],)

            if output_attentions:
                all_self_attns += (layer_outputs[1],)

        hidden_states = self.norm(hidden_states)

        # add hidden states from the last decoder layer
        if output_hidden_states:
            all_hidden_states += (hidden_states,)

        next_cache = next_decoder_cache if use_cache else None
        if not return_dict:
            return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
        return BaseModelOutputWithPast(
            last_hidden_state=hidden_states,
            past_key_values=next_cache,
            hidden_states=all_hidden_states,
            attentions=all_self_attns,
        )


class LlamaForCausalLM(LlamaPreTrainedModel):
    _tied_weights_keys = ["lm_head.weight"]

    def __init__(self, config):
        super().__init__(config)
        self.model = LlamaModel(config)
        self.vocab_size = config.vocab_size
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

        # Initialize weights and apply final processing
        self.post_init()

    def get_input_embeddings(self):
        return self.model.embed_tokens

    def set_input_embeddings(self, value):
        self.model.embed_tokens = value

    def get_output_embeddings(self):
        return self.lm_head

    def set_output_embeddings(self, new_embeddings):
        self.lm_head = new_embeddings

    def set_decoder(self, decoder):
        self.model = decoder

    def get_decoder(self):
        return self.model

    @add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[List[torch.FloatTensor]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        is_padded_inputs: Optional[bool] = None,
    ) -> Union[Tuple, CausalLMOutputWithPast]:
        r"""
        Args:
            labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
                Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
                config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
                (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
        Returns:
        Example:
        ```python
         from transformers import AutoTokenizer, LlamaForCausalLM
         model = LlamaForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS)
         tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER)
         prompt = "Hey, are you conscious? Can you talk to me?"
         inputs = tokenizer(prompt, return_tensors="pt")
         # Generate
         generate_ids = model.generate(inputs.input_ids, max_length=30)
         tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
        "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
        ```"""

        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        is_padded_inputs = (attention_mask is not None) and (not attention_mask.all().item())

        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
        outputs = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            is_padded_inputs=is_padded_inputs,
        )

        hidden_states = outputs[0]
        if self.config.pretraining_tp > 1:
            lm_head_slices = self.lm_head.weight.split(self.vocab_size // self.config.pretraining_tp, dim=0)
            logits = [F.linear(hidden_states, lm_head_slices[i]) for i in range(self.config.pretraining_tp)]
            logits = torch.cat(logits, dim=-1)
        else:
            logits = self.lm_head(hidden_states)
        logits = logits.float()

        loss = None
        if labels is not None:
            # Shift so that tokens < n predict n
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            # Flatten the tokens
            loss_fct = CrossEntropyLoss()
            shift_logits = shift_logits.view(-1, self.config.vocab_size)
            shift_labels = shift_labels.view(-1)
            # Enable model parallelism
            shift_labels = shift_labels.to(shift_logits.device)
            loss = loss_fct(shift_logits, shift_labels)

        if not return_dict:
            output = (logits,) + outputs[1:]
            return (loss,) + output if loss is not None else output

        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )

    def prepare_inputs_for_generation(
        self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
    ):
        if past_key_values:
            input_ids = input_ids[:, -1:]

        position_ids = kwargs.get("position_ids", None)

        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
        if inputs_embeds is not None and past_key_values is None:
            model_inputs = {"inputs_embeds": inputs_embeds}
        else:
            model_inputs = {"input_ids": input_ids}

        model_inputs.update(
            {
                "position_ids": position_ids,
                "past_key_values": past_key_values,
                "use_cache": kwargs.get("use_cache"),
                "attention_mask": attention_mask,
                "is_padded_inputs": (attention_mask is not None) and (not attention_mask.all().item()),
            }
        )
        return model_inputs

    @staticmethod
    def _reorder_cache(past_key_values, beam_idx):
        reordered_past = ()
        for layer_past in past_key_values:
            reordered_past += (
                tuple(
                    (
                        past_state.index_select(0, beam_idx.to(past_state.device))
                        if type(past_state) == torch.Tensor
                        else past_state  # There is an int in the last layer, it is not supposed to be there,
                        # but this hack works to deal with it.
                    )
                    for past_state in layer_past
                ),
            )
        return reordered_past


@add_start_docstrings(
    """
    The LLaMa Model transformer with a sequence classification head on top (linear layer).
    [`LlamaForSequenceClassification`] uses the last token in order to do the classification, as other causal models
    (e.g. GPT-2) do.
    Since it does classification on the last token, it requires to know the position of the last token. If a
    `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
    no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the
    padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
    each row of the batch).
    """,
    LLAMA_START_DOCSTRING,
)
class LlamaForSequenceClassification(LlamaPreTrainedModel):
    def __init__(self, config):
        super().__init__(config)
        self.num_labels = config.num_labels
        self.model = LlamaModel(config)
        self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)

        # Initialize weights and apply final processing
        self.post_init()

    def get_input_embeddings(self):
        return self.model.embed_tokens

    def set_input_embeddings(self, value):
        self.model.embed_tokens = value

    @add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING)
    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[List[torch.FloatTensor]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, SequenceClassifierOutputWithPast]:
        r"""
        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
            Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
        """
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        transformer_outputs = self.model(
            input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        hidden_states = transformer_outputs[0]
        logits = self.score(hidden_states)

        if input_ids is not None:
            batch_size = input_ids.shape[0]
        else:
            batch_size = inputs_embeds.shape[0]

        if self.config.pad_token_id is None and batch_size != 1:
            raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
        if self.config.pad_token_id is None:
            sequence_lengths = -1
        else:
            if input_ids is not None:
                sequence_lengths = (torch.ne(input_ids, self.config.pad_token_id).sum(-1) - 1).to(logits.device)
            else:
                sequence_lengths = -1

        pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]

        loss = None
        if labels is not None:
            labels = labels.to(logits.device)
            if self.config.problem_type is None:
                if self.num_labels == 1:
                    self.config.problem_type = "regression"
                elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
                    self.config.problem_type = "single_label_classification"
                else:
                    self.config.problem_type = "multi_label_classification"

            if self.config.problem_type == "regression":
                loss_fct = MSELoss()
                if self.num_labels == 1:
                    loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
                else:
                    loss = loss_fct(pooled_logits, labels)
            elif self.config.problem_type == "single_label_classification":
                loss_fct = CrossEntropyLoss()
                loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))
            elif self.config.problem_type == "multi_label_classification":
                loss_fct = BCEWithLogitsLoss()
                loss = loss_fct(pooled_logits, labels)
        if not return_dict:
            output = (pooled_logits,) + transformer_outputs[1:]
            return ((loss,) + output) if loss is not None else output

        return SequenceClassifierOutputWithPast(
            loss=loss,
            logits=pooled_logits,
            past_key_values=transformer_outputs.past_key_values,
            hidden_states=transformer_outputs.hidden_states,
            attentions=transformer_outputs.attentions,
        )

src/model/patch_models/modeling_flash_llama.py
hitz-zentroa/GoLLIE
from __future__ import annotations  # To make it not choke over FlashSelfAttention

import warnings
from functools import partial
from typing import Callable, Optional

import torch.nn as nn

import transformers
from src.model.patch_models.patching_llama import llama_forward_with_flash_attn
from src.model.patch_models.patching_neox import neox_forward_with_flash_attn
from transformers import GPTNeoXForCausalLM, GPTNeoXModel, LlamaForCausalLM, LlamaModel


SUPPORTED_MODELS = [
    GPTNeoXModel,
    GPTNeoXForCausalLM,
    LlamaForCausalLM,
    LlamaModel,
]


def _patched_mlp_forward(post_module: nn.Module, module: nn.Module, *args, **kwargs):
    post_module.train(module.training)
    out = module.old_forward(*args, **kwargs)
    out = post_module(out)
    return out


def _patched_attn_forward(post_module: nn.Module, module: nn.Module, *args, **kwargs):
    post_module.train(module.training)
    out = module.old_forward(*args, **kwargs)
    hiddens = post_module(out[0])
    return (hiddens,) + out[1:]


def add_dropout(module: nn.Module, patched_fwd: Callable, p_dropout: float = 0.1):
    dropout = nn.Dropout(p=p_dropout)
    module.old_forward = module.forward
    module.forward = partial(patched_fwd, dropout, module)


def add_flash_attn(module: nn.Module, causal: bool = True):
    """
    Replaces the standard attention implementation with Flash Attention [1].
    Limitations:
      - Only works for fp16 or bf16 inputs
      - Requires inputs to be on CUDA
      - `output_attentions=True` does not work after patching, attention weights will be None
      - Non-contiguous attention masks are not supported (e.g. [1, 1, 0, 1, 1, 0, 0] will just become [1, 1, 1, 1, 1, 0, 0]).

    [1] https://github.com/HazyResearch/flash-attention
    """

    flash_attn = FlashSelfAttention(causal=causal)
    if isinstance(module, transformers.models.llama.modeling_llama.LlamaAttention):
        module.old_forward = module.forward
        module.forward = partial(llama_forward_with_flash_attn, module, flash_attn)
    elif isinstance(module, transformers.models.gpt_neox.modeling_gpt_neox.GPTNeoXAttention):
        if not hasattr(module, "_attn"):
            warnings.warn("Provided module doesn't have a _attn() function to be patched.")
        module._attn = partial(neox_forward_with_flash_attn, module, flash_attn)
    else:
        raise NotImplementedError(f"Flash attention is not implemented for {module.__class__.__name__}.")


def patch_model(
    model: nn.Module,
    resid_pdrop: Optional[float] = 0.1,
    flash_attention: bool = True,
    patch_unsupported: bool = False,
    residual_dropout_lima: bool = False,
):
    """
    Helper function for patching HF language models.
    Currently supports: GPTNeoX-based models

    Limitations:
      - Flash attention requires CUDA and fp16/bf16 training. It also requires contiguous attention masks.
      - Residual dropout does not support multi-GPU training without DeepDpeed.
    """
    global FlashSelfAttention
    if flash_attention:
        try:
            from flash_attn.modules.mha import FlashSelfAttention  # pyright: reportMissingImports=false
        except ModuleNotFoundError:
            warnings.warn(
                """\nmodule flash_attn not found - either install:
  pip3 install flash_attn
or run with:
  --use_flash_attention=false """
            )
            exit(1)
    if (resid_pdrop is None or resid_pdrop == 0.0) and not flash_attention:
        print("Continuing without patching")
        return

    if resid_pdrop is not None and (resid_pdrop < 0 or resid_pdrop > 1.0):
        raise ValueError("Invalid argument: `resid_pdrop` must be between 0.0 and 1.0")

    if not flash_attention and (resid_pdrop is None or resid_pdrop == 0.0):
        return

    if (
        not any(isinstance(model, model_class) for model_class in SUPPORTED_MODELS)
        and model.__class__.__name__ != "RWForCausalLM"
    ):
        if not flash_attention and (resid_pdrop is None or resid_pdrop == 0.0):
            return  # nothing to patch

        if not patch_unsupported:
            warnings.warn(
                "Model patching does not support this model class. No patches will be applied. "
                "If you want to force patch this model, please set `patch_unsupported=True`."
            )
            return

        warnings.warn(
            "Patching residual dropout has only been tested with this model class. "
            f"Please make sure that it also works for `{model.__class__.__name__}`.\n"
            "Or disable flash_attention and residual_dropout with:\n"
            "--use_flash_attention=false  --no-residual_dropout"
        )

    if isinstance(model, LlamaForCausalLM):
        model = model.model

    if model.__class__.__name__ == "RWForCausalLM":
        model = model.base_model

    attention_key_lookup = {
        GPTNeoXModel: "attention",
        LlamaModel: "self_attn",
    }
    mlp_key_lookup = {
        GPTNeoXModel: "mlp",
        LlamaModel: "mlp",
    }
    if model.__class__.__name__ == "RWModel":
        layers = model.h
        attention_key = "self_attention"
        mlp_key = "mlp"
    else:
        layers = model.layers
        attention_key = attention_key_lookup.get(model.__class__, "attention")
        mlp_key = mlp_key_lookup.get(model.__class__, "mlp")
    num_layers = len(layers)
    resid_pdrop_last_layer = resid_pdrop
    for i, layer in enumerate(layers):
        if flash_attention:
            add_flash_attn(getattr(layer, attention_key), causal=True)
        if residual_dropout_lima:
            resid_pdrop = i / (num_layers - 1) * resid_pdrop_last_layer
        if resid_pdrop is not None and resid_pdrop > 0:
            add_dropout(getattr(layer, attention_key), _patched_attn_forward, resid_pdrop)
            add_dropout(getattr(layer, mlp_key), _patched_mlp_forward, resid_pdrop)

src/model/patch_models/patching.py
hitz-zentroa/GoLLIE
import torch
import torch.nn as nn

import transformers
from src.model.patch_models.patching_utils import compute_flash_attention


def neox_forward_with_flash_attn(
    self: transformers.models.gpt_neox.modeling_gpt_neox.GPTNeoXAttention,
    flash_attn: nn.Module,  # flash_attn.modules.mha.FlashSelfAttention
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    attention_mask=None,
    head_mask=None,
):
    # query, key, value: [bs, num_attention_heads, seq_len, attn_head_size]
    if query.shape == key.shape:
        flash_attn.train(self.training)
        out_dtype = value.dtype
        q, k, v = query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2)
        if attention_mask is not None:
            attention_mask = attention_mask[:, 0, 0, :]
        out = compute_flash_attention(flash_attn, q, k, v, attention_mask)
        out = out.transpose(1, 2).to(out_dtype)
        return out, None
    else:
        return self.old_forward(query, key, value, attention_mask, head_mask)

src/model/patch_models/patching_neox.py
hitz-zentroa/GoLLIE
import torch
import torch.nn.functional as F


def compute_flash_attention(flash_attn, q, k, v, attention_mask=None, head_mask=None):
    # q, k, v: [bs, seq_len, num_attention_heads, attn_head_size]
    # attention_mask (float): [bs, seq_len]
    batch_size, max_len = q.size(0), q.size(1)

    qkv = torch.stack([q, k, v], dim=2).to(torch.float16)  # need to truncate in case input is fp32
    cu_seqlens, max_seqlen = None, None

    if attention_mask is None:
        return flash_attn(qkv, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen)
    else:
        # Limitation: non-contiguous attention mask will not be handled correctly
        # model will be able to pay attention between the first and last non-masked token, i.e. left- and right-side padding is supported.
        csums = (attention_mask >= 0).cumsum(dim=1)
        ends = csums.argmax(dim=1) + 1
        starts = ends - csums.max(dim=1).values
        seqlens = ends - starts

        qkv = torch.cat([qkv[i, starts[i] : ends[i]] for i in range(batch_size)], dim=0)
        zero = torch.zeros_like(seqlens[:1])  # torch.tensor([0]) with correct dtype and device
        cu_seqlens = torch.cat([zero, seqlens.cumsum(dim=0)], dim=0).to(torch.int32)
        max_seqlen = seqlens.max().item()

        out = flash_attn(qkv, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen)
        # out: [num_unmasked_tokens, num_attention_heads, attn_head_size]

        seqs = [out[start:end] for start, end in zip(cu_seqlens[:-1], cu_seqlens[1:])]
        # stack and pad sequences together
        padded_seqs = [
            F.pad(seqs[i], (0, 0) * (seqs[i].dim() - 1) + (starts[i], max_len - ends[i]), value=0.0)
            for i in range(batch_size)
        ]
        out = torch.stack(padded_seqs)
        return out


if __name__ == "__main__":
    from flash_attn.modules.mha import FlashSelfAttention

    flash_attn = FlashSelfAttention(causal=True)

    dtype = torch.float16
    device = torch.device("cuda:0")

    batch_size, seq_len, num_heads, head_size = 4, 18, 8, 32
    q = torch.randn(batch_size, seq_len, num_heads, head_size, dtype=dtype, device=device)
    k = torch.randn(batch_size, seq_len, num_heads, head_size, dtype=dtype, device=device)
    v = torch.randn(batch_size, seq_len, num_heads, head_size, dtype=dtype, device=device)

    attn_mask = torch.randn(batch_size, seq_len, dtype=dtype, device=device).abs().cumsum(dim=1)
    attn_mask = ((attn_mask > 3) & (attn_mask < 10)).int().log()

    out = compute_flash_attention(flash_attn, q, k, v, attention_mask=attn_mask)

src/model/patch_models/patching_utils.py
hitz-zentroa/GoLLIE
from dataclasses import dataclass, field
from typing import List


@dataclass
class DataInferenceArguments:
    """
    Arguments pertaining to what data we are going to input our model for paraphrasing.
    """

    datasets: List[str] = field(
        default=None,
        metadata={"help": "The tasks to train on. Can be a list of tasks or a single task."},
    )

    config_template: str = field(
        default="vicuna_v1.1",
        metadata={
            "help": (
                "The config template to use. Available templates: 'one_shot', 'vicuna_v1.1', 'koala_v1', 'dolly_v2',"
                " 'oasst_pythia', 'stablelm', 'baize', 'rwkv', 'openbuddy', 'phoenix', 'chatgpt', 'claude', 'mpt'"
            )
        },
    )

    language: str = field(
        default="en",
        metadata={"help": "The language to do phrase paraphrasing."},
    )

    generation_args_json: str = field(
        default=None,
        metadata={"help": "The generation args json file."},
    )

src/paraphrase/config.py
hitz-zentroa/GoLLIE
import logging
from typing import List

from fastchat.conversation import get_conv_template
from torch.utils.data import Dataset
from tqdm import tqdm

from src.paraphrase.utils import clean_guidelines
from src.tasks import task_id_to_guidelines
from transformers import BatchEncoding, PreTrainedTokenizerBase


def prepare_data(
    example: str,
    tokenizer: PreTrainedTokenizerBase,
    is_encoder_decoder: bool = False,
    max_length: int = 2048,
    conv_template: str = None,
) -> BatchEncoding:
    """
    Prepare data for training or inference.

    Args:
        example (`str`):
            The example to prepare.
        tokenizer (`PreTrainedTokenizerBase`):
            The tokenizer to use.
        is_encoder_decoder (`bool`, optional):
            Whether the model is an encoder-decoder model. Defaults to `False`.
        max_length (`int`, optional):
            The maximum length of the input. Defaults to `2048`.
        conv_template (`str`, optional):
            The conversation template to use. Defaults to `None`. If `None` we will return the prompt.

    Returns:
        `BatchEncoding`: `BatchEncoding` with the prepared data.
    """

    prompt = (
        "Please, generate a paraphrase of the following text. Ensure that no information is lost in the paraphrase. If"
        " possible rearrange the word order. Write a short explanation of how are you paraphrasing the text.\n"
        'Use the following format: "Explanation: <your explanation>". \n "Paraphrase: <your paraphrase>". \n'
        f'Text: "{example}"'
    )

    if conv_template is not None:
        conv = get_conv_template(conv_template)
        conv.append_message(conv.roles[0], prompt)
        conv.append_message(conv.roles[1], None)
        prompt = conv.get_prompt()

    model_inputs = tokenizer(
        text=prompt,
        max_length=max_length,
        truncation=True,
        padding=False,
        return_tensors=None,
        add_special_tokens=True,
    )

    if not is_encoder_decoder:
        # Remove the last token if it is an eos token
        if model_inputs["input_ids"][-1] == tokenizer.eos_token_id:
            model_inputs["input_ids"] = model_inputs["input_ids"][:-1]
            model_inputs["attention_mask"] = model_inputs["attention_mask"][:-1]

    if "token_type_ids" in model_inputs:
        # LLaMa tokenizer adds token type ids, but we don't need them
        model_inputs.pop("token_type_ids")

    return model_inputs


class ParaphraseDataset(Dataset):
    """
    Dataset for Collie.

    Args:
        tokenizer (`PreTrainedTokenizerBase`):
            The tokenizer to use.
        dataset_name (`str`):
            The name of the dataset.
        language (`str`):
            The language to get the data for.
        is_encoder_decoder (`bool`, optional):
            Whether the model is an encoder-decoder model. Defaults to `False`.
        max_length (`int`, optional):
            The maximum length of the input. Defaults to `2048`.
        conv_template (`str`, optional):
            The conversation template to use. Defaults to `None`. If `None` we will return the prompt.
    """

    def __init__(
        self,
        tokenizer: PreTrainedTokenizerBase,
        dataset_name: str,
        language: str,
        is_encoder_decoder: bool = False,
        max_length: int = 2048,
        conv_template: str = None,
    ):
        self.dataset_name = dataset_name
        self.language = language
        self.is_encoder_decoder = is_encoder_decoder
        self.max_length = max_length
        self.conv_template = conv_template
        guidelines = task_id_to_guidelines(dataset_name)
        guidelines = clean_guidelines(guidelines)
        self.dataset: List[BatchEncoding] = []
        for guideline in tqdm(guidelines.values(), desc="Data Tokenization"):
            for text in guideline[language]:
                # rich.print(f"Guideline: {text}")
                self.dataset.append(prepare_data(text, tokenizer, is_encoder_decoder, max_length, conv_template))

        logging.info(f"Dataset {dataset_name} has {len(self.dataset)} guidelines.")

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, idx):
        return self.dataset[idx]

src/paraphrase/dataset.py
hitz-zentroa/GoLLIE
import json
import logging
import os
import sys

from fastchat.conversation import get_conv_template

from src.config import ModelArguments
from src.model.load_model import load_model
from src.paraphrase.config import DataInferenceArguments
from src.paraphrase.dataset import ParaphraseDataset
from src.paraphrase.utils import clean_guidelines, format_guidelines_as_py, update_guidelines
from src.tasks import task_id_to_guidelines
from src.trainer import get_correct_torch_dtype
from transformers import DataCollatorForSeq2Seq, HfArgumentParser, Seq2SeqTrainer, Seq2SeqTrainingArguments


def run_paraphrasing(
    model_args: ModelArguments,
    data_args: DataInferenceArguments,
    training_args: Seq2SeqTrainingArguments,
):
    if not training_args.predict_with_generate:
        raise ValueError("Set `predict_with_generate` to `True` in the config file to run this script.")

    logging.info(
        f"Loading model from {model_args.model_name_or_path}.\n"
        f"   - quantization: {model_args.quantization}\n"
        f"   - lora_weights_name_or_path: {model_args.lora_weights_name_or_path}\n"
    )

    model, tokenizer = load_model(
        inference=True,
        model_weights_name_or_path=model_args.model_name_or_path,
        quantization=model_args.quantization,
        use_lora=model_args.lora_weights_name_or_path is not None,
        lora_weights_name_or_path=model_args.lora_weights_name_or_path,
        force_auto_device_map=model_args.force_auto_device_map,
        torch_dtype=get_correct_torch_dtype(
            quantization=model_args.quantization_inference, model_args=model_args, training_args=training_args
        ),
        use_better_transformer=model_args.use_better_transformer,
        trust_remote_code=model_args.trust_remote_code,
        use_flash_attention=model_args.use_flash_attention,
        max_memory_MB=model_args.max_memory_MB,
    )

    trainer = Seq2SeqTrainer(
        tokenizer=tokenizer,
        model=model,
        args=training_args,
        data_collator=DataCollatorForSeq2Seq(
            tokenizer,
            pad_to_multiple_of=8,
            return_tensors="pt",
            padding=True,
            label_pad_token_id=tokenizer.pad_token_id,
        ),
    )

    with open(data_args.generation_args_json, "r", encoding="utf8") as f:
        gen_kwargs = json.load(f)
        logging.info(f"Generation kwargs: {json.dumps(gen_kwargs, indent=2,ensure_ascii=False)}")

    if gen_kwargs["num_beams"] > 1 and gen_kwargs["do_sample"]:
        logging.warning(
            "You are using `num_beams` > 1 and `do_sample`. Beam-search multinomial sampling from huggingface"
            " transformers is not stable, you may get this error: 'RuntimeError: probability tensor contains either"
            " inf, nan or element < 0'. To avoid this error, set `do_sample` to `False` or use `num_beams` = 1."
        )

    num_return_sequences = gen_kwargs.pop("num_return_sequences", 1)

    for dataset_name in data_args.datasets:
        logging.info(f"Running inference on {dataset_name}...")

        guidelines = task_id_to_guidelines(dataset_name)
        guidelines = clean_guidelines(guidelines)

        test_dataset = ParaphraseDataset(
            tokenizer=tokenizer,
            dataset_name=dataset_name,
            language=data_args.language,
            is_encoder_decoder=model.config.is_encoder_decoder,
            max_length=1024,
            conv_template=data_args.config_template,
        )

        output_path = os.path.join(training_args.output_dir, dataset_name, "guidelines.py")
        output_path_orig_outputs = os.path.join(training_args.output_dir, dataset_name, "original_outputs.txt")

        os.makedirs(os.path.dirname(output_path), exist_ok=True)

        # Trainer.predict does not support multiple return sequences, so we have to do it manually

        with open(output_path_orig_outputs, "w", encoding="utf8") as f:
            for i in range(num_return_sequences):
                predictions = trainer.predict(test_dataset, **gen_kwargs).predictions
                predictions[predictions == -100] = tokenizer.pad_token_id

                try:
                    predictions = tokenizer.batch_decode(predictions, skip_special_tokens=True)
                except OverflowError:
                    raise OverflowError(f"Unable to decode predictions: {predictions}")

                for prediction in predictions:
                    print(prediction, file=f)
                    print("\n====================\n", file=f)

                if data_args.config_template is not None:
                    conv = get_conv_template(data_args.config_template)
                    predictions = [prediction.split(conv.roles[1])[-1].strip() for prediction in predictions]
                # rich.print(predictions)

                predictions = [prediction.strip().strip("\n") for prediction in predictions]
                predictions = [prediction.split("\n")[-1].strip() for prediction in predictions]
                predictions = [":".join(prediction.split(":")[1:]) for prediction in predictions]

                guidelines = update_guidelines(
                    paraphrases=predictions,
                    guidelines=guidelines,
                    language=data_args.language,
                )

                f.flush()

        with open(output_path, "w", encoding="utf8") as f:
            print(format_guidelines_as_py(guidelines), file=f)
            logging.info(f"Saved guidelines to {output_path}")


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    parser = HfArgumentParser((ModelArguments, DataInferenceArguments, Seq2SeqTrainingArguments))
    logging.info(f"Sys args {sys.argv}")

    if len(sys.argv) > 0 and sys.argv[-1].endswith(".json"):
        # If we pass only one argument to the script, and it's the path to a json file,
        # let's parse it to get our arguments.
        logging.info(f"Loading json config {sys.argv[-1]}")
        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[-1]))

    elif len(sys.argv) > 0 and sys.argv[-1].endswith(".yaml"):
        # If we pass only one argument to the script, and it's the path to a yaml file,
        # let's parse it to get our arguments.
        logging.info(f"Loading yaml config {sys.argv[-1]}")
        model_args, data_args, training_args = parser.parse_yaml_file(yaml_file=os.path.abspath(sys.argv[-1]))
    else:
        logging.info("No config file passed, using command line arguments.")
        model_args, data_args, training_args = parser.parse_args_into_dataclasses()

    run_paraphrasing(
        model_args=model_args,
        data_args=data_args,
        training_args=training_args,
    )

src/paraphrase/run_paraphrasing.py
hitz-zentroa/GoLLIE
import json
import math
from typing import Dict, Iterator, List, Sized

import black


def batch(iterable: Sized, n=1) -> Iterator:
    """
    Yield successive n-sized chunks from iterable.

    Args:
        iterable (`Sized`):
            The iterable to split.
        n (`int`, optional):
            The size of the chunks. Defaults to `1`.

    Yields:
        `Iterator`:
            An iterator with the chunks.
    """
    l: int = len(iterable)
    p: int = math.ceil(l / n)
    for ndx in range(0, l, p):
        yield iterable[ndx : min(ndx + p, l)]


def clean_guidelines(guidelines: Dict[str, Dict[str, List[str]]]):
    """
    Clean the guidelines.

    Args:
        guidelines (Dict[str, Dict[str, List[str]]]): The guidelines.

    Returns:
        The cleaned guidelines.
    """

    for guideline in guidelines.values():
        for language in guideline.keys():
            for i in range(len(guideline[language])):
                text = guideline[language][i]
                text = text.replace("\n", " ")
                text = text.replace("\t", " ")
                text = " ".join(text.split())
                guideline[language][i] = text

    return guidelines


def update_guidelines(paraphrases: List[str], guidelines: Dict[str, Dict[str, List[str]]], language: str):
    """
    Update the guidelines for a given task.

    Args:
        paraphrases (List[str]): The paraphrases.
        guidelines (Dict[str, Dict[str, List[str]]]): The guidelines.
        language (str): The language for which the paraphrases were generated.

    Returns:
        The updated guidelines.
    """

    i = 0
    for guideline in guidelines.values():
        guideline[language].append(paraphrases[i])
        i += 1
    return guidelines


def get_num_return_sentences(config_path: str):
    """
    Get the number of sentences to return.

    Args:
        config_path (str): The path to the config json file.

    Returns:
        The number of sentences to return.
    """

    with open(config_path, "r") as f:
        config = json.load(f)
    return config["num_return_sequences"]


def format_guidelines_as_py(guidelines: Dict[str, Dict[str, List[str]]]):
    # Add a \n every 98 characters. Do not split words.
    for guideline in guidelines.values():
        for language in guideline.keys():
            for i in range(len(guideline[language])):
                text = guideline[language][i]

                # Remove \" and \' at the beginning and end of the string
                text = text.strip()
                text = text.strip('"').strip()
                text = text.strip("'").strip()

                words = text.split()
                c: int = 0
                for j in range(len(words)):
                    c += len(words[j])
                    if c > 98:
                        words[j] = "\n" + words[j]
                        c = 0
                text = " ".join(words)
                guideline[language][i] = text

    guidelines_py = json.dumps(guidelines, indent=4)
    guidelines_py = f"GUIDELINES = {guidelines_py}"
    guidelines_py = black.format_str(guidelines_py, mode=black.Mode())
    return guidelines_py

src/paraphrase/utils.py
hitz-zentroa/GoLLIE
import gc
import glob
import json
import logging
import os
import sys

import torch
import torch.utils.data
from datasets import DatasetDict

from src.config import DataTrainingArguments, ModelArguments
from src.dataset.dataset import CollieDataset, DataCollatorForCoLLIE
from src.evaluate import evaluate
from src.model.load_model import load_model, merge_lora_model
from src.trainer import CollieTrainer, ConcatDataset, get_correct_torch_dtype
from transformers import (
    HfArgumentParser,
    Seq2SeqTrainingArguments,
)


def clean_cache():
    """Clean cache to avoid memory leak.
    This fixes this issue: https://github.com/huggingface/transformers/issues/22801"""

    logging.info(f"Cleaning GPU memory. Current memory usage: {torch.cuda.memory_allocated()}")
    torch.cuda.empty_cache()
    gc.collect()
    torch.cuda.empty_cache()
    logging.info(f"GPU memory usage after cleaning: {torch.cuda.memory_allocated()}")


def train_collie(
    model_args: ModelArguments,
    data_args: DataTrainingArguments,
    training_args: Seq2SeqTrainingArguments,
):
    logging.info(f"Loading {model_args.model_name_or_path} model...")

    model, tokenizer = load_model(
        inference=False,
        model_weights_name_or_path=model_args.model_name_or_path,
        quantization=model_args.quantization,
        use_lora=model_args.use_lora,
        lora_r=model_args.lora_r,
        lora_target_modules=model_args.lora_target_modules,
        torch_dtype=get_correct_torch_dtype(
            quantization=model_args.quantization, model_args=model_args, training_args=training_args
        ),
        force_auto_device_map=model_args.force_auto_device_map,
        use_gradient_checkpointing=training_args.gradient_checkpointing,
        use_better_transformer=model_args.use_better_transformer,
        trust_remote_code=model_args.trust_remote_code,
        use_flash_attention=model_args.use_flash_attention,
        fsdp_training=len(training_args.fsdp) > 1 or training_args.fsdp_config is not None,
        max_memory_MB=model_args.max_memory_MB,
    )

    logging.info("Loading datasets...")
    training_datasets_path = [
        f"{os.path.join(data_args.dataset_dir, task)}.train.jsonl" for task in data_args.train_tasks
    ]
    development_datasets_path = [
        f"{os.path.join(data_args.dataset_dir, task)}.dev.jsonl" for task in data_args.validation_tasks
    ]

    logging.info(
        f"We will train CoLLIE on {len(training_datasets_path)} datasets: {', '.join(training_datasets_path)}"
    )

    logging.info(
        f"We will validate CoLLIE on {len(development_datasets_path)} datasets: {', '.join(development_datasets_path)}"
    )

    logging.info(
        "Training dataset will be loaded with. 'ignore_pad_token_for_loss':"
        f" {data_args.ignore_pad_token_for_loss} and 'prompt_loss_weight':"
        f" {data_args.prompt_loss_weight}"
    )

    training_datasets = []
    for train_task in data_args.train_tasks:
        train_path = os.path.join(data_args.dataset_dir, f"{train_task}.train.jsonl")
        train_dataset = CollieDataset(
            tokenizer=tokenizer,
            dataset_path=train_path,
            max_length=data_args.max_seq_length,
            is_encoder_decoder=model.config.is_encoder_decoder,
            inference=False,
            prompt_loss_weight=data_args.prompt_loss_weight,
            max_examples=data_args.max_examples_per_task_train,
        )
        training_datasets.append(train_dataset)

    train_dataset = ConcatDataset(training_datasets)

    dev_datasets = DatasetDict()
    for dev_task in data_args.validation_tasks:
        dev_path = os.path.join(data_args.dataset_dir, f"{dev_task}.dev.jsonl")
        dev_dataset = CollieDataset(
            tokenizer=tokenizer,
            dataset_path=dev_path,
            max_length=data_args.max_seq_length,
            is_encoder_decoder=model.config.is_encoder_decoder,
            inference=False,
            prompt_loss_weight=0.0,
            max_examples=data_args.max_examples_per_task_val,
        )
        dev_datasets[os.path.splitext(os.path.basename(dev_path))[0]] = dev_dataset

    trainer = CollieTrainer(
        model=model,
        tokenizer=tokenizer,
        train_dataset=train_dataset,
        eval_dataset=dev_datasets,
        args=training_args,
        data_collator=DataCollatorForCoLLIE(
            tokenizer,
            pad_to_multiple_of=8,
            return_tensors="pt",
            padding=True,
            label_pad_token_id=(-100 if data_args.ignore_pad_token_for_loss else tokenizer.pad_token_id),
        ),
    )

    trainer.train(resume_from_checkpoint=training_args.resume_from_checkpoint)

    # Save the model
    trainer.save_model()

    # model.save_pretrained(training_args.output_dir)
    # model.config.save_pretrained(training_args.output_dir)
    # tokenizer.save_pretrained(training_args.output_dir)


def inference_collie(
    model_args: ModelArguments,
    data_args: DataTrainingArguments,
    training_args: Seq2SeqTrainingArguments,
    checkpoint_path: str = None,
):
    if not training_args.predict_with_generate:
        logging.warning(
            "You have set predict_with_generate to False. We will only compute the loss"
            " on the test set. If you want to generate predictions, set"
            " predict_with_generate to True."
        )

        if not training_args.prediction_loss_only:
            logging.warning(
                "You have set predict_with_generate to False, so you only "
                "want to compute the loss on the test set. But you have set "
                "prediction_loss_only to False. This is contradictory, please "
                "review you configuration. We will attempt to continue but "
                "you might get unexpected results."
            )

    if training_args.do_train:
        if not checkpoint_path:
            logging.warning(
                "You are doing inference after training a model! We will load the "
                f"pretrained model saved in {training_args.output_dir}."
            )
            if model_args.use_lora:
                model_path = model_args.model_name_or_path
                lora_weights_name_or_path = training_args.output_dir
            else:
                model_path = training_args.output_dir
                lora_weights_name_or_path = None
        else:
            logging.warning(
                "You are doing inference after training a model! We will load the "
                f"pretrained model saved in {checkpoint_path}."
            )
            if model_args.use_lora:
                model_path = model_args.model_name_or_path
                lora_weights_name_or_path = checkpoint_path
            else:
                model_path = checkpoint_path
                lora_weights_name_or_path = None
    else:
        if not checkpoint_path:
            model_path = model_args.model_name_or_path
            lora_weights_name_or_path = model_args.lora_weights_name_or_path
        else:
            if model_args.use_lora:
                model_path = model_args.model_name_or_path
                lora_weights_name_or_path = checkpoint_path
            else:
                model_path = checkpoint_path
                lora_weights_name_or_path = None

    if model_args.use_lora and lora_weights_name_or_path is None:
        logging.warning(
            "You are have specified to use LORA, but have not specified a path to the "
            "LORA weights. We will attempt to load the LORA weights from the same "
            f"path as the model weights: {model_path}."
        )

    delete_merged_model: bool = False
    if model_args.merge_lora_before_eval:
        logging.info("You have specified to merge the LORA weights before evaluation. We will attempt to do so.")
        if model_args.quantization is None:
            logging.warning(
                "You have specified to create a merged model (merge_lora_before_eval=True), but you have not specified"
                " a quantization precision. Model loades without quantization are automatically merged when loaded for"
                " inference, so there is no need to save a merged model and reaload it. This flag is only useful when"
                " you want to merge a model and then load it using 4 bits ot 8 bits quantization or if you plan to"
                " release the merged model."
            )
        if os.path.exists(os.path.join(training_args.output_dir, "merged_model")):
            logging.info(
                f"A merged model already exists at {os.path.join(training_args.output_dir,'merged_model')}. We will"
                " use this model."
            )
            delete_merged_model = False

        else:
            merge_lora_model(
                weights_path=model_path,
                lora_weights_name_or_path=lora_weights_name_or_path,
                torch_dtype=model_args.torch_dtype,
                output_path=os.path.join(training_args.output_dir, "merged_model"),
            )
            delete_merged_model = not model_args.keep_merged_model_after_eval

        model_path = os.path.join(training_args.output_dir, "merged_model")
        lora_weights_name_or_path = None
        clean_cache()  # Ensure that nothing remains in the cache, as we will load the mergen model next.

    model, tokenizer = load_model(
        inference=True,
        model_weights_name_or_path=model_path,
        quantization=model_args.quantization_inference,
        use_lora=model_args.lora_weights_name_or_path is not None,
        lora_weights_name_or_path=lora_weights_name_or_path,
        force_auto_device_map=model_args.force_auto_device_map,
        torch_dtype=get_correct_torch_dtype(
            quantization=model_args.quantization_inference, model_args=model_args, training_args=training_args
        ),
        use_better_transformer=model_args.use_better_transformer,
        trust_remote_code=model_args.trust_remote_code,
        use_flash_attention=model_args.use_flash_attention,
        max_memory_MB=model_args.max_memory_MB,
    )

    trainer = CollieTrainer(
        model=model,
        args=training_args,
        data_collator=DataCollatorForCoLLIE(
            tokenizer,
            pad_to_multiple_of=8,
            return_tensors="pt",
            padding=True,
            label_pad_token_id=tokenizer.pad_token_id,
        ),
    )

    for test_task in data_args.test_tasks:
        test_dataset = os.path.join(
            data_args.dataset_dir,
            f"{test_task}.{'test' if not data_args.use_dev_inference else 'dev'}.jsonl",
        )
        test_dataset = CollieDataset(
            tokenizer=tokenizer,
            dataset_path=test_dataset,
            max_length=data_args.max_seq_length,
            is_encoder_decoder=model.config.is_encoder_decoder,
            inference=True if training_args.predict_with_generate else False,
            prompt_loss_weight=0.0,
            max_examples=data_args.max_examples_per_task_test,
        )

        logging.info(f"Running inference on {test_task}...")
        predictions = trainer.predict(test_dataset)

        if not trainer.is_world_process_zero():
            # In distributed training, we only want one process to write predictions to the file.
            continue

        output_dir = training_args.output_dir if checkpoint_path is None else checkpoint_path
        if training_args.predict_with_generate:
            output_name = f"{os.path.join(output_dir,'predictions',test_task)}.predictions.jsonl"

            os.makedirs(os.path.join(output_dir, "predictions"), exist_ok=True)

            with open(output_name, "w", encoding="utf8") as f:
                logging.info(f"Writing predictions to {output_name}")
                predictions = predictions.predictions
                # Switch all -100 to tokenizer.pad_token_id, so we can decode the predictions
                predictions[predictions == -100] = tokenizer.pad_token_id

                try:
                    predictions = tokenizer.batch_decode(predictions, skip_special_tokens=True)
                except OverflowError:
                    raise OverflowError(f"Unable to decode predictions: {predictions}")

                for prediction in predictions:
                    print(
                        json.dumps({"model_prediction": prediction}, ensure_ascii=False),
                        file=f,
                    )

        else:
            metrics_name = f"{os.path.join(output_dir,test_task)}.metrics.json"
            with open(metrics_name, "w", encoding="utf8") as f:
                logging.info(f"Writing metrics to {metrics_name}")
                json.dump(predictions.metrics, fp=f, ensure_ascii=False, indent=4)

    if training_args.predict_with_generate and trainer.is_world_process_zero():
        task_scores = evaluate(model_args, data_args, training_args, checkpoint_path=checkpoint_path)
        # Add test_ prefix to report test scores
        task_scores = {f"test_{task}": score for task, score in task_scores.items()}
        # Report
        trainer.log(task_scores)

    if delete_merged_model:
        logging.info(f"Deleting merged model at {model_path}")
        import shutil

        try:
            shutil.rmtree(model_path)
        except OSError as e:
            logging.error(
                f"Unable to delete the merged model {model_path} : {e.strerror}\n"
                "You may need to delete the merged model manually."
            )


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)

    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, Seq2SeqTrainingArguments))
    logging.info(f"Sys args {sys.argv}")

    if len(sys.argv) > 0 and sys.argv[-1].endswith(".json"):
        # If we pass only one argument to the script, and it's the path to a json file,
        # let's parse it to get our arguments.
        logging.info(f"Loading json config {sys.argv[-1]}")
        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[-1]))

    elif len(sys.argv) > 0 and sys.argv[-1].endswith(".yaml"):
        # If we pass only one argument to the script, and it's the path to a yaml file,
        # let's parse it to get our arguments.
        logging.info(f"Loading yaml config {sys.argv[-1]}")
        model_args, data_args, training_args = parser.parse_yaml_file(yaml_file=os.path.abspath(sys.argv[-1]))
    else:
        logging.info("No config file passed, using command line arguments.")
        model_args, data_args, training_args = parser.parse_args_into_dataclasses()

    if training_args.do_train and data_args.train_tasks is not None:
        train_collie(
            model_args,
            data_args,
            training_args,
        )
        clean_cache()
        if model_args.use_lora and model_args.merge_lora_after_training:
            merge_lora_model(
                weights_path=model_args.model_name_or_path,
                lora_weights_name_or_path=training_args.output_dir,
                torch_dtype=model_args.torch_dtype,
                output_path=os.path.join(training_args.output_dir, "merged_model"),
            )
            clean_cache()

    if training_args.do_predict and data_args.test_tasks is not None:
        if not data_args.evaluate_all_checkpoints:
            inference_collie(
                model_args,
                data_args,
                training_args,
            )
            clean_cache()
        else:
            # Find all checkpoints in the output directory
            checkpoints = [
                c
                for c in glob.glob(
                    os.path.join(training_args.output_dir, "checkpoint-*"),
                )
                if os.path.isdir(c)
            ]

            # Sort checkpoints by step number
            checkpoints = sorted(checkpoints, key=lambda x: int(x.split("-")[-1]))
            # Evaluate only checkpoints trained for 1000 or more steps, underfit models are very slow to evaluate
            no_eval_checkpoints = [c for c in checkpoints if int(c.split("-")[-1]) < 1000]
            if len(no_eval_checkpoints) > 0:
                logging.warning(
                    f"Found {len(no_eval_checkpoints)} checkpoints in {training_args.output_dir} that will not be"
                    f" evaluated: {no_eval_checkpoints} . We will evaluate only checkpoints trained for 1000 or more"
                    " steps, underfit models are very slow to evaluate."
                )

            checkpoints = [c for c in checkpoints if int(c.split("-")[-1]) >= 1000]

            # Add also the last checkpoint (stored in the output_dir)
            # checkpoints.append(training_args.output_dir)

            logging.info(
                f"Found {len(checkpoints)} checkpoints in {training_args.output_dir}:"
                f" {checkpoints} . We will evaluate each of them."
            )

            # Evaluate each checkpoint
            for checkpoint in checkpoints:
                inference_collie(
                    model_args,
                    data_args,
                    training_args,
                    checkpoint_path=checkpoint,
                )
                clean_cache()

src/run.py
hitz-zentroa/GoLLIE
import argparse
import json
import os
from typing import Dict, List


def compare_class_scores(model_paths: str, output_path: str):
    """
    Compare the class scores of multiple models and save the results to a CSV file.

    Args:
        model_paths (`List[str]`):
            List of paths to the models to compare.
        output_path (`str`):
            Path to the output CSV file.
    """
    class_scores: Dict[str, Dict[str, List[str]]] = {}
    for model_path in model_paths:
        task_scores_json = os.path.join(model_path, "task_scores.json")
        if not os.path.exists(task_scores_json):
            raise FileNotFoundError(f"task_scores.json not found in {model_path}")
        with open(task_scores_json, "r", encoding="utf8") as f:
            task_scores = json.load(f)
        for dataset, values in task_scores.items():
            for task, values in values.items():
                if "class_scores" in values:
                    name = f"{dataset}.{task}"
                    if name not in class_scores:
                        class_scores[name] = {}

                    for label, score in values["class_scores"].items():
                        if label not in class_scores[name]:
                            class_scores[name][label] = []
                        class_scores[name][label].append(str(score["f1-score"]))

    # Convert dict into CSV
    with open(output_path, "w", encoding="utf8") as f:
        print(f"Label,{','.join(model_paths)}", file=f)
        for name, scores in class_scores.items():
            print(name, file=f)
            for label, score in scores.items():
                print(f"{label},{','.join(score)}", file=f)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--model_paths", type=str, nargs="+", required=True, help="Paths to the models to compare.")
    parser.add_argument("--output_path", type=str, required=True, help="Path to the output CSV file.")
    args = parser.parse_args()
    compare_class_scores(args.model_paths, args.output_path)

src/scripts/compare_class_scores.py
hitz-zentroa/GoLLIE
import argparse
import json

from src.generate_data import get_class


def get_top_label_per_class(config, top_k: int = 10):
    """
    Get the top k labels per class of the train set
    Args:
        config (dict): The config dictionary
        top_k (int): The number of top labels to return
    Returns:
        (dict): A dictionary with the top k labels per class
    """
    if config["tasks"] != ["NER"] and "ner" not in config["tasks"][0].lower():
        raise ValueError(f"This function only supports NER task. Got {config['tasks']}")
    dataloader_cls = get_class(config["dataloader_cls"])
    if "train_file" not in config:
        raise ValueError("train_file is not specified in config")
    dataloader = dataloader_cls(config["train_file"], **config)
    ENTITY_TO_CLASS_MAPPING = dataloader.ENTITY_TO_CLASS_MAPPING
    {v: k for k, v in ENTITY_TO_CLASS_MAPPING.items()}

    ENTITY_COUNT = {str(k).split(".")[-1][:-2]: {} for k in ENTITY_TO_CLASS_MAPPING.values()}
    print(ENTITY_COUNT)
    for example in dataloader:
        entities = example["entities"]
        for entity in entities:
            ENTITY_COUNT[str(entity).split("(")[0]][entity.span] = (
                ENTITY_COUNT[str(entity).split("(")[0]].get(entity.span, 0) + 1
            )

    top_k_labels = {}

    for entity_type, entity_count in ENTITY_COUNT.items():
        top_k_labels[entity_type] = sorted(entity_count.items(), key=lambda x: x[1], reverse=True)[:top_k]

    print("Top k labels per class")
    print(top_k_labels)

    print("\n\nTop k labels per class (formatted)\n")
    for entity_type, entity_count in top_k_labels.items():
        print(f"{entity_type}. Examples: {', '.join([f'{k}' for k, v in entity_count])}.\n")

    print("\n\nTop k labels per class (Python List)\n")

    for entity_type, entity_count in top_k_labels.items():
        print(entity_type)
        print([f"{k}" for k, v in entity_count])

    return top_k_labels


def add_examples(config, top_k, lang="en"):
    # Get entity to example_name mapping
    entity_to_example_name = {}
    prompt_template = config["dataloader_cls"]
    prompt_template = prompt_template.split(".")[:3]
    guilines_gold = prompt_template + ["guidelines_gold.py"]
    prompt_template += ["prompts.py"]

    prompt_template = "/".join(prompt_template)

    if config["tasks"] == ["CrossNER_POLITICS"]:
        prompt_template = prompt_template.replace("prompts.py", "prompts_politics.py")
    elif config["tasks"] == ["CrossNER_MUSIC"]:
        prompt_template = prompt_template.replace("prompts.py", "prompts_music.py")
    elif config["tasks"] == ["CrossNER_AI"]:
        prompt_template = prompt_template.replace("prompts.py", "prompts_ai.py")
    elif config["tasks"] == ["CrossNER_NATURAL_SCIENCE"]:
        prompt_template = prompt_template.replace("prompts.py", "prompts_natural_science.py")
    elif config["tasks"] == ["CrossNER_LITERATURE"]:
        prompt_template = prompt_template.replace("prompts.py", "prompts_literature.py")

    guilines_gold = "/".join(guilines_gold)

    with open(prompt_template, "r", encoding="utf8") as f:
        lines = f.readlines()
    i = 0
    while i < len(lines):
        if lines[i].startswith("class"):
            entity = lines[i].split("(")[0].split(" ")[1]
            i += 1
            while not lines[i].strip().startswith("span"):
                i += 1

            name = lines[i].split("{")[1].split("}")[0]
            entity_to_example_name[entity] = name
        i += 1

    print(f"entity_to_example_name:\n{json.dumps(entity_to_example_name, indent=4)}")

    # Get top k labels per class
    top_k_labels = get_top_label_per_class(config, top_k)

    # Add examples to guilines
    examples = {}
    for entity_type, entity_count in top_k_labels.items():
        examples[entity_to_example_name[entity_type]] = {lang: [f"{k}" for k, v in entity_count]}

    print(f"\nEXAMPLES={json.dumps(examples, indent=4)}", file=open(guilines_gold, "a", encoding="utf8"))


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", type=str, required=True, help="The config file")
    parser.add_argument("--top_k", type=int, default=10, help="The number of top labels to return")
    parser.add_argument("--add_to_guidelines", action="store_true", help="Add examples to guidelines_gold.py")
    args = parser.parse_args()
    config = json.load(open(args.config, "r"))
    if args.add_to_guidelines:
        add_examples(config, args.top_k)
    else:
        get_top_label_per_class(config, args.top_k)

src/scripts/get_examples.py
hitz-zentroa/GoLLIE
import json

import numpy as np
import pandas as pd
from tabulate import tabulate


SUP_DATASETS = [
    ("ace05.ner", "entities"),
    ("ace05.re", "relations"),
    ("ace05.ee", "events"),
    ("ace05.eae", "arguments"),
    ("bc5cdr.ner", "entities"),
    ("conll03.ner", "entities"),
    ("diann.ner", "entities"),
    ("ncbidisease.ner", "entities"),
    ("ontonotes5.ner", "entities"),
    ("rams.eae", "arguments"),
    ("tacred.sf", "slots"),
    ("wnut17.ner", "entities"),
]

ZERO_DATASETS = [
    ("broadtwitter.ner", "entities"),
    ("casie.ee", "events"),
    ("casie.eae", "arguments"),
    ("crossner.crossner_ai", "entities"),
    ("crossner.crossner_literature", "entities"),
    ("crossner.crossner_music", "entities"),
    ("crossner.crossner_politics", "entities"),
    ("crossner.crossner_natural_science", "entities"),
    ("e3c.ner", "entities"),
    ("fabner.ner", "entities"),
    ("harveyner.ner", "entities"),
    ("mitmovie.ner", "entities"),
    ("mitrestaurant.ner", "entities"),
    ("multinerd.ner", "entities"),
    ("wikievents.ner", "entities"),
    ("wikievents.ee", "events"),
    ("wikievents.eae", "arguments"),
]

PATHS = {
    "Baseline": "/ikerlariak/osainz006/models/collie/Baseline-7b_CodeLLaMA{seed}/task_scores_summary.json",
    "GoLLIE": "/ikerlariak/osainz006/models/collie/CoLLIE+-7b_CodeLLaMA{seed}/task_scores_summary.json",
    "w/o Candidates": "/ikerlariak/osainz006/models/collie/CoLLIE-7b_CodeLLaMA{seed}/task_scores_summary.json",
    "w/o Masking": (
        "/ikerlariak/osainz006/models/collie/CoLLIE+-7b_CodeLLaMA{seed}_abl_masking/task_scores_summary.json"
    ),
    "w/o Dropout": (
        "/ikerlariak/osainz006/models/collie/CoLLIE+-7b_CodeLLaMA{seed}_abl_dropout/task_scores_summary.json"
    ),
    "13B": "/ikerlariak/osainz006/models/collie/CoLLIE+-13b_CodeLLaMA{seed}/task_scores_summary.json",
    "34B": "/ikerlariak/osainz006/models/collie/CoLLIE+-34b_CodeLLaMA{seed}/task_scores_summary.json",
}

SEEDS = ["", "_2", "_3"]


data = []
for name, path in PATHS.items():
    print()
    print(name)
    sup_results, zero_results = [], []
    i = 0
    for seed in SEEDS:
        print(seed)
        if name == "Baseline" and seed:
            seed += "/1"
        try:
            with open(path.format(seed=seed)) as f:
                results = json.load(f)
        except FileNotFoundError:
            continue
        i += 1
        _sup_results = []
        for dataset, task in SUP_DATASETS:
            _sup_results.append(results[dataset][task]["f1-score"] * 100)
        sup_results.append(_sup_results + [np.mean(_sup_results)])
        _zero_results, _zero_sota_results = [], []
        for dataset, task in ZERO_DATASETS:
            _zero_results.append(results[dataset][task]["f1-score"] * 100)
            if dataset.startswith("crossner") or dataset.startswith("mit") or dataset.startswith("wikievents"):
                _zero_sota_results.append(results[dataset][task]["f1-score"] * 100)
        zero_results.append(_zero_results + [np.mean(_zero_sota_results), np.mean(_zero_results)])

    sup_results.append(np.array(sup_results).mean(0))
    sup_results.append(np.array(sup_results).std(0))

    zero_results.append(np.array(zero_results).mean(0))
    zero_results.append(np.array(zero_results).std(0))

    columns = list(map(str, range(i))) + ["Average", "Std"]

    sup_results = pd.DataFrame(sup_results, columns=[dataset[0] for dataset in SUP_DATASETS] + ["Average"]).T
    sup_results.columns = columns
    print(tabulate(sup_results, headers=columns, floatfmt=".1f"))

    zero_results = pd.DataFrame(
        zero_results, columns=[dataset[0] for dataset in ZERO_DATASETS] + ["Zero Average", "Average"]
    ).T
    zero_results.columns = columns
    print(tabulate(zero_results, headers=columns, floatfmt=".1f"))

    # print(tabulate(sup_results, floatfmt=".1f", headers=[dataset[0] for dataset in SUP_DATASETS]))
    # print(tabulate(zero_results, floatfmt=".1f", headers=[dataset[0] for dataset in ZERO_DATASETS]))

src/scripts/get_result_table.py
hitz-zentroa/GoLLIE
import glob
import json
import logging
import os
import time
from argparse import ArgumentParser

import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns


sns.set_theme(style="whitegrid")


def plot_curves(args):
    results_dict = {}
    for dir_path in glob.glob(os.path.join(args.models_path, args.regexp)):
        logging.info("Processing model output:" + dir_path)

        key = os.path.basename(dir_path)

        for chkpt in glob.glob(os.path.join(dir_path, "checkpoint-*")):
            logging.info("Checkpoint: " + os.path.basename(chkpt))
            scores_file = os.path.join(chkpt, "task_scores.json")
            if not os.path.exists(scores_file):
                continue

            step = int(os.path.basename(chkpt).split("-")[-1])

            trainer_state_file = os.path.join(chkpt, "trainer_state.json")
            if os.path.exists(trainer_state_file):
                with open(trainer_state_file) as f:
                    step_info = json.load(f)["log_history"][-1]
            else:
                step_info = {"step": step}

            with open(scores_file) as f:
                scores = json.load(f)

            for dataset, values in scores.items():
                for task, values in values.items():
                    if task == "prediction_stats":
                        continue
                    if dataset not in results_dict:
                        results_dict[dataset] = {}

                    if task not in results_dict[dataset]:
                        results_dict[dataset][task] = {}

                    if args.hue:
                        _key = key.replace(args.hue, "").replace("__", "_").rstrip("_")

                    if args.hue and _key not in results_dict[dataset][task]:
                        results_dict[dataset][task][_key] = {}
                    elif not args.hue and key not in results_dict[dataset][task]:
                        results_dict[dataset][task][key] = {}

                    if args.hue:
                        if step not in results_dict[dataset][task][_key]:
                            results_dict[dataset][task][_key][step] = {}

                        results_dict[dataset][task][_key][step][args.hue in key] = {**values, **step_info}
                    else:
                        if step not in results_dict[dataset][task][key]:
                            results_dict[dataset][task][key][step] = {}

                        results_dict[dataset][task][key][step] = {**values, **step_info}

    result_list = []
    for dataset, values in results_dict.items():
        for task, values in values.items():
            for model, values in values.items():
                for chkpt, values in values.items():
                    if args.hue:
                        for hue, values in values.items():
                            result_list.append(
                                {
                                    "dataset": dataset,
                                    "task": task,
                                    "model": model,
                                    "steps": chkpt,
                                    f"{args.hue}": hue,
                                    **values,
                                }
                            )
                    else:
                        result_list.append(
                            {
                                "dataset": dataset,
                                "task": task,
                                "model": model,
                                "steps": chkpt,
                                f"{args.hue}": False,
                                **values,
                            }
                        )
    localtime = time.localtime()
    timestamp = (
        f"{localtime.tm_year}-{localtime.tm_mon}-{localtime.tm_mday}-{localtime.tm_hour}-"
        + f"{localtime.tm_min}-{localtime.tm_sec}"
    )

    df = pd.DataFrame(result_list)
    logging.info(f"Available columns: {df.columns}")
    f = plt.figure(figsize=(10, 5))
    sns.lineplot(
        data=df,
        x=args.x,
        y=args.y,
        hue="model",
        size=args.hue,
        size_order=[True, False] if args.hue else None,
        palette="tab10",
        linewidth=2.5,
    )
    plt.tight_layout()
    f.savefig(f"assets/plots/{timestamp}.png", dpi=400)


if __name__ == "__main__":
    """
    Example: python -m src.scripts.plot_f1_curves --regexp "*optim*" --hue "ignore_prompt_loss"
    """
    logging.basicConfig(level=logging.INFO)

    parser = ArgumentParser()

    parser.add_argument("--models_path", type=str, default="/ikerlariak/igarcia945/CoLLIE/pretrained_models")
    parser.add_argument("--regexp", type=str, default="*")
    parser.add_argument("--hue", type=str, default=None)
    parser.add_argument("--x", type=str, default="steps")
    parser.add_argument("--y", type=str, default="f1-score")

    args = parser.parse_args()
    plot_curves(args)

src/scripts/plot_f1_curves.py
hitz-zentroa/GoLLIE
import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns


sns.set_theme(style="whitegrid")

# PromptNER https://arxiv.org/pdf/2305.15444.pdf
GPT3_5 = [
    00.00,  # Movie
    00.00,  # Restaurant
    20.30,  # Politics
    31.30,  # Literature
    24.50,  # Music
    40.70,  # AI
    40.60,  # Science
]

# Instruct UIE https://arxiv.org/pdf/2304.08085.pdf
InstructUIE = [
    63.00,  # Movie
    20.99,  # Restaurant
    49.00,  # Politics
    47.21,  # Literature
    53.16,  # Music
    48.15,  # AI
    49.30,  # Science
]

CoLLIE = [
    63.0,  # Movie
    43.4,  # Restaurant
    57.2,  # Politics
    62.7,  # Literature
    67.8,  # Music
    59.1,  # AI
    55.5,  # Science
]


def main():
    fig, ax = plt.subplots(1, 7, figsize=(12, 4), sharey=True, layout="constrained")

    TASK_NAMES = ["Movie", "Restaurant", "Politics", "Literature", "Music", "AI", "Science"]

    for i, (gpt, iuie, collie, name) in enumerate(zip(GPT3_5, InstructUIE, CoLLIE, TASK_NAMES)):
        rect = ax[i].bar(
            [1],
            [np.round(gpt)],
            width=1.0,
            label="GPT-3",
            # color="#a40e26",
            color=sns.color_palette("crest", 3)[0],
            hatch="//",
        )
        if gpt:
            ax[i].bar_label(rect, padding=3, fontsize=12)

        rect = ax[i].bar(
            [2],
            [np.round(iuie)],
            width=1.0,
            label="Instruct-UIE",
            # color="#6639ba",
            color=sns.color_palette("crest", 3)[1],
            hatch="/",
        )
        ax[i].bar_label(rect, padding=3, fontsize=12)

        rect = ax[i].bar(
            [3],
            [np.round(collie)],
            width=1.0,
            label="CoLLIE",
            # color="#0a3069"
            color=sns.color_palette("crest", 3)[2],
        )
        ax[i].bar_label(rect, padding=3, fontsize=12)

        ax[i].set_xticks([1, 2, 3])
        ax[i].set_yticklabels([])
        ax[i].set_xticklabels(["", name, ""], fontsize=14)
        # ax[i].set_title(name)
        ax[i].grid(False)
        ax[i].spines["top"].set_visible(False)
        ax[i].spines["right"].set_visible(False)
        ax[i].spines["bottom"].set_visible(False)
        ax[i].spines["left"].set_visible(False)

    fig.legend(["GPT-3.5", "SOTA", "GoLLIE"], loc="outside upper center", ncol=3, fontsize=14, frameon=False)
    # ax[3].legend(["GPT-3", "Instruct-UIE", "CoLLIE"], fontsize=12, ncol=3, bbox_to_anchor=(1.00, 1.15), loc="lower center")

    # plt.tight_layout()
    plt.savefig("assets/plots/zero_shot_results.pdf", dpi=300)


if __name__ == "__main__":
    main()

src/scripts/plot_results.py
hitz-zentroa/GoLLIE
import argparse
import logging
import random
import string

import bitsandbytes as bnb
import torch
from accelerate import Accelerator
from fairseq.optim.adafactor import Adafactor
from tabulate import tabulate
from torch.optim import AdamW
from torch.utils.data import DataLoader, Dataset

from src.model.load_model import load_model
from transformers import (
    BatchEncoding,
    DataCollatorForSeq2Seq,
    PreTrainedModel,
    PreTrainedTokenizerBase,
)


def generate_random_sentence(sentence_length: int = 5120) -> str:
    """
    Generates a random string of specific length with ascii characters.

    Args:
        sentence_length (`int`, optional):
            Length of the sentence. Defaults to `5120`.

    Returns:
        str:
            A random sentence of length `sentence_length`.
    """
    sentence = ""
    for i in range(sentence_length):
        # generate a random word of length between 1 and 10 characters
        word = "".join(random.choice(string.ascii_lowercase) for _ in range(random.randint(1, 10)))
        sentence += word + " "
    return sentence


class TestDataset(Dataset):
    """
    A dummy dataset used for testing.

    Args:
        tokenizer (`PreTrainedTokenizerBase`):
            The pre-trained tokenizer.
        seq_len (`int`):
            The length of the sentences.
        data_len (`int`):
            The length of the data.
    """

    def __init__(self, tokenizer: PreTrainedTokenizerBase, seq_len: int, data_len: int):
        self.data = []
        for i in range(data_len):
            # model_inputs = tokenizer(
            #    generate_random_sentence(),
            #    return_tensors=None,
            #    padding="max_length",
            #    truncation=True,
            #    max_length=seq_len,
            #    add_special_tokens=True,
            # )

            # model_inputs["labels"] = model_inputs["input_ids"].copy()

            # self.data.append(model_inputs)

            inputs_ids = random.sample(range(100, len(tokenizer) - 100), seq_len)

            self.data.append(
                BatchEncoding(
                    {
                        "input_ids": inputs_ids,
                        "attention_mask": [1] * seq_len,
                        "labels": inputs_ids.copy(),
                    }
                )
            )

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx]


def get_dataloader(
    tokenizer: PreTrainedTokenizerBase,
    batch_size: int,
    seq_len: int,
) -> DataLoader:
    """
    A function that returns a `DataLoader` instance

    Args:
        tokenizer (`PreTrainedTokenizerBase`):
            The pre-trained tokenizer.
        batch_size (`int`):
            The size of the batch.
        seq_len (`int`):
            The size of the sequence.

    Returns:
        `DataLoader`:
            The dataloader that fits the configuration given.
    """
    dataset = TestDataset(tokenizer=tokenizer, seq_len=seq_len, data_len=batch_size)
    data_collator = DataCollatorForSeq2Seq(tokenizer, padding=True, pad_to_multiple_of=8, return_tensors="pt")
    return DataLoader(
        dataset,
        batch_size=batch_size,
        collate_fn=data_collator,
        shuffle=True,
        pin_memory=True,
    )


def run_training_test(
    model: PreTrainedModel,
    tokenizer: PreTrainedTokenizerBase,
    accelerator: Accelerator,
    seq_len: int,
    batch_size: int,
    use_lora: bool,
    quantization: bool,
    optimizer_name: str,
    learning_rate: float,
):
    """
    Runs the stress test to check if the current hardware allows it.

    Args:
        model (`PreTrainedModel`):
            The model to be tested.
        tokenizer (`PreTrainedTokenizerBase`):
            The tokenizer of the model to be tested.
        accelerator (`Accelerator`):
            The accelerator instance.
        seq_len (`int`):
            The length of the sequence.
        batch_size (`int`):
            The length of the batch.
        use_lora (`bool`):
            Whether to use LoRA or not.
        quantization (`int`):
            Whether to use 4 bits / 8 bits quantization or not.
        optimizer_name (`str`):
            The optimizer to use on the test.
        learning_rate (`float`):
            The learning rate to use on the test.

    Raises:
        `ValueError`:
            raised when `optimizer_name!="adamW"` and `quantization!=None`.
    """
    assert optimizer_name in [
        "adamW",
        "AdaFactor",
    ], f"Optimizer {optimizer_name} not supported. Choose between adamW and AdaFactor."

    model.config.max_length = seq_len

    dataloader = get_dataloader(tokenizer=tokenizer, batch_size=batch_size, seq_len=seq_len)

    if quantization:
        if optimizer_name != "adamW":
            raise ValueError("Quantization is only supported with adamW optimizer.")
        optimizer = bnb.optim.AdamW8bit(
            model.parameters(),
            lr=learning_rate,
            betas=(0.9, 0.995),
        )
    else:
        if optimizer_name == "adamW":
            optimizer = AdamW(model.parameters(), lr=learning_rate, eps=1e-7)
        else:
            optimizer = Adafactor(
                params=model.parameters(),
                scale_parameter=False,
                relative_step=False,
                warmup_init=False,
                lr=learning_rate,
                clip_threshold=1.0,
            )

    optimizer, dataloader = accelerator.prepare(optimizer, dataloader)

    model.train()

    for i, batch in enumerate(dataloader):
        optimizer.zero_grad()
        # print("Data devices:")
        # print(batch.input_ids.device)
        # print(batch.attention_mask.device)
        # print(batch.labels.device)
        # print("Model devices:")
        # print(model.device)
        outputs = model(
            input_ids=batch.input_ids,
            labels=batch.labels,
            attention_mask=batch.attention_mask,
        )
        loss = outputs.loss
        accelerator.backward(loss)
        optimizer.step()
        optimizer.zero_grad()

    accelerator.wait_for_everyone()

    return 1


def get_params():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--model_name_or_path",
        type=str,
        default="/gaueko1/hizkuntza-ereduak/LLaMA/lm/huggingface/65B/",
        help="Model name or path",
    )
    parser.add_argument(
        "--seq_len",
        type=int,
        nargs="+",
        default=[128, 256, 512, 1024, 1536, 2048, 4096],
        help="Sequence lengths to test",
    )
    parser.add_argument(
        "--batch_size",
        type=int,
        nargs="+",
        default=[1, 4, 8, 16, 32, 64, 128],
        help="Batch sizes to test",
    )
    parser.add_argument(
        "--use_lora",
        action="store_true",
        help="Use LORA",
    )
    parser.add_argument(
        "--quantization",
        type=int,
        default=None,
        help="Use 4 bits / 8 bits quantization or not",
    )
    parser.add_argument(
        "--optimizer_name",
        type=str,
        default="adamW",
        choices=["adamW", "AdaFactor"],
        help="Optimizer to use",
    )
    parser.add_argument(
        "--learning_rate",
        type=float,
        default=5e-5,
        help="Learning rate",
    )
    return parser.parse_args()


def main():
    args = get_params()
    seq_lens = args.seq_len
    batch_sizes = args.batch_size

    result_matrix = [[0 for i in range(len(batch_sizes))] for j in range(len(seq_lens))]

    logging.info("***** Experiment parameters *****")
    logging.info(f"Model name or path: {args.model_name_or_path}")
    logging.info(f"Sequence lengths: {seq_lens}")
    logging.info(f"Batch sizes: {batch_sizes}")
    logging.info(f"Use LORA: {args.use_lora}")
    logging.info(f"Use quantization: {args.quantization}")
    logging.info(f"Optimizer: {args.optimizer_name}")
    logging.info(f"Learning rate: {args.learning_rate}")
    logging.info("*********************************")
    logging.info(f"Loading model {args.model_name_or_path}...")
    model, tokenizer = load_model(
        inference=False,
        model_weights_name_or_path=args.model_name_or_path,
        use_lora=args.use_lora,
        quantization=args.quantization,
        use_gradient_checkpointing=True,
        use_flash_attention=True,
    )
    logging.info("Model loaded.")

    accelerator = Accelerator()
    model = accelerator.prepare(model)

    for i, seq_len in enumerate(seq_lens):
        for j, batch_size in enumerate(batch_sizes):
            logging.info(f"Running test with seq_len={seq_len} and batch_size={batch_size}...")
            try:
                result_matrix[i][j] = run_training_test(
                    model=model,
                    tokenizer=tokenizer,
                    accelerator=accelerator,
                    seq_len=seq_len,
                    batch_size=batch_size,
                    use_lora=args.use_lora,
                    quantization=args.quantization,
                    optimizer_name=args.optimizer_name,
                    learning_rate=args.learning_rate,
                )
                logging.info("Test passed.")
            except RuntimeError as e:
                if "out of memory" in str(e):
                    logging.info("| WARNING: ran out of memory, retrying batch")
                    for p in model.parameters():
                        if p.grad is not None:
                            del p.grad  # free some memory
                    torch.cuda.empty_cache()
                    logging.info("Test failed.")
                    continue
                else:
                    raise e

    table = tabulate(tabular_data=result_matrix, headers=batch_sizes, showindex=seq_lens)
    print(table)


if __name__ == "__main__":
    main()

src/scripts/test_context_batch_size.py
hitz-zentroa/GoLLIE
import json
import os
from argparse import ArgumentParser

from src.tasks import TASK_ID_TO_TASKS


def main(args):
    # Create the directory if not exists
    os.makedirs(args.output_dir, exist_ok=True)

    with open(args.input_file) as f:
        lines = [json.loads(line) for line in f]

    basename = os.path.basename(args.input_file).rstrip(".jsonl")
    output_path = os.path.join(args.output_dir, f"{basename}.{args.row}.py")
    with open(output_path, "wt", encoding="utf-8") as f:
        line = lines[args.row]
        imports = TASK_ID_TO_TASKS[line["task_id"]]

        print(f"from {imports} import *", file=f)
        print("from src.tasks.utils_typing import Entity, Value, Relation, Event", file=f)
        print("from dataclasses import dataclass", file=f)
        print(line["text"], file=f)
        print(f"labels = {line['labels']}", file=f)

    os.system(f"black {output_path}")
    os.system(f"ruff check {output_path} --fix")


if __name__ == "__main__":
    parser = ArgumentParser("Visualize examples")

    parser.add_argument("-i", "--input_file", dest="input_file", type=str)
    parser.add_argument("-r", "--row", dest="row", type=int, default=0)
    parser.add_argument("-o", "--output_dir", dest="output_dir", default=".ignore/examples")

    args = parser.parse_args()
    main(args)

src/scripts/visualize_example.py
hitz-zentroa/GoLLIE
from typing import Dict, List

from . import (
    ace,
    bc5cdr,
    broadtwitter,
    casie,
    conll03,
    crossner,
    diann,
    e3c,
    fabner,
    harveyner,
    mitmovie,
    mitrestaurant,
    multinerd,
    ncbidisease,
    ontonotes,
    rams,
    tacred,
    wikievents,
    wnut,
)


TASK_ID_TO_TASKS = {
    "ACE05_NER": "src.tasks.ace",
    "ACE05_VER": "src.tasks.ace",
    "ACE05_RE": "src.tasks.ace",
    "ACE05_RC": "src.tasks.ace",
    "ACE05_EE": "src.tasks.ace",
    "ACE05_EAE": "src.tasks.ace",
    "RAMS_EAE": "src.tasks.rams",
    "CoNLL03_NER": "src.tasks.conll03",
    "CASIE_EE": "src.tasks.casie",
    "CASIE_EAE": "src.tasks.casie",
    "Europarl_NER": "src.tasks.conll03",
    "TACRED_SF": "src.tasks.tacred",
    "OntoNotes5_NER": "src.tasks.ontonotes",
    "NcbiDisease_NER": "src.tasks.ncbidisease",
    "DIANN_NER": "src.tasks.diann",
    "WNUT17_NER": "src.tasks.wnut",
    "MultiNERD_NER": "src.tasks.multinerd",
    "WikiEvents_NER": "src.tasks.wikievents",
    "WikiEvents_EE": "src.tasks.wikievents",
    "WikiEvents_EAE": "src.tasks.wikievents",
    "FabNER_NER": "src.tasks.fabner",
    "E3C_NER": "src.tasks.e3c",
    "BC5CDR_NER": "src.tasks.bc5cdr",
    "BroadTwitter_NER": "src.tasks.broadtwitter",
    "HarveyNER_NER": "src.tasks.harveyner",
    "MITMovie_NER": "src.tasks.mitmovie",
    "MITRestaurant_NER": "src.tasks.mitrestaurant",
    "CrossNER_CrossNER_AI": "src.tasks.crossner",
    "CrossNER_CrossNER_POLITICS": "src.tasks.crossner",
    "CrossNER_CrossNER_NATURAL_SCIENCES": "src.tasks.crossner",
    "CrossNER_CrossNER_LITERATURE": "src.tasks.crossner",
    "CrossNER_CrossNER_MUSIC": "src.tasks.crossner",
    "CrossNER_woMISC_CrossNER_AI": "src.tasks.crossner",
    "CrossNER_woMISC_CrossNER_POLITICS": "src.tasks.crossner",
    "CrossNER_woMISC_CrossNER_NATURAL_SCIENCES": "src.tasks.crossner",
    "CrossNER_woMISC_CrossNER_LITERATURE": "src.tasks.crossner",
    "CrossNER_woMISC_CrossNER_MUSIC": "src.tasks.crossner",
}

__all__ = [
    "ace",
    "rams",
    "conll03",
    "casie",
    "tacred",
    "ontonotes",
    "ncbidisease",
    "bc5cdr",
    "diann",
    "wnut",
    "multinerd",
    "wikievents",
    "fabner",
    "e3c",
    "broadtwitter",
    "harveyner",
    "mitmovie",
    "mitrestaurant",
    "crossner",
    "TASK_ID_TO_TASKS",
    "task_id_to_guidelines",
]


def task_id_to_prompts(task_id: str) -> str:
    """
    Returns the prompts path for a given task.

    Args:
        task_id (str): The task id

    Returns:
        str: The path to the prompts
    """

    if task_id.upper() == "CASIE_EE":
        return "src.tasks.casie.prompts_ed"

    elif task_id.upper() == "CASIE_EAE":
        return "src.tasks.casie.prompts_eae"

    elif task_id.upper() == "CrossNER_CrossNER_POLITICS".upper():
        return "src.tasks.crossner.prompts_politics"
    elif task_id.upper() == "CrossNER_CrossNER_AI".upper():
        return "src.tasks.crossner.prompts_ai"
    elif task_id.upper() == "CrossNER_CrossNER_NATURAL_SCIENCE".upper():
        return "src.tasks.crossner.prompts_natural_science"
    elif task_id.upper() == "CrossNER_CrossNER_LITERATURE".upper():
        return "src.tasks.crossner.prompts_literature"
    elif task_id.upper() == "CrossNER_CrossNER_MUSIC".upper():
        return "src.tasks.crossner.prompts_music"
    elif task_id.upper() == "CrossNER_woMISC_CrossNER_POLITICS".upper():
        return "src.tasks.crossner.prompts_politics"
    elif task_id.upper() == "CrossNER_woMISC_CrossNER_AI".upper():
        return "src.tasks.crossner.prompts_ai"
    elif task_id.upper() == "CrossNER_woMISC_CrossNER_NATURAL_SCIENCE".upper():
        return "src.tasks.crossner.prompts_natural_science"
    elif task_id.upper() == "CrossNER_woMISC_CrossNER_LITERATURE".upper():
        return "src.tasks.crossner.prompts_literature"
    elif task_id.upper() == "CrossNER_woMISC_CrossNER_MUSIC".upper():
        return "src.tasks.crossner.prompts_music"

    # Default case
    else:
        return TASK_ID_TO_TASKS[task_id] + ".prompts"


def task_id_to_guidelines(task_id: str) -> Dict[str, Dict[str, List[str]]]:
    """
    Return the guidelines for a given task.

    Args:
        task_id (str): The task id.

    Returns:
        The guidelines for the task.
    """
    if task_id.lower() == "ace05":
        from src.tasks.ace.guidelines_gold import GUIDELINES

        return GUIDELINES
    elif task_id.lower() == "rams":
        from src.tasks.rams.guidelines_gold import GUIDELINES

        return GUIDELINES
    elif task_id.lower() == "conll03":
        from src.tasks.conll03.guidelines_gold import GUIDELINES

        return GUIDELINES
    elif task_id.lower() == "casie":
        from src.tasks.casie.guidelines_gold import GUIDELINES

        return GUIDELINES
    elif task_id.lower() == "tacred":
        from src.tasks.tacred.guidelines_gold import GUIDELINES

        return GUIDELINES
    elif task_id.lower() == "ontonotes5":
        from src.tasks.ontonotes.guidelines_gold import GUIDELINES

        return GUIDELINES
    elif task_id.lower() == "ncbidisease":
        from src.tasks.ncbidisease.guidelines_gold import GUIDELINES

        return GUIDELINES
    elif task_id.lower() == "bc5cdr":
        from src.tasks.bc5cdr.guidelines_gold import GUIDELINES

        return GUIDELINES
    elif task_id.lower() == "diann":
        from src.tasks.diann.guidelines_gold import GUIDELINES

        return GUIDELINES
    elif task_id.lower() == "wnut17":
        from src.tasks.wnut.guidelines_gold import GUIDELINES

        return GUIDELINES
    elif task_id.lower() == "multinerd":
        from src.tasks.multinerd.guidelines_gold import GUIDELINES

        return GUIDELINES
    elif task_id.lower() == "wikievents":
        from src.tasks.wikievents.guidelines_gold import GUIDELINES

        return GUIDELINES
    elif task_id.lower() == "fabner":
        from src.tasks.fabner.guidelines_gold import GUIDELINES

        return GUIDELINES
    elif task_id.lower() == "e3c":
        from src.tasks.e3c.guidelines_gold import GUIDELINES

        return GUIDELINES
    elif task_id.lower() == "broadtwitter":
        from src.tasks.broadtwitter.guidelines_gold import GUIDELINES

        return GUIDELINES
    elif task_id.lower() == "harveyner":
        from src.tasks.harveyner.guidelines_gold import GUIDELINES

        return GUIDELINES
    elif task_id.lower() == "mitmovie":
        from src.tasks.mitmovie.guidelines_gold import GUIDELINES

        return GUIDELINES
    elif task_id.lower() == "mitrestaurant":
        from src.tasks.mitrestaurant.guidelines_gold import GUIDELINES

        return GUIDELINES
    elif task_id.lower() == "crossner":
        from src.tasks.crossner.guidelines_gold import GUIDELINES

        return GUIDELINES
    else:
        raise ValueError(f"Task {task_id} not supported.")

src/tasks/__init__.py
hitz-zentroa/GoLLIE
from . import data_loader, prompts, scorer


__all__ = ["data_loader", "prompts", "scorer"]

src/tasks/ace/__init__.py
hitz-zentroa/GoLLIE
import inspect
import json
from typing import Tuple, Union

from src.tasks.ace.guidelines import GUIDELINES
from src.tasks.ace.guidelines_gold import EXAMPLES
from src.tasks.ace.prompts import (
    COARSE_EVENT_DEFINITIONS,
    COARSE_RELATION_DEFINITIONS,
    COARSE_TO_FINE_EVENTS,
    COARSE_TO_FINE_RELATIONS,
    ENTITY_DEFINITIONS,
    EVENT_DEFINITIONS,
    FINE_TO_COARSE_EVENTS,
    FINE_TO_COARSE_RELATIONS,
    GPE,
    RELATION_DEFINITIONS,
    VALUE_DEFINITIONS,
    Acquit,
    AgentArtifactRelationRelation,
    Appeal,
    ArrestJail,
    Attack,
    BeBorn,
    Business,
    BusinessEvent,
    ChargeIndict,
    CitizenResidentReligionEthnicity,
    ConflictEvent,
    ContactEvent,
    ContactInfo,
    Convict,
    Crime,
    DeclareBankruptcy,
    Demonstrate,
    Die,
    Divorce,
    Elect,
    Employment,
    EndOrg,
    EndPosition,
    Execute,
    Extradite,
    Facility,
    Family,
    Fine,
    Founder,
    GenAffiliationRelation,
    Geographical,
    Injure,
    InvestorShareholder,
    JobTitle,
    JusticeEvent,
    LastingPersonal,
    LifeEvent,
    Located,
    Location,
    Marry,
    Meet,
    Membership,
    MergeOrg,
    MovementEvent,
    Near,
    Nominate,
    Numeric,
    Organization,
    OrganizationAffiliationRelation,
    OrgLocationOrigin,
    Ownership,
    Pardon,
    PartWholeRelation,
    Person,
    PersonalSocialRelation,
    PersonnelEvent,
    PhoneWrite,
    PhysicalRelation,
    ReleaseParole,
    Sentence,
    SentenceAct,
    SportsAffiliation,
    StartOrg,
    StartPosition,
    StudentAlum,
    Subsidiary,
    Sue,
    Time,
    TransactionEvent,
    TransferMoney,
    TransferOwnership,
    Transport,
    TrialHearing,
    UserOwnerInventorManufacturer,
    Vehicle,
    Weapon,
)

from ..utils_data import DatasetLoader, Sampler


class ACEDatasetLoader(DatasetLoader):
    """
    A `DatasetLoader` for the ACE05 dataset.

    Args:
        path (`str`):
            The location of the dataset directory.
        group_by (`str`, optional):
            Whether to group the texts by sentence or documents. Defaults to "sentence".

    Raises:
        `ValueError`:
            raised when a not defined value found.
    """

    ENTITY_TO_CLASS_MAPPING = {
        "FAC": Facility,
        "GPE": GPE,
        "LOC": Location,
        "ORG": Organization,
        "PER": Person,
        "VEH": Vehicle,
        "WEA": Weapon,
    }
    VALUE_TO_CLASS_MAPPING = {
        "Contact-Info": ContactInfo,
        "Crime": Crime,
        "Job-Title": JobTitle,
        "Numeric": Numeric,
        "Sentence": Sentence,
        "TIME": Time,
    }
    RELATION_TO_CLASS_MAPPING = {
        "ART:User-Owner-Inventor-Manufacturer": (AgentArtifactRelationRelation, UserOwnerInventorManufacturer),
        "GEN-AFF:Citizen-Resident-Religion-Ethnicity": (GenAffiliationRelation, CitizenResidentReligionEthnicity),
        "GEN-AFF:Org-Location": (GenAffiliationRelation, OrgLocationOrigin),
        "ORG-AFF:Employment": (OrganizationAffiliationRelation, Employment),
        "ORG-AFF:Founder": (OrganizationAffiliationRelation, Founder),
        "ORG-AFF:Investor-Shareholder": (OrganizationAffiliationRelation, InvestorShareholder),
        "ORG-AFF:Membership": (OrganizationAffiliationRelation, Membership),
        "ORG-AFF:Ownership": (OrganizationAffiliationRelation, Ownership),
        "ORG-AFF:Sports-Affiliation": (OrganizationAffiliationRelation, SportsAffiliation),
        "ORG-AFF:Student-Alum": (OrganizationAffiliationRelation, StudentAlum),
        "PART-WHOLE:Artifact": (
            PartWholeRelation,
            Geographical,
        ),  # There is no definition for Artifact relation on the guidelines
        "PART-WHOLE:Geographical": (PartWholeRelation, Geographical),
        "PART-WHOLE:Subsidiary": (PartWholeRelation, Subsidiary),
        "PER-SOC:Business": (PersonalSocialRelation, Business),
        "PER-SOC:Family": (PersonalSocialRelation, Family),
        "PER-SOC:Lasting-Personal": (PersonalSocialRelation, LastingPersonal),
        "PHYS:Located": (PhysicalRelation, Located),
        "PHYS:Near": (PhysicalRelation, Near),
    }
    _EVENT_CONSTANTS_MAPPING = {
        "trigger": "mention",
        "Place": "place",
        "Time-After": "time",
        "Time-At-Begginning": "time",
        "Time-At-Beginning": "time",  # A bug on the data
        "Time-At-End": "time",
        "Time-Before": "time",
        "Time-Ending": "time",
        "Time-Holds": "time",
        "Time-Starting": "time",
        "Time-Within": "time",
    }
    EVENT_TO_CLASS_MAPPING = {
        "Business:Declare-Bankruptcy": {
            "coarse": BusinessEvent,
            "class": DeclareBankruptcy,
            "Org": "org",
        },
        "Business:End-Org": {
            "coarse": BusinessEvent,
            "class": EndOrg,
            "Org": "org",
        },
        "Business:Merge-Org": {
            "coarse": BusinessEvent,
            "class": MergeOrg,
            "Org": "org",
        },
        "Business:Start-Org": {"coarse": BusinessEvent, "class": StartOrg, "Agent": "agent", "Org": "org"},
        "Conflict:Attack": {
            "coarse": ConflictEvent,
            "class": Attack,
            "Agent": "attacker",
            "Attacker": "attacker",
            "Instrument": "instrument",
            "Target": "target",
            "Victim": "target",
        },
        "Conflict:Demonstrate": {
            "coarse": ConflictEvent,
            "class": Demonstrate,
            "Entity": "entity",
        },
        "Contact:Meet": {
            "coarse": ContactEvent,
            "class": Meet,
            "Entity": "entity",
        },
        "Contact:Phone-Write": {
            "coarse": ContactEvent,
            "class": PhoneWrite,
            "Entity": "entity",
        },
        "Justice:Acquit": {
            "coarse": JusticeEvent,
            "class": Acquit,
            "Adjudicator": "adjudicator",
            "Crime": "crime",
            "Defendant": "defendant",
        },
        "Justice:Appeal": {
            "coarse": JusticeEvent,
            "class": Appeal,
            "Adjudicator": "adjudicator",
            "Crime": "crime",
            "Plaintiff": "prosecutor",
        },
        "Justice:Arrest-Jail": {
            "coarse": JusticeEvent,
            "class": ArrestJail,
            "Agent": "agent",
            "Crime": "crime",
            "Person": "person",
        },
        "Justice:Charge-Indict": {
            "coarse": JusticeEvent,
            "class": ChargeIndict,
            "Adjudicator": "adjudicator",
            "Crime": "crime",
            "Defendant": "defendant",
            "Prosecutor": "prosecutor",
        },
        "Justice:Convict": {
            "coarse": JusticeEvent,
            "class": Convict,
            "Adjudicator": "adjudicator",
            "Crime": "crime",
            "Defendant": "defendant",
        },
        "Justice:Execute": {
            "coarse": JusticeEvent,
            "class": Execute,
            "Agent": "agent",
            "Crime": "crime",
            "Person": "person",
        },
        "Justice:Extradite": {
            "coarse": JusticeEvent,
            "class": Extradite,
            "Agent": "agent",
            "Destination": "destination",
            "Origin": "origin",
            "Person": "person",
        },
        "Justice:Fine": {
            "coarse": JusticeEvent,
            "class": Fine,
            "Adjudicator": "adjudicator",
            "Crime": "crime",
            "Entity": "entity",
            "Money": "money",
        },
        "Justice:Pardon": {
            "coarse": JusticeEvent,
            "class": Pardon,
            "Adjudicator": "adjudicator",
            "Defendant": "defendant",
            "Crime": "crime",
        },
        "Justice:Release-Parole": {
            "coarse": JusticeEvent,
            "class": ReleaseParole,
            "Crime": "crime",
            "Entity": "entity",
            "Person": "person",
        },
        "Justice:Sentence": {
            "coarse": JusticeEvent,
            "class": SentenceAct,
            "Adjudicator": "adjudicator",
            "Crime": "crime",
            "Defendant": "defendant",
            "Sentence": "sentence",
        },
        "Justice:Sue": {
            "coarse": JusticeEvent,
            "class": Sue,
            "Adjudicator": "adjudicator",
            "Crime": "crime",
            "Defendant": "defendant",
            "Plaintiff": "plaintiff",
        },
        "Justice:Trial-Hearing": {
            "coarse": JusticeEvent,
            "class": TrialHearing,
            "Adjudicator": "adjudicator",
            "Crime": "crime",
            "Defendant": "defendant",
            "Prosecutor": "prosecutor",
        },
        "Life:Be-Born": {
            "coarse": LifeEvent,
            "class": BeBorn,
            "Person": "person",
        },
        "Life:Die": {
            "coarse": LifeEvent,
            "class": Die,
            "Agent": "agent",
            "Instrument": "instrument",
            "Person": "victim",
            "Victim": "victim",
        },
        "Life:Divorce": {
            "coarse": LifeEvent,
            "class": Divorce,
            "Person": "person",
        },
        "Life:Injure": {
            "coarse": LifeEvent,
            "class": Injure,
            "Agent": "agent",
            "Instrument": "instrument",
            "Victim": "victim",
        },
        "Life:Marry": {
            "coarse": LifeEvent,
            "class": Marry,
            "Person": "person",
        },
        "Movement:Transport": {
            "coarse": MovementEvent,
            "class": Transport,
            "Agent": "agent",
            "Artifact": "artifact",
            "Destination": "destination",
            "Origin": "origin",
            "Place": "destination",
            "Vehicle": "vehicle",
            "Victim": "artifact",  # MMMmmm WTF
            "Price": "price",
        },
        "Personnel:Elect": {
            "coarse": PersonnelEvent,
            "class": Elect,
            "Entity": "entity",
            "Person": "person",
            "Position": "position",
        },
        "Personnel:End-Position": {
            "coarse": PersonnelEvent,
            "class": EndPosition,
            "Entity": "entity",
            "Person": "person",
            "Position": "position",
        },
        "Personnel:Nominate": {
            "coarse": PersonnelEvent,
            "class": Nominate,
            "Agent": "agent",
            "Person": "person",
            "Position": "position",
        },
        "Personnel:Start-Position": {
            "coarse": PersonnelEvent,
            "class": StartPosition,
            "Entity": "entity",
            "Person": "person",
            "Position": "position",
        },
        "Transaction:Transfer-Money": {
            "coarse": TransactionEvent,
            "class": TransferMoney,
            "Beneficiary": "beneficiary",
            "Giver": "giver",
            "Money": "money",
            "Recipient": "recipient",
        },
        "Transaction:Transfer-Ownership": {
            "coarse": TransactionEvent,
            "class": TransferOwnership,
            "Artifact": "artifact",
            "Beneficiary": "beneficiary",
            "Buyer": "buyer",
            "Price": "price",
            "Seller": "seller",
        },
    }

    def __init__(self, path: str, group_by: str = "sentence", **kwargs) -> None:
        assert group_by in [
            "sentence",
            "document",
        ], "`group_by` must be either 'sentence' or 'document'."

        self.elements = {}

        with open(path, "rt") as in_f:
            for line in in_f:
                line = json.loads(line.strip())

                key = line["sent_id"] if group_by == "sentence" else line["doc_id"]
                if key not in self.elements:
                    self.elements[key] = {
                        "id": key,
                        "doc_id": line["doc_id"],
                        "text": "",
                        "entities": [],
                        "values": [],
                        "coarse_relations": [],
                        "relations": [],
                        "events": [],
                        "arguments": [],
                        "gold": [],
                    }

                entities = [
                    self.ENTITY_TO_CLASS_MAPPING[entity["entity_type"]](span=entity["text"])
                    for entity in line["entity_mentions"]
                    if entity["entity_type"] in self.ENTITY_TO_CLASS_MAPPING
                ]
                values = [
                    self.VALUE_TO_CLASS_MAPPING[entity["entity_type"]](span=entity["text"])
                    for entity in line["entity_mentions"]
                    if entity["entity_type"] in self.VALUE_TO_CLASS_MAPPING
                ]
                coarse_relations, relations = [], []
                for rel in line["relation_mentions"]:
                    if rel["relation_subtype"] not in self.RELATION_TO_CLASS_MAPPING:
                        continue
                    coarse_relations.append(
                        self.RELATION_TO_CLASS_MAPPING[rel["relation_subtype"]][0](
                            arg1=rel["arguments"][0]["text"],
                            arg2=rel["arguments"][1]["text"],
                        )
                    )
                    relations.append(
                        self.RELATION_TO_CLASS_MAPPING[rel["relation_subtype"]][-1](
                            arg1=rel["arguments"][0]["text"],
                            arg2=rel["arguments"][1]["text"],
                        )
                    )
                events, arguments = [], []
                for event in line["event_mentions"]:
                    if event["event_type"] not in self.EVENT_TO_CLASS_MAPPING:
                        continue
                    info = self.EVENT_TO_CLASS_MAPPING[event["event_type"]]
                    _inst = {param: [] for param in inspect.signature(info["class"]).parameters.keys()}
                    _inst["mention"] = event["trigger"]["text"]
                    for argument in event["arguments"]:
                        if argument["role"] in info:
                            name = info[argument["role"]]
                            _inst[name].append(argument["text"])
                        elif argument["role"] in self._EVENT_CONSTANTS_MAPPING:
                            name = self._EVENT_CONSTANTS_MAPPING[argument["role"]]
                            if name not in _inst:
                                continue
                            _inst[name].append(argument["text"])
                        else:
                            raise ValueError(f"Argument {event['event_type']}:{argument['role']} not found!")

                    events.append(info["coarse"](mention=_inst["mention"]))
                    arguments.append(info["class"](**_inst))

                self.elements[key]["text"] += " " + line["sentence"].strip()
                self.elements[key]["entities"] += entities
                self.elements[key]["values"] += values
                self.elements[key]["coarse_relations"] += coarse_relations
                self.elements[key]["relations"] += relations
                self.elements[key]["events"] += events
                self.elements[key]["arguments"] += arguments
                self.elements[key]["gold"] += entities  # Is not used anyway


class ACESampler(Sampler):
    """
    A data `Sampler` for the ACE05 dataset.

    Args:
        dataset_loader (`ACEDatasetLoader`):
            The dataset loader that contains the data information.
        task (`str`, optional):
            The task to sample. It must be one of the following: NER, VER, RE, EE.
            Defaults to `None`.
        split (`str`, optional):
            The split to sample. It must be one of the following: "train", "dev" or
            "test". Depending on the split the sampling strategy differs. Defaults to
            `"train"`.
        parallel_instances (`Union[int, Tuple[int, int]]`, optional):
            The number of sentences sampled in parallel. Options:

                * **`int`**: The amount of elements that will be sampled in parallel.
                * **`tuple`**: The range of elements that will be sampled in parallel.

            Defaults to 1.
        max_guidelines (`int`, optional):
            The number of guidelines to append to the example at the same time. If `-1`
            is given then all the guidelines are appended. Defaults to `-1`.
        guideline_dropout (`float`, optional):
            The probability to dropout a guideline definition for the given example. This
            is only applied on training. Defaults to `0.0`.
        seed (`float`, optional):
            The seed to sample the examples. Defaults to `0`.
        prompt_template (`str`, optional):
            The path to the prompt template. Defaults to `"templates/prompt.txt"`.
        ensure_positives_on_train (bool, optional):
            Whether to ensure that the guidelines of annotated examples are not removed.
            Defaults to `False`.
        dataset_name (str, optional):
            The name of the dataset. Defaults to `None`.
        scorer (`str`, optional):
           The scorer class import string. Defaults to `None`.
        sample_only_gold_guidelines (`bool`, optional):
            Whether to sample only guidelines of present annotations. Defaults to `False`.
    """

    def __init__(
        self,
        dataset_loader: ACEDatasetLoader,
        task: str = None,
        split: str = "train",
        parallel_instances: Union[int, Tuple[int, int]] = 1,
        max_guidelines: int = -1,
        guideline_dropout: float = 0.0,
        seed: float = 0,
        ensure_positives_on_train: bool = False,
        dataset_name: str = None,
        scorer: str = None,
        sample_only_gold_guidelines: bool = False,
        **kwargs,
    ) -> None:
        assert task in [
            "NER",
            "VER",
            "RE",
            "RC",
            "EE",
            "EAE",
        ], f"{task} must be either 'NER', 'VER', 'RE', 'RC', 'EE', 'EAE'."

        task_definitions, task_target, task_template = {
            "NER": (ENTITY_DEFINITIONS, "entities", "templates/prompt.txt"),
            "VER": (VALUE_DEFINITIONS, "values", "templates/prompt.txt"),
            "RE": (COARSE_RELATION_DEFINITIONS, "coarse_relations", "templates/prompt_ace_re.txt"),
            "RC": (RELATION_DEFINITIONS, "relations", "templates/prompt_ace_rc.txt"),
            "EE": (COARSE_EVENT_DEFINITIONS, "events", "templates/prompt.txt"),
            "EAE": (EVENT_DEFINITIONS, "arguments", "templates/prompt_ace_eae.txt"),
        }[task]

        if task in ["RC", "EAE"]:
            is_coarse_to_fine: bool = True
            COARSE_TO_FINE = COARSE_TO_FINE_EVENTS if task == "EAE" else COARSE_TO_FINE_RELATIONS
            FINE_TO_COARSE = FINE_TO_COARSE_EVENTS if task == "EAE" else FINE_TO_COARSE_RELATIONS
        else:
            is_coarse_to_fine = False
            COARSE_TO_FINE = None
            FINE_TO_COARSE = None

        kwargs.pop("prompt_template")

        super().__init__(
            dataset_loader=dataset_loader,
            task=task,
            split=split,
            parallel_instances=parallel_instances,
            max_guidelines=max_guidelines,
            guideline_dropout=guideline_dropout,
            seed=seed,
            prompt_template=task_template,
            ensure_positives_on_train=ensure_positives_on_train,
            sample_only_gold_guidelines=sample_only_gold_guidelines,
            dataset_name=dataset_name,
            scorer=scorer,
            task_definitions=task_definitions,
            task_target=task_target,
            is_coarse_to_fine=is_coarse_to_fine,
            coarse_to_fine=COARSE_TO_FINE,
            fine_to_coarse=FINE_TO_COARSE,
            definitions=GUIDELINES,
            examples=EXAMPLES,
            **kwargs,
        )

src/tasks/ace/data_loader.py
hitz-zentroa/GoLLIE
GUIDELINES = {
    "ace_gpe": {
        "en": [
            (
                "Geo-Political Entities are composite entities comprised of a population, a government, a physical"
                " location, and a \nnation (or province, state, country, city, etc.)."
            ),
            (
                "Geo-Political Entities are complex systems that consist of a group of people, a governing body, a"
                " specific \nterritory, and a nation (or a province, state, country, city, and so on)."
            ),
            (
                "Geo-Political Entities are comprehensive entities that consist of a population, a governing body, a"
                " specific \nlocation, and a nation (or an equivalent regional division such as a province, state,"
                " country, or city)."
            ),
            (
                "Geo-Political Entities are comprehensive entities that encompass a population, a government, a"
                " specific \nlocation, and a nation (or an equivalent administrative division such as a province,"
                " state, country, or city)."
            ),
            (
                "Geo-Political Entities are complex entities that encompass a population, a government, a specific"
                " location, and a \nnation (or a province, state, country, city, and so on)."
            ),
        ]
    },
    "ace_acquit": {
        "en": [
            (
                "An Acquit Event occurs whenever a trial ends but fails to produce a conviction. This will include"
                " cases where the \ncharges are dropped by the Prosecutor."
            ),
            (
                "A Acquit Event takes place whenever a trial concludes without resulting in a conviction. This"
                " encompasses \nsituations in which the prosecutor chooses to drop the charges."
            ),
            (
                "A Acquit Event transpires whenever a trial concludes without resulting in a conviction. This"
                " encompasses \ninstances in which the prosecutor chooses to drop the charges."
            ),
            (
                "A Dismissal Occurrence takes place whenever a judicial proceeding concludes without resulting in a"
                " conviction. \nThis encompasses situations in which the prosecutor chooses to drop the charges."
            ),
            (
                "A Acquit Event transpires whenever a judicial proceeding concludes without resulting in a conviction."
                " Such \ninstances encompass situations in which the prosecutor chooses to abandon the charges."
            ),
        ]
    },
    "ace_agentartifactrelationrelation": {
        "en": [
            (
                "The AgentArtifact Relation applies when an agent owns an artifact, has possession of an artifact,"
                " uses an artifact, \nor caused an artifact to come into being. Note: if the `arg2` is an"
                " Organization, use OrganizationAffiliation when \n`arg1` is a Person or PartWhole when `arg1` is an"
                " Organization or GPE."
            ),
            (
                "The AgentArtifact Relation is relevant when an agent exercises control over an artifact, holds an"
                " artifact, \nemploys an artifact, or brings an artifact into existence. It's important to note that"
                " when the `arg2` is an Organization, \nutilize the OrganizationAffiliation relation if `arg1` is a"
                " Person, or the PartWhole relationship if `arg1` is an \nOrganization or GPE."
            ),
            (
                "The AgentArtifact Relation is relevant when an agent holds, utilizes, or creates an artifact, or when"
                " the agent is \nlinked to the artifact in some other way. It's important to note that if the `arg2`"
                " is an Organization, the relationship \nshould be categorized as OrganizationAffiliation when `arg1`"
                " is a Person, or PartWhole when `arg1` is an Organization or \nGPE."
            ),
            (
                "The AgentArtifact Relation is applicable when an agent controls an artifact, holds an artifact,"
                " utilizes an \nartifact, or brings an artifact into existence. It's important to note that when the"
                " second argument 'arg2' is associated with \nan Organization, the relationship type should be changed"
                " to OrganizationAffiliation when the first argument 'arg1' \nis a Person, or to PartWhole when 'arg1'"
                " is an Organization or GPE (Gazetteer of Planetary Nomenclature)."
            ),
            (
                "The AgentArtifact Relation is relevant when an agent exercises control over an artifact, holds an"
                " artifact, \nutilizes an artifact, or brings an artifact into existence. It's important to note that"
                " when the `arg2` is an Organization, \nthe OrganizationAffiliation should be used if `arg1` is a"
                " Person, or the PartWhole relationship should be employed if \n`arg1` is an Organization or GPE."
            ),
        ]
    },
    "ace_appeal": {
        "en": [
            "An Appeal Event occurs whenever the decision of a court is taken to a higher court for review.",
            (
                "An Appeal Event happens whenever a judgment rendered by a lower court is brought before a higher"
                " court for \nreevaluation."
            ),
            (
                "An Appeal Event transpires whenever a judicial decision is brought before a higher court for"
                " reevaluation."
            ),
            (
                "An Appeal Event transpires whenever a judicial decision is conveyed to a more elevated court for"
                " reconsideration."
            ),
            (
                "An Appeal Event transpires whenever a judicial decision is conveyed to a superior court for"
                " reevaluation."
            ),
        ]
    },
    "ace_arrestjail": {
        "en": [
            (
                "A Jail Event occurs whenever the movement of a Person is constrained by a state actor (a GPE, its"
                " Organization \nsubparts, or its Person representatives)."
            ),
            (
                "A Jail Event takes place whenever the movement of an individual, defined as a Person, is limited or"
                " restricted by a \nstate representative, which can include a Government, its affiliated"
                " Organizations, or the individuals acting on their \nbehalf."
            ),
            (
                "A Jail Event takes place whenever the movement of an individual, defined as a Person, is restricted"
                " by a \nstate-affiliated entity, which can include a Governmental or Private Entity (GPE), any of its"
                " organizational subdivisions, or its \nrepresentative individuals."
            ),
            (
                "A Jail Event takes place whenever the movement of an individual, defined as a Person, is restricted"
                " by a state agent, \nwho can be any part of the Government, Public Sector, or a related Organization,"
                " or a person acting on their behalf."
            ),
            (
                "A Jail Event takes place whenever the movement of an individual, defined as a Person, is restricted"
                " by a state agent, \nwhich can include a Governmental or Private Entity, any of its organizational"
                " subdivisions, or its representatives who \nare themselves Persons."
            ),
        ]
    },
    "ace_attack": {
        "en": [
            (
                "An Attack Event is defined as a violent physical act causing harm or damage. Attack Events include"
                " any such Event not \ncovered by the Injure or Die subtypes, including Events where there is no"
                " stated agent."
            ),
            (
                "An Assault Incident is characterized as a forceful physical act resulting in injury or destruction."
                " Assault \nIncidents encompass any occurrence not classified under the Injure or Perish subtypes, as"
                " well as instances where the \nresponsible party is not explicitly identified."
            ),
            (
                "An Assault Occurrence is characterized as a hostile physical act resulting in injury or destruction."
                " Assault \nOccurrences encompass all events not classified under the Injure or Perish subtypes, as"
                " well as instances where the instigator \nis not explicitly identified."
            ),
            (
                "An Assault Incident is characterized as a hostile physical act resulting in injury or destruction."
                " Assault \nIncidents encompass all events not encompassed by the Injure or Die subtypes, as well as"
                " instances where the instigator is not \nexplicitly identified."
            ),
            (
                "An Assault Incident is characterized as a hostile physical act resulting in injury or destruction."
                " Assault \nIncidents encompass all events not classified under the Injure or Die subtypes, as well as"
                " instances where the responsible \nparty is not explicitly identified."
            ),
        ]
    },
    "ace_beborn": {
        "en": [
            "A BeBorn Event occurs whenever a Person Entity is given birth to.",
            "A BeBorn Event takes place whenever a new Person Entity comes into existence.",
            "A BeBorn Event takes place whenever a new Person Entity comes into existence.",
            "A BeBorn Event takes place whenever a new Person Entity comes into existence.",
            (
                "A BeBorn Event takes place whenever a new Person Entity comes into existence through the process of"
                " birth."
            ),
        ]
    },
    "ace_business": {
        "en": [
            (
                "The Business Relation captures the connection between two entities in any professional relationship."
                " Both \narguments must be entities of type Person."
            ),
            (
                "The Business Relation highlights the association between two parties in any professional connection."
                " Both \nfactors need to be entities of the Person type."
            ),
            (
                "The Business Relation component encapsulates the association between two entities involved in any"
                " professional \nconnection. Both elements must be of the Person class type."
            ),
            (
                "The Business Relation highlights the association between two parties in any professional connection."
                " Both \nelements need to be entities of the Person type."
            ),
            (
                "The Business Relation highlights the association between two parties in any professional connection."
                " Both \nelements involved must be of the Person entity type."
            ),
        ]
    },
    "ace_businessevent": {
        "en": [
            (
                "A BusinessEvent refers to actions related to Organizations such as: creating, merging, declaring"
                " bankruptcy or \nending organizations (including government agencies)."
            ),
            (
                "A BusinessEvent encompasses actions connected to Organizations, such as their establishment,"
                " consolidation, \ninsolvency proceedings, or cessation (including government entities)."
            ),
            (
                "A BusinessEvent entails actions connected to Organizations, including: initiation, consolidation,"
                " insolvency \npronouncements, and cessation (encompassing government entities)."
            ),
            (
                "A BusinessEvent encompasses activities associated with Organizations, including: establishment,"
                " \nconsolidation, insolvency proceedings, and cessation (encompassing government entities)."
            ),
            (
                "A BusinessEvent encompasses actions involving Organizations, such as their establishment,"
                " consolidation, \ninsolvency proceedings, or cessation (including government entities)."
            ),
        ]
    },
    "ace_chargeindict": {
        "en": [
            (
                "A Charge Event occurs whenever a Person, Organization or GPE is accused of a crime by a state actor"
                " (GPE, an \nOrganization subpart of a GPE or a Person representing a GPE)."
            ),
            (
                "A Charge Event transpires whenever an individual, organization, or government-related entity is"
                " accused of an \noffense by a state representative, encompassing a government-related entity, an"
                " organization subordinate to a \ngovernment-related entity, or an individual acting on behalf of a"
                " government-related entity."
            ),
            (
                "A Charge Event transpires whenever an individual, organization, or government-related entity is"
                " accused of an \noffense by a state representative, such as another government-related entity, an"
                " organization subordinate to a \ngovernment-related entity, or an individual acting on behalf of a"
                " government-related entity."
            ),
            (
                "A Charge Event takes place whenever an individual, organization, or government-led entity is accused"
                " of a crime by a \nstate representative, such as another government-led entity, an organization"
                " subordinate to a government-led entity, \nor an individual acting on behalf of a government-led"
                " entity."
            ),
            (
                "A Charge Event takes place whenever an individual, organization, or government-led entity is accused"
                " of a \nviolation by a state-affiliated party, such as another government-led entity, a subsidiary"
                " organization of that entity, or \nan individual acting on behalf of a government-led entity."
            ),
        ]
    },
    "ace_citizenresidentreligionethnicity": {
        "en": [
            (
                "CitizenResidentReligionEthnicity describes the relation between a Person entity and (1) the GPE in"
                " which they \nhave citizenship, (2) the GPE or Location in which they live, the religious"
                " Organization or Person entity with which they \nhave affiliation and (3) the GPE or PER entity that"
                " indicates their ethnicity."
            ),
            (
                "CitizenResidentReligionEthnicity illustrates the connection between a Person entity and (1) the GPE"
                " in which \nthey possess citizenship, (2) the GPE or Location where they reside, the religious"
                " Organization or Person entity they \nare associated with, and (3) the GPE or PER entity signifying"
                " their ethnicity."
            ),
            (
                "The text describes the connection between a Person entity and four distinct aspects: (1) the GPE"
                " where they possess \ncitizenship, (2) the GPE or Location where they reside, (3) the religious"
                " Organization or Person entity they are associated with, \nand (4) the GPE or PER entity signifying"
                " their ethnicity."
            ),
            (
                "CitizenResidentReligionEthnicity illustrates the connection between a Person entity and (1) the"
                " \nGeographicPointEntity in which they possess citizenship, (2) the GeographicPointEntity or Location"
                " where they reside, the religious \nOrganization or Person entity they are associated with, and (3)"
                " the GeographicPointEntity or Person entity signifying their \nethnicity."
            ),
            (
                "The text describes the connection between a Person entity and four distinct aspects: (1) the"
                " \nGeographicPointEntity (GPE) in which they possess citizenship, (2) the GPE or Location where they"
                " reside, (3) the religious Organization \nor Person entity that they are associated with, and (4) the"
                " GPE or Person entity that signifies their ethnicity."
            ),
        ]
    },
    "ace_conflictevent": {
        "en": [
            (
                "A ConflictEvent refers to either violent physical acts causing harm or damage, but are not covered by"
                " Life events \n(conflicts, clashes, fighting, gunfire, ...) or demonstrations (protests, sit-ins,"
                " strikes, riots, ...)."
            ),
            (
                "A ConflictEvent entails either destructive physical actions that inflict harm or destruction, but it"
                " does not \nencompass events classified as Life events (such as conflicts, clashes, fighting,"
                " gunfire, and so on) or demonstrations \n(including protests, sit-ins, strikes, riots, and the like)."
            ),
            (
                "A ConflictEvent entails either violent physical actions resulting in harm or destruction, but it does"
                " not \nencompass events classified as Life events (such as conflicts, clashes, fighting, gunfire, and"
                " so on) or demonstrations \n(like protests, sit-ins, strikes, riots, and the like)."
            ),
            (
                "A ConflictEvent entails either violent physical actions resulting in injury or destruction, but it"
                " does not \nencompass situations categorized as life events (such as disputes, disagreements,"
                " altercations, gunfire, etc.) or \ndemonstrations (including protests, sit-ins, strikes, riots, and"
                " so on)."
            ),
            (
                "A ConflictEvent entails either destructive physical actions that inflict harm or damage, but it does"
                " not encompass \nsituations categorized as life events (such as disputes, disagreements,"
                " altercations, gunfire, etc.) or demonstrations \n(including protests, sit-ins, strikes, riots, and"
                " so on)."
            ),
        ]
    },
    "ace_contactevent": {
        "en": [
            (
                "A ContactEvent occurs whenever two or more entities (persons or organization's representatives) come"
                " together at \na single location and interact with one another face-to-face or directly enages in"
                " discussion via written or \ntelephone communication."
            ),
            (
                "A ContactEvent takes place whenever two or more individuals or organizational representatives gather"
                " at a single \nplace and communicate with each other in person or through written or telephone"
                " dialogue."
            ),
            (
                "A ContactEvent takes place whenever two or more entities, such as individuals or organization"
                " representatives, \ngather at a single location and interact with each other in person or through"
                " direct communication via written or telephone \nmeans."
            ),
            (
                "A ContactEvent takes place whenever two or more individuals or organizational representatives gather"
                " at a single \nlocation and communicate with each other in person or through written or telephone"
                " communication."
            ),
            (
                "A ContactEvent takes place whenever two or more entities, such as individuals or organizational"
                " representatives, \ngather at a single location and communicate with each other in person or through"
                " direct written or telephone dialogue."
            ),
        ]
    },
    "ace_contactinfo": {
        "en": [
            (
                "A ContactInfo value refers to contact information values such as telephone numbers, emails,"
                " addresses. For \nexample: 'mich...@sumptionandwyland.com', ..."
            ),
            (
                "A ContactInfo entity encompasses various types of contact data, including phone numbers, email"
                " addresses, and \nmailing addresses. For instance,"
                " '[mich...@sumptionandwyland.com](mailto:mich...@sumptionandwyland.com)', and \nadditional examples."
            ),
            (
                "A ContactInfo entity encompasses various contact information entries, like phone numbers, email"
                " addresses, and \nmailing addresses. For instance,"
                " '[mich...@sumptionandwyland.com](mailto:mich...@sumptionandwyland.com)', \namong others."
            ),
            (
                "A ContactInfo object encompasses various contact information entries, such as phone numbers, email"
                " addresses, \nand mailing addresses. For instance,"
                " \n'[mich...@sumptionandwyland.com](mailto:mich...@sumptionandwyland.com)', ..."
            ),
            (
                "A ContactInfo entity encompasses various contact details, including phone numbers, email addresses,"
                " and mailing \naddresses. For instance,"
                " '[mich...@sumptionandwyland.com](mailto:mich...@sumptionandwyland.com)', among others."
            ),
        ]
    },
    "ace_convict": {
        "en": [
            (
                "A Convict Event occurs whenever a Try Event ends with a successful prosecution of the Defendant. In"
                " other words, a \nPerson, Organization or GPE Entity is convicted whenever that Entity has been found"
                " guilty of a Crime."
            ),
            (
                "A Convict Event takes place whenever a Try Event results in a successful prosecution, meaning that a"
                " person, \norganization, or general purpose entity is found guilty of a crime."
            ),
            (
                "A Convict Event takes place whenever a Try Event results in a successful prosecution, meaning that a"
                " Person, \nOrganization, or GPE Entity is convicted whenever that Entity has been judged guilty of a"
                " Crime."
            ),
            (
                "A Convict Event takes place whenever a Try Event results in a successful prosecution, meaning that a"
                " Person, \nOrganization, or GPE Entity is convicted whenever that Entity has been proven guilty of a"
                " crime."
            ),
            (
                "A Convict Event takes place whenever a Try Event reaches a successful conclusion with the prosecution"
                " of a \nDefendant. To put it another way, a Person, Organization, or Governmental/Public Entity is"
                " declared guilty and convicted when \nthey have been found to have committed a Crime."
            ),
        ]
    },
    "ace_crime": {
        "en": [
            (
                "A Crime value refers to the specific reason (crime) that a Person entity can be judged or sentenced"
                " for. For example: \n'raping', 'murder', 'drug', ..."
            ),
            (
                "A Crime value represents the particular criminal act (crime) for which a Person entity may be"
                " penalized or \nconvicted. Some examples include 'rape', 'murder', 'drug-related offense', and so on."
            ),
            (
                "A Crime value denotes the particular reason (crime) for which a Person entity may be subjected to"
                " judgment or \nsentencing. For example, 'rape', 'murder', 'drug-related offense', and so on."
            ),
            (
                "A Crime value represents the particular crime (referred to as 'crime' in the Person entity) that"
                " leads to a judgment \nor sentencing. Some examples include 'rape', 'murder', 'drug-related offense',"
                " and so on."
            ),
            (
                "A Crime value denotes the particular crime (offense) for which a Person entity may be convicted or"
                " penalized. \nExamples of such crimes include 'rape', 'murder', 'drug-related offenses', and so on."
            ),
        ]
    },
    "ace_declarebankruptcy": {
        "en": [
            (
                "A DeclareBankruptcy Event will occur whenever an Entity officially requests legal protection from"
                " debt \ncollection due to an extremely negative balance sheet."
            ),
            (
                "A DeclareBankruptcy Event will transpire whenever an Entity formally applies for legal relief from"
                " debt \ncollection as a result of an extremely unfavorable financial situation."
            ),
            (
                'A "DeclareBankruptcy Event" will take place whenever an "Entity" officially seeks a legal exemption'
                " from debt \ncollection as a result of an overwhelmingly unfavorable financial situation."
            ),
            (
                "A DeclareBankruptcy Event will transpire whenever an Entity formally applies for legal relief from"
                " debt \ncollection as a result of an extremely unfavorable balance sheet."
            ),
            (
                'A "DeclareBankruptcy Event" will take place whenever a "Entity" formally applies for legal relief'
                ' from debt \ncollection as a result of an "extremely negative balance sheet.'
            ),
        ]
    },
    "ace_demonstrate": {
        "en": [
            (
                "A Demonstrate Event occurs whenever a large number of people come together in a public area to"
                " protest or demand some \nsort of official action. Demonstrate Events include, but are not limited"
                " to, protests, sit-ins, strikes, and riots."
            ),
            (
                "A Demonstrate Event takes place whenever a sizable gathering of individuals congregates in a public"
                " space to \nexpress their dissatisfaction or call for governmental intervention. Such events"
                " encompass a wide range of activities, \nsuch as protests, sit-ins, strikes, and riots."
            ),
            (
                "A Demonstrate Event takes place whenever a considerable gathering of individuals congregates in a"
                " public space to \nexpress their dissent or call for governmental action. Demonstrate Events"
                " encompass a variety of events, such as protests, \nsit-ins, strikes, and riots."
            ),
            (
                "A Demonstrate Event takes place whenever a sizable gathering of individuals congregates in a public"
                " space to \nexpress their dissatisfaction or call for specific governmental action. Demonstrate"
                " Events encompass a variety of \nactions, such as protests, sit-ins, strikes, and riots."
            ),
            (
                "A Show of Unity Event takes place whenever a sizable gathering of individuals congregates in a"
                " communal space to \nexpress their dissatisfaction or advocate for specific government action. Such"
                " events encompass a variety of forms, such \nas protests, sit-ins, strikes, and riots."
            ),
        ]
    },
    "ace_die": {
        "en": [
            (
                "A Die Event occurs whenever the life of a Person Entity ends. Die Events can be accidental,"
                " intentional or \nself-inflicted"
            ),
            (
                "A Die Event takes place whenever the existence of a Person Entity comes to an end. These events can"
                " be due to \naccidental, intentional, or self-inflicted causes."
            ),
            (
                "A Die Event transpires whenever the existence of a Person Entity comes to an end. Such events can be"
                " due to accidental, \nintentional, or self-inflicted causes."
            ),
            (
                "A Die Event takes place whenever the life of a Person Entity comes to an end. Such events can be the"
                " result of accidents, \nintentional actions, or self-inflicted incidents."
            ),
            (
                "A Die Event takes place whenever the existence of a Person Entity comes to an end. These events can"
                " be the result of \naccidents, intentional actions, or self-inflicted incidents."
            ),
        ]
    },
    "ace_divorce": {
        "en": [
            (
                "A Divorce Event occurs whenever two people are officially divorced under the legal definition of"
                " divorce. We do not \ninclude separations or church annulments."
            ),
            (
                "A Divorce Event takes place whenever a legal decree of divorce is issued, excluding any instances of"
                " separation or \nreligious annulments."
            ),
            (
                "A divorce event takes place whenever a legal dissolution of marriage occurs between two individuals,"
                " excluding any \ninstances of separation or religious annulments."
            ),
            (
                "A Divorce Event takes place whenever a legal decree of divorce is issued, excluding any instances of"
                " separation or \nreligious annulments."
            ),
            (
                "A Divorce Event takes place whenever a legal decree of divorce is issued, excluding any instances of"
                " separation or \nreligious annulments."
            ),
        ]
    },
    "ace_elect": {
        "en": [
            (
                "An Elect Event occurs whenever a candidate wins an election designed to determine the Person argument"
                " of a \nStartPosition Event."
            ),
            (
                "An Elect Event takes place whenever a candidate emerges victorious in an election that aims to"
                " establish the Person \nrole within a StartPosition Event."
            ),
            (
                "An Elect Event takes place whenever a candidate emerges victorious in an election that aims to"
                " establish the Person \nrole within a StartPosition Event."
            ),
            (
                "An Elect Event takes place whenever a candidate is victorious in an election that is intended to"
                " establish the Person \nrole within a StartPosition Event."
            ),
            (
                "An Elect Event takes place whenever a candidate is victorious in an election that is intended to"
                " establish the Person \nrole within a StartPosition Event."
            ),
        ]
    },
    "ace_employment": {
        "en": [
            (
                "Employment captures the relationship between Persons and their employers. This Relation is only"
                " taggable when it \ncan be reasonably assumed that the Person is paid by the ORG or GPE."
            ),
            (
                "Employment highlights the connection between Individuals and their employers. This association can"
                " only be \nconsidered relevant when it's fairly certain that the Person receives payment from the"
                " Organization or Group."
            ),
            (
                "Employment highlights the connection between Individuals and their employers. This association is"
                " only \nconsidered relevant when it can be fairly presumed that the Person receives remuneration from"
                " the Organization or Business \nEntity."
            ),
            (
                "Employment connections illustrate the association between Individuals and their employers. This"
                " association \ncan only be considered relevant when it's fairly certain that the Individual is"
                " receiving payment from the ORG or GPE."
            ),
            (
                "Employment establishes the connection between Persons and their employers. This association is only"
                " considered \nrelevant when it can be fairly presumed that the Person receives payment from the ORG"
                " or GPE."
            ),
        ]
    },
    "ace_endorg": {
        "en": [
            "An EndOrg Event occurs whenever an Organization ceases to exist (in other words 'goes out of business').",
            (
                "An EndOrg Occurrence transpires whenever an Entity ceases to continue operating (in simpler terms,"
                ' "closes \ndown").'
            ),
            (
                "An EndOrg Event takes place whenever an Organization comes to an end, meaning it has gone out of"
                " business."
            ),
            (
                "An EndOrg Event takes place whenever an Organization comes to an end (in simpler terms, when it shuts"
                " down)."
            ),
            "An EndOrg Event takes place whenever a business entity comes to an end and ceases its operations.",
        ]
    },
    "ace_endposition": {
        "en": [
            (
                "An EndPosition Event occurs whenever a Person Entity stops working for (or changes offices within) an"
                " Organization \nor GPE."
            ),
            (
                "A EndPosition Event is triggered whenever an individual associated with a Person Entity ceases to"
                " work for, or \ntransfers between departments within, an Organization or Governmental/Public Entity."
            ),
            (
                "An EndPosition Event takes place whenever a Person Entity finishes working for (or transitions"
                " between \ndepartments within) an Organization or Governmental/Public Entity."
            ),
            (
                "A EndPosition Event takes place whenever a Person Entity concludes their work for (or transfers"
                " between \ndepartments within) an Organization or Governmental/Public Entity."
            ),
            (
                "A EndPosition Event takes place whenever a Person Entity finishes their work for (or transfers"
                " between departments \nwithin) an Organization or GPE."
            ),
        ]
    },
    "ace_execute": {
        "en": [
            (
                "An Execute Event occurs whenever the life of a Person is taken by a state actor (a GPE, its"
                " Organization subparts, or \nPerson representatives)."
            ),
            (
                "An Execute Event takes place whenever the life of an individual, which is a Person entity, is"
                " terminated by a state \nactor, who is a part of a Government, Public Enterprise, or an Organization"
                " along with its representatives."
            ),
            (
                "An Execute Event takes place whenever the life of an individual, which is governed by a state entity"
                " or one of its \nsubsidiaries, is terminated by a person who holds a position of authority within"
                " that organization."
            ),
            (
                "An Execute Event takes place whenever the life of an individual is terminated by a state"
                " representative, which can \ninclude a Governmental or Political Entity, its affiliated"
                " organizations, or individuals acting on their behalf."
            ),
            (
                "An Execute Event takes place whenever the life of an individual is terminated by a state"
                " representative, \nencompassing entities such as a Governing Power Entity, its affiliated"
                " organizations, or authorized individual \nrepresentatives."
            ),
        ]
    },
    "ace_extradite": {
        "en": [
            (
                "An Extradite Event occurs whenever a Person is sent by a state actor from one Place (normally the GPE"
                " associated with \nthe state actor, but sometimes a Facility under its control) to another place"
                " (Location, GPE or Facility) for the \npurposes of legal proceedings there."
            ),
            (
                "An Extradite Occurrence takes place whenever a person is transferred by a government official from"
                " one specific \nlocation (usually the government-associated location, but occasionally a controlled"
                " facility) to another location (such \nas a place, government-connected location, or controlled"
                " facility) for the purpose of engaging in legal actions \nthere."
            ),
            (
                "An Extradite Occurrence transpires whenever an individual is transported by a government official"
                " from one \nspecific location (usually the territory affiliated with the government agent, but"
                " occasionally a facility managed by \nthem) to another location (Location, GPE, or Facility) for the"
                " intention of engaging in legal processes there."
            ),
            (
                "An Extradite Event takes place whenever a government official transports an individual from one"
                " designated area \n(usually the government-related location, but occasionally a controlled facility)"
                " to another location (such as a court or \ndetention center) for the purpose of legal proceedings."
            ),
            (
                "An Extradite Occurrence takes place whenever an individual is transferred by a governmental entity"
                " from one \nspecific location (usually the territory connected to the governmental entity, but"
                " occasionally a facility under its \njurisdiction) to another location (Location, GPE, or Facility)"
                " for the purpose of engaging in legal proceedings there."
            ),
        ]
    },
    "ace_facility": {
        "en": [
            (
                "A facility is a functional, primarily man-made structure. These include buildings and similar"
                " facilities \ndesigned for human habitation, such as houses, factories, stadiums, office buildings,"
                " ... Roughly speaking, facilities \nare artifacts falling under the domains of architecture and civil"
                " engineering."
            ),
            (
                "A facility is a practical, predominantly human-constructed structure, encompassing buildings and"
                " analogous \nfacilities intended for human inhabitation, such as residential homes, industrial"
                " plants, sports stadiums, and office \ncomplexes. In essence, facilities can be considered"
                " architectural and civil engineering creations."
            ),
            (
                "A facility is a practical, predominantly human-created structure, encompassing buildings and"
                " analogous \nfacilities intended for human inhabitation, such as residential homes, manufacturing"
                " plants, sports stadiums, and office \ncomplexes. In essence, facilities comprise objects within the"
                " realms of architecture and civil engineering."
            ),
            (
                "A facility is a practical, predominantly human-crafted structure, encompassing buildings and"
                " analogous \nfacilities intended for human inhabitation, such as dwellings, manufacturing plants,"
                " sports arenas, office complexes, and \nso forth. In essence, facilities comprise artifacts that"
                " belong to the realms of architecture and civil engineering."
            ),
            (
                "A facility is a practical, predominantly human-created structure, encompassing buildings and"
                " analogous \nhuman-occupied spaces like homes, factories, sports arenas, office buildings, and more."
                " In essence, facilities are creations \nthat belong to the realms of architecture and civil"
                " engineering."
            ),
        ]
    },
    "ace_family": {
        "en": [
            (
                "The Family Relation captures the connection between one entity and another with which it is in any"
                " familial \nrelationship. Both arguments must be entities of type Person."
            ),
            (
                "The Family Relation depicts the association between two entities, one of which is of the Person type,"
                " and the other \nwith which it holds any familial connection."
            ),
            (
                "The Family Relation embodies the association between two entities, one of which is linked through any"
                " familial \nconnection. Both entities must be of the Person type."
            ),
            (
                "The Family Relation encapsulates the association between two entities, either of which is a person"
                " and maintains a \nfamilial connection with the other."
            ),
            (
                "The Family Relation represents the association between two entities, one of which must be a person,"
                " and the other \nwith whom they have a familial connection."
            ),
        ]
    },
    "ace_fine": {
        "en": [
            (
                "A Fine Event takes place whenever a state actor issues a financial punishment to a GPE, Person or"
                " Organization \nEntity, typically as a result of court proceedings."
            ),
            (
                "A Fine Event occurs whenever a state entity imposes a financial penalty on a Goods, People, or Entity"
                " (GPE), usually \nfollowing a legal proceeding."
            ),
            (
                "A Fine Event occurs whenever a state entity imposes a financial penalty on a General Person or"
                " Entity, usually \nfollowing a legal proceeding."
            ),
            (
                "A Fine Event occurs whenever a state actor imposes a financial penalty on a Governmental, Person, or"
                " Organization \nEntity, usually following a court decision."
            ),
            (
                "A Fine Event occurs whenever a state entity imposes a financial penalty on a Goods, Person, or"
                " Organization Entity, \nusually following a court-ordered decision."
            ),
        ]
    },
    "ace_founder": {
        "en": [
            (
                "Founder captures the relationship between an agent (Person, Organization, or GPE) and an Organization"
                " or GPE \nestablished or set up by that agent."
            ),
            (
                "The founder captures the connection between an agent, whether it's a Person, Organization, or Group"
                " of Persons, and \nan Organization or Group of Persons created or established by that agent."
            ),
            (
                "Founder encapsulates the connection between an agent (Person, Organization, or GPE) and an"
                " Organization or GPE \nthat has been conceived or instituted by that agent."
            ),
            (
                "Founder encapsulates the connection between an agent (Person, Organization, or GPE) and the"
                " Organization or GPE \nthat was created or established by that agent."
            ),
            (
                "The founder captures the connection between an agent, whether a Person, Organization, or Group of"
                " Persons, and an \nOrganization or Group of Persons created or established by that agent."
            ),
        ]
    },
    "ace_genaffiliationrelation": {
        "en": [
            (
                "The GenAffiliation Relation describes the citizen, resident, religion or ethnicity relation when the"
                " `arg1` is a \nPerson. When the `arg1` is an Organization, the relation describes where it is"
                " located, based or does business."
            ),
            (
                "The GenAffiliation Relation refers to the connection between an individual and their association with"
                " a \nparticular group, such as citizenship, residency, religious affiliation, or ethnicity, when the"
                " first argument (arg1) is a \nPerson. When the arg1 is an Organization, the relation highlights the"
                " organization's location, base, or areas of business \noperation."
            ),
            (
                "The GenAffiliation Relation refers to the connection between an individual and their associated"
                " \ncharacteristics, such as citizenship, residency, religious affiliation, or ethnicity, when the"
                " first argument (arg1) is a Person. \nWhen the arg1 is an Organization, the relation highlights the"
                " organization's location, base, or areas of operation."
            ),
            (
                "The GenAffiliation Relation conveys the connection between an individual and their association with"
                " regards to \ncitizenship, residency, religious affiliation, or ethnicity when the first argument,"
                " `arg1`, is a Person. However, when `arg1` \nis an Organization, the relation illustrates where it"
                " operates, is situated, or conducts its business."
            ),
            (
                "The GenAffiliation Relation refers to the connection between an individual and their associated"
                " \ncharacteristics, such as citizenship, residency, religious affiliation, or ethnicity, when the"
                " first argument (arg1) is a Person. \nWhen the arg1 is an Organization, the relation specifies the"
                " organization's location, base, or areas of business \nactivity."
            ),
        ]
    },
    "ace_geographical": {
        "en": [
            (
                "The Geographical relation captures the location of a Facility, Location, or GPE in or at or as a part"
                " of another \nFacility, Location, or GPE."
            ),
            (
                "The geographical relationship highlights the placement of a facility, location, or key geopolitical"
                " entity \nwithin or as an integral part of another facility, location, or key geopolitical entity."
            ),
            (
                "The geographical relationship highlights the positioning of a facility, location, or key point of"
                " interest within \nor as an integral component of another facility, location, or key point of"
                " interest."
            ),
            (
                "The geographic relationship highlights the positioning of a facility, location, or geopolitical"
                " entity within or \nas part of another facility, location, or geopolitical entity."
            ),
            (
                "The geographic relationship reveals the positioning of a facility, location, or key point of interest"
                " within or as \npart of another facility, location, or key point of interest."
            ),
        ]
    },
    "ace_injure": {
        "en": [
            (
                "An Injure Event occurs whenever a Person Entity experiences physical harm. Injure Events can be"
                " accidental, \nintentional or self-inflicted."
            ),
            (
                "An Injure Event takes place whenever a Person Entity encounters bodily harm. Injure Events can arise"
                " from \naccidents, intentional actions, or self-infliction."
            ),
            (
                "An Injure Event takes place whenever a Person Entity encounters physical harm. Injure Events can"
                " arise from \nunintentional, purposeful, or self-inflicted actions."
            ),
            (
                "An Injure Event takes place whenever a Person Entity encounters bodily harm. Injure Events can arise"
                " from \nunintentional, purposeful, or self-inflicted actions."
            ),
            (
                "An Injure Event takes place whenever a Person Entity encounters bodily harm. Injure Events can arise"
                " from \naccidental, intentional, or self-inflicted causes."
            ),
        ]
    },
    "ace_investorshareholder": {
        "en": [
            (
                "InvestorShareholder captures the relationship between an agent (Person, Organization, or GPE) and an"
                " \nOrganization in which the agent has invested or in which the agent owns shares/stock. Please note"
                " that agents may invest in GPEs."
            ),
            (
                "InvestorShareholder establishes the connection between a principal (Person, Organization, or GPE) and"
                " a Company \nin which the principal has financially engaged or holds an ownership stake through"
                " shares/stock. It is important to \nrecognize that principals can also invest in GPEs."
            ),
            (
                "InvestorShareholder highlights the connection between a decision-maker (Person, Organization, or GPE)"
                " and a \nCompany in which the decision-maker has put money or holds ownership through shares/stock."
                " Remember that decision-makers \ncan also invest in GPEs."
            ),
            (
                "InvestorShareholder highlights the connection between a decision-maker (an individual, organization,"
                " or group \nof people) and a company in which the decision-maker has financially invested or holds an"
                " ownership stake through \nshares/stock. Be aware that the decision-makers may also invest in groups"
                " of people."
            ),
            (
                "InvestorShareholder establishes the connection between a principal (Person, Organization, or GPE) and"
                " a \ncorporation in which the principal has allocated resources or holds an ownership stake. It is"
                " essential to recognize that agents \nmay invest in GPEs as well."
            ),
        ]
    },
    "ace_jobtitle": {
        "en": [
            (
                "A JobTitle value refers to the name of the job or position of a Person entity in a Organization. For"
                " example: 'co-chief \nexecutive', 'move coordinator', 'interim ED', ..."
            ),
            (
                "A JobTitle field within the Person entity of an Organization represents the specific job or role held"
                " by an \nindividual. For example, 'co-chief executive', 'move coordinator', 'interim ED', and so on."
            ),
            (
                "A JobTitle field in the Person entity of an Organization represents the specific job or role held by"
                " an individual. \nIllustrative examples include 'co-chief executive', 'move coordinator', and"
                " 'interim ED'."
            ),
            (
                "A JobTitle field in the Person entity of an Organization represents the specific job or role held by"
                " an individual. \nIllustrative examples include 'co-chief executive', 'move coordinator', 'interim"
                " ED', and so forth."
            ),
            (
                "A JobTitle field in the Person entity of an Organization represents the title or designation of the"
                " individual's \nrole. For example, titles may include 'Co-Chief Executive Officer', 'Move"
                " Coordinator', and 'Interim Executive \nDirector'."
            ),
        ]
    },
    "ace_justiceevent": {
        "en": [
            (
                "A JusticeEvent refers to any judicial action such as: arresting, jailing, releasing, granting parole,"
                " trial \nstarting, hearing, charging, indicting, suing, convicting, sentencing, fine, executing,"
                " extraditing, adquiting, \nappealing or pardoning a Person entity."
            ),
            (
                "A JusticeEvent is a term that encompasses various legal proceedings involving persons, such as"
                " apprehension, \nincarceration, release, parole granting, case commencement, hearings, charges,"
                " indictments, lawsuits, convictions, \nsentencing, fines, executions, extraditions, acquittals, and"
                " appeals or pardoning of Person entities."
            ),
            (
                "A JusticeEvent encompasses any and all legal proceedings involving persons, including but not limited"
                " to: \napprehension, incarceration, release, parole granting, trial commencement, hearings, charges,"
                " indictments, lawsuits, \nconvictions, sentencing, fines, executions, extraditions, acquittals,"
                " appeals, and pardons."
            ),
            (
                "A JusticeEvent encompasses any and all legal proceedings involving persons, which can include but are"
                " not limited \nto: apprehension, incarceration, release, parole granting, case commencement,"
                " hearings, charges, indictments, \nlawsuits, convictions, sentencing, fines, executions,"
                " extraditions, acquittals, and appeals or pardons of Person \nentities."
            ),
            (
                "A JusticeEvent is associated with any legal proceeding involving persons, including but not limited"
                " to: \napprehension, incarceration, release, parole granting, case commencement, hearings, charges,"
                " indictments, lawsuits, \nconvictions, sentencing, fines, executions, extraditions, acquittals, and"
                " appeals or pardons of Person entities."
            ),
        ]
    },
    "ace_lastingpersonal": {
        "en": [
            (
                "Lasting-Personal captures relationships that meet the following conditions: (1) The relationship must"
                " involve \npersonal contact (or a reasonable assumption thereof). (2) There must be some indication"
                " or expectation that the \nrelationship exists outside of a particular cited interaction. Both"
                " arguments must be entities of type Person."
            ),
            (
                "Lasting-Personal establishes relationships that fulfill the following requirements: (1) the"
                " connection must \ninvolve personal contact or a logical assumption of it, and (2) there must be a"
                " hint or expectation that the relationship \ncontinues beyond a specific interaction. Both parties"
                " involved must be entities of type Person."
            ),
            (
                "Lasting-Personal entails the establishment of connections that fulfill the following criteria: (1)"
                " the \ninvolvement of personal contact (or a logical assumption thereof) is essential, and (2) there"
                " must be a indication or \nexpectation suggesting that the relationship extends beyond a specific"
                " mentioned interaction. Both parties involved must be \nof the Person entity type."
            ),
            (
                "Lasting-Personal establishes relationships that fulfill the following criteria: (1) there must be a"
                " personal \nconnection, either directly or inferred, between the entities, and (2) there should be a"
                " suggestion or anticipation that the \nrelationship continues beyond a specific interaction. Both"
                " parties involved must be of the Person type."
            ),
            (
                "Lasting-Personal establishes relationships that comply with the following criteria: (1) there must be"
                " a personal \nconnection, either directly or inferred, and (2) there should be a indication or"
                " expectation that the relationship continues \nbeyond a specific interaction. Both parties involved"
                " must be of the Person type."
            ),
        ]
    },
    "ace_lifeevent": {
        "en": [
            "A LifeEvent occurs whenever a Person Entity borns, dies, gets married, divorced or gets injured.",
            (
                "A LifeEvent takes place whenever a Person Entity experiences the events of being born, passing away,"
                " getting \nmarried, obtaining a divorce, or suffering an injury."
            ),
            (
                "A LifeEvent takes place whenever a Person Entity experiences a birth, death, wedding, divorce, or"
                " injury."
            ),
            (
                "A LifeEvent takes place whenever a Person Entity experiences a birth, death, marriage, divorce, or"
                " injury."
            ),
            (
                "A LifeEvent takes place whenever a Person Entity experiences a birth, death, marriage, divorce, or"
                " injury."
            ),
        ]
    },
    "ace_located": {
        "en": [
            (
                "The Located relation captures the physical location of an entity. This relation is restricted to"
                " people. In other \nwords, `arg1` in Located relations can only be occupied by mentions of Entities"
                " of type Person."
            ),
            (
                "The Located relation embodies the physical placement of a entity. This connection is specifically"
                " applicable to \nhuman entities. In simpler terms, the first component, `arg1`, in the Located"
                " relation can only be filled by references to \nentities of the Person type."
            ),
            (
                "The Located relation embodies the physical placement of a entity. This connection is specifically"
                " applicable to \npeople. To clarify, the `arg1` position in Located relations can solely contain"
                " references to Entity instances of the \nPerson type."
            ),
            (
                "The Located relation embodies the physical placement of a entity. This connection is specifically"
                " applicable to \nindividuals. To clarify, the first position, 'arg1', in Located relations can only"
                " contain references to Entity instances of the \nPerson type."
            ),
            (
                "The Located relation embodies the physical placement of a entity. This association is specifically"
                " applicable to \npeople. To be more precise, in Located relations, the first position, `arg1`, can"
                " only be filled by references to Entities of \nthe Person type."
            ),
        ]
    },
    "ace_location": {
        "en": [
            (
                "Places defined on a geographical or astronomical basis which are mentioned in a document and do not"
                " constitute a \npolitical entity give rise to Location entities. These include, for example, the"
                " solar system, Mars, the Hudson River, Mt. \nEverest, and Death Valley."
            ),
            (
                "Locations specified by their geographical or astronomical attributes within a document, which do not"
                " represent a \npolitical entity, form Location entities. Examples of such places include the solar"
                " system, Mars, the Hudson River, Mt. \nEverest, and Death Valley."
            ),
            (
                "Geographical or astronomical locations specified in a document that don't represent a political"
                " entity lead to the \ncreation of Location entities. These may encompass examples such as the solar"
                " system, Mars, the Hudson River, Mt. Everest, \nand Death Valley."
            ),
            (
                "Geographical or astronomical locations specified in a document that don't represent a political"
                " entity result in \nthe creation of Location entities. These may encompass examples such as the solar"
                " system, Mars, the Hudson River, Mt. \nEverest, and Death Valley."
            ),
            (
                "Geographical or astronomical locations referred to in a document that do not represent a political"
                " entity are the \nbasis for creating Location entities. Examples of such locations include the solar"
                " system, Mars, the Hudson River, Mt. \nEverest, and Death Valley."
            ),
        ]
    },
    "ace_marry": {
        "en": [
            "Marry Events are official Events, where two people are married under the legal definition.",
            (
                "Marry Events are formal occasions where two individuals are united in marriage according to the legal"
                " criteria."
            ),
            (
                "Marry Events are formal occasions, during which two individuals are joined in matrimony in accordance"
                " with the \nlegal criteria."
            ),
            (
                "Marry Events are events that are considered official, during which two individuals are joined in"
                " marriage \naccording to the legal criteria."
            ),
            (
                "Marry Events are formal occasions wherein two individuals are united in marriage according to the"
                " established \nlegal criteria."
            ),
        ]
    },
    "ace_meet": {
        "en": [
            (
                "A Meet Event occurs whenever two or more Entities come together at a single location and interact"
                " with one another \nface-to-face. Meet Events include talks, summits, conferences, meetings, visits,"
                " and any other Event where two or more parties \nget together at some location."
            ),
            (
                "A Meet Event takes place whenever multiple Entities gather at a single location and engage in"
                " face-to-face \ninteractions. Such events encompass meetings, discussions, conferences, summits, and"
                " any other gatherings where two or more \nparticipants converge at a specific site."
            ),
            (
                "A Meet Event takes place whenever multiple Entities gather at a single venue and engage in"
                " person-to-person \ninteractions. Such events encompass meetings, discussions, assemblies,"
                " conventions, and any other occasions where multiple \nparticipants join at a specific location."
            ),
            (
                "A Meet Event transpires whenever multiple Entities assemble at a solitary location and engage with"
                " each other in \nperson. Meet Events encompass gatherings, conferences, meetings, visits, and any"
                " other occurrence where two or more \nentities convene at a specific place."
            ),
            (
                "A Meet Event takes place whenever multiple Entities gather at a single site and engage with each"
                " other in person. \nThese events encompass meetings, discussions, assemblies, gatherings, and any"
                " other occurrences where multiple \nparticipants join at a specific location."
            ),
        ]
    },
    "ace_membership": {
        "en": [
            (
                "Membership captures the relationship between an agent and an organization of which the agent is a"
                " member. \nOrganizations and GPEs can be members of other Organizations (such as NATO or the UN)."
            ),
            (
                "Membership encapsulates the connection between an agent and the institution to which they belong."
                " Both agents and \ngeneralized pseudo-entities (GPEs) can be members of larger organizations (e.g.,"
                " NATO or the UN)."
            ),
            (
                "Membership encapsulates the connection between an agent and the institution to which they hold"
                " membership. Both \nagents and general-purpose entities (GPEs) can be members of larger"
                " organizations, such as NATO or the United Nations."
            ),
            (
                "Membership reflects the connection between an agent and the organization they belong to. Both agents"
                " and GPEs can be \nmembers of larger organizations like NATO or the UN."
            ),
            (
                "Membership reflects the connection between an agent and the institution they belong to, with"
                " organizations and \ngroupings of entities (GPEs) sometimes being members of other organizations like"
                " NATO or the UN."
            ),
        ]
    },
    "ace_mergeorg": {
        "en": [
            (
                "A MergeOrg Event occurs whenever two or more Organization Entities come together to form a new"
                " Organization Entity. \nThis Event applies to any kind of Organization, including government"
                " agencies. It also includes joint venture."
            ),
            (
                "A MergeOrg Event takes place whenever two or more Organization Entities combine to create a new"
                " Organization \nEntity. This event can involve various types of organizations, such as government"
                " agencies, and also encompasses joint \nventures."
            ),
            (
                "A MergeOrg Event takes place whenever two or more Organization Entities combine to create a new"
                " Organization \nEntity. This event can involve various types of organizations, such as government"
                " agencies, as well as joint ventures."
            ),
            (
                "A MergeOrg Event takes place whenever two or more Organization Entities combine to create a new"
                " Organization \nEntity. This event can involve various types of organizations, such as government"
                " agencies, and also encompasses joint \nventures."
            ),
            (
                "A MergeOrg Event takes place whenever multiple Organization Entities combine to create a new"
                " Organization Entity. \nThis event can involve various types of organizations, such as government"
                " agencies, as well as joint ventures."
            ),
        ]
    },
    "ace_movementevent": {
        "en": [
            (
                "A TransportEvent occurs whenever an Artifact (Weapon or Vehicle) or a Person is moved from one Place"
                " (GPE, Facility, \nLocation) to another. This event requires the explicit mention of the Artifact or"
                " Person."
            ),
            (
                "A TransportEvent takes place whenever an Artifact, such as a Weapon or Vehicle, or a Person is"
                " transferred from one \nPlace, like a GPE, Facility, or Location, to another. This event necessitates"
                " the explicit mention of the Artifact or \nPerson."
            ),
            (
                "A TransportEvent takes place whenever an Artifact, such as a Weapon or Vehicle, or a Person is"
                " relocated from one \nPlace, like a GPE, Facility, or Location, to another. This event necessitates"
                " an explicit reference to the Artifact or \nPerson involved."
            ),
            (
                "A TransportEvent takes place whenever an Artifact, such as a Weapon or Vehicle, or a Person is"
                " relocated from one \nPlace, like a GPE, Facility, or Location, to another. This event necessitates a"
                " clear reference to the Artifact or Person \ninvolved."
            ),
            (
                "A TransportEvent takes place whenever an Artifact, such as a Weapon or Vehicle, or a Person is"
                " relocated from one \nPlace, like a GPE, Facility, or Location, to another. This event necessitates a"
                " clear reference to the Artifact or Person \ninvolved."
            ),
        ]
    },
    "ace_near": {
        "en": [
            (
                "Near indicates that an entity is explicitly near another entity, but neither entity is a part of the"
                " other or located \nin/at the other."
            ),
            (
                "Near implies that two entities are distinctly positioned in close proximity, without either being an"
                " intrinsic \ncomponent of the other or situated inside/within the other."
            ),
            (
                "Near implies that two entities are positioned close together, without either being a component of the"
                " other or \nsituated inside/within the other."
            ),
            (
                "Near implies that two entities are positioned close together, but neither is incorporated into the"
                " other nor \nsituated within/at the other."
            ),
            (
                'Text paraphrase: "Near signifies that two entities are positioned in close proximity to each other,'
                " without either \nbeing incorporated into the other or situated inside/within the other."
            ),
        ]
    },
    "ace_nominate": {
        "en": [
            (
                "A Nominate Event occurs whenever a Person is proposed for a StartPosition Event by the appropriate"
                " Person, through \nofficial channels."
            ),
            (
                "A Nominate Event takes place whenever a person is recommended for a StartPosition Event by the"
                " appropriate \nindividual, using official methods of communication."
            ),
            (
                "A Nominate Event takes place whenever a person is recommended for a StartPosition Event by the"
                " relevant individual, \nusing proper procedures."
            ),
            (
                "A Nominate Event takes place whenever a person is recommended for a StartPosition Event by the"
                " relevant individual, \nusing official methods of communication."
            ),
            (
                "A Nominate Event takes place whenever a person is recommended for a StartPosition Event by the"
                " relevant individual, \nusing proper communication methods."
            ),
        ]
    },
    "ace_numeric": {
        "en": [
            (
                "A Numeric value refers to relevant numbers, amounts, etc. For example: 'billions of dollars', '50"
                " percent', \n'100%', ..."
            ),
            (
                "A numeric value encompasses significant numbers, monetary amounts, and suchlike figures. For"
                " instance: \n'billions of dollars', '50 percent', '100%', and so on."
            ),
            (
                "A Numeric value encompasses significant numbers, monetary amounts, and such. For instance: 'billions"
                " of \ndollars', '50 percent', '100%', and so forth."
            ),
            (
                "A Numeric value encompasses significant figures, amounts, and related quantities. For instance:"
                ' "billions of \ndollars", "50 percent", "100%", and so on.'
            ),
            (
                'A Numeric value encompasses significant numbers, monetary amounts, and proportions, such as "billions'
                ' of \ndollars", "50 percent", and "100%".'
            ),
        ]
    },
    "ace_organization": {
        "en": [
            (
                "Each organization or set of organizations mentioned in a document gives rise to an entity of type"
                " Organization. \nTypical examples are businesses, government units, sports teams, and formally"
                " organized music groups."
            ),
            (
                "Every entity referenced in a document as an organization or group creates an object classified as"
                " Organization. \nCommon examples include commercial enterprises, government departments, sports"
                " clubs, and formally structured \nmusical groups."
            ),
            (
                "In any given document, each instance of an organization or a group of organizations leads to the"
                " creation of an entity \nclassified as Organization. Common illustrations include commercial"
                " enterprises, government departments, sports clubs, \nand formally structured musical groups."
            ),
            (
                "In every document, each organization or group of organizations represented forms an entity classified"
                " as \nOrganization. Common instances encompass businesses, government departments, sports clubs, and"
                " formally structured musical \ngroups."
            ),
            (
                "In every document, each entity that is a part of a company, government agency, sports squad, or an"
                " officially \nstructured musical group constitutes a distinct entity of the Organization type."
            ),
        ]
    },
    "ace_organizationaffiliationrelation": {
        "en": [
            (
                "The OrganizationAffiliation Relation describes the relations between a Person (or other"
                " Organizations) and a \nrelated Organization. This relation includes: employment, ownership, founder,"
                " student or alumn, sport affiliation, \ninverstor or shareholder and membership relations."
            ),
            (
                "The OrganizationAffiliation Relation refers to the connections between an individual or another"
                " organization \nand an associated organization, encompassing various types of relationships such as"
                " employment, ownership, \nfoundership, academic affiliation, sports connections, investment or"
                " shareholding, and membership ties."
            ),
            (
                "The OrganizationAffiliation Relation refers to the connections between an individual (either a Person"
                " or another \ntype of Organization) and a related Organization. This relationship encompasses various"
                " aspects, such as employment, \nownership, foundership, academic affiliation, sports connections,"
                " investment or shareholding, and membership ties."
            ),
            (
                "The OrganizationAffiliation Relation refers to the connections between an individual or another"
                " organization \nand a related organization, encompassing various types of associations such as"
                " employment, ownership, foundership, \nacademic ties as a student or alumnus, sports affiliations,"
                " investment or shareholder roles, and membership links."
            ),
            (
                "The OrganizationAffiliation Relation refers to the connections between an individual (either a Person"
                " or another \nOrganization) and a related Organization. This association encompasses various types of"
                " relationships, such as employment, \nownership, foundership, academic affiliation, sports"
                " connections, investment or shareholding, and membership ties."
            ),
        ]
    },
    "ace_orglocationorigin": {
        "en": [
            (
                "OrgLocationOrigin captures the relationship between an organization and the Location or GPE where it"
                " is located, \nbased, or does business. Note: Subsidiary trumps this relation for government"
                " organizations."
            ),
            (
                "OrgLocationOrigin establishes the connection between an organization and the specific Location or GPE"
                " where it \noperates, is headquartered, or conducts business. It's worth noting that for government"
                " entities, the Subsidiary \nrelationship takes precedence over this connection."
            ),
            (
                "OrgLocationOrigin establishes the connection between an organization and the specific Location or GPE"
                " where it is \nsituated, centered, or conducts its operations. It's worth mentioning that for"
                " government entities, the Subsidiary \nrelation takes precedence over this connection."
            ),
            (
                "OrgLocationOrigin establishes the connection between an organization and the specific Location or GPE"
                " where it is \nsituated, centered, or conducts business. Take note: For government organizations, the"
                " Subsidiary relationship takes \nprecedence over this connection."
            ),
            (
                "OrgLocationOrigin establishes the connection between an organization and the specific Location or GPE"
                " where it is \nsituated, headquartered, or conducts its operations. Keep in mind that for government"
                " entities, the Subsidiary \nrelationship takes precedence."
            ),
        ]
    },
    "ace_ownership": {
        "en": [
            (
                "Ownership captures the relationship between a Person and an Organization owned by that Person. If the"
                " `arg2` is not \nan Organization, use the Agent-Artifact relation."
            ),
            (
                "Paraphrase: The concept of ownership defines the connection between an individual (Person) and a"
                " company \n(Organization) that is possessed by that individual. If the second component (arg2) is not"
                " an organization, it is essential to \nemploy the Agent-Artifact association."
            ),
            (
                'Paraphrase: "The Ownership association illustrates the connection between a Person and an'
                " Organization that the \nPerson possesses. If the `arg2` is not an Organization, utilize the"
                " Agent-Artifact association."
            ),
            (
                'Paraphrase: "The concept of ownership defines the connection between an individual (Person) and a'
                " business entity \n(Organization) that the person possesses. If the second argument (arg2) is not an"
                " organization, you can employ the Agent-Artifact \nassociation."
            ),
            (
                'Paraphrase: "The Ownership association illustrates the connection between an individual Person and'
                " an \nOrganization that the Person possesses. If the second component, 'arg2', is not an"
                " Organization, you can employ the \nAgent-Artifact association."
            ),
        ]
    },
    "ace_pardon": {
        "en": [
            (
                "A Pardon Event occurs whenever a head-of-state or their appointed representative lifts a sentence"
                " imposed by the \njudiciary."
            ),
            (
                "A Pardon Event takes place whenever a nation's highest-ranking leader or their designated"
                " representative \nnullifies a punishment that has been issued by the judicial system."
            ),
            (
                "A Pardon Event takes place whenever a nation's highest-ranking leader or a designated representative"
                " on their \nbehalf nullifies a judicial punishment."
            ),
            (
                "A Pardon Event takes place whenever a nation's top leader or their designated representative"
                " nullifies a \npunishment that has been levied by the judicial system."
            ),
            (
                "A Pardon Event takes place whenever a nation's top leader or their designated representative"
                " nullifies a penalty \nissued by the judicial system."
            ),
        ]
    },
    "ace_partwholerelation": {
        "en": [
            (
                "The PartWhole Relation refers to the semantic relation between two entities that are parts of a"
                " larger whole or vice \nversa. For example, the relation between a country and its states, or between"
                " a company and its subsidiaries, are instances \nof PartWhole relations."
            ),
            (
                "The PartWhole Relation encompasses the semantic connection between two entities that function as"
                " components of a \nmore extensive whole or, conversely, as a larger entity comprising smaller parts."
                " For example, the connection between a \ncountry and its constituent states, or the relationship"
                " between a parent company and its subsidiaries, embody the \nPartWhole relationship."
            ),
            (
                "The PartWhole Relation encompasses the semantic connection between two components that compose a more"
                " extensive \nwhole or, conversely, between a greater entity and its constituent parts. Illustrative"
                " instances of this relation \ninclude the connection between a country and its individual states, as"
                " well as the relationship between a corporation and \nits subsidiaries."
            ),
            (
                "The PartWhole Relation encompasses the semantic connection between two components that compose a more"
                " extensive \nwhole or, conversely, between parts of a larger entity. Illustrative instances of this"
                " relation include the connection \nbetween a country and its constituent states, as well as the"
                " relationship between a parent company and its subordinate \nsubsidiaries."
            ),
            (
                "The PartWhole Relation encompasses the semantic connection between two components that comprise a"
                " more extensive \nentity or conversely, between a larger whole and its constituent parts."
                " Illustrations of PartWhole relations include the \nconnection between a country and its individual"
                " states, as well as the relationship between a company and its subsidiaries."
            ),
        ]
    },
    "ace_person": {
        "en": [
            (
                "Each distinct person or set of people mentioned in a document refers to an entity of type Person. For"
                ' example, people \nmay be specified by name ("John Smith"), occupation ("the butcher"), family'
                ' relation ("dad"), pronoun ("he"), etc., \nor by some combination of these.'
            ),
            (
                "In every document, individuals or groups of individuals mentioned are associated with an entity of"
                ' the Person type. \nThis can include identifying people by their name ("John Smith"), their'
                ' profession ("the butcher"), familial \nrelationships ("dad"), or by using pronouns ("he"), among'
                " other methods, or even a mix of these approaches."
            ),
            (
                "In every document, individuals or groups of individuals mentioned are associated with an entity of"
                ' the Person type. \nThis can include instances where people are identified by their name ("John'
                ' Smith"), profession ("the butcher"), \nfamilial relationship ("dad"), or by the use of pronouns'
                ' ("he"), as well as situations where these methods are combined.'
            ),
            (
                "In every document, individuals or groups of individuals mentioned are associated with an entity of"
                ' the Person type. \nThis can include instances where people are identified by their name ("John'
                ' Smith"), occupation ("the butcher"), \nfamilial relationship ("dad"), or even by pronouns ("he"), as'
                " well as scenarios where these identifiers are used in \ncombination."
            ),
            (
                "In every document, individuals or groups of individuals mentioned are associated with an entity"
                ' classified as \nPerson. These persons could be identified by their given name ("John Smith"), their'
                ' profession ("the butcher"), their \nfamilial relationship ("dad"), or by utilizing pronouns ("he"),'
                " among other possibilities, or even a mix of these \napproaches."
            ),
        ]
    },
    "ace_personalsocialrelation": {
        "en": [
            (
                "The Personal-Social Relation describe the relationship between people. Both arguments must be"
                " entities of type \nPerson. Please note: The arguments of these Relations are not ordered. The"
                " Relations are symmetric."
            ),
            (
                "The Personal-Social Relation highlights the connection between individuals. Both parties involved"
                " must be of the \nPerson entity type. It's essential to note that the arguments in these"
                " relationships are not arranged in a particular \norder, and the relationships are symmetrical."
            ),
            (
                "The Personal-Social Relation encompasses the connections between individuals. In this context, both"
                " parties \nmust be entities of the Person type. It is crucial to highlight that the arguments of"
                " these relationships are not \nsequenced; as a result, these connections are symmetric."
            ),
            (
                "The Personal-Social Relation refers to the connection between individuals. In this context, both"
                " parties must be \nof the Person entity type. It is essential to note that the arguments of these"
                " relations are not arranged in a particular \norder, which indicates that these relationships are"
                " symmetric."
            ),
            (
                "The Personal-Social Relation highlights the connection between individuals. In this context, both"
                " parties must \nbe of the Person entity type. It is important to note that the participants in these"
                " relationships are not arranged in a \nspecific order, and the associations are symmetrical."
            ),
        ]
    },
    "ace_personnelevent": {
        "en": [
            (
                "A PersonnelEvent occurs when a Person entity changes its job position (JobTitle entity) with respect"
                " an \nOrganization entity. It includes when a person starts working, ends working, changes offices"
                " within, gets nominated or is \nelected for a position in a Organization."
            ),
            (
                "A PersonnelEvent transpires when a Person element undergoes a transition in their occupation"
                " (JobTitle entity) in \nrelation to an Organization entity. This encompasses instances of commencing"
                " work, concluding work, shifting roles within \nthe organization, securing a nomination or being"
                " appointed to a position in the organization."
            ),
            (
                "A PersonnelEvent transpires when a Person entity undergoes a transition in their job position"
                " (JobTitle entity) in \nrelation to an Organization entity. This encompasses instances such as the"
                " commencement or termination of employment, \nrelocation within the organization, securing a"
                " nomination or election for a position within the organization."
            ),
            (
                "A PersonnelEvent transpires when a Person entity undergoes a transition in their job position"
                " (JobTitle entity) \nrelative to an Organization entity. This encompasses instances such as"
                " commencing employment, terminating employment, \nrelocating within the organization, securing a"
                " nomination or election for a position within the organization."
            ),
            (
                "A PersonnelEvent takes place whenever a Person entity undergoes a change in their job position"
                " (JobTitle entity) in \nrelation to an Organization entity. This encompasses instances such as when a"
                " person commences employment, terminates \ntheir employment, transfers between offices, receives a"
                " nomination or is elected to a position within the organization."
            ),
        ]
    },
    "ace_phonewrite": {
        "en": [
            (
                "A PhoneWrite Event occurs when two or more people directly engage in discussion which does not take"
                " place \n'face-to-face'. To make this Event less open-ended, we limit it to written or telephone"
                " communication where at least two parties are \nspecified."
            ),
            (
                "A PhoneWrite Event transpires when two or more individuals participate in a dialogue that does not"
                " happen in a \nconventional face-to-face setting. In order to narrow down the scope of this event and"
                " make it more definitive, we restrict it to \nsituations involving written or telephone"
                " communication, with a minimum of two identified participants."
            ),
            (
                "A PhoneWrite Event transpires when two or more individuals engage in a dialogue that does not happen"
                " in a direct, \nface-to-face context. In order to narrow down the scope of this event, we confine it"
                " to situations involving written or telephone \ncommunication, where at least two distinct"
                " participants are involved."
            ),
            (
                "A PhoneWrite Event transpires when two or more individuals participate in a dialogue that does not"
                ' transpire in a \nconventional "face-to-face" manner. In order to make this event more focused, we'
                " restrict it to situations involving written or \ntelephone communication, where at least two"
                " distinct participants are involved."
            ),
            (
                "A PhoneWrite Event transpires when two or more individuals interactively connect through a dialogue"
                " that is not \nconducted in a confrontational manner. In order to narrow down the scope of this"
                " event, we confine it to instances of written or \ntelephonic communication that involve at least two"
                " predefined participants."
            ),
        ]
    },
    "ace_physicalrelation": {
        "en": [
            (
                "The Physical Relation captures the physical location relation of entities such as: a Person entity"
                " located in a \nFacility, Location or GPE; or two entities that are near, but neither entity is a"
                " part of the other or located in/at the other."
            ),
            (
                "The Physical Relation encapsulates the spatial connection between entities, for example: a Person"
                " entity \nsituated within a Facility, Location, or GPE; or the proximity of two entities, without"
                " either being an aspect of the other or \nlocated inside/at the other."
            ),
            (
                "The Physical Relation encapsulates the spatial connection between entities, including scenarios like"
                " a Person \nentity situated within a Facility, Location, or GPE; or when two entities are close to"
                " each other, without either being a \ncomponent of the other or located within/at the other."
            ),
            (
                "The Physical Relation encapsulates the spatial association between entities, such as a Person entity"
                " situated \nwithin a Facility, Location, or GPE; or the proximity of two entities, neither of which"
                " is a component of the other or located \nwithin/at the other."
            ),
            (
                "The Physical Relation embodies the spatial connection between entities, including scenarios where a"
                " Person \nentity is situated within a Facility, Location, or GPE; as well as instances where two"
                " entities are close to each other, \nwithout either being part of the other or located within/at the"
                " other."
            ),
        ]
    },
    "ace_releaseparole": {
        "en": [
            (
                "A Release Event occurs whenever a state actor (GPE, Organization subpart, or Person representative)"
                " ends its \ncustody of a Person Entity."
            ),
            (
                "A Release Incident takes place whenever a key player in the system (whether a Governmental or Private"
                " Entity, an \nOrganization subdivision, or an individual acting on behalf of another person)"
                " relinquishes their responsibility for a Person \nEntity."
            ),
            (
                "A Release Incident takes place whenever a key player in the system (Governmental Person Entity,"
                " Organizational \nSubpart, or Individual acting on behalf of another Individual) relinquishes"
                " responsibility for a Person Entity."
            ),
            (
                "A Release Incident takes place whenever a state entity (Governmental or Private Enterprise,"
                " Organization \nsubdivision, or Individual acting on behalf of another) terminates its supervision of"
                " a Person Entity."
            ),
            (
                "A Release Incident takes place whenever a significant entity (Governmental Payer or Employer,"
                " Organizational \nSubunit, or Individual Rep), which is in control of a Person Entity, terminates"
                " that control."
            ),
        ]
    },
    "ace_sentence": {
        "en": [
            (
                "A Sentence value refers to sentences decided by a court or judge for a specific crime. For example:"
                " '124 years in \nprison', 'a sentence', 'death'..."
            ),
            (
                'A Sentence value represents court-imposed sentences for a particular offense. For instance: "124'
                ' years in \nprison", "a sentence", or "death"....'
            ),
            (
                "A Sentence value encompasses judicial sentences imposed by a court or judge for a particular offense."
                " For instance: \n'124 years incarceration', 'a sentence', or 'capital punishment'."
            ),
            (
                "A Sentence value denotes sentences imposed by a court or judge for a particular offense. For"
                ' instance: "124 years in \nprison", "a sentence", or "death"....'
            ),
            (
                "A Sentence value encompasses judicially determined sentences for a particular offense. For instance:"
                ' "124 years \nin prison", "a sentence", or "death'
            ),
        ]
    },
    "ace_sentenceact": {
        "en": [
            (
                "A SentenceAct Event takes place whenever the punishment (particularly incarceration) for the"
                " Defendant of a Try \nEvent is issued by a state actor (a GPE, an Organization subpart or a Person"
                " representing them)"
            ),
            (
                "A SentenceAct Event occurs whenever the consequences, particularly imprisonment, for the defendant in"
                " a Try Event \nare administered by a state representative, such as a governmental body, an"
                " organization, or an individual acting on \ntheir behalf."
            ),
            (
                "A SentenceAct Event occurs whenever the penalty, specifically imprisonment, is imposed on the"
                " defendant in a Try \nEvent by a state representative, such as a Governmental Party Entity, an"
                " Organization subpart, or an individual acting on \ntheir behalf."
            ),
            (
                "A SentenceAct Event occurs whenever the repercussion, specifically imprisonment, for the defendant in"
                " a Try Event \nis administered by a state representative, such as a Governmental or Private Entity,"
                " an Organization subdivision, or \nan individual acting on their behalf."
            ),
            (
                "A SentenceAct Event occurs whenever the penalty, specifically imprisonment, is imposed on the"
                " Defendant in a Try \nEvent by a state representative, such as a Governmental Body, an Organization"
                " subdivision, or an individual acting on \ntheir behalf."
            ),
        ]
    },
    "ace_sportsaffiliation": {
        "en": [
            (
                "Sports-Affiliation captures the relationship between a player, coach, manager, or assistant and his"
                " or her \naffiliation with a sports organization (including sports leagues or divisions as well as"
                " individual sports teams)."
            ),
            (
                "Sports-Affiliation highlights the connection between an athlete, coach, manager, or assistant and"
                " their \nassociation with a sports institution (encompassing sports leagues, divisions, and both"
                " individual and team sports \naffiliations)."
            ),
            (
                "Sports-Affiliation highlights the connection between an athlete, coach, manager, or assistant and"
                " their \nassociation with a sports institution, encompassing both sports leagues or divisions and"
                " individual sports teams."
            ),
            (
                "Sports-Affiliation highlights the connection between an athlete, coach, manager, or assistant and"
                " their \nassociation with a sports group (encompassing sports leagues, divisions, and individual"
                " sports teams)."
            ),
            (
                "Sports-Affiliation highlights the connection between an athlete, coach, manager, or assistant and"
                " their \nassociation with a sports organization, encompassing both team-based and individual sports"
                " as well as their respective \nleagues or divisions."
            ),
        ]
    },
    "ace_startorg": {
        "en": [
            "A StartOrg Event occurs whenever a new Organization is created.",
            "A StartOrg Event takes place whenever a fresh Organization is established.",
            "A StartOrg Event takes place whenever a fresh Organization is established.",
            "A StartOrg Occurrence takes place whenever a fresh Organization comes into existence.",
            "A StartOrg Occurrence transpires whenever a novel Organization emerges.",
        ]
    },
    "ace_startposition": {
        "en": [
            (
                "A StartPosition Event occurs whenever a Person Entity begins working for (or changes offices within)"
                " an \nOrganization or GPE. This includes government officials starting their terms, whether elected"
                " or appointed."
            ),
            (
                "A StartPosition Event takes place whenever a Person Entity commences employment with or transfers"
                " within an \nOrganization or Government/Public Entity (GPE). This encompasses instances such as"
                " government officials initiating their \ntenures, whether through election or appointment."
            ),
            (
                "A StartPosition Event takes place whenever a Person Entity commences employment with or transfers"
                " within an \nOrganization or Government-Private Entity. This encompasses instances such as government"
                " officials initiating their \ntenures, whether by means of election or appointment."
            ),
            (
                "A StartPosition Event takes place whenever a Person Entity commences employment with or transfers"
                " within an \nOrganization or Government/Public Entity (GPE). This encompasses situations where"
                " government officials embark on their \ntenures, either through election or appointment."
            ),
            (
                "A StartPosition Event takes place whenever a Person Entity commences employment with or transfers"
                " within an \nOrganization or Government/Public Entity (GPE). This encompasses instances when"
                " government officials initiate their \ntenures, irrespective of whether they are elected or"
                " appointed."
            ),
        ]
    },
    "ace_studentalum": {
        "en": [
            (
                "StudentAlum captures the relationship between a Person and an educational institution the Person"
                " attends or \nattended."
            ),
            (
                "StudentAlum entity embodies the connection between an individual and the academic institution they"
                " are enrolled \nin or have previously attended."
            ),
            (
                "StudentAlum entity represents the connection between an individual and the educational institution"
                " they are \nattending or have attended."
            ),
            (
                "StudentAlum entity embodies the connection between an individual and the academic institution they"
                " are enrolled \nin or have previously attended."
            ),
            (
                "StudentAlum relationship connects an individual with the academic institution they are currently"
                " enrolled in or \nhave previously attended."
            ),
        ]
    },
    "ace_subsidiary": {
        "en": [
            (
                "Subsidiary captures the ownership, administrative, and other hierarchical relationships between"
                " \norganizations and between organizations and GPEs."
            ),
            (
                'Paraphrase: "Subsidiary relationships encompass the ownership, managerial, and hierarchical'
                " connections \nbetween businesses and their connections with GPEs."
            ),
            (
                'Paraphrase: "Subsidiary encapsulates the proprietary, managerial, and hierarchical connections'
                " between \nbusinesses and between these entities and GPEs."
            ),
            (
                'Paraphrase: "Subsidiary relationships encompass the ownership, managerial, and hierarchical'
                " connections \nbetween businesses and their connections with GPEs."
            ),
            (
                'Paraphrase: "The concept of subsidiary encompasses the ownership, managerial, and hierarchical'
                " connections \nbetween businesses and between these entities and legal entities of a public nature."
            ),
        ]
    },
    "ace_sue": {
        "en": [
            (
                "A Sue Event occurs whenever a court proceeding has been initiated for the purposes of determining the"
                " liability of a \nPerson, Organization or GPE accused of committing a crime or neglecting a"
                " commitment."
            ),
            (
                "A Sue Event takes place whenever a legal proceeding is commenced to ascertain the accountability of"
                " an individual, \norganization, or government entity alleged to have committed a crime or breached a"
                " duty."
            ),
            (
                "A Sue Event takes place whenever a legal case is brought forth to assess the responsibility of an"
                " individual, \norganization, or government entity that has been accused of either committing a crime"
                " or failing to fulfill a commitment."
            ),
            (
                "A Sue Event takes place whenever a legal proceeding is commenced to assess the accountability of an"
                " individual, \norganization, or government entity alleged to have committed a crime or breached a"
                " duty."
            ),
            (
                "A Sue Event takes place whenever a legal proceeding is commenced to assess the accountability of an"
                " individual, \norganization, or government entity accused of either committing a crime or breaching a"
                " commitment."
            ),
        ]
    },
    "ace_time": {
        "en": [
            (
                "A Time value refers to a specific time frame. Usually known as time expressions. For example: '4"
                " years', 'today', \n'December', 'future', ..."
            ),
            (
                "A Time value represents a particular time duration. Commonly called time expressions, examples"
                " include: '4 \nyears', 'today', 'December', and 'future'."
            ),
            (
                "A Time value represents a particular time duration. Typically called time expressions, examples"
                " include: '4 \nyears', 'today', 'December', and 'future'."
            ),
            (
                "A Time value represents a particular time duration. Typically called time expressions, examples"
                " include: '4 \nyears', 'today', 'December', and 'future'."
            ),
            (
                "A Time value represents a particular time duration. Commonly called time expressions, examples"
                " include: '4 \nyears', 'today', 'December', and 'future'."
            ),
        ]
    },
    "ace_transactionevent": {
        "en": [
            (
                "A TransactionEvent refers to buying, selling, loaning, borrowing, giving, or receving of Artifacts or"
                " \nOrganizations; or giving, receiving, borrowing, or lending Money."
            ),
            (
                "A TransactionEvent involves activities such as acquiring, exchanging, lending, or borrowing Artifacts"
                " or \nOrganizations, as well as engaging in transactions that involve giving, receiving, borrowing,"
                " or lending money."
            ),
            (
                "A TransactionEvent involves activities such as purchasing, selling, lending, borrowing, donating, or"
                " acquiring \nArtifacts or Organizations; or engaging in actions like giving, receiving, borrowing, or"
                " lending Money."
            ),
            (
                "A TransactionEvent involves activities such as purchasing, selling, lending, borrowing, gifting, or"
                " acquiring \nArtifacts or Organizations; or engaging in giving, receiving, loaning, or borrowing"
                " Money."
            ),
            (
                "A TransactionEvent involves activities such as acquiring, disposing of, lending, or borrowing"
                " artifacts or \norganizations, as well as engaging in monetary transactions like giving, receiving,"
                " borrowing, or lending money."
            ),
        ]
    },
    "ace_transfermoney": {
        "en": [
            (
                "TransferMoney Events refer to the giving, receiving, borrowing, or lending money when it is not in"
                " the context of \npurchasing something. The canonical examples are: (1) people giving money to"
                " organizations (and getting nothing tangible in \nreturn); and (2) organizations lending money to"
                " people or other orgs."
            ),
            (
                "TransferMoney Events involve the process of exchanging money for non-commercial purposes, where no"
                " goods or \nservices are obtained in return. Typical scenarios include: (1) individuals donating"
                " money to charitable organizations \nwithout receiving any tangible items; and (2) financial"
                " institutions loaning money to individuals or other \norganizations."
            ),
            (
                "TransferMoney Events involve the process of giving, receiving, borrowing, or lending funds when not"
                " situated \nwithin the context of a purchase. Typical examples include: (1) individuals donating"
                " money to organizations without \nreceiving any tangible item in return; and (2) organizations"
                " providing loans to individuals or other entities."
            ),
            (
                "TransferMoney Events involve the transfer of money in situations that do not involve purchasing goods"
                " or services. \nTypical examples include: (1) individuals donating money to organizations without"
                " receiving any tangible items in \nreturn; and (2) organizations loaning money to individuals or"
                " other entities."
            ),
            (
                "TransferMoney Events involve the process of exchanging money between parties without the context of"
                " purchasing \ngoods or services. Key illustrations include: (1) individuals donating money to"
                " institutions without expecting any \ntangible items in return; and (2) organizations providing loans"
                " to individuals or other entities."
            ),
        ]
    },
    "ace_transferownership": {
        "en": [
            (
                "TransferOwnership Events refer to the buying, selling, loaning, borrowing, giving, or receiving of"
                " artifacts or \norganizations."
            ),
            (
                "TransferOwnership Events involve the processes of acquiring, disposing of, loaning, borrowing,"
                " gifting, or \ntransferring ownership of artifacts or organizations."
            ),
            (
                "TransferOwnership Events involve the acquisition, disposition, lending, borrowing, gifting, or taking"
                " \npossession of items or institutions."
            ),
            (
                "TransferOwnership Events encompass the acquisition, disposition, lending, borrowing, gifting, or"
                " obtaining of \nitems or establishments."
            ),
            (
                "TransferOwnership Events involve the acquisition, disposal, lending, borrowing, gifting, or obtaining"
                " of items \nor institutions."
            ),
        ]
    },
    "ace_transport": {
        "en": [
            (
                "A Transport Event occurs whenever an Artifact (Weapon or Vehicle) or a Person is moved from one Place"
                " (GPE, Facility, \nLocation) to another."
            ),
            "A Mobility Incident happens whenever an item",
            (
                "A Mobility Occurrence takes place whenever an item such as an Artifact (Battleship or Fighter), a"
                " Person, or a \nVehicle is relocated from one Place (Global Province and Economy, Facility, Location)"
                " to another."
            ),
            (
                "A Movement Incident takes place whenever an item, such as a Tool (Firearm or Vehicle) or an"
                " individual, is relocated \nfrom one Setting (Geographical Point of Entry, Facility, Location) to"
                " another."
            ),
            "A Mobility Incident happens whenever an item",
        ]
    },
    "ace_trialhearing": {
        "en": [
            (
                "A Trial Event occurs whenever a court proceeding has been initiated for the purposes of determining"
                " the guilt or \ninnocence of a Person, Organization or GPE accused of committing a crime."
            ),
            (
                "A Trial Event takes place whenever a legal hearing has been instituted with the objective of"
                " ascertaining the \ninnocence or guilt of an individual, organization, or government entity accused"
                " of violating a law."
            ),
            (
                "A Trial Event takes place whenever a legal hearing has been instituted with the objective of"
                " ascertaining the \ninnocence or guilt of an individual, organization, or government entity accused"
                " of violating a law."
            ),
            (
                "A Trial Event takes place whenever a legal hearing has been convened with the objective of"
                " ascertaining the \ninnocence or guilt of an individual, organization, or government entity facing"
                " criminal charges."
            ),
            (
                "A Trial Event takes place whenever a legal hearing has been instituted with the objective of"
                " ascertaining the guilt \nor innocence of an individual, organization, or group suspected of"
                " committing a crime."
            ),
        ]
    },
    "ace_userownerinventormanufacturer": {
        "en": [
            (
                "This Relation applies when an agent owns an artifact, has possession of an artifact, uses an"
                " artifact, or caused an \nartifact to come into being. Note: if `arg2` is an Organization, use"
                " Ownership relation (arg1=PER) or Subsidiary relation \n(arg1=ORG or GPE)."
            ),
            (
                "This principle is relevant when an agent claims ownership of an object, holds possession of an"
                " object, employs an \nobject, or brings an object into existence. Take note: if the second argument"
                " is an organization, use the Ownership relation \nwith the first argument set to PERSON or the"
                " Subsidiary relation with the first argument designated as ORGANIZATION or \nGOVERNMENTAL ENTITY."
            ),
            (
                "This principle is relevant when an agent holds, possesses, utilizes, or brings into existence an"
                " artifact. It's \nimportant to mention that if the second entity is an organization, the relationship"
                " can be categorized as ownership (with arg1 \nas PER) or as a subsidiary relationship (with arg1 as"
                " ORG or GPE)."
            ),
            (
                "This principle is relevant when an agent claims ownership of a product, holds possession of a"
                " product, utilizes a \nproduct, or brings a product into existence. It's important to note that if"
                " the second entity is an organization, the \nrelationship should be classified as Ownership if the"
                " first entity is a Person, or as a Subsidiary relationship if the first entity \nis an Organization"
                " or a Governmental or Public Enterprise."
            ),
            (
                "This principle is relevant when an agent claims ownership of an object, holds possession of an"
                " object, utilizes an \nobject, or brings an object into existence. It's important to note that if the"
                " second entity is an organization, the \nrelationship should be classified as Ownership if the first"
                " entity is a person, or as a Subsidiary relationship if the first entity \nis an organization or a"
                " group of people."
            ),
        ]
    },
    "ace_vehicle": {
        "en": [
            (
                "A Vehicle entity refers to vehicles that are used for transportation. The vehicles can transport"
                " either persons or \nartifacts. For example: 'car', 'plane', 'cabin', ..."
            ),
            (
                "A Vehicle entity encompasses vehicles utilized for transportation purposes, which can convey either"
                " individuals \nor goods. For instance: automobiles, airplanes, ships, and so on."
            ),
            (
                "A Vehicle entity encompasses vehicles utilized for transportation purposes, which can convey either"
                " individuals \nor items. For instance: 'car', 'plane', 'cabin', and so forth."
            ),
            (
                "A Vehicle entity encompasses various means of transport, which are utilized for the conveyance of"
                " individuals or \ngoods. Illustrative instances include automobiles, airplanes, and freight"
                " containers."
            ),
            (
                "A Vehicle entity encompasses vehicles utilized for transportation purposes, which can convey either"
                " individuals \nor items. For instance, 'car', 'plane', 'cabin', and so on."
            ),
        ]
    },
    "ace_weapon": {
        "en": [
            (
                "A Weapon entity refers to instruments that can be used to deal physical damage, destroy something or"
                " kill someone. \nFor example: 'bomb', 'm-16s', 'missile', ..."
            ),
            (
                "A Weapon entity encompasses items designed to inflict physical harm, cause destruction, or result in"
                " the death of a \nliving being. Examples include: 'bomb', 'm-16s', 'missile', and so forth."
            ),
            (
                "A Weapon entity encompasses items designed to inflict physical harm, cause destruction, or result in"
                " the death of a \nliving being. Examples include: 'bomb', 'm-16s', 'missile', and so forth."
            ),
            (
                "A Weapon entity encompasses items designed to inflict physical harm, cause destruction, or result in"
                " the loss of \nlife. Examples include: 'bomb', 'm-16s', 'missile', and so forth."
            ),
            (
                "A Weapon entity encompasses items designed to inflict physical harm, annihilate objects, or cause the"
                " death of \nliving beings, such as 'bombs', 'M-16 assault rifles', and 'missiles'."
            ),
        ]
    },
}

src/tasks/ace/guidelines.py
hitz-zentroa/GoLLIE
GUIDELINES = {
    "ace_gpe": {
        "en": [
            "Geo-Political Entities are composite entities comprised of a population, a government, a physical\n"
            "location, and a nation (or province, state, country, city, etc.).\n"
        ]
    },
    "ace_acquit": {
        "en": [
            "An Acquit Event occurs whenever a trial ends but fails to produce a conviction. This will include cases\n"
            "where the charges are dropped by the Prosecutor.\n"
        ]
    },
    "ace_agentartifactrelationrelation": {
        "en": [
            "The AgentArtifact Relation applies when an agent owns an artifact, has possession of an artifact, uses\n"
            "an artifact, or caused an artifact to come into being. Note: if the `arg2` is an Organization, use\n"
            "OrganizationAffiliation when `arg1` is a Person or PartWhole when `arg1` is an Organization or GPE.\n"
        ]
    },
    "ace_appeal": {
        "en": ["An Appeal Event occurs whenever the decision of a court is taken to a higher court for review."]
    },
    "ace_arrestjail": {
        "en": [
            "A Jail Event occurs whenever the movement of a Person is constrained by a state actor (a GPE, its\n"
            "Organization subparts, or its Person representatives).\n"
        ]
    },
    "ace_attack": {
        "en": [
            "An Attack Event is defined as a violent physical act causing harm or damage. Attack Events include any\n"
            "such Event not covered by the Injure or Die subtypes, including Events where there is no stated agent.\n"
        ]
    },
    "ace_beborn": {"en": ["A BeBorn Event occurs whenever a Person Entity is given birth to."]},
    "ace_business": {
        "en": [
            "The Business Relation captures the connection between two entities in any professional relationship.\n"
            "Both arguments must be entities of type Person.\n"
        ]
    },
    "ace_businessevent": {
        "en": [
            "A BusinessEvent refers to actions related to Organizations such as: creating, merging, declaring\n"
            "bankruptcy or ending organizations (including government agencies).\n"
        ]
    },
    "ace_chargeindict": {
        "en": [
            "A Charge Event occurs whenever a Person, Organization or GPE is accused of a crime by a state actor\n"
            "(GPE, an Organization subpart of a GPE or a Person representing a GPE).\n"
        ]
    },
    "ace_citizenresidentreligionethnicity": {
        "en": [
            "CitizenResidentReligionEthnicity describes the relation between a Person entity and (1) the GPE in\n"
            "which they have citizenship, (2) the GPE or Location in which they live, the religious Organization or\n"
            "Person entity with which they have affiliation and (3) the GPE or PER entity that indicates their\n"
            "ethnicity.\n"
        ]
    },
    "ace_conflictevent": {
        "en": [
            "A ConflictEvent refers to either violent physical acts causing harm or damage, but are not covered by\n"
            "Life events (conflicts, clashes, fighting, gunfire, ...) or demonstrations (protests, sit-ins,\n"
            "strikes, riots, ...).\n"
        ]
    },
    "ace_contactevent": {
        "en": [
            "A ContactEvent occurs whenever two or more entities (persons or organization's representatives) come\n"
            "together at a single location and interact with one another face-to-face or directly enages in\n"
            "discussion via written or telephone communication.\n"
        ]
    },
    "ace_contactinfo": {
        "en": [
            "A ContactInfo value refers to contact information values such as telephone numbers, emails, addresses.\n"
            "For example: 'mich...@sumptionandwyland.com', ...\n"
        ]
    },
    "ace_convict": {
        "en": [
            "A Convict Event occurs whenever a Try Event ends with a successful prosecution of the Defendant. In\n"
            "other words, a Person, Organization or GPE Entity is convicted whenever that Entity has been found\n"
            "guilty of a Crime.\n"
        ]
    },
    "ace_crime": {
        "en": [
            "A Crime value refers to the specific reason (crime) that a Person entity can be judged or sentenced\n"
            "for. For example: 'raping', 'murder', 'drug', ...\n"
        ]
    },
    "ace_declarebankruptcy": {
        "en": [
            "A DeclareBankruptcy Event will occur whenever an Entity officially requests legal protection from debt\n"
            "collection due to an extremely negative balance sheet.\n"
        ]
    },
    "ace_demonstrate": {
        "en": [
            "A Demonstrate Event occurs whenever a large number of people come together in a public area to protest\n"
            "or demand some sort of official action. Demonstrate Events include, but are not limited to,\n"
            "protests, sit-ins, strikes, and riots.\n"
        ]
    },
    "ace_die": {
        "en": [
            "A Die Event occurs whenever the life of a Person Entity ends. Die Events can be accidental, intentional\n"
            "or self-inflicted\n"
        ]
    },
    "ace_divorce": {
        "en": [
            "A Divorce Event occurs whenever two people are officially divorced under the legal definition of\n"
            "divorce. We do not include separations or church annulments.\n"
        ]
    },
    "ace_elect": {
        "en": [
            "An Elect Event occurs whenever a candidate wins an election designed to determine the Person argument\n"
            "of a StartPosition Event.\n"
        ]
    },
    "ace_employment": {
        "en": [
            "Employment captures the relationship between Persons and their employers. This Relation is only\n"
            "taggable when it can be reasonably assumed that the Person is paid by the ORG or GPE.\n"
        ]
    },
    "ace_endorg": {
        "en": [
            "An EndOrg Event occurs whenever an Organization ceases to exist (in other words 'goes out of\n"
            "business').\n"
        ]
    },
    "ace_endposition": {
        "en": [
            "An EndPosition Event occurs whenever a Person Entity stops working for (or changes offices within) an\n"
            "Organization or GPE.\n"
        ]
    },
    "ace_execute": {
        "en": [
            "An Execute Event occurs whenever the life of a Person is taken by a state actor (a GPE, its\n"
            "Organization subparts, or Person representatives).\n"
        ]
    },
    "ace_extradite": {
        "en": [
            "An Extradite Event occurs whenever a Person is sent by a state actor from one Place (normally the GPE\n"
            "associated with the state actor, but sometimes a Facility under its control) to another place\n"
            "(Location, GPE or Facility) for the purposes of legal proceedings there.\n"
        ]
    },
    "ace_facility": {
        "en": [
            "A facility is a functional, primarily man-made structure. These include buildings and similar\n"
            "facilities designed for human habitation, such as houses, factories, stadiums, office buildings, ...\n"
            "Roughly speaking, facilities are artifacts falling under the domains of architecture and civil\n"
            "engineering.\n"
        ]
    },
    "ace_family": {
        "en": [
            "The Family Relation captures the connection between one entity and another with which it is in any\n"
            "familial relationship. Both arguments must be entities of type Person.\n"
        ]
    },
    "ace_fine": {
        "en": [
            "A Fine Event takes place whenever a state actor issues a financial punishment to a GPE, Person or\n"
            "Organization Entity, typically as a result of court proceedings.\n"
        ]
    },
    "ace_founder": {
        "en": [
            "Founder captures the relationship between an agent (Person, Organization, or GPE) and an Organization\n"
            "or GPE established or set up by that agent.\n"
        ]
    },
    "ace_genaffiliationrelation": {
        "en": [
            "The GenAffiliation Relation describes the citizen, resident, religion or ethnicity relation when the\n"
            "`arg1` is a Person. When the `arg1` is an Organization, the relation describes where it is located,\n"
            "based or does business.\n"
        ]
    },
    "ace_geographical": {
        "en": [
            "The Geographical relation captures the location of a Facility, Location, or GPE in or at or as a part\n"
            "of another Facility, Location, or GPE.\n"
        ]
    },
    "ace_injure": {
        "en": [
            "An Injure Event occurs whenever a Person Entity experiences physical harm. Injure Events can be\n"
            "accidental, intentional or self-inflicted.\n"
        ]
    },
    "ace_investorshareholder": {
        "en": [
            "InvestorShareholder captures the relationship between an agent (Person, Organization, or GPE) and an\n"
            "Organization in which the agent has invested or in which the agent owns shares/stock. Please note that\n"
            "agents may invest in GPEs.\n"
        ]
    },
    "ace_jobtitle": {
        "en": [
            "A JobTitle value refers to the name of the job or position of a Person entity in a Organization. For\n"
            "example: 'co-chief executive', 'move coordinator', 'interim ED', ...\n"
        ]
    },
    "ace_justiceevent": {
        "en": [
            "A JusticeEvent refers to any judicial action such as: arresting, jailing, releasing, granting parole,\n"
            "trial starting, hearing, charging, indicting, suing, convicting, sentencing, fine, executing,\n"
            "extraditing, adquiting, appealing or pardoning a Person entity.\n"
        ]
    },
    "ace_lastingpersonal": {
        "en": [
            "Lasting-Personal captures relationships that meet the following conditions: (1) The relationship must\n"
            "involve personal contact (or a reasonable assumption thereof). (2) There must be some indication or\n"
            "expectation that the relationship exists outside of a particular cited interaction. Both arguments\n"
            "must be entities of type Person.\n"
        ]
    },
    "ace_lifeevent": {
        "en": ["A LifeEvent occurs whenever a Person Entity borns, dies, gets married, divorced or gets injured."]
    },
    "ace_located": {
        "en": [
            "The Located relation captures the physical location of an entity. This relation is restricted to\n"
            "people. In other words, `arg1` in Located relations can only be occupied by mentions of Entities of\n"
            "type Person.\n"
        ]
    },
    "ace_location": {
        "en": [
            "Places defined on a geographical or astronomical basis which are mentioned in a document and do not\n"
            "constitute a political entity give rise to Location entities. These include, for example, the solar\n"
            "system, Mars, the Hudson River, Mt. Everest, and Death Valley.\n"
        ]
    },
    "ace_marry": {
        "en": ["Marry Events are official Events, where two people are married under the legal definition."]
    },
    "ace_meet": {
        "en": [
            "A Meet Event occurs whenever two or more Entities come together at a single location and interact with\n"
            "one another face-to-face. Meet Events include talks, summits, conferences, meetings, visits, and any\n"
            "other Event where two or more parties get together at some location.\n"
        ]
    },
    "ace_membership": {
        "en": [
            "Membership captures the relationship between an agent and an organization of which the agent is a\n"
            "member. Organizations and GPEs can be members of other Organizations (such as NATO or the UN).\n"
        ]
    },
    "ace_mergeorg": {
        "en": [
            "A MergeOrg Event occurs whenever two or more Organization Entities come together to form a new\n"
            "Organization Entity. This Event applies to any kind of Organization, including government agencies. It\n"
            "also includes joint venture.\n"
        ]
    },
    "ace_movementevent": {
        "en": [
            "A TransportEvent occurs whenever an Artifact (Weapon or Vehicle) or a Person is moved from one Place\n"
            "(GPE, Facility, Location) to another. This event requires the explicit mention of the Artifact or\n"
            "Person.\n"
        ]
    },
    "ace_near": {
        "en": [
            "Near indicates that an entity is explicitly near another entity, but neither entity is a part of the\n"
            "other or located in/at the other.\n"
        ]
    },
    "ace_nominate": {
        "en": [
            "A Nominate Event occurs whenever a Person is proposed for a StartPosition Event by the appropriate\n"
            "Person, through official channels.\n"
        ]
    },
    "ace_numeric": {
        "en": [
            "A Numeric value refers to relevant numbers, amounts, etc. For example: 'billions of dollars', '50\n"
            "percent', '100%', ...\n"
        ]
    },
    "ace_organization": {
        "en": [
            "Each organization or set of organizations mentioned in a document gives rise to an entity of type\n"
            "Organization. Typical examples are businesses, government units, sports teams, and formally organized\n"
            "music groups.\n"
        ]
    },
    "ace_organizationaffiliationrelation": {
        "en": [
            "The OrganizationAffiliation Relation describes the relations between a Person (or other Organizations)\n"
            "and a related Organization. This relation includes: employment, ownership, founder, student or alumn,\n"
            "sport affiliation, inverstor or shareholder and membership relations.\n"
        ]
    },
    "ace_orglocationorigin": {
        "en": [
            "OrgLocationOrigin captures the relationship between an organization and the Location or GPE where it is\n"
            "located, based, or does business. Note: Subsidiary trumps this relation for government organizations.\n"
        ]
    },
    "ace_ownership": {
        "en": [
            "Ownership captures the relationship between a Person and an Organization owned by that Person. If the\n"
            "`arg2` is not an Organization, use the Agent-Artifact relation.\n"
        ]
    },
    "ace_pardon": {
        "en": [
            "A Pardon Event occurs whenever a head-of-state or their appointed representative lifts a sentence\n"
            "imposed by the judiciary.\n"
        ]
    },
    "ace_partwholerelation": {
        "en": [
            "The PartWhole Relation refers to the semantic relation between two entities that are parts of a larger\n"
            "whole or vice versa. For example, the relation between a country and its states, or between a company\n"
            "and its subsidiaries, are instances of PartWhole relations.\n"
        ]
    },
    "ace_person": {
        "en": [
            "Each distinct person or set of people mentioned in a document refers to an entity of type Person. For\n"
            'example, people may be specified by name ("John Smith"), occupation ("the butcher"), family relation\n'
            '("dad"), pronoun ("he"), etc., or by some combination of these.'
        ]
    },
    "ace_personalsocialrelation": {
        "en": [
            "The Personal-Social Relation describe the relationship between people. Both arguments must be entities\n"
            "of type Person. Please note: The arguments of these Relations are not ordered. The Relations are\n"
            "symmetric.\n"
        ]
    },
    "ace_personnelevent": {
        "en": [
            "A PersonnelEvent occurs when a Person entity changes its job position (JobTitle entity) with respect an\n"
            "Organization entity. It includes when a person starts working, ends working, changes offices within,\n"
            "gets nominated or is elected for a position in a Organization.\n"
        ]
    },
    "ace_phonewrite": {
        "en": [
            "A PhoneWrite Event occurs when two or more people directly engage in discussion which does not take\n"
            "place 'face-to-face'. To make this Event less open-ended, we limit it to written or telephone\n"
            "communication where at least two parties are specified.\n"
        ]
    },
    "ace_physicalrelation": {
        "en": [
            "The Physical Relation captures the physical location relation of entities such as: a Person entity\n"
            "located in a Facility, Location or GPE; or two entities that are near, but neither entity is a part of\n"
            "the other or located in/at the other.\n"
        ]
    },
    "ace_releaseparole": {
        "en": [
            "A Release Event occurs whenever a state actor (GPE, Organization subpart, or Person representative)\n"
            "ends its custody of a Person Entity.\n"
        ]
    },
    "ace_sentence": {
        "en": [
            "A Sentence value refers to sentences decided by a court or judge for a specific crime. For example:\n"
            "'124 years in prison', 'a sentence', 'death'...\n"
        ]
    },
    "ace_sentenceact": {
        "en": [
            "A SentenceAct Event takes place whenever the punishment (particularly incarceration) for the Defendant\n"
            "of a Try Event is issued by a state actor (a GPE, an Organization subpart or a Person representing\n"
            "them)\n"
        ]
    },
    "ace_sportsaffiliation": {
        "en": [
            "Sports-Affiliation captures the relationship between a player, coach, manager, or assistant and his or\n"
            "her affiliation with a sports organization (including sports leagues or divisions as well as\n"
            "individual sports teams).\n"
        ]
    },
    "ace_startorg": {"en": ["A StartOrg Event occurs whenever a new Organization is created."]},
    "ace_startposition": {
        "en": [
            "A StartPosition Event occurs whenever a Person Entity begins working for (or changes offices within) an\n"
            "Organization or GPE. This includes government officials starting their terms, whether elected or\n"
            "appointed.\n"
        ]
    },
    "ace_studentalum": {
        "en": [
            "StudentAlum captures the relationship between a Person and an educational institution the Person\n"
            "attends or attended.\n"
        ]
    },
    "ace_subsidiary": {
        "en": [
            "Subsidiary captures the ownership, administrative, and other hierarchical relationships between\n"
            "organizations and between organizations and GPEs.\n"
        ]
    },
    "ace_sue": {
        "en": [
            "A Sue Event occurs whenever a court proceeding has been initiated for the purposes of determining the\n"
            "liability of a Person, Organization or GPE accused of committing a crime or neglecting a commitment.\n"
        ]
    },
    "ace_time": {
        "en": [
            "A Time value refers to a specific time frame. Usually known as time expressions. For example: '4\n"
            "years', 'today', 'December', 'future', ...\n"
        ]
    },
    "ace_transactionevent": {
        "en": [
            "A TransactionEvent refers to buying, selling, loaning, borrowing, giving, or receving of Artifacts or\n"
            "Organizations; or giving, receiving, borrowing, or lending Money.\n"
        ]
    },
    "ace_transfermoney": {
        "en": [
            "TransferMoney Events refer to the giving, receiving, borrowing, or lending money when it is not in the\n"
            "context of purchasing something. The canonical examples are: (1) people giving money to organizations\n"
            "(and getting nothing tangible in return); and (2) organizations lending money to people or other orgs.\n"
        ]
    },
    "ace_transferownership": {
        "en": [
            "TransferOwnership Events refer to the buying, selling, loaning, borrowing, giving, or receiving of\n"
            "artifacts or organizations.\n"
        ]
    },
    "ace_transport": {
        "en": [
            "A Transport Event occurs whenever an Artifact (Weapon or Vehicle) or a Person is moved from one Place\n"
            "(GPE, Facility, Location) to another.\n"
        ]
    },
    "ace_trialhearing": {
        "en": [
            "A Trial Event occurs whenever a court proceeding has been initiated for the purposes of determining the\n"
            "guilt or innocence of a Person, Organization or GPE accused of committing a crime.\n"
        ]
    },
    "ace_userownerinventormanufacturer": {
        "en": [
            "This Relation applies when an agent owns an artifact, has possession of an artifact, uses an artifact,\n"
            "or caused an artifact to come into being. Note: if `arg2` is an Organization, use Ownership relation\n"
            "(arg1=PER) or Subsidiary relation (arg1=ORG or GPE).\n"
        ]
    },
    "ace_vehicle": {
        "en": [
            "A Vehicle entity refers to vehicles that are used for transportation. The vehicles can transport either\n"
            "persons or artifacts. For example: 'car', 'plane', 'cabin', ...\n"
        ]
    },
    "ace_weapon": {
        "en": [
            "A Weapon entity refers to instruments that can be used to deal physical damage, destroy something or\n"
            "kill someone. For example: 'bomb', 'm-16s', 'missile', ...\n"
        ]
    },
}

EXAMPLES = {
    "ace_lifeevent_examples": {
        "en": [
            "wounded",
            "divorce",
            "birth",
            "born",
            "marriage",
            "married",
            "suicide",
            "killing",
            "injured",
            "died",
            "killed",
        ]
    },
    "ace_movementevent_examples": {
        "en": ["travel", "arrived", "going", "moving", "take", "went", "get", "trip", "come", "go"]
    },
    "ace_transactionevent_examples": {
        "en": [
            "donate",
            "received",
            "paying",
            "seize",
            "contributions",
            "donations",
            "captured",
            "bought",
            "paid",
            "pay",
        ]
    },
    "ace_businessevent_examples": {
        "en": [
            "started",
            "open",
            "create",
            "closing",
            "merged",
            "merging",
            "set up",
            "founded",
            "merger",
            "bankruptcy",
        ]
    },
    "ace_conflictevent_examples": {
        "en": ["terrorism", "combat", "hit", "fight", "bombing", "fire", "attacks", "fighting", "attack", "war"]
    },
    "ace_contactevent_examples": {
        "en": ["meetings", "conference", "talked", "met", "letters", "summit", "call", "meet", "talks", "meeting"]
    },
    "ace_personnelevent_examples": {
        "en": [
            "won",
            "appoint",
            "retired",
            "fired",
            "appointed",
            "resigned",
            "elected",
            "former",
            "elections",
            "election",
        ]
    },
    "ace_justiceevent_examples": {
        "en": [
            "parole",
            "sued",
            "sentenced",
            "appeal",
            "charged",
            "sentence",
            "arrested",
            "convicted",
            "charges",
            "trial",
        ]
    },
    "ace_person_examples": {
        "en": ["I", "you", "they", "troops", "forces", "people", "who", "officials", "we", "them"]
    },
    "ace_organization_examples": {
        "en": [
            "they",
            "military",
            "company",
            "CNN",
            "U.N.",
            "companies",
            "police",
            "administration",
            "court",
            "government",
        ]
    },
    "ace_gpe_examples": {
        "en": ["Iraq", "U.S.", "they", "country", "Baghdad", "city", "US", "coalition", "Israel", "British"]
    },
    "ace_location_examples": {
        "en": ["world", "area", "there", "where", "here", "mars", "areas", "border", "planet", "region"]
    },
    "ace_facility_examples": {
        "en": ["home", "airport", "base", "there", "hospital", "house", "prison", "room", "bridge", "road"]
    },
    "ace_weapon_examples": {
        "en": ["weapons", "gun", "missiles", "nuclear", "artillery", "bomb", "arms", "granades", "explosive", "knife"]
    },
    "ace_vehicle_examples": {
        "en": ["ship", "car", "aircraft", "tank", "plane", "bus", "vehicle", "boat", "helicopter", "truck"]
    },
    "ace_time_examples": {
        "en": [
            "Tuesday",
            "Monday",
            "tonight",
            "new",
            "this morning",
            "yesterday",
            "right now",
            "former",
            "today",
            "now",
        ]
    },
    "ace_numeric_examples": {
        "en": [
            "$800 million",
            "$30 million",
            "60 percent",
            "$2 million",
            "$500 billion",
            "billions of dollars",
            "40%",
            "$90 million",
            "100%",
            "$250,000",
        ]
    },
    "ace_jobtitle_examples": {
        "en": [
            "a pastor",
            "a counselor",
            "Presidential",
            "Vice President",
            "Governor of uh New Jersey",
            "head of Homeland Security",
            "president",
            "governor",
            "the Shah of Iran",
            "the CEO",
        ]
    },
    "ace_crime_examples": {
        "en": [
            "corruption",
            "theft",
            "receiving bribes",
            "killing",
            "insider trading",
            "attempted murder",
            "treason",
            "obstruction of justice",
            "sodomy",
            "murder",
        ]
    },
    "ace_sentence_examples": {
        "en": [
            "prison for life",
            "the maximum sentence",
            "a death penalty",
            "more than seven years in prison",
            "death",
            "five years probation",
            "100 lashes",
            "life",
            "life in prison",
            "the death penalty",
        ]
    },
    "ace_contactinfo_examples": {
        "en": [
            "911",
            "www.werkheiserfordelegate.com",
            "770-239-X X X X",
            "DKo...@hotmail.com",
            "a...@yahoo.com",
            "cnn.com/wolf",
            "www.vitac.com",
            "am@cnn.com",
        ]
    },
}

src/tasks/ace/guidelines_gold.py
hitz-zentroa/GoLLIE
"""
Adapted from OneIE: http://blender.cs.illinois.edu/software/oneie/
This script extracts IE annotations from ACE2005 (LDC2006T06).
"""

import glob
import json
import os
import re
from argparse import ArgumentParser
from dataclasses import dataclass
from typing import Any, Dict, List, Tuple

import tqdm
from bs4 import BeautifulSoup
from nltk import (
    sent_tokenize as sent_tokenize_,
)
from nltk import (
    wordpunct_tokenize as wordpunct_tokenize_,
)


TAG_PATTERN = re.compile("<[^<>]+>", re.MULTILINE)

DOCS_TO_REVISE_SENT = {
    "CNN_ENG_20030529_130011.6": [
        (209, 254),
        (461, 504),
        (668, 859),
        (984, 1074),
        (1577, 1632),
    ],
    "CNN_ENG_20030626_203133.11": [(1497, 1527)],
    "CNN_ENG_20030526_180540.6": [(67, 99)],
    "CNNHL_ENG_20030523_221118.14": [(136, 174)],
    "BACONSREBELLION_20050127.1017": [(2659, 2663), (4381, 4405), (410, 458)],
    "misc.legal.moderated_20050129.2225": [(4118, 4127), (4710, 4794)],
    "alt.vacation.las-vegas_20050109.0133": [(1201, 1248)],
    "alt.obituaries_20041121.1339": [(1947, 2044), (1731, 1737)],
    "APW_ENG_20030326.0190": [(638, 739)],
    "APW_ENG_20030403.0862": [(729, 781)],
    "CNN_IP_20030405.1600.02": [(699, 705)],
    "CNN_IP_20030403.1600.00-1": [(2392, 2399)],
    "CNN_IP_20030409.1600.04": [(1039, 1050)],
    "CNN_IP_20030412.1600.03": [(741, 772)],
    "CNN_IP_20030402.1600.02-1": [(885, 892)],
    "CNN_IP_20030329.1600.02": [(3229, 3235)],
    "CNN_IP_20030409.1600.02": [(477, 498)],
    "CNN_CF_20030304.1900.04": [(522, 575), (5193, 5210), (5461, 5542)],
    "CNN_IP_20030403.1600.00-3": [(1487, 1493)],
    "soc.history.war.world-war-ii_20050127.2403": [(414, 441)],
}

# Inconsistency between data and annotation guideline argument names
arg_name_mapping = {
    "ATTACK": {"Victim": "Target", "Agent": "Attacker"},
    "APPEAL": {"Plaintiff": "Prosecutor"},
    "PHONE-WRITE": {"Place": None},
}


def mask_escape(text: str) -> str:
    """Replaces escaped characters with rare sequences.

    Args:
        text (str): text to mask.

    Returns:
        str: masked string.
    """
    return text.replace("&amp;", "ҪҪҪҪҪ").replace("&lt;", "ҚҚҚҚ").replace("&gt;", "ҺҺҺҺ")


def unmask_escape(text: str) -> str:
    """Replaces masking sequences with the original escaped characters.

    Args:
        text (str): masked string.

    Returns:
        str: unmasked string.
    """
    return text.replace("ҪҪҪҪҪ", "&amp;").replace("ҚҚҚҚ", "&lt;").replace("ҺҺҺҺ", "&gt;")


def recover_escape(text: str) -> str:
    """Converts named character references in the given string to the corresponding
    Unicode characters. I didn't notice any numeric character references in this
    dataset.

    Args:
        text (str): text to unescape.

    Returns:
        str: unescaped string.
    """
    return text.replace("&amp;", "&").replace("&lt;", "<").replace("&gt;", ">")


def sent_tokenize(text: Tuple[str, int, int], language: str = "english") -> List[Tuple[str, int, int]]:
    """Performs sentence tokenization. For English, it uses NLTK's sent_tokenize
    function. For Chinese, it uses split_chinese_sentence, a simple sentence
    tokenizer implemented by myself.

    Args:
        text (Tuple[str, int, int]): a tuple of three elements, text to split
            into sentences, start offset, and end offset.
        language (str): available options: english, chinese.

    Returns:
        List[Tuple[str, int, int]]: a list of sentences.
    """
    text, start, end = text
    if language == "chinese":
        sentences = split_chinese_sentence(text)
    else:
        sentences = sent_tokenize_(text, language=language)

    last = 0
    sentences_ = []
    for sent in sentences:
        index = text[last:].find(sent)
        if index == -1:
            print(text, sent)
        else:
            sentences_.append((sent, last + index + start, last + index + len(sent) + start))
        last += index + len(sent)
    return sentences_


def wordpunct_tokenize(text: str, language: str = "english") -> List[str]:
    """Performs word tokenization. For English, it uses NLTK's
    wordpunct_tokenize function. For Chinese, it simply splits the sentence into
    characters.

    Args:
        text (str): text to split into words.
        language (str): available options: english, chinese.

    Returns:
        List[str]: a list of words.
    """
    if language == "chinese":
        return [c for c in text if c.strip()]
    return wordpunct_tokenize_(text)


def split_chinese_sentence(text: str) -> List[str]:
    """Performs sentence tokenization for Chinese.

    Args:
        text (str): text to split into sentences.

    Returns:
        List[str]: a list of sentences.
    """
    sentences = []
    quote_mark_count = 0
    sentence = ""
    for i, c in enumerate(text):
        sentence += c
        if c in {"”", "」"}:
            sentences.append(sentence)
            sentence = ""
        elif c in {"。", "!", "?", "！", "？"}:
            if i < len(text) - 1 and text[i + 1] not in {"”", '"', "」"}:
                sentences.append(sentence)
                sentence = ""
        elif c == '"':
            quote_mark_count += 1
            if quote_mark_count % 2 == 0 and len(sentence) > 2 and sentence[-2] in {"？", "！", "。", "?", "!"}:
                sentences.append(sentence)
                sentence = ""
    if sentence:
        sentences.append(sentence)
    return sentences


@dataclass
class Span:
    start: int
    end: int
    text: str

    def __post_init__(self):
        self.start = int(self.start)
        self.end = int(self.end)
        self.text = self.text.replace("\n", " ")

    def char_offsets_to_token_offsets(self, tokens: List[Tuple[int, int, str]]):
        """Converts self.start and self.end from character offsets to token
        offsets.

        Args:
            tokens (List[int, int, str]): a list of token tuples. Each item in
                the list is a triple (start_offset, end_offset, text).
        """
        start_ = end_ = -1
        for i, (s, e, _) in enumerate(tokens):
            if s == self.start:
                start_ = i
            if e == self.end:
                end_ = i + 1
        if start_ == -1 or end_ == -1 or start_ > end_:
            raise ValueError(
                "Failed to update offsets for {}-{}:{} in {}".format(self.start, self.end, self.text, tokens)
            )
        self.start, self.end = start_, end_

    def to_dict(self) -> Dict[str, Any]:
        """Converts instance variables to a dict.

        Returns:
            dict: a dict of instance variables.
        """
        return {"text": recover_escape(self.text), "start": self.start, "end": self.end}

    def remove_space(self):
        """Removes heading and trailing spaces in the span text."""
        # heading spaces
        text = self.text.lstrip(" ")
        self.start += len(self.text) - len(text)
        # trailing spaces
        text = text.rstrip(" ")
        self.text = text
        self.end = self.start + len(text)

    def copy(self):
        """Makes a copy of itself.

        Returns:
            Span: a copy of itself."""
        return Span(self.start, self.end, self.text)


@dataclass
class Entity(Span):
    entity_id: str
    mention_id: str
    entity_type: str
    entity_subtype: str
    mention_type: str
    value: str = None

    def to_dict(self) -> Dict[str, Any]:
        """Converts instance variables to a dict.

        Returns:
            Dict: a dict of instance variables.
        """
        entity_dict = {
            "text": recover_escape(self.text),
            "entity_id": self.entity_id,
            "mention_id": self.mention_id,
            "start": self.start,
            "end": self.end,
            "entity_type": self.entity_type,
            "entity_subtype": self.entity_subtype,
            "mention_type": self.mention_type,
        }
        if self.value:
            entity_dict["value"] = self.value
        return entity_dict


@dataclass
class RelationArgument:
    mention_id: str
    role: str
    text: str

    def to_dict(self) -> Dict[str, Any]:
        """Converts instance variables to a dict.

        Returns:
            Dict[str, Any]: a dict of instance variables.
        """
        return {
            "mention_id": self.mention_id,
            "role": self.role,
            "text": recover_escape(self.text),
        }


@dataclass
class Relation:
    relation_id: str
    relation_type: str
    relation_subtype: str
    arg1: RelationArgument
    arg2: RelationArgument

    def to_dict(self) -> Dict[str, Any]:
        """Converts instance variables to a dict.

        Returns:
            Dict[str, Any]: a dict of instance variables.
        """
        return {
            "relation_id": self.relation_id,
            "relation_type": self.relation_type,
            "relation_subtype": self.relation_subtype,
            "arg1": self.arg1.to_dict(),
            "arg2": self.arg2.to_dict(),
        }


@dataclass
class EventArgument:
    mention_id: str
    role: str
    text: str

    def to_dict(self) -> Dict[str, Any]:
        """Converts instance variables to a dict.

        Returns:
            Dict[str, Any]: a dict of instance variables.
        """
        return {
            "mention_id": self.mention_id,
            "role": self.role,
            "text": recover_escape(self.text),
        }


@dataclass
class Event:
    event_id: str
    mention_id: str
    event_type: str
    event_subtype: str
    trigger: Span
    arguments: List[EventArgument]

    def to_dict(self) -> Dict[str, Any]:
        """Converts instance variables to a dict.

        Returns:
            Dict[str, Any]: a dict of instance variables.
        """
        return {
            "event_id": self.event_id,
            "mention_id": self.mention_id,
            "event_type": self.event_type,
            "event_subtype": self.event_subtype,
            "trigger": self.trigger.to_dict(),
            "arguments": [arg.to_dict() for arg in self.arguments],
        }


@dataclass
class Sentence(Span):
    sent_id: str
    tokens: List[str]
    entities: List[Entity]
    relations: List[Relation]
    events: List[Event]

    def to_dict(self) -> Dict[str, Any]:
        """Converts instance variables to a dict.

        Returns:
            Dict[str, Any]: a dict of instance variables.
        """
        return {
            "sent_id": self.sent_id,
            "tokens": [recover_escape(t) for t in self.tokens],
            "entities": [entity.to_dict() for entity in self.entities],
            "relations": [relation.to_dict() for relation in self.relations],
            "events": [event.to_dict() for event in self.events],
            "start": self.start,
            "end": self.end,
            "text": recover_escape(self.text).replace("\t", " "),
        }


@dataclass
class Document:
    doc_id: str
    sentences: List[Sentence]

    def to_dict(self) -> Dict[str, Any]:
        """Converts instance variables to a dict.

        Returns:
            Dict[str, Any]: a dict of instance variables.
        """
        return {
            "doc_id": self.doc_id,
            "sentences": [sent.to_dict() for sent in self.sentences],
        }


def revise_sentences(sentences: List[Tuple[str, int, int]], doc_id: str) -> List[Tuple[int, int, str]]:
    """Automatic sentence tokenization may have errors for a few documents.

    Args:
        sentences (List[Tuple[str, int, int]]): a list of sentence tuples.
        doc_id (str): document ID.

    Returns:
        List[Tuple[str, int, int]]: a list of revised sentence tuples.
    """
    sentences_ = []

    offset_list = DOCS_TO_REVISE_SENT[doc_id]
    first_part_offsets = {offset for offset, _ in offset_list}
    second_part_offsets = {offset for _, offset in offset_list}

    for sentence_idx, (text, start, end) in enumerate(sentences):
        if start in first_part_offsets:
            next_text, next_start, next_end = sentences[sentence_idx + 1]
            space = " " * (next_start - end)
            sentences_.append((text + space + next_text, start, next_end))
        elif start in second_part_offsets:
            continue
        else:
            sentences_.append((text, start, end))

    return sentences_


def read_sgm_file(path: str, language: str = "english") -> List[Tuple[str, int, int]]:
    """Reads a SGM text file.

    Args:
        path (str): path to the input file.
        language (str): document language. Valid values: "english" or "chinese".

    Returns:
        List[Tuple[str, int, int]]: a list of sentences. Each item in the list
            is a tuple of three elements, sentence text, start offset, and end
            offset.
    """
    data = open(path, "r", encoding="utf-8").read()

    # Chunk the document
    chunks = TAG_PATTERN.sub("⁑", data).split("⁑")

    # Get the offset of <TEXT>
    data = data.replace("<TEXT>", "⁂")
    data = TAG_PATTERN.sub("", data)
    min_offset = max(0, data.find("⁂"))
    data = data.replace("⁂", "")

    # Extract sentences from chunks
    chunk_offset = 0
    sentences = []
    for chunk in chunks:
        lines = chunk.split("\n")
        current_sentence = []
        start = offset = 0
        for line in lines:
            offset += len(line) + 1
            if line.strip():
                current_sentence.append(line)
            else:
                # empty line
                if current_sentence:
                    sentence = " ".join(current_sentence)
                    if start + chunk_offset >= min_offset:
                        sentences.append(
                            (
                                sentence,
                                start + chunk_offset,
                                start + chunk_offset + len(sentence),
                            )
                        )
                    current_sentence = []
                start = offset
        if current_sentence:
            sentence = " ".join(current_sentence)
            if start + chunk_offset >= min_offset:
                sentences.append(
                    (
                        sentence,
                        start + chunk_offset,
                        start + chunk_offset + len(sentence),
                    )
                )
        chunk_offset += len(chunk)

    # Re-tokenize sentences
    sentences = [s for sent in sentences for s in sent_tokenize(sent, language=language)]

    return sentences


def read_apf_file(path: str, time_and_val: bool = False) -> Tuple[str, str, List[Entity], List[Relation], List[Event]]:
    """Reads an APF file.

    Args:
        path (str): path to the input file.
        time_and_val (bool): extract times and values or not.

    Returns:
        doc_id (str): document ID.
        source (str): document source.
        entity_list (List[Entity]): a list of Entity instances.
        relation_list (List[Relation]): a list of Relation instances.
        event_list (List[Event]): a list of Events instances.
    """
    data = open(path, "r", encoding="utf-8").read()
    soup = BeautifulSoup(data, "lxml-xml")

    # metadata
    root = soup.find("source_file")
    source = root["SOURCE"]
    doc = root.find("document")
    doc_id = doc["DOCID"]

    entity_list, relation_list, event_list = [], [], []

    # entities: nam, nom, pro
    for entity in doc.find_all("entity"):
        entity_id = entity["ID"]
        entity_type = entity["TYPE"]
        entity_subtype = entity["SUBTYPE"]
        for entity_mention in entity.find_all("entity_mention"):
            mention_id = entity_mention["ID"]
            mention_type = entity_mention["TYPE"]
            head = entity_mention.find("head").find("charseq")
            start, end, text = int(head["START"]), int(head["END"]), head.text
            entity_list.append(
                Entity(
                    start,
                    end,
                    text,
                    entity_id,
                    mention_id,
                    entity_type,
                    entity_subtype,
                    mention_type,
                )
            )

    if time_and_val:
        # entities: value
        for entity in doc.find_all("value"):
            entity_id = entity["ID"]
            entity_type = entity["TYPE"]
            entity_subtype = entity.get("SUBTYPE", None)
            for entity_mention in entity.find_all("value_mention"):
                mention_id = entity_mention["ID"]
                mention_type = "VALUE"
                extent = entity_mention.find("extent").find("charseq")
                start, end, text = int(extent["START"]), int(extent["END"]), extent.text
                entity_list.append(
                    Entity(
                        start,
                        end,
                        text,
                        entity_id,
                        mention_id,
                        entity_type,
                        entity_subtype,
                        mention_type,
                    )
                )

        # entities: timex
        for entity in doc.find_all("timex2"):
            entity_id = entity["ID"]
            entity_subtype = "TIME"
            value = entity.get("VAL", None)
            for entity_mention in entity.find_all("timex2_mention"):
                mention_id = entity_mention["ID"]
                mention_type = "TIME"
                extent = entity_mention.find("extent").find("charseq")
                start, end, text = int(extent["START"]), int(extent["END"]), extent.text
                entity_list.append(
                    Entity(
                        start,
                        end,
                        text,
                        entity_id,
                        mention_id,
                        entity_type,
                        entity_subtype,
                        mention_type,
                        value=value,
                    )
                )

    # relations
    for relation in doc.find_all("relation"):
        relation["ID"]
        relation_type = relation["TYPE"]
        if relation_type == "METONYMY":
            continue
        relation_subtype = relation["SUBTYPE"]
        for relation_mention in relation.find_all("relation_mention"):
            mention_id = relation_mention["ID"]
            arg1 = arg2 = None
            for arg in relation_mention.find_all("relation_mention_argument"):
                arg_mention_id = arg["REFID"]
                arg_role = arg["ROLE"]
                arg_text = arg.find("extent").find("charseq").text
                if arg_role == "Arg-1":
                    arg1 = RelationArgument(arg_mention_id, arg_role, arg_text)
                elif arg_role == "Arg-2":
                    arg2 = RelationArgument(arg_mention_id, arg_role, arg_text)
            if arg1 and arg2:
                relation_list.append(Relation(mention_id, relation_type, relation_subtype, arg1, arg2))

    # events
    for event in doc.find_all("event"):
        event_id = event["ID"]
        event_type = event["TYPE"]
        event_subtype = event["SUBTYPE"]
        for event_mention in event.find_all("event_mention"):
            mention_id = event_mention["ID"]
            trigger = event_mention.find("anchor").find("charseq")
            trigger_start, trigger_end = int(trigger["START"]), int(trigger["END"])
            trigger_text = trigger.text
            event_args = []
            for arg in event_mention.find_all("event_mention_argument"):
                arg_mention_id = arg["REFID"]
                arg_role = arg["ROLE"]
                arg_text = arg.find("extent").find("charseq").text
                event_args.append(EventArgument(arg_mention_id, arg_role, arg_text))
            event_list.append(
                Event(
                    event_id,
                    mention_id,
                    event_type,
                    event_subtype,
                    Span(trigger_start, trigger_end + 1, trigger_text),
                    event_args,
                )
            )

    # remove heading/tailing spaces
    for entity in entity_list:
        entity.remove_space()
    for event in event_list:
        event.trigger.remove_space()

    return doc_id, source, entity_list, relation_list, event_list


def process_entities(entities: List[Entity], sentences: List[Tuple[str, int, int]]) -> List[List[Entity]]:
    """Cleans entities and splits them into lists

    Args:
        entities (List[Entity]): a list of Entity instances.
        sentences (List[Tuple[str, int, int]]): a list of sentences.

    Returns:
        List[List[Entity]]: a list of sentence entity lists.
    """
    sentence_entities = [[] for _ in range(len(sentences))]

    # assign each entity to the sentence where it appears
    for entity in entities:
        start, end = entity.start, entity.end
        for i, (_, s, e) in enumerate(sentences):
            if start >= s and end <= e:
                sentence_entities[i].append(entity)
                break

    # remove overlapping entities
    sentence_entities_cleaned = [[] for _ in range(len(sentences))]
    for i, entities in enumerate(sentence_entities):
        if not entities:
            continue
        # prefer longer entities
        entities.sort(key=lambda x: (x.end - x.start), reverse=True)
        chars = [0] * max([x.end for x in entities])
        for entity in entities:
            overlap = False
            for j in range(entity.start, entity.end):
                if chars[j] == 1:
                    overlap = True
                    break
            if not overlap:
                chars[entity.start : entity.end] = [1] * (entity.end - entity.start)
                sentence_entities_cleaned[i].append(entity)
        sentence_entities_cleaned[i].sort(key=lambda x: x.start)

    return sentence_entities_cleaned


def process_events(
    events: List[Event],
    sentence_entities: List[List[Entity]],
    sentences: List[Tuple[str, int, int]],
) -> List[List[Event]]:
    """Cleans and assigns events.

    Args:
        events (List[Event]): A list of Event objects
        entence_entities (List[List[Entity]]): A list of sentence entity lists.
        sentences (List[Tuple[str, int, int]]): A list of sentences.

    Returns:
        List[List[Event]]: a list of sentence event lists.
    """
    sentence_events = [[] for _ in range(len(sentences))]
    # assign each event mention to the sentence where it appears
    for event in events:
        start, end = event.trigger.start, event.trigger.end
        for i, (_, s, e) in enumerate(sentences):
            sent_entities = sentence_entities[i]
            if start >= s and end <= e:
                # clean the argument list
                arguments = []
                for argument in event.arguments:
                    # entity_id = argument.entity_id
                    mention_id = argument.mention_id
                    for entity in sent_entities:
                        if entity.mention_id == mention_id:
                            arguments.append(argument)
                            break
                event_cleaned = Event(
                    event.event_id,
                    event.mention_id,
                    event.event_type,
                    event.event_subtype,
                    trigger=event.trigger.copy(),
                    arguments=arguments,
                )
                sentence_events[i].append(event_cleaned)

    # remove overlapping events
    sentence_events_cleaned = [[] for _ in range(len(sentences))]
    for i, events in enumerate(sentence_events):
        if not events:
            continue
        events.sort(key=lambda x: (x.trigger.end - x.trigger.start), reverse=True)
        chars = [0] * max([x.trigger.end for x in events])
        for event in events:
            overlap = False
            for j in range(event.trigger.start, event.trigger.end):
                if chars[j] == 1:
                    overlap = True
                    break
            if not overlap:
                chars[event.trigger.start : event.trigger.end] = [1] * (event.trigger.end - event.trigger.start)
                sentence_events_cleaned[i].append(event)
        sentence_events_cleaned[i].sort(key=lambda x: x.trigger.start)

    return sentence_events_cleaned


def process_relation(
    relations: List[Relation],
    sentence_entities: List[List[Entity]],
    sentences: List[Tuple[str, int, int]],
) -> List[List[Relation]]:
    """Cleans and assigns relations

    Args:
        relations (List[Relation]): a list of Relation instances.
        sentence_entities (List[List[Entity]]): a list of sentence entity lists.
        sentences (List[Tuple[str, int, int]]): a list of sentences.

    Returns:
        List[List[Relation]]: a list of sentence relation lists.
    """
    sentence_relations = [[] for _ in range(len(sentences))]
    for relation in relations:
        mention_id1 = relation.arg1.mention_id
        mention_id2 = relation.arg2.mention_id
        for i, entities in enumerate(sentence_entities):
            arg1_in_sent = any(mention_id1 == e.mention_id for e in entities)
            arg2_in_sent = any(mention_id2 == e.mention_id for e in entities)
            if arg1_in_sent and arg2_in_sent:
                sentence_relations[i].append(relation)
                break
            elif arg1_in_sent != arg2_in_sent:
                break
    return sentence_relations


def tokenize(
    sentence: Tuple[str, int, int],
    entities: List[Entity],
    events: List[Event],
    language: str = "english",
) -> List[Tuple[int, int, str]]:
    """Tokenizes a sentence.
    Each sentence is first split into chunks that are entity/event spans or words
    between two spans. After that, word tokenization is performed on each chunk.

    Args:
        sentence (Tuple[str, int, int]): Sentence tuple (text, start, end)
        entities (List[Entity]): A list of Entity instances.
        events (List[Event]): A list of Event instances.

    Returns:
        List[Tuple[int, int, str]]: a list of token tuples. Each tuple consists
        of three elements, start offset, end offset, and token text.
    """
    text, start, end = sentence
    text = mask_escape(text)

    # split the sentence into chunks
    splits = {0, len(text)}
    for entity in entities:
        splits.add(entity.start - start)
        splits.add(entity.end - start)
    for event in events:
        splits.add(event.trigger.start - start)
        splits.add(event.trigger.end - start)
    splits = sorted(splits)
    chunks = [(splits[i], splits[i + 1], text[splits[i] : splits[i + 1]]) for i in range(len(splits) - 1)]

    # tokenize each chunk
    chunks = [(s, e, t, wordpunct_tokenize(t, language=language)) for s, e, t in chunks]

    # merge chunks and add word offsets
    tokens = []
    for chunk_start, chunk_end, chunk_text, chunk_tokens in chunks:
        last = 0
        chunk_tokens_ = []
        for token in chunk_tokens:
            token_start = chunk_text[last:].find(token)
            if token_start == -1:
                raise ValueError("Cannot find token {} in {}".format(token, text))
            token_end = token_start + len(token)
            chunk_tokens_.append(
                (
                    token_start + start + last + chunk_start,
                    token_end + start + last + chunk_start,
                    unmask_escape(token),
                )
            )
            last += token_end
        tokens.extend(chunk_tokens_)
    return tokens


def convert(sgm_file: str, apf_file: str, time_and_val: bool = False, language: str = "english") -> Document:
    """Converts a document.

    Args:
        sgm_file (str): path to a SGM file.
        apf_file (str): path to a APF file.
        time_and_val (bool, optional): extracts times and values or not.
            Defaults to False.
        language (str, optional): document language. Available options: english,
            chinese. Defaults to 'english'.

    Returns:
        Document: a Document instance.
    """
    sentences = read_sgm_file(sgm_file, language=language)
    doc_id, source, entities, relations, events = read_apf_file(apf_file, time_and_val=time_and_val)

    # Reivse sentences
    if doc_id in DOCS_TO_REVISE_SENT:
        sentences = revise_sentences(sentences, doc_id)

    # Process entities, relations, and events
    sentence_entities = process_entities(entities, sentences)
    sentence_relations = process_relation(relations, sentence_entities, sentences)
    sentence_events = process_events(events, sentence_entities, sentences)

    # Tokenization
    sentence_tokens = [
        tokenize(s, ent, evt, language=language) for s, ent, evt in zip(sentences, sentence_entities, sentence_events)
    ]

    # Convert span character offsets to token indices
    sentence_objs = []
    for i, (toks, ents, evts, rels, sent) in enumerate(
        zip(
            sentence_tokens,
            sentence_entities,
            sentence_events,
            sentence_relations,
            sentences,
        )
    ):
        for entity in ents:
            entity.char_offsets_to_token_offsets(toks)
        for event in evts:
            event.trigger.char_offsets_to_token_offsets(toks)
        sent_id = "{}-{}".format(doc_id, i)
        sentence_objs.append(
            Sentence(
                start=sent[1],
                end=sent[2],
                text=sent[0],
                sent_id=sent_id,
                tokens=[t for _, _, t in toks],
                entities=ents,
                relations=rels,
                events=evts,
            )
        )
    return Document(doc_id, sentence_objs)


def convert_batch(
    input_path: str,
    output_path: str,
    time_and_val: bool = False,
    language: str = "english",
):
    """Converts a batch of documents.

    Args:
        input_path (str): path to the input directory. Usually, it is the path
            to the LDC2006T06/data/English or LDC2006T06/data/Chinese folder.
        output_path (str): path to the output JSON file.
        time_and_val (bool, optional): extracts times and values or not.
            Defaults to False.
        language (str, optional): document language. Available options: english,
            chinese. Defaults to 'english'.
    """
    if language == "english":
        sgm_files = glob.glob(os.path.join(input_path, "**", "timex2norm", "*.sgm"))
    elif language == "chinese":
        sgm_files = glob.glob(os.path.join(input_path, "**", "adj", "*.sgm"))
    else:
        raise ValueError("Unknown language: {}".format(language))
    print("Converting the dataset to JSON format")
    print("#SGM files: {}".format(len(sgm_files)))
    progress = tqdm.tqdm(total=len(sgm_files))

    with open(output_path, "w", encoding="utf-8") as w:
        for sgm_file in sgm_files:
            progress.update(1)
            apf_file = sgm_file.replace(".sgm", ".apf.xml")
            doc = convert(sgm_file, apf_file, time_and_val=time_and_val, language=language)
            w.write(json.dumps(doc.to_dict()) + "\n")
    progress.close()


def convert_to_oneie(input_path: str, output_path: str):
    """Converts files to OneIE format.

    Args:
        input_path (str): path to the input file.
        output_path (str): path to the output file.
    """
    print("Converting the dataset to OneIE format")
    with open(input_path, "r", encoding="utf-8") as r, open(output_path, "w", encoding="utf-8") as w:
        for line in r:
            doc = json.loads(line)
            for sentence in doc["sentences"]:
                tokens = sentence["tokens"]
                # pieces = [tokenizer.tokenize(t) for t in tokens]
                # token_lens = [len(x) for x in pieces]
                # if 0 in token_lens:
                #     skip_num += 1
                #     continue
                # pieces = [p for ps in pieces for p in ps]
                # if len(pieces) == 0:
                #     skip_num += 1
                #     continue

                entity_text = {e["mention_id"]: e["text"] for e in sentence["entities"]}
                # update argument text
                for relation in sentence["relations"]:
                    arg1, arg2 = relation["arg1"], relation["arg2"]
                    arg1["text"] = entity_text[arg1["mention_id"]]
                    arg2["text"] = entity_text[arg2["mention_id"]]
                for event in sentence["events"]:
                    for arg in event["arguments"]:
                        arg["text"] = entity_text[arg["mention_id"]]

                # entities
                entities = []
                for entity in sentence["entities"]:
                    entities.append(
                        {
                            "id": entity["mention_id"],
                            "text": entity["text"],
                            "entity_type": entity["entity_type"] if entity["mention_type"] != "TIME" else "TIME",
                            "mention_type": entity["mention_type"],
                            "entity_subtype": entity["entity_subtype"],
                            "start": entity["start"],
                            "end": entity["end"],
                        }
                    )
                # relations
                relations = []
                for relation in sentence["relations"]:
                    relations.append(
                        {
                            "id": relation["relation_id"],
                            "relation_type": relation["relation_type"],
                            "relation_subtype": "{}:{}".format(
                                relation["relation_type"], relation["relation_subtype"]
                            ),
                            "arguments": [
                                {
                                    "entity_id": relation["arg1"]["mention_id"],
                                    "text": relation["arg1"]["text"],
                                    "role": relation["arg1"]["role"],
                                },
                                {
                                    "entity_id": relation["arg2"]["mention_id"],
                                    "text": relation["arg2"]["text"],
                                    "role": relation["arg2"]["role"],
                                },
                            ],
                        }
                    )

                # events
                events = []
                for event in sentence["events"]:
                    events.append(
                        {
                            "id": event["mention_id"],
                            "event_type": "{}:{}".format(event["event_type"], event["event_subtype"]),
                            "trigger": event["trigger"],
                            "arguments": [
                                {
                                    "entity_id": arg["mention_id"],
                                    "text": arg["text"],
                                    "role": arg["role"],
                                }
                                for arg in event["arguments"]
                            ],
                        }
                    )

                sent_obj = {
                    "doc_id": doc["doc_id"],
                    "sent_id": sentence["sent_id"],
                    "tokens": tokens,
                    "sentence": sentence["text"],
                    "entity_mentions": entities,
                    "relation_mentions": relations,
                    "event_mentions": events,
                }
                w.write(json.dumps(sent_obj) + "\n")


def convert_to_event_only(input_path: str, output_path: str):
    """Converts files to event-only format.

    Args:
        input_path (str): path to the input file.
        output_path (str): path to the output file.
        tokenizer (PreTrainedTokenizer): wordpiece tokenizer.
    """
    print("Converting the dataset to event-only format...")
    skip_num = 0
    with open(input_path, "r", encoding="utf-8") as r, open(output_path, "w", encoding="utf-8") as w:
        for line in r:
            doc = json.loads(line)
            for sentence in doc["sentences"]:
                tokens = sentence["tokens"]

                entity_text = {e["mention_id"]: e["text"] for e in sentence["entities"]}
                entity_span = {e["mention_id"]: {"start": e["start"], "end": e["end"]} for e in sentence["entities"]}

                ## update argument text

                for event in sentence["events"]:
                    event_type = event["event_subtype"].upper()
                    cleaned_arguments = []

                    for arg in event["arguments"]:
                        arg["text"] = entity_text[arg["mention_id"]]
                        arg["start"] = entity_span[arg["mention_id"]]["start"]
                        arg["end"] = entity_span[arg["mention_id"]]["end"]
                        arg["role"] = simplify_arg_role_name(arg["role"])

                        if event_type in arg_name_mapping:
                            if arg["role"] in arg_name_mapping[event_type]:
                                new_role = arg_name_mapping[event_type][arg["role"]]
                                if new_role is None:  # arg type isn't in ontology at all
                                    continue  # delete it from data
                                else:  # arg type is in ontology, but misnamed in data
                                    arg["role"] = new_role  # update its name

                        cleaned_arguments.append(
                            {
                                "text": arg["text"],
                                "role": arg["role"],
                                "start": arg["start"],
                                "end": arg["end"],
                            }
                        )
                    event["arguments"] = cleaned_arguments

                # events

                events = []
                for event in sentence["events"]:
                    events.append(
                        {
                            "event_type": "{}:{}".format(event["event_type"], event["event_subtype"]),
                            "trigger": event["trigger"],
                            "arguments": event["arguments"],
                        }
                    )

                sent_obj = {
                    "doc_id": doc["doc_id"],
                    "sent_id": sentence["sent_id"],
                    "tokens": tokens,
                    "sentence": sentence["text"],
                    "event_mentions": events,
                }
                w.write(json.dumps(sent_obj) + "\n")
    print("skip num: {}".format(skip_num))


def simplify_arg_role_name(role_name: str):
    """
    Converts the name of the argument role to the simple form (e.g. "Time-Within" to
    "Time").
    """
    cols = role_name.split("-")
    if len(cols) > 1 and cols[0] != "Time":
        print(cols)
    return cols[0]


def split_data(input_file: str, output_dir: str, split_path: str):
    """Splits the input file into train/dev/test sets.

    Args:
        input_file (str): path to the input file.
        output_dir (str): path to the output directory.
        split_path (str): path to the split directory that contains three files,
            train.doc.txt, dev.doc.txt, and test.doc.txt . Each line in these
            files is a document ID.
    """
    print("Splitting the dataset into train/dev/test sets")
    train_docs, dev_docs, test_docs = set(), set(), set()
    # load doc ids
    with open(os.path.join(split_path, "train.doc.txt")) as r:
        train_docs.update(r.read().strip("\n").split("\n"))
    with open(os.path.join(split_path, "dev.doc.txt")) as r:
        dev_docs.update(r.read().strip("\n").split("\n"))
    with open(os.path.join(split_path, "test.doc.txt")) as r:
        test_docs.update(r.read().strip("\n").split("\n"))

    # split the dataset
    with open(input_file, "r", encoding="utf-8") as r, open(
        os.path.join(output_dir, "train.sentence.json"), "w"
    ) as w_train, open(os.path.join(output_dir, "dev.sentence.json"), "w") as w_dev, open(
        os.path.join(output_dir, "test.sentence.json"), "w"
    ) as w_test:
        for line in r:
            inst = json.loads(line)
            doc_id = inst["doc_id"]
            if doc_id in train_docs:
                w_train.write(line)
            elif doc_id in dev_docs:
                w_dev.write(line)
            elif doc_id in test_docs:
                w_test.write(line)


if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument("-i", "--input", help="Path to the input folder")
    parser.add_argument("-o", "--output", help="Path to the output folder")
    parser.add_argument("-s", "--split", default=None, help="Path to the split folder")
    parser.add_argument("-l", "--lang", default="english", help="Document language")
    parser.add_argument("--time_and_val", action="store_true", help="Extracts times and values")

    args = parser.parse_args()
    if args.lang not in ["chinese", "english"]:
        raise ValueError("Unsupported language: {}".format(args.lang))
    input_dir = os.path.join(args.input, args.lang.title())

    # Convert to doc-level JSON format
    json_path = os.path.join(args.output, "{}.json".format(args.lang))
    convert_batch(input_dir, json_path, time_and_val=args.time_and_val, language=args.lang)

    # Convert to event-only format
    sentence_path = os.path.join(args.output, "{}.sentence.json".format(args.lang))
    convert_to_oneie(json_path, sentence_path)

    # Split the data
    if args.split:
        split_data(sentence_path, args.output, args.split)

src/tasks/ace/preprocess_ace.py
hitz-zentroa/GoLLIE
from typing import Dict, List, Type

from ..utils_typing import Entity, Event, Relation, dataclass


"""Entity definitions

The entity definitions are derived from the official ACE guidelines:
https://www.ldc.upenn.edu/sites/www.ldc.upenn.edu/files/english-entities-guidelines-v6.6.pdf
"""


@dataclass
class Person(Entity):
    """{ace_person}"""

    span: str  # {ace_person_examples}


@dataclass
class Organization(Entity):
    """{ace_organization}"""

    span: str  # {ace_organization_examples}


@dataclass
class GPE(Entity):
    """{ace_gpe}"""

    span: str  # {ace_gpe_examples}


@dataclass
class Location(Entity):
    """{ace_location}"""

    span: str  # {ace_location_examples}


@dataclass
class Facility(Entity):
    """{ace_facility}"""

    span: str  # {ace_facility_examples}


@dataclass
class Weapon(Entity):
    """{ace_weapon}"""

    span: str  # {ace_weapon_examples}


@dataclass
class Vehicle(Entity):
    """{ace_vehicle}"""

    span: str  # {ace_vehicle_examples}


ENTITY_DEFINITIONS: List[Type] = [
    Person,
    Organization,
    GPE,
    Location,
    Facility,
    Weapon,
    Vehicle,
]


"""Value definitions

The following "entities" are not defined on the ACE schema, however, they are
annotated as time expressions or value/numeric data.
"""


@dataclass
class Time(Entity):
    """{ace_time}"""

    span: str  # {ace_time_examples}


@dataclass
class Numeric(Entity):
    """{ace_numeric}"""

    span: str  # {ace_numeric_examples}


@dataclass
class JobTitle(Entity):
    """{ace_jobtitle}"""

    span: str  # {ace_jobtitle_examples}


@dataclass
class Crime(Entity):
    """{ace_crime}"""

    span: str  # {ace_crime_examples}


@dataclass
class Sentence(Entity):
    """{ace_sentence}"""

    span: str  # {ace_sentence_examples}


@dataclass
class ContactInfo(Entity):
    """{ace_contactinfo}"""

    span: str  # {ace_contactinfo_examples}


VALUE_DEFINITIONS: List[Type] = [
    Time,
    Numeric,
    JobTitle,
    Crime,
    Sentence,
    ContactInfo,
]


"""Relation definitions

The relations definitions are derived from the oficial ACE guidelines:
https://www.ldc.upenn.edu/sites/www.ldc.upenn.edu/files/english-relations-guidelines-v6.2.pdf
"""


@dataclass
class PhysicalRelation(Relation):
    """{ace_physicalrelation}"""

    arg1: str
    arg2: str


@dataclass
class PartWholeRelation(Relation):
    """{ace_partwholerelation}"""

    arg1: str
    arg2: str


@dataclass
class PersonalSocialRelation(Relation):
    """{ace_personalsocialrelation}"""

    arg1: str
    arg2: str


@dataclass
class OrganizationAffiliationRelation(Relation):
    """{ace_organizationaffiliationrelation}"""

    arg1: str
    arg2: str


@dataclass
class AgentArtifactRelationRelation(Relation):
    """{ace_agentartifactrelationrelation}"""

    arg1: str
    arg2: str


@dataclass
class GenAffiliationRelation(Relation):
    """{ace_genaffiliationrelation}"""

    arg1: str
    arg2: str


COARSE_RELATION_DEFINITIONS: List[Type] = [
    PhysicalRelation,
    PartWholeRelation,
    PersonalSocialRelation,
    OrganizationAffiliationRelation,
    AgentArtifactRelationRelation,
    GenAffiliationRelation,
]


@dataclass
class Located(PhysicalRelation):
    """{ace_located}"""

    arg1: str
    arg2: str


@dataclass
class Near(PhysicalRelation):
    """{ace_near}"""

    arg1: str
    arg2: str


@dataclass
class Geographical(PartWholeRelation):
    """{ace_geographical}"""

    arg1: str
    arg2: str


@dataclass
class Subsidiary(PartWholeRelation):
    """{ace_subsidiary}"""

    arg1: str
    arg2: str


@dataclass
class Business(PersonalSocialRelation):
    """{ace_business}"""

    arg1: str
    arg2: str


@dataclass
class Family(PersonalSocialRelation):
    """{ace_family}"""

    arg1: str
    arg2: str


@dataclass
class LastingPersonal(PersonalSocialRelation):
    """{ace_lastingpersonal}"""

    arg1: str
    arg2: str


@dataclass
class Employment(OrganizationAffiliationRelation):
    """{ace_employment}"""

    arg1: str
    arg2: str


@dataclass
class Ownership(OrganizationAffiliationRelation):
    """{ace_ownership}"""

    arg1: str
    arg2: str


@dataclass
class Founder(OrganizationAffiliationRelation):
    """{ace_founder}"""

    arg1: str
    arg2: str


@dataclass
class StudentAlum(OrganizationAffiliationRelation):
    """{ace_studentalum}"""

    arg1: str
    arg2: str


@dataclass
class SportsAffiliation(OrganizationAffiliationRelation):
    """{ace_sportsaffiliation}"""

    arg1: str
    arg2: str


@dataclass
class InvestorShareholder(OrganizationAffiliationRelation):
    """{ace_investorshareholder}"""

    arg1: str
    arg2: str


@dataclass
class Membership(OrganizationAffiliationRelation):
    """{ace_membership}"""

    arg1: str
    arg2: str


@dataclass
class UserOwnerInventorManufacturer(AgentArtifactRelationRelation):
    """{ace_userownerinventormanufacturer}"""

    arg1: str
    arg2: str


@dataclass
class CitizenResidentReligionEthnicity(GenAffiliationRelation):
    """{ace_citizenresidentreligionethnicity}"""

    arg1: str
    arg2: str


@dataclass
class OrgLocationOrigin(GenAffiliationRelation):
    """{ace_orglocationorigin}"""

    arg1: str
    arg2: str


RELATION_DEFINITIONS: List[Type] = [
    Located,
    Near,
    Geographical,
    Subsidiary,
    Business,
    Family,
    LastingPersonal,
    Employment,
    Ownership,
    Founder,
    StudentAlum,
    SportsAffiliation,
    InvestorShareholder,
    Membership,
    UserOwnerInventorManufacturer,
    CitizenResidentReligionEthnicity,
    OrgLocationOrigin,
]

FINE_TO_COARSE_RELATIONS: Dict[Type, Type] = {_def: _def.__base__ for _def in RELATION_DEFINITIONS}
COARSE_TO_FINE_RELATIONS: Dict[Type, List[Type]] = {}
for fine, coarse in FINE_TO_COARSE_RELATIONS.items():
    if coarse not in COARSE_TO_FINE_RELATIONS:
        COARSE_TO_FINE_RELATIONS[coarse] = []
    COARSE_TO_FINE_RELATIONS[coarse].append(fine)

"""Event definitions

The events definitions are derived from the oficial ACE guidelines:
https://www.ldc.upenn.edu/sites/www.ldc.upenn.edu/files/english-events-guidelines-v5.4.3.pdf
"""


@dataclass
class LifeEvent(Event):
    """{ace_lifeevent}"""

    mention: str
    """The text span that most clearly expresses the event.
        {ace_lifeevent_examples}
    """


@dataclass
class MovementEvent(Event):
    """{ace_movementevent}"""

    mention: str
    """The text span that most clearly expresses the event.
        {ace_movementevent_examples}
    """


@dataclass
class TransactionEvent(Event):
    """{ace_transactionevent}"""

    mention: str
    """The text span that most clearly expresses the event.
        {ace_transactionevent_examples}
    """


@dataclass
class BusinessEvent(Event):
    """{ace_businessevent}"""

    mention: str
    """The text span that most clearly expresses the event.
        {ace_businessevent_examples}
    """


@dataclass
class ConflictEvent(Event):
    """{ace_conflictevent}"""

    mention: str
    """The text span that most clearly expresses the event.
        {ace_conflictevent_examples}
    """


@dataclass
class ContactEvent(Event):
    """{ace_contactevent}"""

    mention: str
    """The text span that most clearly expresses the event.
        {ace_contactevent_examples}
    """


@dataclass
class PersonnelEvent(Event):
    """{ace_personnelevent}"""

    mention: str
    """The text span that most clearly expresses the event.
        {ace_personnelevent_examples}
    """


@dataclass
class JusticeEvent(Event):
    """{ace_justiceevent}"""

    mention: str
    """The text span that most clearly expresses the event.
        {ace_justiceevent_examples}
    """


COARSE_EVENT_DEFINITIONS: List[Type] = [
    LifeEvent,
    MovementEvent,
    TransactionEvent,
    BusinessEvent,
    ConflictEvent,
    ContactEvent,
    PersonnelEvent,
    JusticeEvent,
]


@dataclass
class BeBorn(LifeEvent):
    """A BeBorn Event occurs whenever a Person Entity is given birth to."""

    mention: str  # The text span that most clearly expresses (triggers) the event
    person: List[str]  # The person who is born
    time: List[str]  # When the birth takes place
    place: List[str]  # Where the birth takes place


@dataclass
class Marry(LifeEvent):
    """{ace_marry}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    person: List[str]  # The people who are married
    time: List[str]  # When the marriage takes place
    place: List[str]  # Where the marriage takes place


@dataclass
class Divorce(LifeEvent):
    """{ace_divorce}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    person: List[str]  # The people who are divorced
    time: List[str]  # When the divorce takes place
    place: List[str]  # Where the divorce takes place


@dataclass
class Injure(LifeEvent):
    """{ace_injure}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    agent: List[str]  # The attacking agent / The one that enacts the harm
    victim: List[str]  # The harmed person(s)
    instrument: List[str]  # The device used to inflict the harm
    time: List[str]  # When the injuring takes place
    place: List[str]


@dataclass
class Die(LifeEvent):
    """{ace_die}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    agent: List[str]  # (Optional) The attacking agent / The killer
    victim: List[str]  # The person(s) who died
    instrument: List[str]  # The device used to kill
    time: List[str]  # When the death takes place
    place: List[str]  # Where the death takes place


@dataclass
class Transport(MovementEvent):
    """{ace_transport}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    agent: List[str]  # The agent responsible for the transport Event
    artifact: List[str]  # The person doing the traveling or the artifact being traveled
    vehicle: List[str]  # The vehicle used to transport the person or artifact
    price: List[str]  # The price of transporting the person or artifact
    origin: List[str]  # Where the transporting originated
    destination: List[str]  # Where the transporting is directed
    time: List[str]  # When the transporting takes place


@dataclass
class TransferOwnership(TransactionEvent):
    """{ace_transferownership}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    buyer: List[str]  # The buying agent
    seller: List[str]  # The selling agent
    beneficiary: List[str]  # The agent that benefits from the transaction
    artifact: List[str]  # The item or Organization that was bought or sold
    price: List[str]  # The sale price of the artifact
    time: List[str]  # When the sale takes place
    place: List[str]  # Where the sale takes place


@dataclass
class TransferMoney(TransactionEvent):
    """{ace_transfermoney}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    giver: List[str]  # The donating agent
    recipient: List[str]  # The recipient agent
    beneficiary: List[str]  # The agent that benefits from the transfer
    money: List[str]  # The amount given, donated or loaned
    time: List[str]  # When the amount is transferred
    place: List[str]  # Where the transation takes place


@dataclass
class StartOrg(BusinessEvent):
    """{ace_startorg}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    agent: List[str]  # The agent responsible for the StattOrg Event (the founder)
    org: List[str]  # The organization that is started
    time: List[str]  # When the Event takes place
    place: List[str]  # Where the Event takes place


@dataclass
class MergeOrg(BusinessEvent):
    """{ace_mergeorg}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    org: List[str]  # The organization(s) that are merged
    time: List[str]  # When the merger takes place
    place: List[str]  # Where the merger takes place


@dataclass
class DeclareBankruptcy(BusinessEvent):
    """{ace_declarebankruptcy}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    org: List[str]  # The Organization declaring bankruptcy
    time: List[str]  # When the bankruptcy is declared
    place: List[str]  # Where the declaration takes place


@dataclass
class EndOrg(BusinessEvent):
    """{ace_endorg}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    org: List[str]  # The Organization that is ended
    time: List[str]  # When the Event takes place
    place: List[str]  # Where the Event takes place


@dataclass
class Attack(ConflictEvent):
    """{ace_attack}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    attacker: List[str]  # The attacking/instigating agent
    target: List[str]  # The target of the attack (including unintended targets)
    instrument: List[str]  # The instrument used in the attack
    time: List[str]  # When the attack takes place
    place: List[str]  # Where the attack takes place


@dataclass
class Demonstrate(ConflictEvent):
    """{ace_demonstrate}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    entity: List[str]  # The demonstrating agent
    time: List[str]  # When the demonstration takes place
    place: List[str]  # Where the demonstration takes place


@dataclass
class Meet(ContactEvent):
    """{ace_meet}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    entity: List[str]  # The agents who are meeting
    time: List[str]  # When the meeting takes place
    place: List[str]  # Where the meeting takes place


@dataclass
class PhoneWrite(ContactEvent):
    """{ace_phonewrite}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    entity: List[str]  # The communicating agents
    time: List[str]  # When the communication takes place


@dataclass
class StartPosition(PersonnelEvent):
    """{ace_startposition}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    person: List[str]  # The employee
    entity: List[str]  # The employer
    position: List[str]  # The JobTitle for the position being started
    time: List[str]  # When the employment relationship begins
    place: List[str]  # Where the employment relationship begins


@dataclass
class EndPosition(PersonnelEvent):
    """{ace_endposition}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    person: List[str]  # The employee
    entity: List[str]  # The employer
    position: List[str]  # The JobTitle for the position being ended
    time: List[str]  # When the employment relationship ends
    place: List[str]  # Where the employment relationship ends


@dataclass
class Nominate(PersonnelEvent):
    """{ace_nominate}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    person: List[str]  # The person(s) nominated
    agent: List[str]  # The nominating agent
    position: List[str]  # The JobTitle for the position being nominated to
    time: List[str]  # When the nomination takes place
    place: List[str]  # Where the nomination takes place


@dataclass
class Elect(PersonnelEvent):
    """{ace_elect}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    person: List[str]  # The person elected
    entity: List[str]  # The voting agent(s)
    position: List[str]  # The JobTitle for the position being nominated to
    time: List[str]  # When the election takes place
    place: List[str]  # Where the election takes place


@dataclass
class ArrestJail(JusticeEvent):
    """{ace_arrestjail}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    person: List[str]  # The person who is jailed or arrested
    agent: List[str]  # The jailer or the arresting agent
    crime: List[str]  # The Crime for which the Person is being jailed or arrested
    time: List[str]  # When the person is arrested or sent to jail
    place: List[str]  # Where the person is arrested or where they are in jail


@dataclass
class ReleaseParole(JusticeEvent):
    """{ace_releaseparole}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    person: List[str]  # The person who is released
    entity: List[str]  # The former captor agent(s)
    crime: List[str]  # The Crime for which the released Person was being held
    time: List[str]  # When the release takes place
    place: List[str]  # Where the release takes place


@dataclass
class TrialHearing(JusticeEvent):
    """{ace_trialhearing}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    defendant: List[str]  # The agent on trial
    prosecutor: List[str]  # The prosecuting agent
    adjudicator: List[str]  # The judge or court
    crime: List[str]  # The Crime for which the Defendant is being tried
    time: List[str]  # When the trial takes place
    place: List[str]  # Where the trial takes place


@dataclass
class ChargeIndict(JusticeEvent):
    """{ace_chargeindict}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    defendant: List[str]  # The agent that is indicted
    prosecutor: List[str]  # The agent bringing charges or executing the indictment
    adjudicator: List[str]  # The judge our court
    crime: List[str]  # The Crime for which the Defendant is being indicted
    time: List[str]  # When the indictment takes place
    place: List[str]  # When the indictment takes place


@dataclass
class Sue(JusticeEvent):
    """{ace_sue}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    plaintiff: List[str]  # The suing agent
    defendant: List[str]  # The agent being sued
    adjudicator: List[str]  # The judge or court
    crime: List[str]  # The Crime (or offense) for which the Defendant is being sued
    time: List[str]  # When the suit takes place
    place: List[str]  # Where the suit takes place


@dataclass
class Convict(JusticeEvent):
    """{ace_convict}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    defendant: List[str]  # The convicted agent(s)
    adjudicator: List[str]  # The judge or court
    crime: List[str]  # The Crime for which the Defendant has been convicted
    time: List[str]  # When the conviction takes place
    place: List[str]  # Where the conviction takes place


@dataclass
class SentenceAct(JusticeEvent):
    """{ace_sentenceact}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    defendant: List[str]  # The agent who is sentenced
    adjudicator: List[str]  # The judge or court
    crime: List[str]  # The Crime for which the Person is being sentenced
    sentence: List[str]  # The sentence
    time: List[str]  # The time of the sentencing Event
    place: List[str]  # Where the sentencing takes place


@dataclass
class Fine(JusticeEvent):
    """{ace_fine}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    entity: List[str]  # The Entity that was fined
    adjudicator: List[str]  # The Entity doing the finding
    money: List[str]  # The amount of the fine
    crime: List[str]  # The Crime (or offence) for which the Entity is being fined
    time: List[str]  # When the fining Event takes place
    place: List[str]  # Where the fining Event takes place


@dataclass
class Execute(JusticeEvent):
    """{ace_execute}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    person: List[str]  # The person executed
    agent: List[str]  # The agent responsible for carrying out the execution
    crime: List[str]  # The Crime for which the Person is being executed
    time: List[str]  # When te execution takes place
    place: List[str]  # Where the execution takes place


@dataclass
class Extradite(JusticeEvent):
    """{ace_extradite}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    agent: List[str]  # The extraditing agent
    person: List[str]  # The person being extradited
    destination: List[str]  # Where the person is extradited to, the destination
    origin: List[str]  # The original location of the person being extradited
    crime: List[str]  # The Crime for which the Person is being extradited
    time: List[str]  # When the extradition takes place


@dataclass
class Acquit(JusticeEvent):
    """{ace_acquit}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    defendant: List[str]  # The agent being acquitted
    adjudicator: List[str]  # The judge or court
    crime: List[str]  # The Crime of which the Defendant is being acquitted
    time: List[str]  # When the acquittal takes place
    place: List[str]  # Where the acquittal takes place


@dataclass
class Pardon(JusticeEvent):
    """{ace_pardon}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    defendant: List[str]  # The agent being pardoned
    adjudicator: List[str]  # The state official who does the pardoning
    crime: List[str]  # The Crime of which the Defendant is being pardoned
    time: List[str]  # When the pardon takes place
    place: List[str]  # Where the pardon takes place


@dataclass
class Appeal(JusticeEvent):
    """{ace_appeal}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    defendant: List[str]  # The defendant
    prosecutor: List[str]  # The prosecuting agent
    adjudicator: List[str]  # The judge or court
    crime: List[str]  # The Crime which is the subject of the appeal
    time: List[str]  # When the appeal takes place
    place: List[str]  # Where the appeal takes place


EVENT_DEFINITIONS: List[Type] = [
    BeBorn,
    Marry,
    Divorce,
    Injure,
    Die,
    Transport,
    TransferOwnership,
    TransferMoney,
    StartOrg,
    MergeOrg,
    DeclareBankruptcy,
    EndOrg,
    Attack,
    Demonstrate,
    Meet,
    PhoneWrite,
    StartPosition,
    EndPosition,
    Nominate,
    Elect,
    ArrestJail,
    ReleaseParole,
    TrialHearing,
    ChargeIndict,
    Sue,
    Convict,
    SentenceAct,
    Fine,
    Execute,
    Extradite,
    Acquit,
    Pardon,
    Appeal,
]

FINE_TO_COARSE_EVENTS: Dict[Type, Type] = {_def: _def.__base__ for _def in EVENT_DEFINITIONS}
COARSE_TO_FINE_EVENTS: Dict[Type, List[Type]] = {}
for fine, coarse in FINE_TO_COARSE_EVENTS.items():
    if coarse not in COARSE_TO_FINE_EVENTS:
        COARSE_TO_FINE_EVENTS[coarse] = []
    COARSE_TO_FINE_EVENTS[coarse].append(fine)

# __all__ = list(map(str, [*ENTITY_DEFINITIONS, *RELATION_DEFINITIONS, *EVENT_DEFINITIONS]))

src/tasks/ace/prompts.py
hitz-zentroa/GoLLIE
from typing import Dict, List, Type

from src.tasks.ace.prompts import (
    COARSE_EVENT_DEFINITIONS,
    COARSE_RELATION_DEFINITIONS,
    ENTITY_DEFINITIONS,
    EVENT_DEFINITIONS,
    RELATION_DEFINITIONS,
    VALUE_DEFINITIONS,
)
from src.tasks.utils_scorer import EventScorer, RelationScorer, SpanScorer
from src.tasks.utils_typing import Entity, Value


class ACEEntityScorer(SpanScorer):
    """ACE Entity identification and classification scorer."""

    valid_types: List[Type] = ENTITY_DEFINITIONS

    def __call__(self, reference: List[Entity], predictions: List[Entity]) -> Dict[str, Dict[str, float]]:
        output = super().__call__(reference, predictions)
        return {"entities": output["spans"]}


class ACEValueScorer(SpanScorer):
    """ACE Values identification and classification scorer."""

    valid_types: List[Type] = VALUE_DEFINITIONS

    def __call__(self, reference: List[Value], predictions: List[Value]) -> Dict[str, Dict[str, float]]:
        output = super().__call__(reference, predictions)
        return {"values": output["spans"]}


class ACECoarseRelationScorer(RelationScorer):
    """ACE Relation identification scorer."""

    valid_types: List[Type] = COARSE_RELATION_DEFINITIONS


class ACERelationScorer(RelationScorer):
    """ACE Relation identification scorer."""

    valid_types: List[Type] = RELATION_DEFINITIONS


class ACEEventScorer(EventScorer):
    """ACE Event and argument classification scorer."""

    valid_types: List[Type] = COARSE_EVENT_DEFINITIONS


class ACEEventArgumentScorer(EventScorer):
    """ACE Event and argument classification scorer."""

    valid_types: List[Type] = EVENT_DEFINITIONS

src/tasks/ace/scorer.py
hitz-zentroa/GoLLIE
from . import data_loader, prompts, scorer


__all__ = ["data_loader", "prompts", "scorer"]

src/tasks/bc5cdr/__init__.py
hitz-zentroa/GoLLIE
from typing import Dict, List, Tuple, Type, Union

from src.tasks.bc5cdr.guidelines import GUIDELINES
from src.tasks.bc5cdr.guidelines_gold import EXAMPLES
from src.tasks.bc5cdr.prompts import ENTITY_DEFINITIONS, Chemical, Disease
from src.tasks.label_encoding import rewrite_labels

from ..utils_data import DatasetLoader, Sampler
from ..utils_typing import Entity


def get_bc5cdr_hf(
    split: str,
    ENTITY_TO_CLASS_MAPPING: Dict[str, Type[Entity]],
) -> Tuple[List[List[str]], List[List[Entity]]]:
    """
    Get the Bc5cdr dataset from the huggingface datasets library
    Args:
        split (str): The path_or_split to load. Can be one of `train`, `validation` or `test`.
    Returns:
        (List[str],List[Entity]): The text and the entities
    """
    from datasets import load_dataset

    dataset = load_dataset("ghadeermobasher/BC5CDR-Chemical-Disease")
    dataset["train"].features["ner_tags"].feature.names

    id2label = dict(enumerate(dataset["train"].features["ner_tags"].feature.names))

    dataset_sentences: List[List[str]] = []
    dataset_entities: List[List[Entity]] = []

    for example in dataset[split]:
        words = example["tokens"]
        # Ensure IOB2 encoding
        labels = rewrite_labels(labels=[id2label[label] for label in example["ner_tags"]], encoding="iob2")

        # Get labeled word spans
        spans = []
        for i, label in enumerate(labels):
            if label == "O":
                continue
            elif label.startswith("B-"):
                spans.append([label[2:], i, i + 1])
            elif label.startswith("I-"):
                spans[-1][2] += 1
            else:
                raise ValueError(f"Found an unexpected label: {label}")

        # Get entities
        entities = []
        for label, start, end in spans:
            entities.append(ENTITY_TO_CLASS_MAPPING[label](span=" ".join(words[start:end])))

        dataset_sentences.append(words)
        dataset_entities.append(entities)

    return dataset_sentences, dataset_entities


class Bc5cdrDatasetLoader(DatasetLoader):
    """
    A `DatasetLoader` for the Bc5cdr dataset.

    Args:
        split (`str`):
            The split to load. Can be one of `train`, `validation` or `test`.

    Raises:
        `ValueError`:
            raised when a not defined value found.
    """

    ENTITY_TO_CLASS_MAPPING = None

    def __init__(self, path_or_split: str, **kwargs) -> None:
        self.ENTITY_TO_CLASS_MAPPING = {
            "Disease": Disease,
            "Chemical": Chemical,
        }

        self.elements = {}

        dataset_words, dataset_entities = get_bc5cdr_hf(
            split=path_or_split,
            ENTITY_TO_CLASS_MAPPING=self.ENTITY_TO_CLASS_MAPPING,
        )

        for id, (words, entities) in enumerate(zip(dataset_words, dataset_entities)):
            self.elements[id] = {
                "id": id,
                "doc_id": id,
                "text": " ".join(words),
                "entities": entities,
                "gold": entities,
            }


class Bc5cdrSampler(Sampler):
    """
    A data `Sampler` for the BC5cdr dataset.

    Args:
        dataset_loader (`Bc5cdrDatasetLoader`):
            The dataset loader that contains the data information.
        task (`str`, optional):
            The task to sample. It must be one of the following: NER, VER, RE, EE.
            Defaults to `None`.
        split (`str`, optional):
            The path_or_split to sample. It must be one of the following: "train", "dev" or
            "test". Depending on the path_or_split the sampling strategy differs. Defaults to
            `"train"`.
        parallel_instances (`Union[int, Tuple[int, int]]`, optional):
            The number of sentences sampled in parallel. Options:

                * **`int`**: The amount of elements that will be sampled in parallel.
                * **`tuple`**: The range of elements that will be sampled in parallel.

            Defaults to 1.
        max_guidelines (`int`, optional):
            The number of guidelines to append to the example at the same time. If `-1`
            is given then all the guidelines are appended. Defaults to `-1`.
        guideline_dropout (`float`, optional):
            The probability to dropout a guideline definition for the given example. This
            is only applied on training. Defaults to `0.0`.
        seed (`float`, optional):
            The seed to sample the examples. Defaults to `0`.
        prompt_template (`str`, optional):
            The path to the prompt template. Defaults to `"templates/prompt.txt"`.
        ensure_positives_on_train (bool, optional):
            Whether to ensure that the guidelines of annotated examples are not removed.
            Defaults to `False`.
        dataset_name (str, optional):
            The name of the dataset. Defaults to `None`.
        scorer (`str`, optional):
           The scorer class import string. Defaults to `None`.
        sample_only_gold_guidelines (`bool`, optional):
            Whether to sample only guidelines of present annotations. Defaults to `False`.
    """

    def __init__(
        self,
        dataset_loader: Bc5cdrDatasetLoader,
        task: str = None,
        split: str = "train",
        parallel_instances: Union[int, Tuple[int, int]] = 1,
        max_guidelines: int = -1,
        guideline_dropout: float = 0.0,
        seed: float = 0,
        prompt_template: str = "templates/prompt.txt",
        ensure_positives_on_train: bool = False,
        dataset_name: str = None,
        scorer: str = None,
        sample_only_gold_guidelines: bool = False,
        **kwargs,
    ) -> None:
        assert task in [
            "NER",
        ], f"BC5cdr only supports NER task. {task} is not supported."

        task_definitions, task_target = {
            "NER": (ENTITY_DEFINITIONS, "entities"),
        }[task]

        super().__init__(
            dataset_loader=dataset_loader,
            task=task,
            split=split,
            parallel_instances=parallel_instances,
            max_guidelines=max_guidelines,
            guideline_dropout=guideline_dropout,
            seed=seed,
            prompt_template=prompt_template,
            ensure_positives_on_train=ensure_positives_on_train,
            sample_only_gold_guidelines=sample_only_gold_guidelines,
            dataset_name=dataset_name,
            scorer=scorer,
            task_definitions=task_definitions,
            task_target=task_target,
            definitions=GUIDELINES,
            examples=EXAMPLES,
            **kwargs,
        )

src/tasks/bc5cdr/data_loader.py
hitz-zentroa/GoLLIE
GUIDELINES = {
    "bc5cdr_disease": {
        "en": [
            (
                "A disease is a particular abnormal condition that negatively affects the structure or function of all"
                " or part of an \norganism, and that is not immediately due to any external injury. Disease mentions,"
                " include Specific Disease (i.e., \nclear-cell renal cell carcinoma), Disease Class (i.e., cystic"
                " kidney diseases), Composite mentions (i.e., prostatic, \npancreas, skin, and lung cancer), and"
                " Modifier (i.e., hereditary breast cancer families)."
            ),
            (
                "A disease refers to a distinct irregular condition that detrimentally impacts the structure or"
                " function of an \nentire organism or a portion of it, and is not directly caused by any external"
                " injury. Disease mentions encompass Specific \nDisease (for example, clear-cell renal cell"
                " carcinoma), Disease Class (such as cystic kidney diseases), Composite \nmentions (like prostatic,"
                " pancreas, skin, and lung cancer), and Modifier (for instance, hereditary breast cancer \nfamilies)."
            ),
            (
                "A disease refers to a distinct anomalous state that unfavorably impacts the structure or function of"
                " an entire \norganism or part of it, and isn't immediately attributable to any external injury."
                " Disease-related mentions encompass \nSpecific Disease (for example, clear-cell renal cell"
                " carcinoma), Disease Class (for instance, cystic kidney diseases), \nComposite mentions (such as"
                " prostatic, pancreas, skin, and lung cancer), and Modifier (like hereditary breast cancer"
                " \nfamilies)."
            ),
            (
                "A disease refers to a distinct irregular condition that unfavorably impacts the structure or function"
                " of an entire \norganism or a portion of it, and is not directly caused by any external injury."
                " Disease mentions encompass Specific Disease \n(for example, clear-cell renal cell carcinoma),"
                " Disease Class (for instance, cystic kidney diseases), Composite \nmentions (such as prostatic,"
                " pancreas, skin, and lung cancer), and Modifier (like hereditary breast cancer families)."
            ),
            (
                "A disease refers to a specific, unusual condition that unfavorably impacts the structure or function"
                " of all or part \nof an organism, and is not the result of any immediate external injury. Disease"
                " mentions encompass Specific Disease \n(for example, clear-cell renal cell carcinoma), Disease Class"
                " (for instance, cystic kidney diseases), Composite \nmentions (such as prostatic, pancreas, skin, and"
                " lung cancer), and Modifier (like hereditary breast cancer families)."
            ),
        ]
    },
    "bc5cdr_chemical": {
        "en": [
            (
                "The basic rule for chemical entity annotation is that the chemical should have a specific structure."
                " Chemicals \nwhich should be annotated are listed as follows: 1) Chemical Nouns convertible to: - A"
                " single chemical structure diagram: \nsingle atoms, ions, isotopes, pure elements and molecules such"
                " as: Calcium (Ca), Iron (Fe), Lithium (Li), Potassium (K), \nOxygen (O2), - A general Markush diagram"
                " with R groups such as: Amino acids 2) General class names where the definition of the \nclass"
                " includes information on some structural or elemental composition such as: steroids, sugars, fatty"
                " acids, \nsaturated fatty acids \u2026 3) Small Biochemicals - Monosaccharides, disaccharides and"
                " trisaccharides: Glucose, Sucrose \u2026 - \nPeptides and proteins with less than 15 aminoacids:"
                " Angiotensin II \u2026 - Monomers, dimmers, trimmers of nucleotides: e.g. ATP, \ncAMP \u2026 - Fatty"
                " acids and their derivatives excluding polymeric structures. e.g. Cholesterol, glycerol,"
                " prostaglandin \nE1 \u2026 4) Synthetic Polymers such as: Polyethylene glycol 5) Special chemicals"
                " having well-defined chemical \ncompositions. E.g. \u201cethanolic extract of Daucus carota seeds"
                " (DCE)\u201d in 16755009; \u201cgrape seed proanthocyanidin extract\u201d in \n11334364."
            ),
            (
                "The fundamental principle for annotating chemical entities is that they must possess a distinct"
                " structure. \nChemicals that need to be annotated include: 1) Chemical nouns convertible to: - A"
                " single chemical structure diagram: \nfeaturing individual atoms, ions, isotopes, pure elements, and"
                " molecules like Calcium (Ca), Iron (Fe), Lithium (Li), \nPotassium (K), Oxygen (O2), and - A general"
                " Markush diagram with R groups such as Amino acids. 2) General class names where the \nclass"
                " definition encompasses information on structural or elemental composition, for instance, steroids,"
                " sugars, \nfatty acids, saturated fatty acids, and so on. 3) Small biochemicals, such as:"
                " Monosaccharides, disaccharides, and \ntrisaccharides like Glucose and Sucrose; Peptides and proteins"
                " containing up to 15 amino acids, for example, Angiotensin II; \nMonomers, dimers, and trimers of"
                " nucleotides, e.g., ATP and cAMP; Fatty acids and their derivatives, excluding polymeric"
                " \nstructures, e.g., Cholesterol, glycerol, and prostaglandin E1. 4) Synthetic polymers, like"
                " Polyethylene glycol. 5) \nExceptional chemicals featuring well-defined chemical compositions, for"
                ' example, "ethanolic extract of Daucus carota seeds \n(DCE)" in 16755009 or "grape seed'
                ' proanthocyanidin extract" in 11334364.'
            ),
            (
                "The fundamental guideline for annotating chemical entities is that the substance must possess a"
                " distinct \nstructure. Chemicals that need to be annotated include: 1) Chemical Nouns convertible to:"
                " - A single chemical structure \ndiagram: covering individual atoms, ions, isotopes, pure elements,"
                " and molecules, such as: Calcium (Ca), Iron (Fe), \nLithium (Li), Potassium (K), Oxygen (O2), - A"
                " general Markush diagram with R groups, like: Amino acids 2) General class names \nwhere the"
                " definition of the class encompasses information on structural or elemental composition, for example:"
                " \nsteroids, sugars, fatty acids, saturated fatty acids, and more. 3) Small Biochemicals -"
                " Monosaccharides, disaccharides, \nand trisaccharides: Glucose, Sucrose, and so on - Peptides and"
                " proteins with fewer than 15 amino acids: Angiotensin II, \netc. - Monomers, dimers, and trimers of"
                " nucleotides: e.g., ATP, cAMP, etc. - Fatty acids and their derivatives, \nexcluding polymeric"
                " structures, for instance: Cholesterol, glycerol, prostaglandin E1, etc. 4) Synthetic Polymers, such"
                " \nas: Polyethylene glycol 5) Special chemicals featuring well-defined chemical compositions, e.g.,"
                ' "ethanolic \nextract of Daucus carota seeds (DCE)" in 16755009, "grape seed proanthocyanidin'
                ' extract" in 11334364.'
            ),
            (
                "The fundamental principle for annotating chemical entities is that they must possess a distinct"
                " structure. \nChemicals that should be annotated include: 1) Chemical Nouns convertible to: - A"
                " single chemical structure diagram: single \natoms, ions, isotopes, pure elements, and molecules such"
                " as calcium (Ca), iron (Fe), lithium (Li), potassium (K), oxygen \n(O2), - A general Markush diagram"
                " with R groups such as amino acids 2) General class names where the definition of the class"
                " \nincorporates information on some structural or elemental composition, for example: steroids,"
                " sugars, fatty acids, saturated \nfatty acids, and so on 3) Small Biochemicals: - Monosaccharides,"
                " disaccharides, and trisaccharides: glucose, sucrose, \nand more - Peptides and proteins with less"
                " than 15 amino acids: angiotensin II, for instance - Monomers, dimmers, and \ntrimmers of"
                " nucleotides: e.g., ATP, cAMP - Fatty acids and their derivatives, excluding polymeric structures:"
                " e.g., \ncholesterol, glycerol, prostaglandin E1, and so on 4) Synthetic Polymers, such as:"
                " polyethylene glycol 5) Special chemicals \nwith well-defined chemical compositions, for example:"
                ' "ethanolic extract of Daucus carota seeds (DCE)" in 16755009, \n"grape seed proanthocyanidin'
                ' extract" in 11334364.'
            ),
            (
                "The fundamental guideline for annotating chemical entities is that the substance must possess a"
                " distinct \nstructure. Chemicals that need to be annotated include: 1) Chemical Nouns convertible to:"
                " - A single chemical structure \ndiagram: individual atoms, ions, isotopes, pure elements, and"
                " molecules such as calcium (Ca), iron (Fe), lithium (Li), \npotassium (K), oxygen (O2), - A general"
                " Markush diagram with R groups such as amino acids 2) General class names that encompass"
                " \ninformation about some structural or elemental composition, for example: steroids, sugars, fatty"
                " acids, saturated fatty \nacids, and so on 3) Small Biochemicals: Monosaccharides, disaccharides, and"
                " trisaccharides like glucose and sucrose; \npeptides and proteins with less than 15 amino acids, for"
                " instance, angiotensin II; monomers, dimmers, and trimmers of \nnucleotides such as ATP and cAMP;"
                " fatty acids and their derivatives, excluding polymeric structures like cholesterol, \nglycerol, and"
                " prostaglandin E1 4) Synthetic Polymers like polyethylene glycol 5) Special chemicals featuring"
                ' well-defined \nchemical compositions, e.g., "ethanolic extract of Daucus carota seeds (DCE)" in'
                ' 16755009, or "grape seed \nproanthocyanidin extract" in 11334364.'
            ),
        ]
    },
}

src/tasks/bc5cdr/guidelines.py
hitz-zentroa/GoLLIE
GUIDELINES = {
    "bc5cdr_disease": {
        "en": [
            "A disease is a particular abnormal condition that negatively affects the structure or function of all or"
            " part of an organism, and that is not immediately due to any external injury. Disease mentions, include"
            " Specific Disease (i.e., clear-cell renal cell carcinoma), Disease Class (i.e., cystic kidney diseases),"
            " Composite mentions (i.e., prostatic, pancreas, skin, and lung cancer), and Modifier (i.e., hereditary"
            " breast cancer families). "
        ]
    },
    "bc5cdr_chemical": {
        "en": [
            (
                "The basic rule for chemical entity annotation is that the chemical should have a specific structure."
                " Chemicals which should be annotated are listed as follows: 1) Chemical Nouns convertible to: - A"
                " single chemical structure diagram: single atoms, ions, isotopes, pure elements and molecules such"
                " as: Calcium (Ca), Iron (Fe), Lithium (Li), Potassium (K), Oxygen (O2), - A general Markush diagram"
                " with R groups such as: Amino acids 2) General class names where the definition of the class includes"
                " information on some structural or elemental composition such as: steroids, sugars, fatty acids,"
                " saturated fatty acids … 3) Small Biochemicals - Monosaccharides, disaccharides and trisaccharides:"
                " Glucose, Sucrose … - Peptides and proteins with less than 15 aminoacids: Angiotensin II … -"
                " Monomers, dimmers, trimmers of nucleotides: e.g. ATP, cAMP … - Fatty acids and their derivatives"
                " excluding polymeric structures. e.g. Cholesterol, glycerol, prostaglandin E1 … 4) Synthetic Polymers"
                " such as: Polyethylene glycol 5) Special chemicals having well-defined chemical compositions. E.g."
                " “ethanolic extract of Daucus carota seeds (DCE)” in 16755009; “grape seed proanthocyanidin extract”"
                " in 11334364."
            ),
        ]
    },
}

EXAMPLES = {
    "bc5cdr_disease_examples": {
        "en": [
            "toxicity",
            "pain",
            "hypotension",
            "proteinuria",
            "seizures",
            "hypertension",
            "seizure",
            "myocardial infarction",
            "hepatitis",
            "bradycardia",
        ]
    },
    "bc5cdr_chemical_examples": {
        "en": [
            "cocaine",
            "dopamine",
            "morphine",
            "nicotine",
            "lithium",
            "haloperidol",
            "clonidine",
            "creatinine",
            "cisplatin",
            "lidocaine",
        ]
    },
}

src/tasks/bc5cdr/guidelines_gold.py
hitz-zentroa/GoLLIE
from typing import List

from ..utils_typing import Entity, dataclass


"""Entity definitions

The entity definitions are derived from the official BC5DR guidelines:
https://biocreative.bioinformatics.udel.edu/media/store/files/2015/bc5_CDR_data_guidelines.pdf


"""


@dataclass
class Disease(Entity):
    """{bc5cdr_disease}"""

    span: str  # {bc5cdr_disease_examples}


@dataclass
class Chemical(Entity):
    """{bc5cdr_chemical}"""

    span: str  # {bc5cdr_chemical_examples}


ENTITY_DEFINITIONS: List[Entity] = [Disease, Chemical]

# __all__ = list(map(str, [*ENTITY_DEFINITIONS]))

src/tasks/bc5cdr/prompts.py
hitz-zentroa/GoLLIE
from typing import Dict, List, Type

from src.tasks.bc5cdr.prompts import ENTITY_DEFINITIONS
from src.tasks.utils_scorer import SpanScorer
from src.tasks.utils_typing import Entity


class Bc5cdrEntityScorer(SpanScorer):
    """CoNLL03 Entity identification and classification scorer."""

    valid_types: List[Type] = ENTITY_DEFINITIONS

    def __call__(self, reference: List[Entity], predictions: List[Entity]) -> Dict[str, Dict[str, float]]:
        output = super().__call__(reference, predictions)
        return {"entities": output["spans"]}

src/tasks/bc5cdr/scorer.py
hitz-zentroa/GoLLIE
from . import data_loader, prompts, scorer


__all__ = ["data_loader", "prompts", "scorer"]

src/tasks/broadtwitter/__init__.py
hitz-zentroa/GoLLIE
from typing import Dict, List, Tuple, Type, Union

from src.tasks.broadtwitter.guidelines import GUIDELINES
from src.tasks.broadtwitter.guidelines_gold import EXAMPLES
from src.tasks.broadtwitter.prompts import ENTITY_DEFINITIONS, Location, Organization, Person
from src.tasks.label_encoding import rewrite_labels

from ..utils_data import DatasetLoader, Sampler
from ..utils_typing import Entity


def get_broadtwitter_hf(
    split: str,
    ENTITY_TO_CLASS_MAPPING: Dict[str, Type[Entity]],
) -> Tuple[List[List[str]], List[List[Entity]]]:
    """
    Get the Broad Twitter Corpus dataset from the huggingface datasets library
    Args:
        split (str): The path_or_split to load. Can be one of `train`, `validation` or `test`.
    Returns:
        (List[str],List[Entity]): The text and the entities
    """
    from datasets import load_dataset

    dataset = load_dataset("strombergnlp/broad_twitter_corpus")
    dataset["train"].features["ner_tags"].feature.names

    id2label = dict(enumerate(dataset["train"].features["ner_tags"].feature.names))

    dataset_sentences: List[List[str]] = []
    dataset_entities: List[List[Entity]] = []

    for example in dataset[split]:
        words = example["tokens"]
        # Ensure IOB2 encoding
        labels = rewrite_labels(labels=[id2label[label] for label in example["ner_tags"]], encoding="iob2")
        # FIX @ labels https://github.com/GateNLP/broad_twitter_corpus/issues/15
        for i, word in enumerate(words):
            if word.strip() == "@":
                if i + 1 < len(labels):
                    if labels[i + 1].startswith("B") and labels[i].startswith("B"):
                        p, l = labels[i + 1].split("-")
                        labels[i + 1] = f"I-{l}"

        # Ensure IOB2 encoding again, just in case
        labels = rewrite_labels(labels=labels, encoding="iob2")

        # Get labeled word spans
        spans = []
        for i, label in enumerate(labels):
            if label == "O":
                continue
            elif label.startswith("B-"):
                spans.append([label[2:], i, i + 1])
            elif label.startswith("I-"):
                spans[-1][2] += 1
            else:
                raise ValueError(f"Found an unexpected label: {label}")

        # Get entities
        entities = []
        for label, start, end in spans:
            entities.append(ENTITY_TO_CLASS_MAPPING[label](span=" ".join(words[start:end])))

        dataset_sentences.append(words)
        dataset_entities.append(entities)

    return dataset_sentences, dataset_entities


class BroadTwitterDatasetLoader(DatasetLoader):
    """
    A `DatasetLoader` for the Broad Twitter Corpus dataset.

    Args:
        split (`str`):
            The split to load. Can be one of `train`, `validation` or `test`.

    Raises:
        `ValueError`:
            raised when a not defined value found.
    """

    ENTITY_TO_CLASS_MAPPING = None

    def __init__(self, path_or_split: str, **kwargs) -> None:
        self.ENTITY_TO_CLASS_MAPPING = {
            "PER": Person,
            "LOC": Location,
            "ORG": Organization,
        }

        self.elements = {}

        dataset_words, dataset_entities = get_broadtwitter_hf(
            split=path_or_split,
            ENTITY_TO_CLASS_MAPPING=self.ENTITY_TO_CLASS_MAPPING,
        )

        for id, (words, entities) in enumerate(zip(dataset_words, dataset_entities)):
            self.elements[id] = {
                "id": id,
                "doc_id": id,
                "text": " ".join(words),
                "entities": entities,
                "gold": entities,
            }


class BroadTwitterSampler(Sampler):
    """
    A data `Sampler` for the Broad Twitter Corpus dataset.

    Args:
        dataset_loader (`BroadTwitterDatasetLoader`):
            The dataset loader that contains the data information.
        task (`str`, optional):
            The task to sample. It must be one of the following: NER, VER, RE, EE.
            Defaults to `None`.
        split (`str`, optional):
            The path_or_split to sample. It must be one of the following: "train", "dev" or
            "test". Depending on the path_or_split the sampling strategy differs. Defaults to
            `"train"`.
        parallel_instances (`Union[int, Tuple[int, int]]`, optional):
            The number of sentences sampled in parallel. Options:

                * **`int`**: The amount of elements that will be sampled in parallel.
                * **`tuple`**: The range of elements that will be sampled in parallel.

            Defaults to 1.
        max_guidelines (`int`, optional):
            The number of guidelines to append to the example at the same time. If `-1`
            is given then all the guidelines are appended. Defaults to `-1`.
        guideline_dropout (`float`, optional):
            The probability to dropout a guideline definition for the given example. This
            is only applied on training. Defaults to `0.0`.
        seed (`float`, optional):
            The seed to sample the examples. Defaults to `0`.
        prompt_template (`str`, optional):
            The path to the prompt template. Defaults to `"templates/prompt.txt"`.
        ensure_positives_on_train (bool, optional):
            Whether to ensure that the guidelines of annotated examples are not removed.
            Defaults to `False`.
        dataset_name (str, optional):
            The name of the dataset. Defaults to `None`.
        scorer (`str`, optional):
           The scorer class import string. Defaults to `None`.
        sample_only_gold_guidelines (`bool`, optional):
            Whether to sample only guidelines of present annotations. Defaults to `False`.
    """

    def __init__(
        self,
        dataset_loader: BroadTwitterDatasetLoader,
        task: str = None,
        split: str = "train",
        parallel_instances: Union[int, Tuple[int, int]] = 1,
        max_guidelines: int = -1,
        guideline_dropout: float = 0.0,
        seed: float = 0,
        prompt_template: str = "templates/prompt.txt",
        ensure_positives_on_train: bool = False,
        dataset_name: str = None,
        scorer: str = None,
        sample_only_gold_guidelines: bool = False,
        **kwargs,
    ) -> None:
        assert task in [
            "NER",
        ], f"Broad Twitter Corpus only supports NER task. {task} is not supported."

        task_definitions, task_target = {
            "NER": (ENTITY_DEFINITIONS, "entities"),
        }[task]

        super().__init__(
            dataset_loader=dataset_loader,
            task=task,
            split=split,
            parallel_instances=parallel_instances,
            max_guidelines=max_guidelines,
            guideline_dropout=guideline_dropout,
            seed=seed,
            prompt_template=prompt_template,
            ensure_positives_on_train=ensure_positives_on_train,
            sample_only_gold_guidelines=sample_only_gold_guidelines,
            dataset_name=dataset_name,
            scorer=scorer,
            task_definitions=task_definitions,
            task_target=task_target,
            definitions=GUIDELINES,
            examples=EXAMPLES,
            **kwargs,
        )

src/tasks/broadtwitter/data_loader.py
hitz-zentroa/GoLLIE
GUIDELINES = {
    "ner_person": {"en": ["first, middle and last names of people, animals and fictional characters aliases."]},
    "ner_organization": {
        "en": [
            "Companies (press agencies, studios, banks, stock\n"
            "markets, manufacturers, cooperatives) subdivisions of companies (newsrooms)\n"
            "brands political movements (political parties, terrorist organisations)\n"
            "government bodies (ministries, councils, courts, political unions of countries\n"
            "(e.g. the {\\it U.N.})) publications (magazines, newspapers, journals)\n"
            "musical companies (bands, choirs, opera companies, orchestras\n"
            "public organisations (schools, universities, charities other collections\n"
            "of people (sports clubs, sports teams, associations, theaters companies,\n"
            "religious orders, youth organisations.\n"
        ]
    },
    "ner_location": {
        "en": [
            "Roads (streets, motorways) trajectories regions (villages, towns, cities, provinces,\n"
            "countries, continents, dioceses, parishes) structures (bridges, ports, dams) natural locations\n"
            "(mountains, mountain ranges, woods, rivers, wells, fields, valleys, gardens, nature reserves,\n"
            "allotments, beaches, national parks) public places (squares, opera houses, museums, schools, markets,\n"
            "airports, stations, swimming pools, hospitals, sports facilities, youth centers, parks, town halls,\n"
            "theaters, cinemas, galleries, camping grounds, NASA launch pads, club houses, universities, libraries,\n"
            "churches, medical centers, parking lots, playgrounds, cemeteries) commercial places (chemists, pubs,\n"
            "restaurants, depots, hostels, hotels, industrial parks, nightclubs, music venues) assorted buildings\n"
            "(houses, monasteries, creches, mills, army barracks, castles, retirement homes, towers, halls, rooms,\n"
            "vicarages, courtyards) abstract ``places'' (e.g. {\\it the free world})\n"
        ]
    },
}

src/tasks/broadtwitter/guidelines.py
hitz-zentroa/GoLLIE
GUIDELINES = {
    "ner_person": {"en": ["first, middle and last names of people, animals and fictional characters aliases."]},
    "ner_organization": {
        "en": [
            "Companies (press agencies, studios, banks, stock\n"
            "markets, manufacturers, cooperatives) subdivisions of companies (newsrooms)\n"
            "brands political movements (political parties, terrorist organisations)\n"
            "government bodies (ministries, councils, courts, political unions of countries\n"
            "(e.g. the {\\it U.N.})) publications (magazines, newspapers, journals)\n"
            "musical companies (bands, choirs, opera companies, orchestras\n"
            "public organisations (schools, universities, charities other collections\n"
            "of people (sports clubs, sports teams, associations, theaters companies,\n"
            "religious orders, youth organisations.\n"
        ]
    },
    "ner_location": {
        "en": [
            "Roads (streets, motorways) trajectories regions (villages, towns, cities, provinces,\n"
            "countries, continents, dioceses, parishes) structures (bridges, ports, dams) natural locations\n"
            "(mountains, mountain ranges, woods, rivers, wells, fields, valleys, gardens, nature reserves,\n"
            "allotments, beaches, national parks) public places (squares, opera houses, museums, schools, markets,\n"
            "airports, stations, swimming pools, hospitals, sports facilities, youth centers, parks, town halls,\n"
            "theaters, cinemas, galleries, camping grounds, NASA launch pads, club houses, universities, libraries,\n"
            "churches, medical centers, parking lots, playgrounds, cemeteries) commercial places (chemists, pubs,\n"
            "restaurants, depots, hostels, hotels, industrial parks, nightclubs, music venues) assorted buildings\n"
            "(houses, monasteries, creches, mills, army barracks, castles, retirement homes, towers, halls, rooms,\n"
            "vicarages, courtyards) abstract ``places'' (e.g. {\\it the free world})\n"
        ]
    },
}
EXAMPLES = {
    "ner_person_examples": {
        "en": [
            "Obama",
            "President Obama",
            "@ SimoLove",
            "Kate Middleton",
            "Putin",
            "Prince William",
            "Cameron",
            "Harper",
            "@ RossMarowits",
            "James Foley",
        ]
    },
    "ner_location_examples": {
        "en": ["UK", "Ukraine", "US", "U . S .", "Iraq", "Canada", "London", "Russia", "Australia", "Ontario"]
    },
    "ner_organization_examples": {
        "en": [
            "@ Independent",
            "Irish News",
            "Malaysia Airlines",
            "Twitter",
            "BBC",
            "Apple",
            "Getty",
            "Isis",
            "@ BBCBreaking",
            "Liverpool",
        ]
    },
}

src/tasks/broadtwitter/guidelines_gold.py
hitz-zentroa/GoLLIE
from typing import List

from ..utils_typing import Entity, dataclass


"""Entity definitions

The entity definitions are derived from the official ConLL2003 guidelines:
https://www.clips.uantwerpen.be/conll2003/ner/
Based on: Nancy Chinchor, Erica Brown, Lisa Ferro, Patty Robinson,
           "1999 Named Entity Task Definition". MITRE and SAIC, 1999.
"""


@dataclass
class Person(Entity):
    """{ner_person}"""

    span: str  # {ner_person_examples}


@dataclass
class Organization(Entity):
    """{ner_organization}"""

    span: str  # {ner_organization_examples}


@dataclass
class Location(Entity):
    """{ner_location}"""

    span: str  # {ner_location_examples}


@dataclass
class Miscellaneous(Entity):
    """{ner_miscellaneous}"""

    span: str  # {ner_miscellaneous_examples}


ENTITY_DEFINITIONS: List[Entity] = [
    Person,
    Organization,
    Location,
]


# __all__ = list(map(str, [*ENTITY_DEFINITIONS]))

src/tasks/broadtwitter/prompts.py
hitz-zentroa/GoLLIE
from typing import Dict, List, Type

from src.tasks.broadtwitter.prompts import ENTITY_DEFINITIONS
from src.tasks.utils_scorer import SpanScorer
from src.tasks.utils_typing import Entity


class BroadTwitterEntityScorer(SpanScorer):
    """CoNLL03 Entity identification and classification scorer."""

    valid_types: List[Type] = ENTITY_DEFINITIONS

    def __call__(self, reference: List[Entity], predictions: List[Entity]) -> Dict[str, Dict[str, float]]:
        output = super().__call__(reference, predictions)
        return {"entities": output["spans"]}

src/tasks/broadtwitter/scorer.py
hitz-zentroa/GoLLIE
from . import data_loader, prompts_eae, prompts_ed, scorer


__all__ = ["data_loader", "prompts_ed", "prompts_eae", "scorer"]

src/tasks/casie/__init__.py
hitz-zentroa/GoLLIE
import inspect
import json
from typing import Tuple, Union

from src.tasks.casie.guidelines import GUIDELINES
from src.tasks.casie.guidelines_gold import EXAMPLES
from src.tasks.casie.prompts_eae import (
    EAE_EVENT_DEFINITIONS,
)
from src.tasks.casie.prompts_eae import (
    DatabreachAttack as EAEDatabreachAttack,
)
from src.tasks.casie.prompts_eae import PhisingAttack as EAEPhisingAttack
from src.tasks.casie.prompts_eae import RansomAttack as EAERansomAttack
from src.tasks.casie.prompts_eae import (
    VulnerabilityDiscover as EAEVulnerabilityDiscover,
)
from src.tasks.casie.prompts_eae import (
    VulnerabilityPatch as EAEVulnerabilityPatch,
)
from src.tasks.casie.prompts_ed import (
    ED_EVENT_DEFINITIONS,
    DatabreachAttack,
    PhisingAttack,
    RansomAttack,
    VulnerabilityDiscover,
    VulnerabilityPatch,
)
from src.tasks.utils_data import DatasetLoader, Sampler


class CASIEDatasetLoader(DatasetLoader):
    _EVENT_CONSTANTS_MAPPING = {
        "trigger": "mention",
        "Time": "time",
        "Place": "place",
        "Attacker": "attacker",
        "Tool": "tool",
        "Victim": "victim",
        "Damage-Amount": "damage_amount",
        "Vulnerability": "vulnerability",
        "Vulnerable_System": "vulnerable_system",
        "Vulnerable_System_Owner": "system_owner",
        "Vulnerable_System_Version": "system_version",
        "Supported_Platform": "supported_platform",
    }
    EVENT_TO_CLASS_MAPPING = {
        "Attack:Databreach": {
            "ed_class": DatabreachAttack,
            "eae_class": EAEDatabreachAttack,
            "Attack-Pattern": "attack_pattern",
            "Compromised-Data": "compromised_data",
            "Number-of-Data": "number_of_data",
            "Number-of-Victim": "number_of_victim",
            "Purpose": "purpose",
        },
        "Attack:Phising": {
            "ed_class": PhisingAttack,
            "eae_class": EAEPhisingAttack,
            "Attack-Pattern": "pattern",
            "Purpose": "purpose",
            "Trusted-Entity": "trusted_entity",
        },
        "Attack:Ransom": {
            "ed_class": RansomAttack,
            "eae_class": EAERansomAttack,
            "Attack-Pattern": "pattern",
            "Payment-Method": "payment_method",
            "Price": "price",
        },
        "Vulnerability:Discover": {
            "ed_class": VulnerabilityDiscover,
            "eae_class": EAEVulnerabilityDiscover,
            "CVE": "cve",
            "Capabilities": "used_for",
            "Discoverer": "discoverer",
            "Supported_Platform": "supported_platform",
        },
        "Vulnerability:Patch": {
            "ed_class": VulnerabilityPatch,
            "eae_class": EAEVulnerabilityPatch,
            "CVE": "cve",
            "Issues-Addressed": "issues_addressed",
            "Patch": "patch",
            "Patch-Number": "patch_number",
            "Releaser": "releaser",
        },
    }

    def __init__(self, path: str, **kwargs) -> None:
        self.elements = {}

        with open(path, "rt") as in_f:
            for line in in_f:
                line = json.loads(line.strip())

                key = line["sent_id"]
                if key not in self.elements:
                    self.elements[key] = {
                        "id": key,
                        "doc_id": key.split(":")[0],
                        "text": "",
                        "events": [],
                        "arguments": [],
                        "gold": [],
                    }

                events, arguments = [], []
                for event in sorted(line["events"], key=lambda x: x["trigger"]["start"]):
                    if event["type"] not in self.EVENT_TO_CLASS_MAPPING:
                        continue
                    info = self.EVENT_TO_CLASS_MAPPING[event["type"]]
                    _inst = {param: [] for param in inspect.signature(info["eae_class"]).parameters.keys()}
                    _inst["mention"] = event["trigger"]["text"]
                    for argument in event["arguments"]:
                        if argument["role"] in info:
                            name = info[argument["role"]]
                            _inst[name].append(argument["text"])
                        elif argument["role"] in self._EVENT_CONSTANTS_MAPPING:
                            name = self._EVENT_CONSTANTS_MAPPING[argument["role"]]
                            if name not in _inst:
                                continue
                            _inst[name].append(argument["text"])
                        else:
                            raise ValueError(f"Argument {event['type']}:{argument['role']} not found!")

                    events.append(info["ed_class"](mention=_inst["mention"]))
                    arguments.append(info["eae_class"](**_inst))

                self.elements[key]["text"] += " " + line["text"].strip()
                self.elements[key]["events"] += events
                self.elements[key]["arguments"] += arguments
                self.elements[key]["gold"] += events


class CASIESampler(Sampler):
    """
    A data `Sampler` for the CASIE dataset.

    Args:
        dataset_loader (`CASIEDatasetLoader`):
            The dataset loader that contains the data information.
        task (`str`, optional):
            The task to sample. It must be one of the following: NER, VER, RE, EE.
            Defaults to `None`.
        split (`str`, optional):
            The split to sample. It must be one of the following: "train", "dev" or
            "test". Depending on the split the sampling strategy differs. Defaults to
            `"train"`.
        parallel_instances (`Union[int, Tuple[int, int]]`, optional):
            The number of sentences sampled in parallel. Options:

                * **`int`**: The amount of elements that will be sampled in parallel.
                * **`tuple`**: The range of elements that will be sampled in parallel.

            Defaults to 1.
        max_guidelines (`int`, optional):
            The number of guidelines to append to the example at the same time. If `-1`
            is given then all the guidelines are appended. Defaults to `-1`.
        guideline_dropout (`float`, optional):
            The probability to dropout a guideline definition for the given example. This
            is only applied on training. Defaults to `0.0`.
        seed (`float`, optional):
            The seed to sample the examples. Defaults to `0`.
        prompt_template (`str`, optional):
            The path to the prompt template. Defaults to `"templates/prompt.txt"`.
        ensure_positives_on_train (bool, optional):
            Whether to ensure that the guidelines of annotated examples are not removed.
            Defaults to `False`.
        dataset_name (str, optional):
            The name of the dataset. Defaults to `None`.
        scorer (`str`, optional):
           The scorer class import string. Defaults to `None`.
        sample_only_gold_guidelines (`bool`, optional):
            Whether to sample only guidelines of present annotations. Defaults to `False`.
    """

    def __init__(
        self,
        dataset_loader: CASIEDatasetLoader,
        task: str = None,
        split: str = "train",
        parallel_instances: Union[int, Tuple[int, int]] = 1,
        max_guidelines: int = -1,
        guideline_dropout: float = 0.0,
        seed: float = 0,
        ensure_positives_on_train: bool = False,
        dataset_name: str = None,
        scorer: str = None,
        sample_only_gold_guidelines: bool = False,
        **kwargs,
    ) -> None:
        assert task in [
            "EE",
            "EAE",
        ], f"{task} must be either 'EE', 'EAE'."

        task_definitions, task_target, task_template = {
            "EE": (ED_EVENT_DEFINITIONS, "events", "templates/prompt.txt"),
            "EAE": (EAE_EVENT_DEFINITIONS, "arguments", "templates/prompt_ace_eae.txt"),
        }[task]

        kwargs.pop("prompt_template")

        super().__init__(
            dataset_loader=dataset_loader,
            task=task,
            split=split,
            parallel_instances=parallel_instances,
            max_guidelines=max_guidelines,
            guideline_dropout=guideline_dropout,
            seed=seed,
            prompt_template=task_template,
            ensure_positives_on_train=ensure_positives_on_train,
            sample_only_gold_guidelines=sample_only_gold_guidelines,
            dataset_name=dataset_name,
            scorer=scorer,
            task_definitions=task_definitions,
            task_target=task_target,
            is_coarse_to_fine=False,
            coarse_to_fine=None,
            fine_to_coarse=None,
            definitions=GUIDELINES,
            examples=EXAMPLES,
            **kwargs,
        )

src/tasks/casie/data_loader.py
hitz-zentroa/GoLLIE
GUIDELINES = {
    "databreach_attack_main": {
        "en": [
            "An DatabreachAttack Event happens when an attacker compromises a system\n"
            " to later remove or expose the data, e.g., to sell, publish or make it accessible."
        ]
    },
    "databreach_attack_mention": {
        "en": ["The text span that triggers the event, such as:\n   'attack', 'expose', 'publish', 'steal', ..."]
    },
    "phising_attack_main": {
        "en": [
            "A PhisingAttack Event happens when an attacker imitates another entity, in\n"
            " an attempt to get a victim to access malicious materials, such as a website or\n"
            " attachments."
        ]
    },
    "phising_attack_mention": {
        "en": [
            "The text span that triggers the event, such as:\n"
            "    'attack', 'purports to be', 'dupe', ...\n"
            "    'masquerading as', 'pretending to be', 'scam', ..."
        ]
    },
    "ransom_attack_main": {
        "en": [
            "A RansomAttack Event happens when n attacker breaks into a system and\n"
            " encrypts data, and will only decrypt the data for a ransom payment."
        ]
    },
    "ransom_attack_mention": {
        "en": ["The text span that triggers the event, such as:\n    'attack', ransomware', 'selling', 'ransom', ..."]
    },
    "vulnerability_discover_main": {
        "en": [
            "A VulnerabilityDiscover Event happens when a security expert or other entity,\n"
            " like a company, finds a software vulnerability."
        ]
    },
    "vulnerability_discover_mention": {
        "en": [
            "The text span that triggers the event, such as:\n    'attack', 'found', 'exploit', 'vulnerability', ..."
        ]
    },
    "vulnerability_patch_main": {
        "en": [
            "A VulnerabiltyPatch Event happens when  software company addresses a known\n"
            " vulnerability by releasing or describing an appropriate update."
        ]
    },
    "vulnerability_patch_mention": {
        "en": [
            "The text span that triggers the event, such as:\n"
            "    'patch', 'fixed', 'addresses', 'implemented',\n"
            "    'released', ..."
        ]
    },
}

src/tasks/casie/guidelines.py
hitz-zentroa/GoLLIE
GUIDELINES = {
    "databreach_attack_main": {
        "en": [
            "An DatabreachAttack Event happens when an attacker compromises a system\n"
            " to later remove or expose the data, e.g., to sell, publish or make it accessible."
        ]
    },
    "databreach_attack_mention": {
        "en": ["The text span that triggers the event, such as:\n   'attack', 'expose', 'publish', 'steal', ..."]
    },
    "phising_attack_main": {
        "en": [
            "A PhisingAttack Event happens when an attacker imitates another entity, in\n"
            " an attempt to get a victim to access malicious materials, such as a website or\n"
            " attachments."
        ]
    },
    "phising_attack_mention": {
        "en": [
            "The text span that triggers the event, such as:\n"
            "    'attack', 'purports to be', 'dupe', ...\n"
            "    'masquerading as', 'pretending to be', 'scam', ..."
        ]
    },
    "ransom_attack_main": {
        "en": [
            "A RansomAttack Event happens when n attacker breaks into a system and\n"
            " encrypts data, and will only decrypt the data for a ransom payment."
        ]
    },
    "ransom_attack_mention": {
        "en": ["The text span that triggers the event, such as:\n    'attack', ransomware', 'selling', 'ransom', ..."]
    },
    "vulnerability_discover_main": {
        "en": [
            "A VulnerabilityDiscover Event happens when a security expert or other entity,\n"
            " like a company, finds a software vulnerability."
        ]
    },
    "vulnerability_discover_mention": {
        "en": [
            "The text span that triggers the event, such as:\n    'attack', 'found', 'exploit', 'vulnerability', ..."
        ]
    },
    "vulnerability_patch_main": {
        "en": [
            "A VulnerabiltyPatch Event happens when  software company addresses a known\n"
            " vulnerability by releasing or describing an appropriate update."
        ]
    },
    "vulnerability_patch_mention": {
        "en": [
            "The text span that triggers the event, such as:\n"
            "    'patch', 'fixed', 'addresses', 'implemented',\n"
            "    'released', ..."
        ]
    },
}

EXAMPLES = {
    "databreach_attack_examples": {"en": ["attack", "stole", "publish", "steal"]},
    "phising_attack_examples": {
        "en": ["attack", "purports to be", "dupe", "masquerading as", "pretending to be", "scam"]
    },
    "ransom_attack_examples": {"en": ["attack", "ransomware", "selling", "ransom"]},
    "vulnerability_discover_examples": {"en": ["attack", "found", "exploit", "vulnerability"]},
    "vulnerability_patch_examples": {"en": ["patch", "fixed", "addresses", "implemented", "released"]},
}

src/tasks/casie/guidelines_gold.py
hitz-zentroa/GoLLIE
from __future__ import annotations

import glob
import json
import os
from argparse import ArgumentParser
from collections import defaultdict, namedtuple
from typing import Dict, Iterable, List, Tuple


class CasieProcessor:
    trigger = namedtuple("Trigger", ["start", "end", "text"])
    argument = namedtuple("Argument", ["start", "end", "text", "role"])

    EVENT_MAPPING: Dict[str, str] = {
        "Vulnerability-related_DiscoverVulnerability": "Vulnerability:Discover",
        "Attack_Phishing": "Attack:Phising",
        "Vulnerability-related_PatchVulnerability": "Vulnerability:Patch",
        "Attack_Databreach": "Attack:Databreach",
        "Attack_Ransom": "Attack:Ransom",
    }

    def __init__(self) -> CasieProcessor:
        self._nlp = None

    def _sent_tokenize(self: CasieProcessor, text: str, start_pos: int) -> Iterable[Tuple[int, str]]:
        if not self._nlp:
            import spacy

            self._nlp = spacy.load("en_core_web_sm")

        for sent in self._nlp(text).sents:
            yield [sent[0].idx + start_pos, sent.text]

    @staticmethod
    def get_sent_id(span: Tuple[int, int], sentences: List[Tuple[int, str]]) -> int:
        for idx, sent in enumerate(sentences):
            if sent[0] <= span[0] and sent[0] + len(sent[1]) >= span[1]:
                return idx
        raise ValueError("No sentence found")

    def __call__(self, input_path: str) -> Iterable[dict]:
        annotations = defaultdict(list)
        with open(input_path, "r", encoding="utf-8") as f:
            try:
                data = json.load(f)
            except json.decoder.JSONDecodeError:
                return []
            doc_name = os.path.basename(input_path).replace(".json", "")
            text = data["content"]
            sentences = list(self._sent_tokenize(text, 0))

            events = [event for idx in data["cyberevent"]["hopper"] for event in idx["events"]]
            for event in events:
                _type = self.EVENT_MAPPING[event["type"] + "_" + event["subtype"]]
                try:
                    sentence_id = self.get_sent_id(
                        (event["nugget"]["startOffset"], event["nugget"]["endOffset"]), sentences
                    )
                except ValueError:
                    continue
                start_idx, sentence = sentences[sentence_id]

                # Normalize the trigger
                if (
                    event["nugget"]["text"]
                    != sentence[event["nugget"]["startOffset"] - start_idx : event["nugget"]["endOffset"] - start_idx]
                ):
                    event["nugget"]["startOffset"] = sentence.find(event["nugget"]["text"]) + start_idx
                    event["nugget"]["endOffset"] = event["nugget"]["startOffset"] + len(event["nugget"]["text"])
                trigger = self.trigger(
                    event["nugget"]["startOffset"] - start_idx,
                    event["nugget"]["endOffset"] - start_idx,
                    event["nugget"]["text"],
                )
                assert trigger.text == sentence[trigger.start : trigger.end]

                arguments = []
                if "argument" not in event:
                    event["argument"] = []
                for argument in event["argument"]:
                    try:
                        arg_sent_id = self.get_sent_id((argument["startOffset"], argument["endOffset"]), sentences)
                    except ValueError:
                        continue
                    if sentence_id != arg_sent_id:
                        print("skip")
                        continue

                    if (
                        argument["text"]
                        != sentence[argument["startOffset"] - start_idx : argument["endOffset"] - start_idx]
                    ):
                        argument["startOffset"] = sentence.find(argument["text"]) + start_idx
                        argument["endOffset"] = argument["startOffset"] + len(argument["text"])

                    argument = self.argument(
                        argument["startOffset"] - start_idx,
                        argument["endOffset"] - start_idx,
                        argument["text"],
                        argument["role"]["type"],
                    )

                    assert argument.text == sentence[argument.start : argument.end]
                    arguments.append(argument)

                annotations[sentence_id].append(
                    {"type": _type, "trigger": trigger._asdict(), "arguments": [arg._asdict() for arg in arguments]}
                )

        for sent_id, anns in annotations.items():
            yield {"sent_id": f"{doc_name}:{sent_id}", "text": sentences[sent_id][-1], "events": anns}


def main(args):
    processor = CasieProcessor()
    with open(os.path.join(args.path, "data.jsonl"), "wt", encoding="utf-8") as f:
        for _file in glob.glob(os.path.join(args.path, "annotation", "*.json")):
            for anns in processor(_file):
                f.write(json.dumps(anns, ensure_ascii=False) + "\n")


if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument("--path", type=str, default="data/casie")
    args = parser.parse_args()
    main(args)

src/tasks/casie/preprocess_casie.py
hitz-zentroa/GoLLIE
from typing import List, Type

from ..utils_typing import Event, dataclass


"""Event definitions

The events definitions are derived from the original paper:
https://ojs.aaai.org/index.php/AAAI/article/view/6401


Attack:Databreach:Attack-Pattern -> 'attack_pattern'
Attack:Databreach:Attacker -> 'attacker'
Attack:Databreach:Compromised-Data -> 'compromised_data'
Attack:Databreach:Damage-Amount -> 'damage_amount'
Attack:Databreach:Number-of-Data -> 'number_of_data'
Attack:Databreach:Number-of-Victim -> 'number_of_victim'
Attack:Databreach:Place -> 'place'
Attack:Databreach:Purpose -> 'purpose'
Attack:Databreach:Time -> 'time'
Attack:Databreach:Tool -> 'tool'
Attack:Databreach:Victim -> 'victim'

Attack:Phising:Attack-Pattern -> 'pattern'
Attack:Phising:Attacker -> 'attacker'
Attack:Phising:Damage-Amount -> 'damage_amount'
Attack:Phising:Place -> 'place'
Attack:Phising:Purpose -> 'purpose'
Attack:Phising:Time -> 'time'
Attack:Phising:Tool -> 'tool'
Attack:Phising:Trusted-Entity -> 'trusted_entity'
Attack:Phising:Victim -> 'victim'

Attack:Ransom:Attack-Pattern -> 'pattern'
Attack:Ransom:Attacker -> 'attacker'
Attack:Ransom:Damage-Amount ->  'damage_amount'
Attack:Ransom:Payment-Method ->  'payment_method'
Attack:Ransom:Place ->  'place'
Attack:Ransom:Price ->  'price'
Attack:Ransom:Time ->  'time'
Attack:Ransom:Tool ->  'tool'
Attack:Ransom:Victim ->  'victim'

Vulnerability:Discover:CVE -> 'cve'
Vulnerability:Discover:Capabilities -> 'used_for'
Vulnerability:Discover:Discoverer -> 'discoverer
Vulnerability:Discover:Supported_Platform -> 'supported_platform
Vulnerability:Discover:Time -> 'time'
Vulnerability:Discover:Vulnerability -> 'vulnerability'
Vulnerability:Discover:Vulnerable_System -> 'vulnerable_system'
Vulnerability:Discover:Vulnerable_System_Owner -> 'system_owner'
Vulnerability:Discover:Vulnerable_System_Version -> 'system_version'

Vulnerability:Patch:CVE -> 'cve'
Vulnerability:Patch:Issues-Addressed -> 'issues_addressed'
Vulnerability:Patch:Patch -> 'patch'
Vulnerability:Patch:Patch-Number -> 'patch_number'
Vulnerability:Patch:Releaser -> 'releaser'
Vulnerability:Patch:Supported_Platform -> 'supported_platform'
Vulnerability:Patch:Time -> 'time'
Vulnerability:Patch:Vulnerability -> 'vulnerability'
Vulnerability:Patch:Vulnerable_System -> 'vulnerable_system'
Vulnerability:Patch:Vulnerable_System_Version -> 'system_version'
"""


@dataclass
class DatabreachAttack(Event):
    """{databreach_attack_main}"""

    mention: str
    """The text span that triggers the event.
    {databreach_attack_examples}
    """
    attacker: List[str]  # The agent (person or organization) of the attack
    attack_pattern: List[str]  # How the attack is done
    victim: List[str]  # The device, organization, person, product or website victim of the attack
    number_of_victim: List[str]  # The number of victims affected by the attack
    compromised_data: List[str]  # The data being compromised: 'information', 'data', ...
    number_of_data: List[str]  # The amount of compromised data
    damage_amount: List[str]  # The amount of damage done to the victim
    tool: List[str]  # The file, malware or website used to attack
    purpose: List[str]  # The reason or purpose behind the attack
    place: List[str]  # Where the attack occurred
    time: List[str]  # When the attack occurred


@dataclass
class PhisingAttack(Event):
    """{phising_attack_main}"""

    mention: str
    """The text span that triggers the event.
    {phising_attack_examples}
    """
    pattern: List[str]  # How was the attack triggered, such as 'opening something' or 'clicking somewhere'
    attacker: List[str]  # The person or organization behind the attack
    victim: List[str]  # The victim of the attack
    damage_amount: List[str]  # The amount of damage done to the victim
    tool: List[str]  # The tool used to send the attack, such as 'email', 'website', 'file'
    trusted_entity: List[str]  # The bait, i.e., what the tool pretended (purported) to be
    purpose: List[str]  # What wants to steal the attacker, such as 'information'
    place: List[str]  # Where the attack occurred
    time: List[str]  # When did the attack occurred, such as 'today', 'tomorrow', ...


@dataclass
class RansomAttack(Event):
    """{ransom_attack_main}"""

    mention: str
    """The text span that triggers the event.
    {ransom_attack_examples}
    """
    pattern: List[str]  # What does the attack do until demands are met.
    attacker: List[str]  # Who performed the attack
    victim: List[str]  # The victim of the attack
    damage_amount: List[str]  # The amount of damage the attack produced
    tool: List[str]  # The tool used
    price: List[str]  # The price of the payment
    payment_method: List[str]  # How the payment have to be done, such as 'a webpage'
    time: List[str]  # When the attack took place
    place: List[str]  # Where the attack took place


@dataclass
class VulnerabilityDiscover(Event):
    """{vulnerability_discover_main}"""

    mention: str
    """The text span that triggers the event.
    {vulnerability_discover_examples}
    """
    cve: List[str]  # The vulnerability identifier: such 'CVE-2018-5003'
    used_for: List[str]  # What is the vulnerability used for such as 'allow to take control'
    discoverer: List[str]  # The entity that reported the vulnerability
    supported_platform: List[str]  # The platforms that support the vulnerability
    vulnerability: List[str]  # The vulnerabilities, such as 'vulnerability'
    vulnerable_system: List[str]  # The systems vulnerable to the vulnerability
    system_owner: List[str]  # The owners of the vulnerable system
    system_version: List[str]  # The version of the vulnerable system
    time: List[str]  # When was the vulnerability discovered


@dataclass
class VulnerabilityPatch(Event):
    """{vulnerability_patch_main}"""

    mention: str
    """The text span that triggers the event.
    {vulnerability_patch_examples}
    """
    cve: List[str]  # The vulnerability identifier: such 'CVE-2018-5003'
    issues_addressed: List[str]  # What did the patch fixed
    supported_platform: List[str]  # The platforms that support the vulnerability
    vulnerability: List[str]  # The vulnerability, such as 'vulnerability'
    vulnerable_system: List[str]  # The affected systems, such as 'infraestructures'
    releaser: List[str]  # The entity releasing the patch
    patch: List[str]  # What was the patch about
    patch_number: List[str]  # Nunber or name of the patch
    system_version: List[str]  # The version of the vulnerable system
    time: List[str]  # When was the patch implemented, the date


EAE_EVENT_DEFINITIONS: List[Type] = [
    DatabreachAttack,
    PhisingAttack,
    RansomAttack,
    VulnerabilityDiscover,
    VulnerabilityPatch,
]

src/tasks/casie/prompts_eae.py
hitz-zentroa/GoLLIE
from typing import List, Type

from ..utils_typing import Event, dataclass


"""Event definitions

The events definitions are derived from the original paper:
https://ojs.aaai.org/index.php/AAAI/article/view/6401


Attack:Databreach:Attack-Pattern -> 'attack_pattern'
Attack:Databreach:Attacker -> 'attacker'
Attack:Databreach:Compromised-Data -> 'compromised_data'
Attack:Databreach:Damage-Amount -> 'damage_amount'
Attack:Databreach:Number-of-Data -> 'number_of_data'
Attack:Databreach:Number-of-Victim -> 'number_of_victim'
Attack:Databreach:Place -> 'place'
Attack:Databreach:Purpose -> 'purpose'
Attack:Databreach:Time -> 'time'
Attack:Databreach:Tool -> 'tool'
Attack:Databreach:Victim -> 'victim'

Attack:Phising:Attack-Pattern -> 'pattern'
Attack:Phising:Attacker -> 'attacker'
Attack:Phising:Damage-Amount -> 'damage_amount'
Attack:Phising:Place -> 'place'
Attack:Phising:Purpose -> 'purpose'
Attack:Phising:Time -> 'time'
Attack:Phising:Tool -> 'tool'
Attack:Phising:Trusted-Entity -> 'trusted_entity'
Attack:Phising:Victim -> 'victim'

Attack:Ransom:Attack-Pattern -> 'pattern'
Attack:Ransom:Attacker -> 'attacker'
Attack:Ransom:Damage-Amount ->  'damage_amount'
Attack:Ransom:Payment-Method ->  'payment_method'
Attack:Ransom:Place ->  'place'
Attack:Ransom:Price ->  'price'
Attack:Ransom:Time ->  'time'
Attack:Ransom:Tool ->  'tool'
Attack:Ransom:Victim ->  'victim'

Vulnerability:Discover:CVE -> 'cve'
Vulnerability:Discover:Capabilities -> 'used_for'
Vulnerability:Discover:Discoverer -> 'discoverer
Vulnerability:Discover:Supported_Platform -> 'supported_platform
Vulnerability:Discover:Time -> 'time'
Vulnerability:Discover:Vulnerability -> 'vulnerability'
Vulnerability:Discover:Vulnerable_System -> 'vulnerable_system'
Vulnerability:Discover:Vulnerable_System_Owner -> 'system_owner'
Vulnerability:Discover:Vulnerable_System_Version -> 'system_version'

Vulnerability:Patch:CVE -> 'cve'
Vulnerability:Patch:Issues-Addressed -> 'issues_addressed'
Vulnerability:Patch:Patch -> 'patch'
Vulnerability:Patch:Patch-Number -> 'patch_number'
Vulnerability:Patch:Releaser -> 'releaser'
Vulnerability:Patch:Supported_Platform -> 'supported_platform'
Vulnerability:Patch:Time -> 'time'
Vulnerability:Patch:Vulnerability -> 'vulnerability'
Vulnerability:Patch:Vulnerable_System -> 'vulnerable_system'
Vulnerability:Patch:Vulnerable_System_Version -> 'system_version'
"""


@dataclass
class DatabreachAttack(Event):
    """{databreach_attack_main}"""

    mention: str
    """The text span that triggers the event.
    {databreach_attack_examples}
    """


@dataclass
class PhisingAttack(Event):
    """{phising_attack_main}"""

    mention: str
    """The text span that triggers the event.
    {phising_attack_examples}
    """


@dataclass
class RansomAttack(Event):
    """{ransom_attack_main}"""

    mention: str
    """The text span that triggers the event.
    {ransom_attack_examples}
    """


@dataclass
class VulnerabilityDiscover(Event):
    """{vulnerability_discover_main}"""

    mention: str
    """The text span that triggers the event.
    {vulnerability_discover_examples}
    """


@dataclass
class VulnerabilityPatch(Event):
    """{vulnerability_patch_main}"""

    mention: str
    """The text span that triggers the event.
    {vulnerability_patch_examples}
    """


ED_EVENT_DEFINITIONS: List[Type] = [
    DatabreachAttack,
    PhisingAttack,
    RansomAttack,
    VulnerabilityDiscover,
    VulnerabilityPatch,
]

src/tasks/casie/prompts_ed.py
hitz-zentroa/GoLLIE
from typing import Any, Dict, List, Type

from typing_extensions import override

from src.tasks.casie.prompts_eae import EAE_EVENT_DEFINITIONS
from src.tasks.casie.prompts_ed import ED_EVENT_DEFINITIONS
from src.tasks.utils_scorer import EventScorer


class CASIEEventScorer(EventScorer):
    """CASIE Event and argument classification scorer."""

    valid_types: List[Type] = ED_EVENT_DEFINITIONS

    def __init__(self, allow_partial_match: bool = True) -> None:
        super().__init__()

        self.allow_partial_match: bool = allow_partial_match

    @override
    def __call__(self, reference: Any, predictions: Any) -> Dict[str, Dict[str, float]]:
        if not len(reference) or (len(reference) and not isinstance(reference[0], list)):
            reference = [reference]
        if not len(predictions) or (len(predictions) and not isinstance(predictions[0], list)):
            predictions = [predictions]

        for ref in reference:
            for ref in ref:
                ref._allow_partial_match = self.allow_partial_match
        for pred in predictions:
            for pred in pred:
                pred._allow_partial_match = self.allow_partial_match

        return super().__call__(reference, predictions)


class CASIEEventArgumentScorer(EventScorer):
    """CASIE Event and argument classification scorer."""

    valid_types: List[Type] = EAE_EVENT_DEFINITIONS

    def __init__(self, allow_partial_match: bool = True) -> None:
        super().__init__()

        self.allow_partial_match: bool = allow_partial_match

    @override
    def __call__(self, reference: Any, predictions: Any) -> Dict[str, Dict[str, float]]:
        if not len(reference) or (len(reference) and not isinstance(reference[0], list)):
            reference = [reference]
        if not len(predictions) or (len(predictions) and not isinstance(predictions[0], list)):
            predictions = [predictions]

        for ref in reference:
            for ref in ref:
                ref._allow_partial_match = self.allow_partial_match
        for pred in predictions:
            for pred in pred:
                pred._allow_partial_match = self.allow_partial_match

        return super().__call__(reference, predictions)

src/tasks/casie/scorer.py
hitz-zentroa/GoLLIE
from . import data_loader, prompts, scorer


__all__ = ["data_loader", "prompts", "scorer"]

src/tasks/conll03/__init__.py
hitz-zentroa/GoLLIE
from typing import Dict, List, Tuple, Type, Union

from src.tasks.conll03.guidelines import GUIDELINES
from src.tasks.conll03.guidelines_gold import EXAMPLES
from src.tasks.conll03.prompts import (
    ENTITY_DEFINITIONS,
    ENTITY_DEFINITIONS_woMISC,
    Location,
    Miscellaneous,
    Organization,
    Person,
)
from src.tasks.label_encoding import rewrite_labels

from ..utils_data import DatasetLoader, Sampler
from ..utils_typing import Entity


def get_conll_hf(
    split: str,
    include_misc: bool,
    ENTITY_TO_CLASS_MAPPING: Dict[str, Type[Entity]],
) -> Tuple[List[List[str]], List[List[Entity]]]:
    """
    Get the conll dataset from the huggingface datasets library
    Args:
        split (str): The path_or_split to load. Can be one of `train`, `validation` or `test`.
        include_misc (bool): Whether to include the MISC entity type. Defaults to `True`.
    Returns:
        (List[str],List[Union[Location,Organization,Person,Miscellaneous]]): The text and the entities
    """
    from datasets import load_dataset

    dataset = load_dataset("conll2003")
    id2label = dict(enumerate(dataset["train"].features["ner_tags"].feature.names))
    dataset_sentences: List[List[str]] = []
    dataset_entities: List[List[Entity]] = []

    for example in dataset[split]:
        words = example["tokens"]
        # Some of the CoNLL02-03 datasets are in IOB1 format instead of IOB2,
        # we convert them to IOB2, so we don't have to deal with this later.
        labels = rewrite_labels(labels=[id2label[label] for label in example["ner_tags"]], encoding="iob2")

        # Get labeled word spans
        spans = []
        for i, label in enumerate(labels):
            if label == "O":
                continue
            elif label.startswith("B-"):
                spans.append([label[2:], i, i + 1])
            elif label.startswith("I-"):
                spans[-1][2] += 1
            else:
                raise ValueError(f"Found an unexpected label: {label}")

        # Get entities
        entities = []
        for label, start, end in spans:
            if include_misc or label != "MISC":
                entities.append(ENTITY_TO_CLASS_MAPPING[label](span=" ".join(words[start:end])))

        dataset_sentences.append(words)
        dataset_entities.append(entities)

    return dataset_sentences, dataset_entities


def read_tsv(filepath) -> Tuple[List[List[str]], List[List[str]]]:
    """
    READ tsv file in conll format
    Args:
        filepath: Path to the file
    Returns: List of words, List of labels
    """

    dataset_words: List[List[str]] = []
    dataset_labels: List[List[str]] = []
    with open(filepath, "r", encoding="utf-8") as f:
        words = []
        labels = []
        for line in f:
            line = line.strip()
            if line.startswith("-DOCSTART-") or line == "" or line == "\n":
                if words:
                    dataset_words.append(words)
                    dataset_labels.append(labels)
                    words = []
                    labels = []
            else:
                try:
                    word, label = line.split()
                except ValueError:
                    try:
                        word, label, _ = line.split()
                    except ValueError:
                        raise ValueError(f"Cannot path_or_split line: {line}")
                if word:
                    words.append(word)
                    labels.append(label)
        if words:
            dataset_words.append(words)
            dataset_labels.append(labels)

    print(f"Read {len(dataset_words)} sentences from {filepath}")

    dataset_labels = [rewrite_labels(labels, encoding="iob2") for labels in dataset_labels]

    return dataset_words, dataset_labels


def load_conll_tsv(
    path: str,
    include_misc: bool,
    ENTITY_TO_CLASS_MAPPING: Dict[str, Type[Entity]],
) -> Tuple[List[List[str]], List[List[Entity]]]:
    """
    Load the conll dataset from a tsv file
    Args:
        path (str): The path to the tsv file
        include_misc (bool): Whether to include the MISC entity type. Defaults to `True`.
    Returns:
        (List[str],List[Union[Location,Organization,Person,Miscellaneous]]): The text and the entities
    """
    dataset_sentences: List[List[str]] = []
    dataset_entities: List[List[Entity]] = []

    dataset_words, dataset_labels = read_tsv(path)

    for words, labels in zip(dataset_words, dataset_labels):
        # Some of the CoNLL02-03 datasets are in IOB1 format instead of IOB2,
        # we convert them to IOB2, so we don't have to deal with this later.
        labels = rewrite_labels(labels=labels, encoding="iob2")
        # Get labeled word spans
        spans = []
        for i, label in enumerate(labels):
            if label == "O":
                continue
            elif label.startswith("B-"):
                spans.append([label[2:], i, i + 1])
            elif label.startswith("I-"):
                spans[-1][2] += 1
            else:
                raise ValueError(f"Found an unexpected label: {label}")

        # Get entities
        entities = []
        for label, start, end in spans:
            if include_misc or label.lower() != "misc":
                entities.append(ENTITY_TO_CLASS_MAPPING[label](span=" ".join(words[start:end])))

        dataset_sentences.append(words)
        dataset_entities.append(entities)

    return dataset_sentences, dataset_entities


class CoNLLDatasetLoader(DatasetLoader):
    """
    A `DatasetLoader` for the ConLL03 dataset.

    Args:
        path_or_split (`str`):
            The split to load. Can be one of `train`, `validation` or `test`. Or the path to a
            tsv file in conll format.
        include_misc (`str`, optional):
            Whether to include the MISC entity type. Defaults to `True`.

    Raises:
        `ValueError`:
            raised when a not defined value found.
    """

    ENTITY_TO_CLASS_MAPPING = None

    def __init__(self, path_or_split: str, include_misc: bool = True, **kwargs) -> None:
        self.ENTITY_TO_CLASS_MAPPING = (
            {
                "LOC": Location,
                "ORG": Organization,
                "PER": Person,
                "MISC": Miscellaneous,
            }
            if include_misc
            else {
                "LOC": Location,
                "ORG": Organization,
                "PER": Person,
            }
        )

        self.elements = {}

        if path_or_split in ["train", "validation", "test"]:
            dataset_words, dataset_entities = get_conll_hf(
                split=path_or_split,
                include_misc=include_misc,
                ENTITY_TO_CLASS_MAPPING=self.ENTITY_TO_CLASS_MAPPING,
            )
        else:
            dataset_words, dataset_entities = load_conll_tsv(
                path=path_or_split,
                include_misc=include_misc,
                ENTITY_TO_CLASS_MAPPING=self.ENTITY_TO_CLASS_MAPPING,
            )

        for id, (words, entities) in enumerate(zip(dataset_words, dataset_entities)):
            self.elements[id] = {
                "id": id,
                "doc_id": id,
                "text": " ".join(words),
                "entities": entities,
                "gold": entities,
            }


class CoNLL03Sampler(Sampler):
    """
    A data `Sampler` for the CONLL03 dataset.

    Args:
        dataset_loader (`CoNLLDatasetLoader`):
            The dataset loader that contains the data information.
        task (`str`, optional):
            The task to sample. It must be one of the following: NER, VER, RE, EE.
            Defaults to `None`.
        split (`str`, optional):
            The path_or_split to sample. It must be one of the following: "train", "dev" or
            "test". Depending on the path_or_split the sampling strategy differs. Defaults to
            `"train"`.
        parallel_instances (`Union[int, Tuple[int, int]]`, optional):
            The number of sentences sampled in parallel. Options:

                * **`int`**: The amount of elements that will be sampled in parallel.
                * **`tuple`**: The range of elements that will be sampled in parallel.

            Defaults to 1.
        max_guidelines (`int`, optional):
            The number of guidelines to append to the example at the same time. If `-1`
            is given then all the guidelines are appended. Defaults to `-1`.
        guideline_dropout (`float`, optional):
            The probability to dropout a guideline definition for the given example. This
            is only applied on training. Defaults to `0.0`.
        seed (`float`, optional):
            The seed to sample the examples. Defaults to `0`.
        prompt_template (`str`, optional):
            The path to the prompt template. Defaults to `"templates/prompt.txt"`.
        ensure_positives_on_train (bool, optional):
            Whether to ensure that the guidelines of annotated examples are not removed.
            Defaults to `False`.
        dataset_name (str, optional):
            The name of the dataset. Defaults to `None`.
        scorer (`str`, optional):
           The scorer class import string. Defaults to `None`.
        sample_only_gold_guidelines (`bool`, optional):
            Whether to sample only guidelines of present annotations. Defaults to `False`.
    """

    def __init__(
        self,
        dataset_loader: CoNLLDatasetLoader,
        task: str = None,
        split: str = "train",
        parallel_instances: Union[int, Tuple[int, int]] = 1,
        max_guidelines: int = -1,
        guideline_dropout: float = 0.0,
        seed: float = 0,
        prompt_template: str = "templates/prompt.txt",
        ensure_positives_on_train: bool = False,
        dataset_name: str = None,
        scorer: str = None,
        sample_only_gold_guidelines: bool = False,
        **kwargs,
    ) -> None:
        assert task in [
            "NER",
        ], f"CoNLL03 only supports NER task. {task} is not supported."

        task_definitions, task_target = {
            "NER": (ENTITY_DEFINITIONS if kwargs["include_misc"] else ENTITY_DEFINITIONS_woMISC, "entities"),
        }[task]

        super().__init__(
            dataset_loader=dataset_loader,
            task=task,
            split=split,
            parallel_instances=parallel_instances,
            max_guidelines=max_guidelines,
            guideline_dropout=guideline_dropout,
            seed=seed,
            prompt_template=prompt_template,
            ensure_positives_on_train=ensure_positives_on_train,
            sample_only_gold_guidelines=sample_only_gold_guidelines,
            dataset_name=dataset_name,
            scorer=scorer,
            task_definitions=task_definitions,
            task_target=task_target,
            definitions=GUIDELINES,
            examples=EXAMPLES,
            **kwargs,
        )

src/tasks/conll03/data_loader.py
hitz-zentroa/GoLLIE
GUIDELINES = {
    "ner_person": {
        "en": [
            "first, middle and last names of people, animals and fictional characters aliases.",
            (
                'Paraphrase: "The initial, middle, and last names of individuals, creatures, and imaginary'
                " personalities, along \nwith their pseudonyms."
            ),
            (
                'Paraphrase: "Initially, focus on the full names of individuals, animals, and imaginary personalities,'
                " as well as \ntheir pseudonyms."
            ),
            (
                'Paraphrase: "Initially, the study focused on the names of individuals, including their first, middle,'
                " and last \nnames, as well as the names of animals and imaginary characters, along with any aliases"
                " they might have."
            ),
            (
                'Original text: "first, middle and last names of people, animals and fictional characters aliases."'
                ' Paraphrased \ntext: "Initially, individuals consider the full names of humans, animals, and'
                " imaginary personalities, along with \ntheir nicknames or pseudonyms."
            ),
        ]
    },
    "ner_organization": {
        "en": [
            (
                "Companies (press agencies, studios, banks, stock markets, manufacturers, cooperatives) subdivisions"
                " of \ncompanies (newsrooms) brands political movements (political parties, terrorist organisations)"
                " government bodies \n(ministries, councils, courts, political unions of countries (e.g. the {\\it"
                " U.N.})) publications (magazines, newspapers, \njournals) musical companies (bands, choirs, opera"
                " companies, orchestras public organisations (schools, universities, \ncharities other collections of"
                " people (sports clubs, sports teams, associations, theaters companies, religious orders, \nyouth"
                " organisations."
            ),
            (
                "Organizations such as news agencies, film studios, banks, stock exchanges, manufacturers,"
                " cooperatives, as well \nas their subdivisions like newsrooms, and various entities like political"
                " movements (political parties, terrorist \ngroups), government bodies (ministries, councils, courts),"
                " political unions of countries (e.g., the United Nations), \npublications (magazines, newspapers,"
                " journals), musical companies (bands, choirs, opera companies, orchestras), public \norganizations"
                " (schools, universities, charities), and other assortments of people (sports clubs, sports teams,"
                " associations, \ntheaters, companies, religious orders, youth organizations) are examples of groups"
                " that can be analyzed using network \nanalysis."
            ),
            (
                "Organizations such as news agencies, film studios, banks, stock exchanges, manufacturers, and"
                " cooperatives, as \nwell as subdivisions of companies like newsrooms, and distinct brands encompass"
                " political movements like political \nparties and terrorist groups, government entities including"
                " ministries, councils, courts, and political unions of \ncountries like the United Nations, various"
                " publications like magazines, newspapers, and journals, musical companies like \nbands, choirs, opera"
                " companies, and orchestras, public organizations like schools, universities, and charities, and"
                " \nother assortments of people such as sports clubs, sports teams, associations, theaters, and"
                " religious orders, as well as \nyouth organizations."
            ),
            (
                "Organizations such as news agencies, film studios, banks, stock exchanges, manufacturers, and"
                " cooperatives, as \nwell as subdivisions of these companies like newsrooms, and brands associated"
                " with political movements like political \nparties and terrorist groups, government entities"
                " including ministries, councils, courts, and political unions of \ncountries like the United Nations,"
                " various publications like magazines, newspapers, and journals, musical companies \nencompassing"
                " bands, choirs, opera companies, and orchestras, public organizations such as schools, universities,"
                " and \ncharities, and finally other assortments of people like sports clubs, sports teams,"
                " associations, theaters, companies, \nreligious orders, and youth organizations."
            ),
            (
                "Organizations such as news agencies, film studios, banks, stock exchanges, manufacturers,"
                " cooperatives, as well \nas subdivisions of companies like newsrooms, brands, political movements"
                " like political parties and terrorist \norganizations, government entities like ministries, councils,"
                " courts, and political unions of countries like the United \nNations, various publications like"
                " magazines, newspapers, and journals, musical companies like bands, choirs, opera \ncompanies, and"
                " orchestras, public organizations like schools, universities, and charities, and other groups of"
                " people like \nsports clubs, sports teams, associations, theaters, companies, religious orders, and"
                " youth organizations."
            ),
        ]
    },
    "ner_location": {
        "en": [
            (
                "Roads (streets, motorways) trajectories regions (villages, towns, cities, provinces, countries,"
                " continents, \ndioceses, parishes) structures (bridges, ports, dams) natural locations (mountains,"
                " mountain ranges, woods, rivers, \nwells, fields, valleys, gardens, nature reserves, allotments,"
                " beaches, national parks) public places (squares, opera \nhouses, museums, schools, markets,"
                " airports, stations, swimming pools, hospitals, sports facilities, youth centers, \nparks, town"
                " halls, theaters, cinemas, galleries, camping grounds, NASA launch pads, club houses, universities,"
                " \nlibraries, churches, medical centers, parking lots, playgrounds, cemeteries) commercial places"
                " (chemists, pubs, \nrestaurants, depots, hostels, hotels, industrial parks, nightclubs, music venues)"
                " assorted buildings (houses, monasteries, \ncreches, mills, army barracks, castles, retirement homes,"
                " towers, halls, rooms, vicarages, courtyards) abstract \n``places'' (e.g. {\\it the free world})"
            ),
            (
                "Roads (including streets and motorways) connect various regions (such as villages, towns, cities,"
                " provinces, \ncountries, continents, dioceses, and parishes). They also provide access to structures"
                " (like bridges, ports, and dams) and \nnatural locations (like mountains, mountain ranges, woods,"
                " rivers, wells, fields, valleys, gardens, nature reserves, \nallotments, beaches, and national"
                " parks). These roads lead to public places (like squares, opera houses, museums, schools, \nmarkets,"
                " airports, stations, swimming pools, hospitals, sports facilities, youth centers, parks, town halls,"
                " theaters, \ncinemas, galleries, camping grounds, NASA launch pads, club houses, universities,"
                " libraries, churches, medical centers, \nand parking lots) and commercial places (such as chemists,"
                " pubs, restaurants, depots, hostels, hotels, industrial \nparks, nightclubs, and music venues)."
                " Additionally, they pass by various types of buildings (including houses, \nmonasteries, creches,"
                " mills, army barracks, castles, retirement homes, towers, halls, rooms, vicarages, and courtyards),"
                ' as \nwell as more abstract "places" (e.g., "the free world").'
            ),
            (
                "Roads (including streets, motorways) lead to various regions (such as villages, towns, cities,"
                " provinces, \ncountries, continents, dioceses, parishes) and are connected to different structures"
                " (like bridges, ports, dams). They also \ntraverse natural locations (like mountains, mountain"
                " ranges, woods, rivers, wells, fields, valleys, gardens, nature \nreserves, allotments, beaches, and"
                " national parks) as well as public places (encompassing squares, opera houses, museums, \nschools,"
                " markets, airports, stations, swimming pools, hospitals, sports facilities, youth centers, parks,"
                " town halls, \ntheaters, cinemas, galleries, camping grounds, NASA launch pads, club houses,"
                " universities, libraries, churches, medical \ncenters, parking lots, playgrounds, and cemeteries)."
                " Furthermore, they pass by commercial places (such as chemists, pubs, \nrestaurants, depots, hostels,"
                " hotels, industrial parks, nightclubs, and music venues) and various buildings (like houses,"
                " \nmonasteries, creches, mills, army barracks, castles, retirement homes, towers, halls, rooms,"
                ' vicarages, and courtyards). \nLastly, they can also encompass abstract "places" (e.g., "the free'
                ' world").'
            ),
            (
                "Roads (including streets, motorways) lead to various regions (encompassing villages, towns, cities,"
                " provinces, \ncountries, continents, dioceses, parishes), and they connect different structural"
                " elements (such as bridges, ports, dams). \nThey also traverse natural locations (like mountains,"
                " mountain ranges, woods, rivers, wells, fields, valleys, \ngardens, nature reserves, allotments,"
                " beaches, and national parks). Furthermore, they provide access to public places \n(like squares,"
                " opera houses, museums, schools, markets, airports, stations, swimming pools, hospitals, sports"
                " \nfacilities, youth centers, parks, town halls, theaters, cinemas, galleries, camping grounds, NASA"
                " launch pads, club houses, \nuniversities, libraries, churches, medical centers, and parking lots)"
                " and commercial establishments (like chemists, pubs, \nrestaurants, depots, hostels, hotels,"
                " industrial parks, nightclubs, and music venues). Additionally, they serve a variety of \nspecific"
                " buildings (including houses, monasteries, creches, mills, army barracks, castles, retirement homes,"
                ' towers, \nhalls, rooms, vicarages, and courtyards), as well as more abstract "places" (e.g., "the'
                ' free world").'
            ),
            (
                "Roads (including streets, motorways) lead to various regions (encompassing villages, towns, cities,"
                " provinces, \ncountries, continents, dioceses, parishes), and they connect diverse structures (such"
                " as bridges, ports, dams). They also \ntraverse natural locations (like mountains, mountain ranges,"
                " woods, rivers, wells, fields, valleys, gardens, nature \nreserves, allotments, beaches, and national"
                " parks). Furthermore, they access public places (like squares, opera houses, \nmuseums, schools,"
                " markets, airports, stations, swimming pools, hospitals, sports facilities, youth centers, parks,"
                " town \nhalls, theaters, cinemas, galleries, camping grounds, NASA launch pads, club houses,"
                " universities, libraries, \nchurches, medical centers, and parking lots), as well as commercial"
                " places (chemists, pubs, restaurants, depots, hostels, \nhotels, industrial parks, nightclubs, music"
                " venues). Additionally, they pass by a variety of buildings (such as houses, \nmonasteries, creches,"
                " mills, army barracks, castles, retirement homes, towers, halls, rooms, vicarages, and courtyards),"
                ' and \nthey even touch abstract "places" (e.g., "the free world").'
            ),
        ]
    },
    "ner_miscellaneous": {
        "en": [
            (
                "Words of which one part is a location, organisation, miscellaneous, or person adjectives and other"
                " words derived \nfrom a word which is location, organisation, miscellaneous, or person religions"
                " political ideologies nationalities \nlanguages programs events (conferences, festivals, sports"
                " competitions, forums, parties, concerts) wars sports related \nnames (league tables, leagues, cups"
                " titles (books, songs, films, stories, albums, musicals, TV programs) slogans eras \nin time types"
                " (not brands) of objects (car types, planes, motorbikes)"
            ),
            (
                "Phrases that include a location, organization, or person descriptor, along with terms derived from"
                " words \npertaining to religions, political ideologies, nationalities, languages, events, programs,"
                " wars, sports-related names, \ntitles, slogans, and types of objects, such as cars, planes, and"
                " motorbikes."
            ),
            (
                "Phrases that include a term denoting a location, organization, abstract concept, or person, as well"
                " as words \nderived from these categories, such as religions, political ideologies, nationalities,"
                " languages, various types of \nevents (like conferences, festivals, sports competitions, forums,"
                " parties, concerts), wars, sports-related \ndesignations (league tables, leagues, cups, titles of"
                " books, songs, films, stories, albums, musicals, TV programs), slogans, \nand eras in time associated"
                " with specific types of objects (like car types, planes, motorbikes)."
            ),
            (
                'Text: "Terms consisting of a component indicating a place, institution, abstract concept, or person,'
                " as well as \nderived words from such terms that encompass religious beliefs, political doctrines,"
                " nationalities, languages, events \n(such as conferences, festivals, sports competitions, forums,"
                " parties, concerts), wars, sports-related \nappellations (like league tables, leagues, cups, and"
                " titles of books, songs, films, stories, albums, musicals, TV programs), \nslogans, and epochs in"
                " time, as well as types (not brands) of objects (like car types, planes, motorbikes)."
            ),
            (
                "Terms consisting of a component that indicates a place, institution, abstract notion, or individual,"
                " as well as \nwords derived from a word that signifies a location, organization, abstract concept, or"
                " nationality, encompassing \nvarious religions, political ideologies, languages, events (including"
                " conferences, festivals, sports competitions, \nforums, parties, and concerts), wars, sports-related"
                " appellations (such as league tables, leagues, and cup titles in \nbooks, songs, films, stories,"
                " albums, musicals, and TV programs), slogans, and epochs in time, while also covering types \n(as"
                " opposed to specific brands) of objects (like car models, planes, and motorbikes)."
            ),
        ]
    },
}

src/tasks/conll03/guidelines.py
hitz-zentroa/GoLLIE
GUIDELINES = {
    "ner_person": {"en": ["first, middle and last names of people, animals and fictional characters aliases."]},
    "ner_organization": {
        "en": [
            "Companies (press agencies, studios, banks, stock\n"
            "markets, manufacturers, cooperatives) subdivisions of companies (newsrooms)\n"
            "brands political movements (political parties, terrorist organisations)\n"
            "government bodies (ministries, councils, courts, political unions of countries\n"
            "(e.g. the {\\it U.N.})) publications (magazines, newspapers, journals)\n"
            "musical companies (bands, choirs, opera companies, orchestras\n"
            "public organisations (schools, universities, charities other collections\n"
            "of people (sports clubs, sports teams, associations, theaters companies,\n"
            "religious orders, youth organisations.\n"
        ]
    },
    "ner_location": {
        "en": [
            "Roads (streets, motorways) trajectories regions (villages, towns, cities, provinces,\n"
            "countries, continents, dioceses, parishes) structures (bridges, ports, dams) natural locations\n"
            "(mountains, mountain ranges, woods, rivers, wells, fields, valleys, gardens, nature reserves,\n"
            "allotments, beaches, national parks) public places (squares, opera houses, museums, schools, markets,\n"
            "airports, stations, swimming pools, hospitals, sports facilities, youth centers, parks, town halls,\n"
            "theaters, cinemas, galleries, camping grounds, NASA launch pads, club houses, universities, libraries,\n"
            "churches, medical centers, parking lots, playgrounds, cemeteries) commercial places (chemists, pubs,\n"
            "restaurants, depots, hostels, hotels, industrial parks, nightclubs, music venues) assorted buildings\n"
            "(houses, monasteries, creches, mills, army barracks, castles, retirement homes, towers, halls, rooms,\n"
            "vicarages, courtyards) abstract ``places'' (e.g. {\\it the free world})\n"
        ]
    },
    "ner_miscellaneous": {
        "en": [
            "Words of which one part is a location, organisation, miscellaneous, or person adjectives\n"
            "and other words derived from a word which is location, organisation, miscellaneous, or person\n"
            "religions political ideologies nationalities languages programs events (conferences, festivals,\n"
            "sports competitions, forums, parties, concerts) wars sports related names (league tables, leagues,\n"
            "cups titles (books, songs, films, stories, albums, musicals, TV programs) slogans eras in time types\n"
            "(not brands) of objects (car types, planes, motorbikes)\n"
        ]
    },
}

EXAMPLES = {
    "ner_location_examples": {
        "en": ["U.S.", "Germany", "Britain", "Australia", "England", "France", "Spain", "Italy", "LONDON", "China"]
    },
    "ner_organization_examples": {
        "en": ["Reuters", "U.N.", "NEW YORK", "CHICAGO", "PUK", "OSCE", "EU", "NATO", "European Union", "Honda"]
    },
    "ner_person_examples": {
        "en": [
            "Clinton",
            "Dole",
            "Arafat",
            "Yeltsin",
            "Lebed",
            "Dutroux",
            "Wasim Akram",
            "Waqar Younis",
            "Mushtaq Ahmed",
            "Mother Teresa",
        ]
    },
    "ner_miscellaneous_examples": {
        "en": ["Russian", "German", "British", "French", "Dutch", "GMT", "Israeli", "English", "Iraqi", "European"]
    },
}

src/tasks/conll03/guidelines_gold.py
hitz-zentroa/GoLLIE
from typing import List

from ..utils_typing import Entity, dataclass


"""Entity definitions

The entity definitions are derived from the official ConLL2003 guidelines:
https://www.clips.uantwerpen.be/conll2003/ner/
Based on: Nancy Chinchor, Erica Brown, Lisa Ferro, Patty Robinson,
           "1999 Named Entity Task Definition". MITRE and SAIC, 1999.
"""


@dataclass
class Person(Entity):
    """{ner_person}"""

    span: str  # {ner_person_examples}


@dataclass
class Organization(Entity):
    """{ner_organization}"""

    span: str  # {ner_organization_examples}


@dataclass
class Location(Entity):
    """{ner_location}"""

    span: str  # {ner_location_examples}


@dataclass
class Miscellaneous(Entity):
    """{ner_miscellaneous}"""

    span: str  # {ner_miscellaneous_examples}


ENTITY_DEFINITIONS: List[Entity] = [
    Person,
    Organization,
    Location,
    Miscellaneous,
]

ENTITY_DEFINITIONS_woMISC: List[Entity] = [
    Person,
    Organization,
    Location,
]


# __all__ = list(map(str, [*ENTITY_DEFINITIONS]))

src/tasks/conll03/prompts.py
hitz-zentroa/GoLLIE
from typing import Dict, List, Type

from src.tasks.conll03.prompts import ENTITY_DEFINITIONS, ENTITY_DEFINITIONS_woMISC
from src.tasks.utils_scorer import SpanScorer
from src.tasks.utils_typing import Entity


class CoNLL03EntityScorer(SpanScorer):
    """CoNLL03 Entity identification and classification scorer."""

    valid_types: List[Type] = ENTITY_DEFINITIONS

    def __call__(self, reference: List[Entity], predictions: List[Entity]) -> Dict[str, Dict[str, float]]:
        output = super().__call__(reference, predictions)
        return {"entities": output["spans"]}


class CoNLL03EntityScorerNoMisc(SpanScorer):
    """CoNLL03 Entity identification and classification scorer."""

    valid_types: List[Type] = ENTITY_DEFINITIONS_woMISC

    def __call__(self, reference: List[Entity], predictions: List[Entity]) -> Dict[str, Dict[str, float]]:
        output = super().__call__(reference, predictions)
        return {"entities": output["spans"]}

src/tasks/conll03/scorer.py
hitz-zentroa/GoLLIE
from . import (
    data_loader,
    prompts_ai,
    prompts_literature,
    prompts_music,
    prompts_natural_science,
    prompts_politics,
    scorer,
)


__all__ = [
    "data_loader",
    "prompts_politics",
    "prompts_natural_science",
    "prompts_music",
    "prompts_literature",
    "prompts_ai",
    "scorer",
]

src/tasks/crossner/__init__.py
hitz-zentroa/GoLLIE
from typing import Tuple, Union

from src.tasks.conll03.data_loader import load_conll_tsv
from src.tasks.crossner.guidelines import GUIDELINES
from src.tasks.crossner.guidelines_gold import EXAMPLES
from src.tasks.crossner.prompts_ai import ENTITY_DEFINITIONS_AI, ENTITY_DEFINITIONS_AI_woMISC
from src.tasks.crossner.prompts_literature import ENTITY_DEFINITIONS_LITERATURE, ENTITY_DEFINITIONS_LITERATURE_woMISC
from src.tasks.crossner.prompts_music import ENTITY_DEFINITIONS_MUSIC, ENTITY_DEFINITIONS_MUSIC_woMISC
from src.tasks.crossner.prompts_natural_science import (
    ENTITY_DEFINITIONS_NATURAL_SCIENCE,
    ENTITY_DEFINITIONS_NATURAL_SCIENCE_woMISC,
)
from src.tasks.crossner.prompts_politics import ENTITY_DEFINITIONS_POLITICS, ENTITY_DEFINITIONS_POLITICS_woMISC

from ..utils_data import DatasetLoader, Sampler


class CrossNERDatasetLoader(DatasetLoader):
    """
    A `DatasetLoader` for the CrossNER dataset.

    Args:
        split (`str`):
            The split to load. Can be one of `train`, `validation` or `test`.
        task (`str`):
            The task to load. Can be one of `politics`, `music`, `literature`, `ai` or `natural_science`.

    Raises:
        `ValueError`:
            raised when a not defined value found.
    """

    ENTITY_TO_CLASS_MAPPING = None

    def __init__(self, path_or_split: str, tasks: str, include_misc: bool = True, **kwargs) -> None:
        if len(tasks) > 1:
            raise ValueError(
                "CrossNER only supports one task at a time. Please specify only one task in the config file. You"
                f" specified {tasks}"
            )
        task = tasks[0].lower().replace("crossner_", "")

        if task == "politics":
            from src.tasks.crossner.prompts_politics import (
                Country,
                Election,
                Event,
                Location,
                Organization,
                Other,
                Person,
                PoliticalParty,
                Politician,
            )

            self.ENTITY_TO_CLASS_MAPPING = {
                "person": Person,
                "organisation": Organization,
                "location": Location,
                "politician": Politician,
                "politicalparty": PoliticalParty,
                "election": Election,
                "event": Event,
                "country": Country,
                "misc": Other,
            }

        elif task == "music":
            from src.tasks.crossner.prompts_music import (
                Album,
                Award,
                Band,
                Country,
                Event,
                Location,
                MusicalArtist,
                MusicalInstrument,
                MusicGenre,
                Organization,
                Other,
                Person,
                Song,
            )

            self.ENTITY_TO_CLASS_MAPPING = {
                "musicgenre": MusicGenre,
                "song": Song,
                "band": Band,
                "album": Album,
                "musicalartist": MusicalArtist,
                "musicalinstrument": MusicalInstrument,
                "award": Award,
                "event": Event,
                "country": Country,
                "location": Location,
                "organisation": Organization,
                "person": Person,
                "misc": Other,
            }

        elif task == "literature":
            from src.tasks.crossner.prompts_literature import (
                Award,
                Book,
                Country,
                Event,
                LiteraryGenre,
                Location,
                Magazine,
                Organization,
                Other,
                Person,
                Poem,
                Writer,
            )

            self.ENTITY_TO_CLASS_MAPPING = {
                "book": Book,
                "writer": Writer,
                "award": Award,
                "poem": Poem,
                "event": Event,
                "magazine": Magazine,
                "literarygenre": LiteraryGenre,
                "person": Person,
                "location": Location,
                "organisation": Organization,
                "country": Country,
                "misc": Other,
            }

        elif task == "ai":
            from src.tasks.crossner.prompts_ai import (
                Algorithm,
                Conference,
                Country,
                Field,
                Location,
                Metric,
                Organization,
                Other,
                Person,
                Product,
                ProgrammingLanguage,
                Researcher,
                Task,
                University,
            )

            self.ENTITY_TO_CLASS_MAPPING = {
                "field": Field,
                "task": Task,
                "product": Product,
                "algorithm": Algorithm,
                "researcher": Researcher,
                "metrics": Metric,
                "university": University,
                "country": Country,
                "person": Person,
                "organisation": Organization,
                "location": Location,
                "programlang": ProgrammingLanguage,
                "conference": Conference,
                "misc": Other,
            }
        elif task == "natural_science":
            from src.tasks.crossner.prompts_natural_science import (
                AcademicJournal,
                AstronomicalObject,
                Award,
                ChemicalCompound,
                ChemicalElement,
                Country,
                Discipline,
                Enzyme,
                Event,
                Location,
                Organization,
                Other,
                Person,
                Protein,
                Scientist,
                Theory,
                University,
            )

            self.ENTITY_TO_CLASS_MAPPING = {
                "scientist": Scientist,
                "person": Person,
                "university": University,
                "organisation": Organization,
                "country": Country,
                "location": Location,
                "discipline": Discipline,
                "enzyme": Enzyme,
                "protein": Protein,
                "chemicalelement": ChemicalElement,
                "chemicalcompound": ChemicalCompound,
                "astronomicalobject": AstronomicalObject,
                "academicjournal": AcademicJournal,
                "event": Event,
                "theory": Theory,
                "award": Award,
                "misc": Other,
            }
        else:
            raise ValueError(
                f"Task {task} not defined. Please choose one of the following: politics, music, literature, ai or"
                " natural_science"
            )

        if not include_misc:
            self.ENTITY_TO_CLASS_MAPPING.pop("misc")

        self.elements = {}

        dataset_words, dataset_entities = load_conll_tsv(
            path=path_or_split,
            include_misc=include_misc,
            ENTITY_TO_CLASS_MAPPING=self.ENTITY_TO_CLASS_MAPPING,
        )

        for id, (words, entities) in enumerate(zip(dataset_words, dataset_entities)):
            self.elements[id] = {
                "id": id,
                "doc_id": id,
                "text": " ".join(words),
                "entities": entities,
                "gold": entities,
            }


class CrossNERSampler(Sampler):
    """
    A data `Sampler` for the CrossNER dataset.

    Args:
        dataset_loader (`CrossNERDatasetLoader`):
            The dataset loader that contains the data information.
        task (`str`, optional):
            The task to sample. It must be one of the following: NER, VER, RE, EE.
            Defaults to `None`.
        split (`str`, optional):
            The path_or_split to sample. It must be one of the following: "train", "dev" or
            "test". Depending on the path_or_split the sampling strategy differs. Defaults to
            `"train"`.
        parallel_instances (`Union[int, Tuple[int, int]]`, optional):
            The number of sentences sampled in parallel. Options:

                * **`int`**: The amount of elements that will be sampled in parallel.
                * **`tuple`**: The range of elements that will be sampled in parallel.

            Defaults to 1.
        max_guidelines (`int`, optional):
            The number of guidelines to append to the example at the same time. If `-1`
            is given then all the guidelines are appended. Defaults to `-1`.
        guideline_dropout (`float`, optional):
            The probability to dropout a guideline definition for the given example. This
            is only applied on training. Defaults to `0.0`.
        seed (`float`, optional):
            The seed to sample the examples. Defaults to `0`.
        prompt_template (`str`, optional):
            The path to the prompt template. Defaults to `"templates/prompt.txt"`.
        ensure_positives_on_train (bool, optional):
            Whether to ensure that the guidelines of annotated examples are not removed.
            Defaults to `False`.
        dataset_name (str, optional):
            The name of the dataset. Defaults to `None`.
        scorer (`str`, optional):
           The scorer class import string. Defaults to `None`.
        sample_only_gold_guidelines (`bool`, optional):
            Whether to sample only guidelines of present annotations. Defaults to `False`.
    """

    def __init__(
        self,
        dataset_loader: CrossNERDatasetLoader,
        task: str = None,
        split: str = "train",
        parallel_instances: Union[int, Tuple[int, int]] = 1,
        max_guidelines: int = -1,
        guideline_dropout: float = 0.0,
        seed: float = 0,
        prompt_template: str = "templates/prompt.txt",
        ensure_positives_on_train: bool = False,
        dataset_name: str = None,
        scorer: str = None,
        sample_only_gold_guidelines: bool = False,
        **kwargs,
    ) -> None:
        assert task in [
            "CrossNER_POLITICS",
            "CrossNER_AI",
            "CrossNER_NATURAL_SCIENCE",
            "CrossNER_LITERATURE",
            "CrossNER_MUSIC",
        ], (
            "CrossNER only supports NER_POLITICS, NER_AI, NER_NATURAL_SCIENCE, NER_LITERATURE, NER_MUSIC. Task"
            f" {task} is not supported."
        )

        include_misc = kwargs["include_misc"]

        task_definitions, task_target = {
            "CrossNER_POLITICS": (
                ENTITY_DEFINITIONS_POLITICS if include_misc else ENTITY_DEFINITIONS_POLITICS_woMISC,
                "entities",
            ),
            "CrossNER_AI": (ENTITY_DEFINITIONS_AI if include_misc else ENTITY_DEFINITIONS_AI_woMISC, "entities"),
            "CrossNER_NATURAL_SCIENCE": (
                ENTITY_DEFINITIONS_NATURAL_SCIENCE if include_misc else ENTITY_DEFINITIONS_NATURAL_SCIENCE_woMISC,
                "entities",
            ),
            "CrossNER_LITERATURE": (
                ENTITY_DEFINITIONS_LITERATURE if include_misc else ENTITY_DEFINITIONS_LITERATURE_woMISC,
                "entities",
            ),
            "CrossNER_MUSIC": (
                ENTITY_DEFINITIONS_MUSIC if include_misc else ENTITY_DEFINITIONS_MUSIC_woMISC,
                "entities",
            ),
        }[task]

        super().__init__(
            dataset_loader=dataset_loader,
            task=task,
            split=split,
            parallel_instances=parallel_instances,
            max_guidelines=max_guidelines,
            guideline_dropout=guideline_dropout,
            seed=seed,
            prompt_template=prompt_template,
            ensure_positives_on_train=ensure_positives_on_train,
            sample_only_gold_guidelines=sample_only_gold_guidelines,
            dataset_name=dataset_name,
            scorer=scorer,
            task_definitions=task_definitions,
            task_target=task_target,
            definitions=GUIDELINES,
            examples=EXAMPLES,
            **kwargs,
        )

src/tasks/crossner/data_loader.py
hitz-zentroa/GoLLIE
GUIDELINES = {
    "crossner_politics_person": {"en": ["Refers to an individual's name that is not a politician."]},
    "crossner_politics_organization": {
        "en": ["Refers to a structured group, institution, company, or association that is not a political party."]
    },
    "crossner_politics_location": {
        "en": [
            "Refers to a specific geographical or structural location. This includes but is not limited"
            " to: places (e.g., parks, landmarks), bridges, cities, towns, villages, areas and other distinct regions."
        ]
    },
    "crossner_politics_politician": {
        "en": [
            "Refers to a person who is actively engaged in politics, holding a public office, involved in political"
            " activities or part of a political party. "
            "If a person entity is identified as a politician, annotate that person's name as a politician entity."
        ]
    },
    "crossner_politics_politicalparty": {
        "en": [
            "A political party is an organization that coordinates candidates to compete in a particular country's"
            " elections. Also annotate acronyms of political parties."
        ]
    },
    "crossner_politics_election": {
        "en": [
            "Refers to an organized process in which individuals vote to choose a candidate or make a"
            " decision. If an event mentioned in the text pertains to an election, annotate it as an election entity."
        ]
    },
    "crossner_politics_event": {
        "en": [
            "Refers to a significant occurrence or happening. Such as wars, conferences, summits, conventions, etc."
        ]
    },
    "crossner_politics_country": {"en": ["Refers to a sovereign nation."]},
    "crossner_politics_miscellaneous": {
        "en": ["Refers to named entities that are not included in any other category."]
    },
    "crossner_naturalscience_scientist": {
        "en": [
            "Refers to a person who is studying or has expert knowledge of one or more of the natural or physical"
            " sciences."
        ]
    },
    "crossner_naturalscience_person": {"en": ["Refers to an individual's name that is not a scientist."]},
    "crossner_naturalscience_university": {
        "en": [
            "Refers to an educational institution of higher learning. Label organizations that are universities as"
            " 'university' entities."
        ]
    },
    "crossner_naturalscience_organization": {
        "en": [
            "Refers to a structured group, institution, company, or association. This category covers a"
            " diverse range of organizations, including businesses, non-profits, educational institutions, and"
            " government agencies."
        ]
    },
    "crossner_naturalscience_country": {"en": ["Refers to a sovereign nation."]},
    "crossner_naturalscience_location": {
        "en": [
            "Refers to a specific geographical or structural location. This includes but is not limited"
            " to: places (e.g., parks, landmarks), bridges, cities, towns, villages, areas and other distinct regions."
        ]
    },
    "crossner_naturalscience_discipline": {
        "en": [
            "Refers to a specialized field of study or research within the domains of natural science. Also label"
            " words such as biology, chemistry, astrophysics, mathematics or physics. This also includes specific"
            " branches such as quantum chemistry."
        ]
    },
    "crossner_naturalscience_enzyme": {
        "en": [
            " Refers to a specialized protein that acts as a biological catalyst, facilitating chemical reactions"
            " within organisms. If a protein is an enzyme, annotate it as an enzyme entity."
        ]
    },
    "crossner_naturalscience_protein": {
        "en": [
            "Refers to a complex biological molecule composed of amino acids, essential for various biological"
            " functions."
        ]
    },
    "crossner_naturalscience_chemicalelement": {
        "en": [
            "Refers to a fundamental substance characterized by its atomic number and unique properties, as found in"
            " the periodic table."
        ]
    },
    "crossner_naturalscience_chemicalcompound": {
        "en": [
            "Refers to a substance formed by combining two or more elements in fixed proportions, with distinct"
            " chemical properties."
        ]
    },
    "crossner_naturalscience_astronomicalobject": {
        "en": ["Refers to a celestial entity such as a star, planet, galaxy, or nebula, present in space."]
    },
    "crossner_naturalscience_academicjournal": {
        "en": [
            "Refers to a periodical publication containing scholarly articles and research findings related to natural"
            " science disciplines."
        ]
    },
    "crossner_naturalscience_event": {
        "en": ["Refers to a significant occurrence or happening. Such as festivals, sports events, conferences, etc."]
    },
    "crossner_naturalscience_award": {
        "en": [
            "Refers to a recognition or distinction presented to individuals or organizations for significant"
            " contributions within the natural science domain."
        ]
    },
    "crossner_naturalscience_theory": {
        "en": [
            "Refers to a scientifically established principle or explanation that describes natural phenomena within"
            " the scope of biology, chemistry, or astrophysics, such as orbital resonance. This encompasses laws and"
            " theories like the Ptolemaic planetary theories."
        ]
    },
    "crossner_naturalscience_miscellaneous": {
        "en": ["Refers to named entities that are not included in any other category."]
    },
    "crossner_music_musicgenre": {
        "en": [
            "Refers to a distinct category or style of music, characterized by its unique characteristics, themes, and"
            " instrumentation. Include regional musics typical of a certain place."
        ]
    },
    "crossner_music_song": {
        "en": [
            "Refers to a distinct category or style of music, characterized by its unique characteristics, themes, and"
            " instrumentation."
        ]
    },
    "crossner_music_band": {
        "en": [
            "Refers to a distinct category or style of music, characterized by its unique characteristics, themes, and"
            " instrumentation."
        ]
    },
    "crossner_music_album": {
        "en": ["Refers to a collection of songs or musical compositions released as a single package"]
    },
    "crossner_music_musicalartist": {
        "en": [
            "Refers to an individual actively involved in the creation, performance, or production of music. If a"
            " person is a singer, composer, or songwriter, annotate them as a musical artist entity instead of a"
            " person entity"
        ]
    },
    "crossner_music_musicalinstrument": {
        "en": ["Refers to a device or tool used to produce musical sounds, such as a piano, guitar, or drum."]
    },
    "crossner_music_award": {
        "en": ["Refers to a device or tool used to produce musical sounds, such as a piano, guitar, or drum."]
    },
    "crossner_music_event": {
        "en": ["Refers to a music-related occurrence or happening, such as a concert, festival, or competition."]
    },
    "crossner_music_country": {"en": ["Refers to a sovereign nation."]},
    "crossner_music_location": {
        "en": [
            "Refers to a specific geographical or structural location. This includes but is not limited"
            " to: places (e.g., parks, landmarks), bridges, cities, towns, villages, areas and other distinct regions."
        ]
    },
    "crossner_music_organization": {
        "en": [
            "Refers to a structured group, institution, company, or association. This category covers a"
            " diverse range of organizations, including businesses, non-profits, educational institutions, and"
            " government agencies."
        ]
    },
    "crossner_music_person": {"en": ["Refers to an individual's name that is not a musical artist."]},
    "crossner_music_miscellaneous": {"en": ["Refers to named entities that are not included in any other category."]},
    "crossner_literature_book": {
        "en": [
            "Refers to a published work of literature, typically consisting of written content bound in a physical or"
            " digital format."
        ]
    },
    "crossner_literature_writer": {
        "en": [
            " Refers to an individual actively engaged in the creation of literary works, including writers,"
            " novelists, scriptwriters, poets, and other literary artists. If a person is involved in"
            " literature-related roles, annotate them as a writer entity instead of a person entity."
        ]
    },
    "crossner_literature_award": {
        "en": [
            "Refers to a recognition or honor bestowed upon individuals or works of literature for notable"
            " achievements. "
        ]
    },
    "crossner_literature_poem": {
        "en": [
            "Refers to a literary composition characterized by its structured form, rhythm, and often heightened"
            " language. "
        ]
    },
    "crossner_literature_event": {
        "en": [
            "Refers to a significant occurrence or happening. Such as wars, conferences, festivals, sport events, etc."
        ]
    },
    "crossner_literature_magazine": {
        "en": ["Refers to a periodical publication that features articles, stories, and various literary works."]
    },
    "crossner_literature_literarygenre": {
        "en": [
            "Refers to categories in literature defined by unique artistic techniques, themes, content, and lengths."
            " Label words such as poetry, essay or novel as well as content categories like Science Fiction, Fantasy,"
            " and Fiction. It also encompasses culturally influenced genres, such as Nordic countries poetry."
        ]
    },
    "crossner_literature_person": {"en": ["Refers to an individual's name that is not a writer."]},
    "crossner_literature_location": {
        "en": [
            "Refers to a specific geographical or structural location. This includes but is not limited"
            " to: places (e.g., parks, landmarks), bridges, cities, towns, villages, areas and other distinct regions."
        ]
    },
    "crossner_literature_organization": {
        "en": [
            "Refers to a structured group, institution, company, or association. This category covers a"
            " diverse range of organizations, including businesses, non-profits, educational institutions, and"
            " government agencies."
        ]
    },
    "crossner_literature_country": {"en": ["Refers to a sovereign nation."]},
    "crossner_literature_miscellaneous": {
        "en": ["Refers to named entities that are not included in any other category."]
    },
    "crossner_ai_field": {
        "en": [
            "Refers to a specific research area or subfield within Artificial Intelligence. Also annotate acronyms"
            " such as NLP."
        ]
    },
    "crossner_ai_task": {
        "en": [
            "Refers to a particular research task or problem within a specific AI research field. Annotate the name of"
            " the specific task, such as machine translation or object detection."
        ]
    },
    "crossner_ai_product": {
        "en": [
            "Refers to a product, system, or toolkit related to Artificial Intelligence. This includes specific"
            " AI-enabled products (e.g., robots like Pepper), systems (e.g., facial recognition systems), and toolkits"
            " (e.g., Tensorflow and PyTorch)."
        ]
    },
    "crossner_ai_algorithm": {
        "en": [
            "Refers to an algorithmic procedure or computational model used in Artificial Intelligence research. This"
            " category includes algorithms (e.g., decision trees) and models (e.g., CNN and LSTM)."
        ]
    },
    "crossner_ai_researcher": {
        "en": [
            "Refers to an individual engaged in research activities within the field of Artificial Intelligence (AI),"
            " including professors, Ph.D. students, and researchers in academia, research institutions, and companies."
            " If a person is involved in AI research, label them as a researcher entity instead of a person entity."
        ]
    },
    "crossner_ai_metric": {
        "en": [
            "Refers to evaluation metrics used to assess the performance of AI models and algorithms. Annotate"
            " specific metrics like F1-score."
        ]
    },
    "crossner_ai_university": {
        "en": [
            "Refers to an educational institution of higher learning. Label organizations that are universities as"
            " 'university'' entities."
        ]
    },
    "crossner_ai_country": {"en": ["Refers to a sovereign nation."]},
    "crossner_ai_person": {"en": ["Refers to an individual's name that is not a researcher."]},
    "crossner_ai_organization": {
        "en": [
            "Refers to a structured group, institution, company, or association. This category covers a"
            " diverse range of organizations, including businesses, non-profits, educational institutions, and"
            " government agencies."
        ]
    },
    "crossner_ai_location": {
        "en": [
            "Refers to a specific geographical or structural location. This includes but is not limited"
            " to: places (e.g., parks, landmarks), bridges, cities, towns, villages, areas and other distinct regions."
        ]
    },
    "crossner_ai_programminglanguage": {
        "en": [
            "Refers to a programming language used in the development of AI applications and research. Annotate the"
            " name of the programming language, such as Java and Python."
        ]
    },
    "crossner_ai_conference": {
        "en": [
            "Refers to a research conference or journal in the field of Artificial Intelligence. This category"
            " includes conferences where AI research is presented and published."
        ]
    },
    "crossner_ai_miscellaneous": {"en": ["Refers to named entities that are not included in any other category."]},
}

src/tasks/crossner/guidelines.py
hitz-zentroa/GoLLIE
GUIDELINES = {
    "crossner_politics_person": {"en": ["Refers to an individual's name that is not a politician."]},
    "crossner_politics_organization": {
        "en": ["Refers to a structured group, institution, company, or association that is not a political party."]
    },
    "crossner_politics_location": {
        "en": [
            "Refers to a specific geographical or structural location. This includes but is not limited"
            " to: places (e.g., parks, landmarks), bridges, cities, towns, villages, areas and other distinct regions."
        ]
    },
    "crossner_politics_politician": {
        "en": [
            "Refers to a person who is actively engaged in politics, holding a public office, involved in political"
            " activities or part of a political party. "
            "If a person entity is identified as a politician, annotate that person's name as a politician entity."
        ]
    },
    "crossner_politics_politicalparty": {
        "en": [
            "A political party is an organization that coordinates candidates to compete in a particular country's"
            " elections. Also annotate acronyms of political parties."
        ]
    },
    "crossner_politics_election": {
        "en": [
            "Refers to an organized process in which individuals vote to choose a candidate or make a"
            " decision. If an event mentioned in the text pertains to an election, annotate it as an election entity."
        ]
    },
    "crossner_politics_event": {
        "en": [
            "Refers to a significant occurrence or happening. Such as wars, conferences, summits, conventions, etc."
        ]
    },
    "crossner_politics_country": {"en": ["Refers to a sovereign nation."]},
    "crossner_politics_miscellaneous": {
        "en": ["Refers to named entities that are not included in any other category."]
    },
    "crossner_naturalscience_scientist": {
        "en": [
            "Refers to a person who is studying or has expert knowledge of one or more of the natural or physical"
            " sciences."
        ]
    },
    "crossner_naturalscience_person": {"en": ["Refers to an individual's name that is not a scientist."]},
    "crossner_naturalscience_university": {
        "en": [
            "Refers to an educational institution of higher learning. Label organizations that are universities as"
            " 'university' entities."
        ]
    },
    "crossner_naturalscience_organization": {
        "en": [
            "Refers to a structured group, institution, company, or association. This category covers a"
            " diverse range of organizations, including businesses, non-profits, educational institutions, and"
            " government agencies."
        ]
    },
    "crossner_naturalscience_country": {"en": ["Refers to a sovereign nation."]},
    "crossner_naturalscience_location": {
        "en": [
            "Refers to a specific geographical or structural location. This includes but is not limited"
            " to: places (e.g., parks, landmarks), bridges, cities, towns, villages, areas and other distinct regions."
        ]
    },
    "crossner_naturalscience_discipline": {
        "en": [
            "Refers to a specialized field of study or research within the domains of natural science. Also label"
            " words such as biology, chemistry, astrophysics, mathematics or physics. This also includes specific"
            " branches such as quantum chemistry."
        ]
    },
    "crossner_naturalscience_enzyme": {
        "en": [
            " Refers to a specialized protein that acts as a biological catalyst, facilitating chemical reactions"
            " within organisms. If a protein is an enzyme, annotate it as an enzyme entity."
        ]
    },
    "crossner_naturalscience_protein": {
        "en": [
            "Refers to a complex biological molecule composed of amino acids, essential for various biological"
            " functions."
        ]
    },
    "crossner_naturalscience_chemicalelement": {
        "en": [
            "Refers to a fundamental substance characterized by its atomic number and unique properties, as found in"
            " the periodic table."
        ]
    },
    "crossner_naturalscience_chemicalcompound": {
        "en": [
            "Refers to a substance formed by combining two or more elements in fixed proportions, with distinct"
            " chemical properties."
        ]
    },
    "crossner_naturalscience_astronomicalobject": {
        "en": ["Refers to a celestial entity such as a star, planet, galaxy, or nebula, present in space."]
    },
    "crossner_naturalscience_academicjournal": {
        "en": [
            "Refers to a periodical publication containing scholarly articles and research findings related to natural"
            " science disciplines."
        ]
    },
    "crossner_naturalscience_event": {
        "en": ["Refers to a significant occurrence or happening. Such as festivals, sports events, conferences, etc."]
    },
    "crossner_naturalscience_award": {
        "en": [
            "Refers to a recognition or distinction presented to individuals or organizations for significant"
            " contributions within the natural science domain."
        ]
    },
    "crossner_naturalscience_theory": {
        "en": [
            "Refers to a scientifically established principle or explanation that describes natural phenomena within"
            " the scope of biology, chemistry, or astrophysics, such as orbital resonance. This encompasses laws and"
            " theories like the Ptolemaic planetary theories."
        ]
    },
    "crossner_naturalscience_miscellaneous": {
        "en": ["Refers to named entities that are not included in any other category."]
    },
    "crossner_music_musicgenre": {
        "en": [
            "Refers to a distinct category or style of music, characterized by its unique characteristics, themes, and"
            " instrumentation. Include regional musics typical of a certain place."
        ]
    },
    "crossner_music_song": {
        "en": [
            "Refers to a distinct category or style of music, characterized by its unique characteristics, themes, and"
            " instrumentation."
        ]
    },
    "crossner_music_band": {
        "en": [
            "Refers to a distinct category or style of music, characterized by its unique characteristics, themes, and"
            " instrumentation."
        ]
    },
    "crossner_music_album": {
        "en": ["Refers to a collection of songs or musical compositions released as a single package"]
    },
    "crossner_music_musicalartist": {
        "en": [
            "Refers to an individual actively involved in the creation, performance, or production of music. If a"
            " person is a singer, composer, or songwriter, annotate them as a musical artist entity instead of a"
            " person entity"
        ]
    },
    "crossner_music_musicalinstrument": {
        "en": ["Refers to a device or tool used to produce musical sounds, such as a piano, guitar, or drum."]
    },
    "crossner_music_award": {
        "en": ["Refers to a device or tool used to produce musical sounds, such as a piano, guitar, or drum."]
    },
    "crossner_music_event": {
        "en": ["Refers to a music-related occurrence or happening, such as a concert, festival, or competition."]
    },
    "crossner_music_country": {"en": ["Refers to a sovereign nation."]},
    "crossner_music_location": {
        "en": [
            "Refers to a specific geographical or structural location. This includes but is not limited"
            " to: places (e.g., parks, landmarks), bridges, cities, towns, villages, areas and other distinct regions."
        ]
    },
    "crossner_music_organization": {
        "en": [
            "Refers to a structured group, institution, company, or association. This category covers a"
            " diverse range of organizations, including businesses, non-profits, educational institutions, and"
            " government agencies."
        ]
    },
    "crossner_music_person": {"en": ["Refers to an individual's name that is not a musical artist."]},
    "crossner_music_miscellaneous": {"en": ["Refers to named entities that are not included in any other category."]},
    "crossner_literature_book": {
        "en": [
            "Refers to a published work of literature, typically consisting of written content bound in a physical or"
            " digital format."
        ]
    },
    "crossner_literature_writer": {
        "en": [
            " Refers to an individual actively engaged in the creation of literary works, including writers,"
            " novelists, scriptwriters, poets, and other literary artists. If a person is involved in"
            " literature-related roles, annotate them as a writer entity instead of a person entity."
        ]
    },
    "crossner_literature_award": {
        "en": [
            "Refers to a recognition or honor bestowed upon individuals or works of literature for notable"
            " achievements. "
        ]
    },
    "crossner_literature_poem": {
        "en": [
            "Refers to a literary composition characterized by its structured form, rhythm, and often heightened"
            " language. "
        ]
    },
    "crossner_literature_event": {
        "en": [
            "Refers to a significant occurrence or happening. Such as wars, conferences, festivals, sport events, etc."
        ]
    },
    "crossner_literature_magazine": {
        "en": ["Refers to a periodical publication that features articles, stories, and various literary works."]
    },
    "crossner_literature_literarygenre": {
        "en": [
            "Refers to categories in literature defined by unique artistic techniques, themes, content, and lengths."
            " Label words such as poetry, essay or novel as well as content categories like Science Fiction, Fantasy,"
            " and Fiction. It also encompasses culturally influenced genres, such as Nordic countries poetry."
        ]
    },
    "crossner_literature_person": {"en": ["Refers to an individual's name that is not a writer."]},
    "crossner_literature_location": {
        "en": [
            "Refers to a specific geographical or structural location. This includes but is not limited"
            " to: places (e.g., parks, landmarks), bridges, cities, towns, villages, areas and other distinct regions."
        ]
    },
    "crossner_literature_organization": {
        "en": [
            "Refers to a structured group, institution, company, or association. This category covers a"
            " diverse range of organizations, including businesses, non-profits, educational institutions, and"
            " government agencies."
        ]
    },
    "crossner_literature_country": {"en": ["Refers to a sovereign nation."]},
    "crossner_literature_miscellaneous": {
        "en": ["Refers to named entities that are not included in any other category."]
    },
    "crossner_ai_field": {
        "en": [
            "Refers to a specific research area or subfield within Artificial Intelligence. Also annotate acronyms"
            " such as NLP."
        ]
    },
    "crossner_ai_task": {
        "en": [
            "Refers to a particular research task or problem within a specific AI research field. Annotate the name of"
            " the specific task, such as machine translation or object detection."
        ]
    },
    "crossner_ai_product": {
        "en": [
            "Refers to a product, system, or toolkit related to Artificial Intelligence. This includes specific"
            " AI-enabled products (e.g., robots like Pepper), systems (e.g., facial recognition systems), and toolkits"
            " (e.g., Tensorflow and PyTorch)."
        ]
    },
    "crossner_ai_algorithm": {
        "en": [
            "Refers to an algorithmic procedure or computational model used in Artificial Intelligence research. This"
            " category includes algorithms (e.g., decision trees) and models (e.g., CNN and LSTM)."
        ]
    },
    "crossner_ai_researcher": {
        "en": [
            "Refers to an individual engaged in research activities within the field of Artificial Intelligence (AI),"
            " including professors, Ph.D. students, and researchers in academia, research institutions, and companies."
            " If a person is involved in AI research, label them as a researcher entity instead of a person entity."
        ]
    },
    "crossner_ai_metric": {
        "en": [
            "Refers to evaluation metrics used to assess the performance of AI models and algorithms. Annotate"
            " specific metrics like F1-score."
        ]
    },
    "crossner_ai_university": {
        "en": [
            "Refers to an educational institution of higher learning. Label organizations that are universities as"
            " 'university'' entities."
        ]
    },
    "crossner_ai_country": {"en": ["Refers to a sovereign nation."]},
    "crossner_ai_person": {"en": ["Refers to an individual's name that is not a researcher."]},
    "crossner_ai_organization": {
        "en": [
            "Refers to a structured group, institution, company, or association. This category covers a"
            " diverse range of organizations, including businesses, non-profits, educational institutions, and"
            " government agencies."
        ]
    },
    "crossner_ai_location": {
        "en": [
            "Refers to a specific geographical or structural location. This includes but is not limited"
            " to: places (e.g., parks, landmarks), bridges, cities, towns, villages, areas and other distinct regions."
        ]
    },
    "crossner_ai_programminglanguage": {
        "en": [
            "Refers to a programming language used in the development of AI applications and research. Annotate the"
            " name of the programming language, such as Java and Python."
        ]
    },
    "crossner_ai_conference": {
        "en": [
            "Refers to a research conference or journal in the field of Artificial Intelligence. This category"
            " includes conferences where AI research is presented and published."
        ]
    },
    "crossner_ai_miscellaneous": {"en": ["Refers to named entities that are not included in any other category."]},
}

EXAMPLES = {
    "crossner_politics_person_examples": {
        "en": [
            "Bennelong",
            "Zappa",
            "Tower",
            "Andrea Fischer",
            "Wyden",
            "Johann Karl Thilo",
            "Rowse",
            "Brooks",
            "Anawrahta",
            "Fidel",
        ]
    },
    "crossner_politics_organization_examples": {
        "en": [
            "NARAL Pro-Choice America",
            "Planned Parenthood",
            "National Right to Life Committee",
            "United Daughters of the Confederacy",
            "Sons of Confederate Veterans",
            "Parliament",
            "Amnesty International",
            "Human Rights Watch",
            "American Enterprise Institute",
            "American Academy of Arts and Sciences",
        ]
    },
    "crossner_politics_location_examples": {
        "en": ["Paris", "London", "Berlin", "Rome", "Dublin", "Cairo", "Amsterdam", "Colombo", "Montreal", "Texas"]
    },
    "crossner_politics_politician_examples": {
        "en": [
            "Bush",
            "Ronald Reagan",
            "Al Gore",
            "Joseph Stalin",
            "Richard Nixon",
            "Thomas Jefferson",
            "Johnson",
            "Genscher",
            "Kennedy",
            "Franklin D. Roosevelt",
        ]
    },
    "crossner_politics_politicalparty_examples": {
        "en": [
            "SPD",
            "National Party of Australia",
            "Sinn F\u00e9in",
            "FDP",
            "Liberal Party of Australia",
            "Liberal",
            "Progressive Conservative Party of Canada",
            "Serbian Radical Party",
            "National Alliance",
            "Forza Italia",
        ]
    },
    "crossner_politics_election_examples": {
        "en": [
            "1997 Irish general election",
            "1992 United States presidential election",
            "1920 United States presidential election",
            "2000 United States presidential election",
            "2008 United States presidential election",
            "1983 United Kingdom general election",
            "1987 United Kingdom general election",
            "1934 Australian federal election",
            "2019 Australian federal election",
            "2007 Australian federal election",
        ]
    },
    "crossner_politics_event_examples": {
        "en": [
            "Decembrist revolt",
            "1952 Republican National Convention",
            "1943 Tehran Conference",
            "World War II",
            "1983 Commonwealth heads of Government summit",
            "TM movement",
            "Tydings-Butler race",
            "War of Liberation",
            "War of the Sixth Coalition",
            "Civil War",
        ]
    },
    "crossner_politics_country_examples": {
        "en": [
            "U.S.",
            "India",
            "China",
            "Germany",
            "United States",
            "Soviet Union",
            "Egypt",
            "Ghana",
            "Australia",
            "Serbia",
        ]
    },
    "crossner_politics_miscellaneous_examples": {
        "en": [
            "Republican",
            "German",
            "Democrat",
            "Soviet",
            "Republicans",
            "Polish",
            "British",
            "Texas Republicans",
            "Australian",
            "Uyghur",
        ]
    },
    "crossner_music_musicgenre_examples": {
        "en": [
            "Hip hop music",
            "Pop music",
            "reggae",
            "Punk rock",
            "Celtic rock",
            "Soul music",
            "Samba",
            "rock",
            "Electronic music",
            "Sacred Harp",
        ]
    },
    "crossner_music_song_examples": {
        "en": [
            "Is That Love",
            "Tempted",
            "Clubland",
            "New Lace Sleeves",
            "The Lunatics ( Have Taken Over the Asylum )",
            "East Coast Killer",
            "West Coast Killer",
            "Layla",
            "Cross Road Blues",
            "Spoonman",
        ]
    },
    "crossner_music_band_examples": {
        "en": [
            "Backstreet Boys",
            "Spice Girls",
            "NSYNC",
            "LFO",
            "O-Town",
            "US5",
            "Black Sabbath",
            "Take That",
            "U2",
            "PUSA",
        ]
    },
    "crossner_music_album_examples": {
        "en": [
            "American Idiot",
            "Dookie",
            "Sun Ra Visits Planet Earth",
            "Interstellar Low Ways",
            "Super-Sonic Jazz",
            "We Travel the Space Ways",
            "The Nubians of Plutonia",
            "Jazz In Silhouette",
            "Tea for the Tillerman",
            "Teaser and the Firecat",
        ]
    },
    "crossner_music_musicalartist_examples": {
        "en": [
            "Elvis Costello",
            "Nas",
            "Madonna",
            "Debbie Gibson",
            "Tiffany",
            "Sun Ra",
            "Stevens",
            "Verni",
            "Skates",
            "Ellsworth",
        ]
    },
    "crossner_music_musicalinstrument_examples": {
        "en": [
            "piano accordion",
            "six-string guitar banjo",
            "guitar",
            "keyboard",
            "bass guitar",
            "drums",
            "bass",
        ]
    },
    "crossner_music_award_examples": {
        "en": [
            "Grammy Award",
            "Brit Awards",
            "MTV Video Music Award",
            "Academy Awards",
            "Grammy Award for Best Alternative Music Album",
            "Academy Award for Best Director",
            "Academy Award for Best Picture",
            "Grammy Award s",
            "Top New Male Vocalist",
            "American Music Awards",
        ]
    },
    "crossner_music_event_examples": {
        "en": [
            "2003 World Championships in Athletics",
            "Vocal Event of the Year",
            "World War II",
            "Saint Patrick 's Day",
            "the first wave of glam metal",
            "panto season",
            "G\u0175yl Gobaith Music Festival",
            "UAAP Season 71",
            "1973 Summer Universiade",
            "1986 Goodwill Games",
        ]
    },
    "crossner_music_country_examples": {
        "en": [
            "U.S.",
            "United States",
            "United Kingdom",
            "Australia",
            "Germany",
            "Poland",
            "Canada",
            "New Zealand",
            "Japan",
            "Ireland",
        ]
    },
    "crossner_music_location_examples": {
        "en": [
            "Europe",
            "Stade de France",
            "Paris",
            "Th\u00e9\u00e2tre de la Ville",
            "Bay Area",
            "Grand Ole Opry",
            "AccorHotels Arena",
            "Bataclan",
            "Casino de Paris",
            "Maple Leaf Gardens",
        ]
    },
    "crossner_music_organization_examples": {
        "en": [
            "Recording Industry Association of America",
            "Billboard",
            "BBC News",
            "Academy of Country Music",
            "Country Music Association",
            "Op\u00e9ra-Comique",
            "United Way of Canada",
            "Official Charts Company",
            "British Phonographic Industry",
            "ARIA",
        ]
    },
    "crossner_music_person_examples": {
        "en": [
            "Barney Bubbles",
            "St\u00e9phane Lafarge",
            "Bubbles",
            "Wolf-R\u00fcdiger M\u00fchlmann",
            "Glen Campbell",
            "Bobbie Gentry",
            "John Denver",
            "Olivia Newton-John",
            "Lou Pearlman",
            "Nichols",
        ]
    },
    "crossner_music_miscellaneous_examples": {
        "en": [
            "Irish-American",
            "European",
            "Les Troyens \u00e0 Carthage",
            "Didon",
            "\u00c9n\u00e9e",
            "Billboard",
            "British duo",
            "King of the Bai\u00e3o",
            "funk",
            "Billboard 200",
        ]
    },
    "crossner_literature_book_examples": {
        "en": [
            "The Plague of Doves",
            "A History of Warfare",
            "The Color Purple",
            "Rabbit at Rest",
            "The Fountainhead",
            "Atlas Shrugged",
            "The Little Prince",
            "Anabasis Alexandri",
            "Pro Archia Poeta",
            "Pantoum of the Great Depression",
        ]
    },
    "crossner_literature_writer_examples": {
        "en": [
            "John Keegan",
            "Leo Tolstoy",
            "John Steinbeck",
            "Yeats",
            "Jacobs",
            "Isaac Bashevis Singer",
            "Capp",
            "Kong Huan",
            "\u5b54\u6d63",
            "Kong Shao",
        ]
    },
    "crossner_literature_award_examples": {
        "en": [
            "Pulitzer Prize",
            "Nebula Award",
            "Hugo Award",
            "Academy Awards",
            "Nobel Prize in Literature",
            "National Book Critics Circle Award",
            "Grand Prix Sp\u00e9cial du Jury",
            "FIPRESCI prize",
            "Palme d 'Or",
            "Silver Bear for Best Actor",
        ]
    },
    "crossner_literature_poem_examples": {
        "en": [
            "Barbara Frietchie",
            "The Barefoot Boy",
            "Maud Muller",
            "Snow-Bound",
            "Shishupala Vadha",
            "Naishadha Charita",
            "Nai\u1e63adh\u012bya-carita",
            "Prick of Conscience",
            "At a Calvary near the Ancre",
            "There Will Come Soft Rains",
        ]
    },
    "crossner_literature_event_examples": {
        "en": [
            "World War II",
            "Cannes Film Festival",
            "47th Berlin International Film Festival",
            "Luge at the 1976 Winter Olympics",
            "Equestrian at the 1976 Summer Olympics",
            "Trojan War",
            "Black Arts Movement",
            "Ether 7 Festival",
            "Cold War",
        ]
    },
    "crossner_literature_magazine_examples": {
        "en": [
            "The New Yorker",
            "Simbolul",
            "Private Eye",
            "The Atlantic",
            "The Crisis",
            "The New York Times Magazine",
            "T Living",
            "Neue Rundschau",
            "The Magazine of Fantasy & Science Fiction",
            "New Statesman",
        ]
    },
    "crossner_literature_literarygenre_examples": {
        "en": [
            "poem",
            "poems",
            "novels",
            "novel",
            "science fiction",
            "anti-slavery writings",
            "Symbolism",
            "experimental poetry",
            "science fiction short stories",
            "Negro primitivism",
        ]
    },
    "crossner_literature_person_examples": {
        "en": [
            "Lindbergh",
            "Giacomo di Benincasa",
            "Baron Cohen",
            "Emperor Toghon Tem\u00fcr",
            "Princess Noguk",
            "Gongmin",
            "Marcel Janco",
            "Jim Broadbent",
            "Major",
            "Norman",
        ]
    },
    "crossner_literature_location_examples": {
        "en": [
            "London",
            "Chalcedon",
            "Charles Dickens Museum",
            "Charles Dickens Birthplace Museum",
            "Portsmouth",
            "Lincoln Center",
            "Beethovensaal",
            "Berlin",
            "Bristol",
            "The Thekla",
        ]
    },
    "crossner_literature_organization_examples": {
        "en": [
            "Science Fiction and Fantasy Writers of America",
            "The New York Times",
            "New York Times",
            "Food52",
            "Amnesty International",
            "S. Fischer Verlag",
            "Project Gutenberg",
            "Hollywood",
            "Royal Society of Literature",
            "National Space Society",
        ]
    },
    "crossner_literature_country_examples": {
        "en": [
            "Germany",
            "Han dynasty",
            "England",
            "Italy",
            "Yuan dynasty",
            "China",
            "Korea",
            "Mauretania",
            "United States",
            "Denmark",
        ]
    },
    "crossner_literature_miscellaneous_examples": {
        "en": [
            "American",
            "British",
            "English",
            "French",
            "Goryeo",
            "Mongolian-born",
            "Rabbit novel",
            "Objectivism",
            "Nazi",
            "Middle English",
        ]
    },
    "crossner_ai_field_examples": {
        "en": [
            "deep learning",
            "pattern recognition",
            "image processing",
            "reinforcement learning",
            "natural language processing",
            "machine learning",
            "unsupervised learning",
            "AI",
            "computer vision",
            "text mining",
        ]
    },
    "crossner_ai_task_examples": {
        "en": [
            "speech synthesis",
            "information retrieval",
            "Feature extraction",
            "dimension reduction",
            "speech recognition",
            "sentiment analysis",
            "Multimodal sentiment analysis",
            "face recognition",
            "handwriting recognition",
            "lip reading",
        ]
    },
    "crossner_ai_product_examples": {
        "en": [
            "MATLAB",
            "Programmable Universal Machine for Assembly",
            "industrial robot",
            "opinion-based recommender system",
            "Octave",
            "Google Translate",
            "SYSTRAN system",
            "BabelFish",
            "Babelfish",
            "RapidMiner",
        ]
    },
    "crossner_ai_algorithm_examples": {
        "en": [
            "principal component analysis",
            "linear discriminant analysis",
            "gradient descent",
            "Support vector machine",
            "recurrent neural network",
            "LSTM",
            "PCA",
            "LDA",
            "canonical correlation analysis",
            "CCA",
        ]
    },
    "crossner_ai_researcher_examples": {
        "en": [
            "J\u00fcrgen Schmidhuber",
            "Seymour Papert",
            "Victor Scheinman",
            "Scheinman",
            "X.Y. Feng",
            "H. Zhang",
            "Y.J. Ren",
            "P.H. Shang",
            "Y. Zhu",
            "Y.C. Liang",
            "Oscar Sainz",
            "Iker García-Ferrero",
        ]
    },
    "crossner_ai_metric_examples": {
        "en": [
            "mean squared error",
            "DCG",
            "maximum likelihood",
            "Recall-Oriented Understudy for Gisting Evaluation",
            "MSE",
            "noise floor measurement",
            "ROUGE",
            "Hinge loss",
            "hinge loss",
            "Sigmoid function Cross entropy loss",
        ]
    },
    "crossner_ai_university_examples": {
        "en": [
            "University of Toronto",
            "Cambridge",
            "University of Groningen",
            "MIT",
            "Brown University",
            "Carnegie Mellon University",
            "MPI Saarbruecken",
            "Stanford University",
            "University of California , San Diego",
            "\u00c9cole Centrale Paris",
            "University of the Basque Country",
        ]
    },
    "crossner_ai_country_examples": {
        "en": ["Netherlands", "Japan", "Germany", "Canada", "Australia", "Brazil", "China", "India", "Italy", "Korea"]
    },
    "crossner_ai_person_examples": {
        "en": ["Francis Ford Coppola", "Michael Jackson", "John Wayne", "Rita Hayworth", "Dean Martin", "Jerry Lewis"]
    },
    "crossner_ai_organization_examples": {
        "en": [
            "Unimation",
            "IAPR",
            "Audio Engineering Society",
            "National Science Foundation",
            "National Aeronautics and Space Administration",
            "NASA",
            "US Department of Energy",
            "US Department of Commerce NIST",
            "US Department of Defense",
            "Defense Advanced Research Projects Agency",
            "HiTZ: Basque Center for Language Technology",
            "IXA NLP Group",
        ]
    },
    "crossner_ai_location_examples": {
        "en": [
            "Chi\u0219in\u0103u",
            "Paris",
            "Montreal",
            "Scotiabank Theatre Toronto",
            "TIFF Bell Lightbox",
            "Moldavian SSR",
        ]
    },
    "crossner_ai_programminglanguage_examples": {
        "en": ["Java", "R", "CLIPS", "Python", "C + +", "GNU Octave", "Java 9", "java", "Perl", "ActiveX"]
    },
    "crossner_ai_conference_examples": {
        "en": [
            "AAAI",
            "1982 Association for the Advancement of Artificial Intelligence",
            "SIGGRAPH",
            "Symposium on Geometry Processing",
            "International Journal of Computer Vision",
            "IJCV",
            "IEEE Computer Society Conference on Computer Vision and Pattern Recognition",
            "CVPR",
            "International Conference on Machine Learning 2011 & 2012",
            "NIST ' s annual Document Understanding Conferences",
        ]
    },
    "crossner_ai_miscellaneous_examples": {
        "en": [
            "unsupervised methods",
            "topological properties",
            "audio signal",
            "eigenfaces",
            "intelligent agents",
            "graphical user interfaces",
            "Heuretics : Theoretical and Study of Heuristic Rules",
            "Best Paper award",
            "Johann Bernoulli Chair",
            "Toshiba Endowed Chair",
        ]
    },
    "crossner_naturalscience_scientist_examples": {
        "en": [
            "Werner Heisenberg",
            "Robert Boyle",
            "Hans Kramers",
            "Isaac Newton",
            "Galileo Galilei",
            "Carl Friedrich von Weizs\u00e4cker",
            "Walther Bothe",
            "August Kopff",
            "Wolf",
            "Edgar H. Booth",
        ]
    },
    "crossner_naturalscience_person_examples": {
        "en": [
            "Shenton",
            "Russell T Davies",
            "Aristotle",
            "Hall",
            "NeNe Leakes",
            "Nicki Minaj",
            "Rihanna",
            "Solange Knowles",
            "Beyonc\u00e9",
            "Sharkeisha",
        ]
    },
    "crossner_naturalscience_university_examples": {
        "en": [
            "University of Sydney",
            "Oxford",
            "Uppsala University",
            "University of Calcutta",
            "Odessa University",
            "Cambridge",
            "University of T\u00fcbingen",
            "University of G\u00f6ttingen",
            "Oak Ridge Associated Universities",
            "Heidelberg University",
        ]
    },
    "crossner_naturalscience_organization_examples": {
        "en": [
            "Royal Society",
            "American Academy of Arts and Sciences",
            "National Academy of Sciences",
            "International Union of Geodesy and Geophysics",
            "American Philosophical Society",
            "Cavendish Laboratory",
            "Royal Swedish Academy of Sciences",
            "American Association for the Advancement of Science",
            "University of Maryland Medical Center",
            "R Adams Cowley Shock Trauma Center",
        ]
    },
    "crossner_naturalscience_country_examples": {
        "en": [
            "Newfoundland Colony",
            "France",
            "United States",
            "Russia",
            "Japan",
            "UK",
            "England",
            "Italy",
            "Britain",
        ]
    },
    "crossner_naturalscience_location_examples": {
        "en": [
            "Cape of Good Hope",
            "Mumbai",
            "Heidelberg",
            "Royal Prince Alfred Hospital",
            "Sydney",
            "Annapolis",
            "MD",
            "Little Bay Copper Mine",
            "Thimble Tickle Bay",
            "Notre Dame Bay",
        ]
    },
    "crossner_naturalscience_discipline_examples": {
        "en": [
            "philosophy",
            "chemistry",
            "medicine",
            "electrodynamics",
            "spectroscopy",
            "geodesy",
            "geophysics",
            "Arts",
            "Literature of Belgium",
            "Electrical Engineering",
        ]
    },
    "crossner_naturalscience_enzyme_examples": {
        "en": [
            "RNA polymerase",
            "DNA methyltransferase",
            "DNA polymerase",
            "Spermidine synthase",
            "Alkaline phosphatase",
            "T7 RNA polymerase",
            "ribulose-1,5-bisphosphate carboxylase oxygenase",
            "RuBisCO",
            "ATP synthase",
            "Histone deacetylase",
        ]
    },
    "crossner_naturalscience_protein_examples": {
        "en": [
            "Argonaute",
            "Histone H3",
            "CYP1A1",
            "Green fluorescent protein",
            "glutamate binding protein",
            "GluBP",
            "Gamma-glutamyltransferase",
            "RNase inhibitor",
            "Histone",
            "nuclear proteins",
        ]
    },
    "crossner_naturalscience_chemicalelement_examples": {
        "en": [
            "sulfur",
            "copper",
            "hydrogen",
            "helium",
            "Nihonium",
            "americium",
            "europium",
            "curium",
            "gadolinium",
            "lead",
        ]
    },
    "crossner_naturalscience_chemicalcompound_examples": {
        "en": [
            "carbon dioxide",
            "Adenosine triphosphate",
            "Nitric oxide",
            "ROS",
            "cytosine",
            "Vespel",
            "neoprene",
            "nylon",
            "Corian",
            "Polytetrafluoroethylene",
        ]
    },
    "crossner_naturalscience_astronomicalobject_examples": {
        "en": ["Saturn", "Jupiter", "Venus", "Earth", "Uranus", "Mercury", "Neptune", "Mars", "Sun", "5145 Pholus"]
    },
    "crossner_naturalscience_academicjournal_examples": {
        "en": [
            "Angewandte Chemie",
            "The Astrophysical Journal",
            "Nature",
            "Chemical Reviews",
            "Accounts of Chemical Research",
            "The Astronomical Journal",
            "Cell",
            "Proceedings of the National Academy of Sciences of the United States of America",
            "Journal of Molecular Biology",
            "Annales de chimie et de physique",
        ]
    },
    "crossner_naturalscience_event_examples": {
        "en": [
            "2017 World Championships in Athletics",
            "Tomorrowland",
            "A State Of Trance",
            "Ultra Music Festival",
            "Electric Daisy Carnival",
            "Electric Zoo",
            "Gatecrasher",
            "Beyond Wonderland",
            "The men 's + 100 kg judo event",
            "2015 European Games",
        ]
    },
    "crossner_naturalscience_theory_examples": {
        "en": [
            "Protein-DNA interaction",
            "circular dichroism",
            "stem cell niche theory",
            "Nicolaus Copernicus heliocentric theory",
            "metaphysics",
            "methodology",
            "theory of evolution",
            "higher-order consciousness",
        ]
    },
    "crossner_naturalscience_award_examples": {
        "en": [
            "Nobel Prize in Physics",
            "Nobel Prize in Chemistry",
            "Fellow of the American Physical Society",
            "fellow of the Royal Society",
            "Copley Medal of the Royal Society",
            "Copley Medal",
            "Baronet MD Royal Society",
            "fellow of the Linnean Society of London",
            "Palit Professor of Physics",
            "Max Planck Medal",
        ]
    },
    "crossner_naturalscience_miscellaneous_examples": {
        "en": [
            "DNA",
            "Greeks",
            "Aryl hydrocarbon receptor repressor",
            "CRISPR",
            "Japanese",
            "American",
            "Arabidopsis thaliana",
            "DNA mismatch repair",
            "keratins",
            "Athletics",
        ]
    },
}

src/tasks/crossner/guidelines_gold.py
hitz-zentroa/GoLLIE
from typing import List

from ..utils_typing import Entity, dataclass


"""Entity definitions

The entity definitions are derived from the official CrossNER corpus guidelines:
https://arxiv.org/pdf/2012.04373.pdf

"""


@dataclass
class Field(Entity):
    """{crossner_ai_field}"""

    span: str  # {crossner_ai_field_examples}


@dataclass
class Task(Entity):
    """{crossner_ai_task}"""

    span: str  # {crossner_ai_task_examples}


@dataclass
class Product(Entity):
    """{crossner_ai_product}"""

    span: str  # {crossner_ai_product_examples}


@dataclass
class Algorithm(Entity):
    """{crossner_ai_algorithm}"""

    span: str  # {crossner_ai_algorithm_examples}


@dataclass
class Researcher(Entity):
    """{crossner_ai_researcher}"""

    span: str  # {crossner_ai_researcher_examples}


@dataclass
class Metric(Entity):
    """{crossner_ai_metric}"""

    span: str  # {crossner_ai_metric_examples}


@dataclass
class University(Entity):
    """{crossner_ai_university}"""

    span: str  # {crossner_ai_university_examples}


@dataclass
class Country(Entity):
    """{crossner_ai_country}"""

    span: str  # {crossner_ai_country_examples}


@dataclass
class Person(Entity):
    """{crossner_ai_person}"""

    span: str  # {crossner_ai_person_examples}


@dataclass
class Organization(Entity):
    """{crossner_ai_organization}"""

    span: str  # {crossner_ai_organization_examples}


@dataclass
class Location(Entity):
    """{crossner_ai_location}"""

    span: str  # {crossner_ai_location_examples}


@dataclass
class ProgrammingLanguage(Entity):
    """{crossner_ai_programminglanguage}"""

    span: str  # {crossner_ai_programminglanguage_examples}


@dataclass
class Conference(Entity):
    """{crossner_ai_conference}"""

    span: str  # {crossner_ai_conference_examples}


@dataclass
class Other(Entity):
    """{crossner_ai_miscellaneous}"""

    span: str  # {crossner_ai_miscellaneous_examples}


ENTITY_DEFINITIONS_AI: List[Entity] = [
    Field,
    Task,
    Product,
    Algorithm,
    Researcher,
    Metric,
    University,
    Country,
    Person,
    Organization,
    Location,
    ProgrammingLanguage,
    Conference,
    Other,
]

ENTITY_DEFINITIONS_AI_woMISC: List[Entity] = [
    Field,
    Task,
    Product,
    Algorithm,
    Researcher,
    Metric,
    University,
    Country,
    Person,
    Organization,
    Location,
    ProgrammingLanguage,
    Conference,
]


# __all__ = list(map(str, [*ENTITY_DEFINITIONS]))

src/tasks/crossner/prompts_ai.py
hitz-zentroa/GoLLIE
from typing import List

from ..utils_typing import Entity, dataclass


"""Entity definitions

The entity definitions are derived from the official CrossNER corpus guidelines:
https://arxiv.org/pdf/2012.04373.pdf

"""


@dataclass
class Book(Entity):
    """{crossner_literature_book}"""

    span: str  # {crossner_literature_book_examples}


@dataclass
class Writer(Entity):
    """{crossner_literature_writer}"""

    span: str  # {crossner_literature_writer_examples}


@dataclass
class Award(Entity):
    """{crossner_literature_award}"""

    span: str  # {crossner_literature_award_examples}


@dataclass
class Poem(Entity):
    """{crossner_literature_poem}"""

    span: str  # {crossner_literature_poem_examples}


@dataclass
class Event(Entity):
    """{crossner_literature_event}"""

    span: str  # {crossner_literature_event_examples}


@dataclass
class Magazine(Entity):
    """{crossner_literature_magazine}"""

    span: str  # {crossner_literature_magazine_examples}


@dataclass
class LiteraryGenre(Entity):
    """{crossner_literature_literarygenre}"""

    span: str  # {crossner_literature_literarygenre_examples}


@dataclass
class Person(Entity):
    """{crossner_literature_person}"""

    span: str  # {crossner_literature_person_examples}


@dataclass
class Location(Entity):
    """{crossner_literature_location}"""

    span: str  # {crossner_literature_location_examples}


@dataclass
class Organization(Entity):
    """{crossner_literature_organization}"""

    span: str  # {crossner_literature_organization_examples}


@dataclass
class Country(Entity):
    """{crossner_literature_country}"""

    span: str  # {crossner_literature_country_examples}


@dataclass
class Other(Entity):
    """{crossner_literature_miscellaneous}"""

    span: str  # {crossner_literature_miscellaneous_examples}


ENTITY_DEFINITIONS_LITERATURE: List[Entity] = [
    Book,
    Writer,
    Award,
    Poem,
    Event,
    Magazine,
    LiteraryGenre,
    Person,
    Location,
    Organization,
    Country,
    Other,
]

ENTITY_DEFINITIONS_LITERATURE_woMISC: List[Entity] = [
    Book,
    Writer,
    Award,
    Poem,
    Event,
    Magazine,
    LiteraryGenre,
    Person,
    Location,
    Organization,
    Country,
]

# __all__ = list(map(str, [*ENTITY_DEFINITIONS]))

src/tasks/crossner/prompts_literature.py
hitz-zentroa/GoLLIE
from typing import List

from ..utils_typing import Entity, dataclass


"""Entity definitions

The entity definitions are derived from the official CrossNER corpus guidelines:
https://arxiv.org/pdf/2012.04373.pdf

"""


@dataclass
class MusicGenre(Entity):
    """{crossner_music_musicgenre}"""

    span: str  # {crossner_music_musicgenre_examples}


@dataclass
class Song(Entity):
    """{crossner_music_song}"""

    span: str  # {crossner_music_song_examples}


@dataclass
class Band(Entity):
    """{crossner_music_band}"""

    span: str  # {crossner_music_band_examples}


@dataclass
class Album(Entity):
    """{crossner_music_album}"""

    span: str  # {crossner_music_album_examples}


@dataclass
class MusicalArtist(Entity):
    """{crossner_music_musicalartist}"""

    span: str  # {crossner_music_musicalartist_examples}


@dataclass
class MusicalInstrument(Entity):
    """{crossner_music_musicalinstrument}"""

    span: str  # {crossner_music_musicalinstrument_examples}


@dataclass
class Award(Entity):
    """{crossner_music_award}"""

    span: str  # {crossner_music_award_examples}


@dataclass
class Event(Entity):
    """{crossner_music_event}"""

    span: str  # {crossner_music_event_examples}


@dataclass
class Country(Entity):
    """{crossner_music_country}"""

    span: str  # {crossner_music_country_examples}


@dataclass
class Location(Entity):
    """{crossner_music_location}"""

    span: str  # {crossner_music_location_examples}


@dataclass
class Organization(Entity):
    """{crossner_music_organization}"""

    span: str  # {crossner_music_organization_examples}


@dataclass
class Person(Entity):
    """{crossner_music_person}"""

    span: str  # {crossner_music_person_examples}


@dataclass
class Other(Entity):
    """{crossner_music_miscellaneous}"""

    span: str  # {crossner_music_miscellaneous_examples}


ENTITY_DEFINITIONS_MUSIC: List[Entity] = [
    MusicGenre,
    Song,
    Band,
    Album,
    MusicalArtist,
    MusicalInstrument,
    Award,
    Event,
    Country,
    Location,
    Organization,
    Person,
    Other,
]

ENTITY_DEFINITIONS_MUSIC_woMISC: List[Entity] = [
    MusicGenre,
    Song,
    Band,
    Album,
    MusicalArtist,
    MusicalInstrument,
    Award,
    Event,
    Country,
    Location,
    Organization,
    Person,
]

# __all__ = list(map(str, [*ENTITY_DEFINITIONS]))

src/tasks/crossner/prompts_music.py
hitz-zentroa/GoLLIE
from typing import List

from ..utils_typing import Entity, dataclass


"""Entity definitions

The entity definitions are derived from the official CrossNER corpus guidelines:
https://arxiv.org/pdf/2012.04373.pdf

"""


@dataclass
class Scientist(Entity):
    """{crossner_naturalscience_scientist}"""

    span: str  # {crossner_naturalscience_scientist_examples}


@dataclass
class Person(Entity):
    """{crossner_naturalscience_person}"""

    span: str  # {crossner_naturalscience_person_examples}


@dataclass
class University(Entity):
    """{crossner_naturalscience_university}"""

    span: str  # {crossner_naturalscience_university_examples}


@dataclass
class Organization(Entity):
    """{crossner_naturalscience_organization}"""

    span: str  # {crossner_naturalscience_organization_examples}


@dataclass
class Country(Entity):
    """{crossner_naturalscience_country}"""

    span: str  # {crossner_naturalscience_country_examples}


@dataclass
class Location(Entity):
    """{crossner_naturalscience_location}"""

    span: str  # {crossner_naturalscience_location_examples}


@dataclass
class Discipline(Entity):
    """{crossner_naturalscience_discipline}"""

    span: str  # {crossner_naturalscience_discipline_examples}


@dataclass
class Enzyme(Entity):
    """{crossner_naturalscience_enzyme}"""

    span: str  # {crossner_naturalscience_enzyme_examples}


@dataclass
class Protein(Entity):
    """{crossner_naturalscience_protein}"""

    span: str  # {crossner_naturalscience_protein_examples}


@dataclass
class ChemicalElement(Entity):
    """{crossner_naturalscience_chemicalelement}"""

    span: str  # {crossner_naturalscience_chemicalelement_examples}


@dataclass
class ChemicalCompound(Entity):
    """{crossner_naturalscience_chemicalcompound}"""

    span: str  # {crossner_naturalscience_chemicalcompound_examples}


@dataclass
class AstronomicalObject(Entity):
    """{crossner_naturalscience_astronomicalobject}"""

    span: str  # {crossner_naturalscience_astronomicalobject_examples}


@dataclass
class AcademicJournal(Entity):
    """{crossner_naturalscience_academicjournal}"""

    span: str  # {crossner_naturalscience_academicjournal_examples}


@dataclass
class Event(Entity):
    """{crossner_naturalscience_event}"""

    span: str  # {crossner_naturalscience_event_examples}


@dataclass
class Theory(Entity):
    """{crossner_naturalscience_theory}"""

    span: str  # {crossner_naturalscience_theory_examples}


@dataclass
class Award(Entity):
    """{crossner_naturalscience_award}"""

    span: str  # {crossner_naturalscience_award_examples}


@dataclass
class Other(Entity):
    """{crossner_naturalscience_miscellaneous}"""

    span: str  # {crossner_naturalscience_miscellaneous_examples}


ENTITY_DEFINITIONS_NATURAL_SCIENCE: List[Entity] = [
    Scientist,
    Person,
    University,
    Organization,
    Country,
    Location,
    Discipline,
    Enzyme,
    Protein,
    ChemicalElement,
    ChemicalCompound,
    AstronomicalObject,
    AcademicJournal,
    Event,
    Theory,
    Award,
    Other,
]

ENTITY_DEFINITIONS_NATURAL_SCIENCE_woMISC: List[Entity] = [
    Scientist,
    Person,
    University,
    Organization,
    Country,
    Location,
    Discipline,
    Enzyme,
    Protein,
    ChemicalElement,
    ChemicalCompound,
    AstronomicalObject,
    AcademicJournal,
    Event,
    Theory,
    Award,
]

# __all__ = list(map(str, [*ENTITY_DEFINITIONS]))

src/tasks/crossner/prompts_natural_science.py
hitz-zentroa/GoLLIE
from typing import List

from ..utils_typing import Entity, dataclass


"""Entity definitions

The entity definitions are derived from the official CrossNER corpus guidelines:
https://arxiv.org/pdf/2012.04373.pdf

"""


@dataclass
class Person(Entity):
    """{crossner_politics_person}"""

    span: str  # {crossner_politics_person_examples}


@dataclass
class Organization(Entity):
    """{crossner_politics_organization}"""

    span: str  # {crossner_politics_organization_examples}


@dataclass
class Location(Entity):
    """{crossner_politics_location}"""

    span: str  # {crossner_politics_location_examples}


@dataclass
class Politician(Entity):
    """{crossner_politics_politician}"""

    span: str  # {crossner_politics_politician_examples}


@dataclass
class PoliticalParty(Entity):
    """{crossner_politics_politicalparty}"""

    span: str  # {crossner_politics_politicalparty_examples}


@dataclass
class Election(Entity):
    """{crossner_politics_election}"""

    span: str  # {crossner_politics_election_examples}


@dataclass
class Event(Entity):
    """{crossner_politics_event}"""

    span: str  # {crossner_politics_event_examples}


@dataclass
class Country(Entity):
    """{crossner_politics_country}"""

    span: str  # {crossner_politics_country_examples}


@dataclass
class Other(Entity):
    """{crossner_politics_miscellaneous}"""

    span: str  # {crossner_politics_miscellaneous_examples}


ENTITY_DEFINITIONS_POLITICS: List[Entity] = [
    Person,
    Organization,
    Location,
    Politician,
    PoliticalParty,
    Election,
    Event,
    Country,
    Other,
]

ENTITY_DEFINITIONS_POLITICS_woMISC: List[Entity] = [
    Person,
    Organization,
    Location,
    Politician,
    PoliticalParty,
    Election,
    Event,
    Country,
]

# __all__ = list(map(str, [*ENTITY_DEFINITIONS]))

src/tasks/crossner/prompts_politics.py
hitz-zentroa/GoLLIE
from typing import Dict, List, Type

from src.tasks.crossner.prompts_ai import ENTITY_DEFINITIONS_AI, ENTITY_DEFINITIONS_AI_woMISC
from src.tasks.crossner.prompts_literature import ENTITY_DEFINITIONS_LITERATURE, ENTITY_DEFINITIONS_LITERATURE_woMISC
from src.tasks.crossner.prompts_music import ENTITY_DEFINITIONS_MUSIC, ENTITY_DEFINITIONS_MUSIC_woMISC
from src.tasks.crossner.prompts_natural_science import (
    ENTITY_DEFINITIONS_NATURAL_SCIENCE,
    ENTITY_DEFINITIONS_NATURAL_SCIENCE_woMISC,
)
from src.tasks.crossner.prompts_politics import ENTITY_DEFINITIONS_POLITICS, ENTITY_DEFINITIONS_POLITICS_woMISC
from src.tasks.utils_scorer import SpanScorer
from src.tasks.utils_typing import Entity


class CrossNERPoliticsEntityScorer(SpanScorer):
    """CoNLL03 Entity identification and classification scorer."""

    valid_types: List[Type] = ENTITY_DEFINITIONS_POLITICS

    def __call__(self, reference: List[Entity], predictions: List[Entity]) -> Dict[str, Dict[str, float]]:
        output = super().__call__(reference, predictions)
        return {"entities": output["spans"]}


class CrossNERMusicEntityScorer(SpanScorer):
    """CoNLL03 Entity identification and classification scorer."""

    valid_types: List[Type] = ENTITY_DEFINITIONS_MUSIC

    def __call__(self, reference: List[Entity], predictions: List[Entity]) -> Dict[str, Dict[str, float]]:
        output = super().__call__(reference, predictions)
        return {"entities": output["spans"]}


class CrossNERLiteratureEntityScorer(SpanScorer):
    """CoNLL03 Entity identification and classification scorer."""

    valid_types: List[Type] = ENTITY_DEFINITIONS_LITERATURE

    def __call__(self, reference: List[Entity], predictions: List[Entity]) -> Dict[str, Dict[str, float]]:
        output = super().__call__(reference, predictions)
        return {"entities": output["spans"]}


class CrossNERAIEntityScorer(SpanScorer):
    """CoNLL03 Entity identification and classification scorer."""

    valid_types: List[Type] = ENTITY_DEFINITIONS_AI

    def __call__(self, reference: List[Entity], predictions: List[Entity]) -> Dict[str, Dict[str, float]]:
        output = super().__call__(reference, predictions)
        return {"entities": output["spans"]}


class CrossNERNaturalScienceEntityScorer(SpanScorer):
    """CoNLL03 Entity identification and classification scorer."""

    valid_types: List[Type] = ENTITY_DEFINITIONS_NATURAL_SCIENCE

    def __call__(self, reference: List[Entity], predictions: List[Entity]) -> Dict[str, Dict[str, float]]:
        output = super().__call__(reference, predictions)
        return {"entities": output["spans"]}


class CrossNERPoliticsEntityScorer_woMISC(SpanScorer):
    """CoNLL03 Entity identification and classification scorer."""

    valid_types: List[Type] = ENTITY_DEFINITIONS_POLITICS_woMISC

    def __call__(self, reference: List[Entity], predictions: List[Entity]) -> Dict[str, Dict[str, float]]:
        output = super().__call__(reference, predictions)
        return {"entities": output["spans"]}


class CrossNERMusicEntityScorer_woMISC(SpanScorer):
    """CoNLL03 Entity identification and classification scorer."""

    valid_types: List[Type] = ENTITY_DEFINITIONS_MUSIC_woMISC

    def __call__(self, reference: List[Entity], predictions: List[Entity]) -> Dict[str, Dict[str, float]]:
        output = super().__call__(reference, predictions)
        return {"entities": output["spans"]}


class CrossNERLiteratureEntityScorer_woMISC(SpanScorer):
    """CoNLL03 Entity identification and classification scorer."""

    valid_types: List[Type] = ENTITY_DEFINITIONS_LITERATURE_woMISC

    def __call__(self, reference: List[Entity], predictions: List[Entity]) -> Dict[str, Dict[str, float]]:
        output = super().__call__(reference, predictions)
        return {"entities": output["spans"]}


class CrossNERAIEntityScorer_woMISC(SpanScorer):
    """CoNLL03 Entity identification and classification scorer."""

    valid_types: List[Type] = ENTITY_DEFINITIONS_AI_woMISC

    def __call__(self, reference: List[Entity], predictions: List[Entity]) -> Dict[str, Dict[str, float]]:
        output = super().__call__(reference, predictions)
        return {"entities": output["spans"]}


class CrossNERNaturalScienceEntityScorer_woMISC(SpanScorer):
    """CoNLL03 Entity identification and classification scorer."""

    valid_types: List[Type] = ENTITY_DEFINITIONS_NATURAL_SCIENCE_woMISC

    def __call__(self, reference: List[Entity], predictions: List[Entity]) -> Dict[str, Dict[str, float]]:
        output = super().__call__(reference, predictions)
        return {"entities": output["spans"]}

src/tasks/crossner/scorer.py
hitz-zentroa/GoLLIE
from . import data_loader, prompts, scorer


__all__ = ["data_loader", "prompts", "scorer"]

src/tasks/diann/__init__.py
hitz-zentroa/GoLLIE
from typing import Dict, List, Tuple, Type, Union

from src.tasks.conll03.data_loader import read_tsv
from src.tasks.diann.guidelines import GUIDELINES
from src.tasks.diann.guidelines_gold import EXAMPLES
from src.tasks.diann.prompts import ENTITY_DEFINITIONS, Disability  # , Negation
from src.tasks.label_encoding import rewrite_labels

from ..utils_data import DatasetLoader, Sampler
from ..utils_typing import Entity


def load_diann_tsv(
    path: str,
    include_misc: bool,
    ENTITY_TO_CLASS_MAPPING: Dict[str, Type[Entity]],
) -> Tuple[List[List[str]], List[List[Entity]]]:
    """
    Load the conll dataset from a tsv file
    REMOVE THE NEGATION LABELS
    Args:
        path (str): The path to the tsv file
        include_misc (bool): Whether to include the MISC entity type. Defaults to `True`.
    Returns:
        (List[str],List[Union[Location,Organization,Person,Miscellaneous]]): The text and the entities
    """
    dataset_sentences: List[List[str]] = []
    dataset_entities: List[List[Entity]] = []

    dataset_words, dataset_labels = read_tsv(path)

    for words, labels in zip(dataset_words, dataset_labels):
        # Some of the CoNLL02-03 datasets are in IOB1 format instead of IOB2,
        # we convert them to IOB2, so we don't have to deal with this later.
        labels = rewrite_labels(labels=labels, encoding="iob2")
        # Remove Negation labels
        for i in range(len(labels)):
            if labels[i].endswith("Neg"):
                labels[i] = "O"

        # Rewrite just in case
        labels = rewrite_labels(labels=labels, encoding="iob2")

        # Get labeled word spans
        spans = []
        for i, label in enumerate(labels):
            if label == "O":
                continue
            elif label.startswith("B-"):
                spans.append([label[2:], i, i + 1])
            elif label.startswith("I-"):
                spans[-1][2] += 1
            else:
                raise ValueError(f"Found an unexpected label: {label}")

        # Get entities
        entities = []
        for label, start, end in spans:
            if include_misc or label.lower() != "misc":
                entities.append(ENTITY_TO_CLASS_MAPPING[label](span=" ".join(words[start:end])))

        dataset_sentences.append(words)
        dataset_entities.append(entities)

    return dataset_sentences, dataset_entities


class DiannDatasetLoader(DatasetLoader):
    """
    A `DatasetLoader` for the diann dataset.

    Args:
        path (`str`):
            The path to the dataset to load.

    Raises:
        `ValueError`:
            raised when a not defined value found.
    """

    ENTITY_TO_CLASS_MAPPING = None

    def __init__(self, path_or_split: str, **kwargs) -> None:
        self.ENTITY_TO_CLASS_MAPPING = {
            "Dis": Disability,
            # "Neg": Negation,
        }

        self.elements = {}

        dataset_words, dataset_entities = load_diann_tsv(
            path=path_or_split,
            include_misc=True,
            ENTITY_TO_CLASS_MAPPING=self.ENTITY_TO_CLASS_MAPPING,
        )

        for id, (words, entities) in enumerate(zip(dataset_words, dataset_entities)):
            self.elements[id] = {
                "id": id,
                "doc_id": id,
                "text": " ".join(words),
                "entities": entities,
                "gold": entities,
            }


class DiannSampler(Sampler):
    """
    A data `Sampler` for the Diann dataset.

    Args:
        dataset_loader (`DiannDiseaseDatasetLoader`):
            The dataset loader that contains the data information.
        task (`str`, optional):
            The task to sample. It must be one of the following: NER, VER, RE, EE.
            Defaults to `None`.
        split (`str`, optional):
            The path_or_split to sample. It must be one of the following: "train", "dev" or
            "test". Depending on the path_or_split the sampling strategy differs. Defaults to
            `"train"`.
        parallel_instances (`Union[int, Tuple[int, int]]`, optional):
            The number of sentences sampled in parallel. Options:

                * **`int`**: The amount of elements that will be sampled in parallel.
                * **`tuple`**: The range of elements that will be sampled in parallel.

            Defaults to 1.
        max_guidelines (`int`, optional):
            The number of guidelines to append to the example at the same time. If `-1`
            is given then all the guidelines are appended. Defaults to `-1`.
        guideline_dropout (`float`, optional):
            The probability to dropout a guideline definition for the given example. This
            is only applied on training. Defaults to `0.0`.
        seed (`float`, optional):
            The seed to sample the examples. Defaults to `0`.
        prompt_template (`str`, optional):
            The path to the prompt template. Defaults to `"templates/prompt.txt"`.
        ensure_positives_on_train (bool, optional):
            Whether to ensure that the guidelines of annotated examples are not removed.
            Defaults to `False`.
        dataset_name (str, optional):
            The name of the dataset. Defaults to `None`.
        scorer (`str`, optional):
           The scorer class import string. Defaults to `None`.
        sample_only_gold_guidelines (`bool`, optional):
            Whether to sample only guidelines of present annotations. Defaults to `False`.
    """

    def __init__(
        self,
        dataset_loader: DiannDatasetLoader,
        task: str = None,
        split: str = "train",
        parallel_instances: Union[int, Tuple[int, int]] = 1,
        max_guidelines: int = -1,
        guideline_dropout: float = 0.0,
        seed: float = 0,
        prompt_template: str = "templates/prompt.txt",
        ensure_positives_on_train: bool = False,
        dataset_name: str = None,
        scorer: str = None,
        sample_only_gold_guidelines: bool = False,
        **kwargs,
    ) -> None:
        assert task in [
            "NER",
        ], f"Diann Disease only supports NER task. {task} is not supported."

        task_definitions, task_target = {
            "NER": (ENTITY_DEFINITIONS, "entities"),
        }[task]

        super().__init__(
            dataset_loader=dataset_loader,
            task=task,
            split=split,
            parallel_instances=parallel_instances,
            max_guidelines=max_guidelines,
            guideline_dropout=guideline_dropout,
            seed=seed,
            prompt_template=prompt_template,
            ensure_positives_on_train=ensure_positives_on_train,
            sample_only_gold_guidelines=sample_only_gold_guidelines,
            dataset_name=dataset_name,
            scorer=scorer,
            task_definitions=task_definitions,
            task_target=task_target,
            definitions=GUIDELINES,
            examples=EXAMPLES,
            **kwargs,
        )

src/tasks/diann/data_loader.py
hitz-zentroa/GoLLIE
GUIDELINES = {
    "diann_disability": {
        "en": [
            (
                "The term disability refers to an umbrella term covering impairments, limitations of activities and"
                " restrictions \non participation. Lists of disability terms and lists of functions whose absence or"
                " limitation has been considered a \ndisability: The Attention, Autonomy, Cognition Communication,"
                " Behaviour, Day-to-day living activities Development, \nEmotions, Executive functioning Feeding,"
                " Functional capacity, Gait Hearing, Language, Learning Mental capabilities, \nMobility, Perception"
                " Psychological capabilities, Sensory, Sight Sleep, Social cognition, Speech Swallowing,"
                " \nOccupational functioning. The modifiers affecting the disability, such as 'isolated aphasia',"
                " 'advanced dementia' and \nAcronyms referring to a disability such as Mild Cognitive Impairment (MCI)"
                " are included in the disability term."
            ),
            (
                "The concept of disability encompasses a range of conditions, including impairments, limitations in"
                " daily \nactivities, and restrictions on social participation. Various lists of disability terms and"
                " functions that may be impacted, \nincluding but not limited to: Attention, Autonomy, Cognition,"
                " Communication, Behavior, Daily living activities, \nDevelopment, Emotions, Executive functioning,"
                " Feeding, Functional capacity, Gait, Hearing, Language, Learning, Mental \ncapabilities, Mobility,"
                " Perception, Psychological capabilities, Sensory, Sight, Sleep, Social cognition, Speech,"
                " \nSwallowing, and Occupational functioning. Additionally, this term includes modifiers that"
                ' influence the disability, such as \n"isolated aphasia" and "advanced dementia," as well as acronyms'
                " denoting specific disabilities like Mild Cognitive \nImpairment (MCI)."
            ),
            (
                "The concept of disability encompasses a range of conditions, including impairments, limitations in"
                " daily \nactivities, and restrictions in social participation. Disability can be characterized by"
                " various factors, such as \ndifficulties with attention, autonomy, cognition, communication,"
                " behavior, and adaptive functioning. Additionally, it can \nbe influenced by aspects like emotional"
                " and psychological capabilities, sensory functioning, sleep patterns, and \nsocial cognition. Lastly,"
                ' disability modifiers, like "isolated aphasia" and "advanced dementia," as well as acronyms'
                " \nreferring to specific disabilities (e.g., Mild Cognitive Impairment or MCI) are also considered"
                " part of this umbrella term."
            ),
            (
                "The concept of disability encompasses a range of conditions, including impairments, limited"
                " activities, and \nrestricted participation. Disability terms and functions that may be impacted,"
                " including attention, autonomy, cognition, \ncommunication, behavior, daily living activities,"
                " development, emotions, executive functioning, feeding, functional \ncapacity, gait, hearing,"
                " language, learning, mental capabilities, mobility, perception, psychological capabilities,"
                " \nsensory, sight, sleep, social cognition, speech, swallowing, and occupational functioning."
                ' Modifiers such as "isolated \naphasia" and "advanced dementia," as well as acronyms like Mild'
                " Cognitive Impairment (MCI), are also considered part of the \ndisability term."
            ),
            (
                "The concept of disability encompasses a range of conditions, including impairments, limitations in"
                " daily \nactivities, and restrictions on social participation. Various lists of disability terms and"
                " functions notably absent or \nrestricted include areas like attention, autonomy, cognition,"
                " communication, behavior, daily living activities, \ndevelopment, emotions, executive functioning,"
                " feeding, functional capacity, gait, hearing, language, learning, mental \ncapabilities, mobility,"
                " perception, psychological capabilities, sensory, sight, sleep, social cognition, speech,"
                ' \nswallowing, and occupational functioning. Additionally, modifiers such as "isolated aphasia" and'
                ' "advanced dementia" are \npart of the disability term, as well as acronyms like Mild Cognitive'
                " Impairment (MCI) that refer to specific \ndisabilities."
            ),
        ]
    },
    "diann_negation": {
        "en": [
            "Negation triggers are annotated when it affects one or more disabilities.",
            "Negation triggers are indicated when they impact one or more disabilities.",
            "Negation markers are indicated when they apply to one or more disabilities.",
            "Negation triggers are indicated when they impact one or multiple disabilities.",
            "Negation triggers are indicated when they impact one or multiple disabilities.",
        ]
    },
}

src/tasks/diann/guidelines.py
hitz-zentroa/GoLLIE
GUIDELINES = {
    "diann_disability": {
        "en": [
            "The term disability refers to an umbrella term covering impairments, limitations of activities and"
            " restrictions on participation. Lists of disability terms and lists of functions whose absence or"
            " limitation has been considered a disability: The Attention, Autonomy, Cognition Communication,"
            " Behaviour, Day-to-day living activities Development, Emotions, Executive functioning Feeding, Functional"
            " capacity, Gait Hearing, Language, Learning Mental capabilities, Mobility, Perception Psychological"
            " capabilities, Sensory, Sight Sleep, Social cognition, Speech Swallowing, Occupational functioning. The"
            " modifiers affecting the disability, such as 'isolated aphasia', 'advanced dementia' and  Acronyms"
            " referring to a disability such as Mild Cognitive Impairment (MCI) are included in the disability"
            " term."
        ]
    },
    "diann_negation": {
        "en": ["Negation triggers are annotated when it affects one or more disabilities."],
    },
}

EXAMPLES = {
    "diann_disability_examples": {
        "en": [
            "cognitive impairment",
            "hearing loss",
            "dementia",
            "mental disorders",
            "mental retardation",
            "sensorineural hearing loss",
            "mild cognitive impairment",
            "functional impairment",
            "mental disorder",
            "dysarthria",
        ]
    },
    "diann_negation_examples": {"en": []},
}

src/tasks/diann/guidelines_gold.py
hitz-zentroa/GoLLIE
from typing import List

from ..utils_typing import Entity, dataclass


"""Entity definitions

The entity definitions are derived from the official Diann guidelines:
http://nlp.uned.es/diann/#data
"""


@dataclass
class Disability(Entity):
    """{diann_disability}"""

    span: str  # {diann_disability_examples}


"""
@dataclass
class Negation(Entity):
    \"""{diann_negation}\"""

    span: str  # {diann_negation_examples}
"""

ENTITY_DEFINITIONS: List[Entity] = [Disability]

# __all__ = list(map(str, [*ENTITY_DEFINITIONS]))

src/tasks/diann/prompts.py
hitz-zentroa/GoLLIE
from typing import Dict, List, Type

from src.tasks.diann.prompts import ENTITY_DEFINITIONS
from src.tasks.utils_scorer import SpanScorer
from src.tasks.utils_typing import Entity


class DiannDiseaseEntityScorer(SpanScorer):
    """CoNLL03 Entity identification and classification scorer."""

    valid_types: List[Type] = ENTITY_DEFINITIONS

    def __call__(self, reference: List[Entity], predictions: List[Entity]) -> Dict[str, Dict[str, float]]:
        output = super().__call__(reference, predictions)
        return {"entities": output["spans"]}

src/tasks/diann/scorer.py
hitz-zentroa/GoLLIE
from . import data_loader, prompts, scorer


__all__ = ["data_loader", "prompts", "scorer"]

src/tasks/e3c/__init__.py
hitz-zentroa/GoLLIE
from typing import Tuple, Union

from src.tasks.conll03.data_loader import load_conll_tsv
from src.tasks.e3c.guidelines import GUIDELINES
from src.tasks.e3c.guidelines_gold import EXAMPLES
from src.tasks.e3c.prompts import (
    ENTITY_DEFINITIONS,
    ClinicalEntity,
)

from ..utils_data import DatasetLoader, Sampler


class E3CDatasetLoader(DatasetLoader):
    """
    A `DatasetLoader` for the E3C dataset.

    Args:
        split (`str`):
            The split to load. Can be one of `train`, `validation` or `test`.

    Raises:
        `ValueError`:
            raised when a not defined value found.
    """

    ENTITY_TO_CLASS_MAPPING = None

    def __init__(self, path_or_split: str, **kwargs) -> None:
        self.ENTITY_TO_CLASS_MAPPING = {
            "ety": ClinicalEntity,
        }

        self.elements = {}

        dataset_words, dataset_entities = load_conll_tsv(
            path=path_or_split,
            include_misc=True,
            ENTITY_TO_CLASS_MAPPING=self.ENTITY_TO_CLASS_MAPPING,
        )

        for id, (words, entities) in enumerate(zip(dataset_words, dataset_entities)):
            self.elements[id] = {
                "id": id,
                "doc_id": id,
                "text": " ".join(words),
                "entities": entities,
                "gold": entities,
            }


class E3CSampler(Sampler):
    """
    A data `Sampler` for the E3C dataset.

    Args:
        dataset_loader (`NcbiDiseaseDatasetLoader`):
            The dataset loader that contains the data information.
        task (`str`, optional):
            The task to sample. It must be one of the following: NER, VER, RE, EE.
            Defaults to `None`.
        split (`str`, optional):
            The path_or_split to sample. It must be one of the following: "train", "dev" or
            "test". Depending on the path_or_split the sampling strategy differs. Defaults to
            `"train"`.
        parallel_instances (`Union[int, Tuple[int, int]]`, optional):
            The number of sentences sampled in parallel. Options:

                * **`int`**: The amount of elements that will be sampled in parallel.
                * **`tuple`**: The range of elements that will be sampled in parallel.

            Defaults to 1.
        max_guidelines (`int`, optional):
            The number of guidelines to append to the example at the same time. If `-1`
            is given then all the guidelines are appended. Defaults to `-1`.
        guideline_dropout (`float`, optional):
            The probability to dropout a guideline definition for the given example. This
            is only applied on training. Defaults to `0.0`.
        seed (`float`, optional):
            The seed to sample the examples. Defaults to `0`.
        prompt_template (`str`, optional):
            The path to the prompt template. Defaults to `"templates/prompt.txt"`.
        ensure_positives_on_train (bool, optional):
            Whether to ensure that the guidelines of annotated examples are not removed.
            Defaults to `False`.
        dataset_name (str, optional):
            The name of the dataset. Defaults to `None`.
        scorer (`str`, optional):
           The scorer class import string. Defaults to `None`.
        sample_only_gold_guidelines (`bool`, optional):
            Whether to sample only guidelines of present annotations. Defaults to `False`.
    """

    def __init__(
        self,
        dataset_loader: E3CDatasetLoader,
        task: str = None,
        split: str = "train",
        parallel_instances: Union[int, Tuple[int, int]] = 1,
        max_guidelines: int = -1,
        guideline_dropout: float = 0.0,
        seed: float = 0,
        prompt_template: str = "templates/prompt.txt",
        ensure_positives_on_train: bool = False,
        dataset_name: str = None,
        scorer: str = None,
        sample_only_gold_guidelines: bool = False,
        **kwargs,
    ) -> None:
        assert task in [
            "NER",
        ], f"E3C Disease only supports NER task. {task} is not supported."

        task_definitions, task_target = {
            "NER": (ENTITY_DEFINITIONS, "entities"),
        }[task]

        super().__init__(
            dataset_loader=dataset_loader,
            task=task,
            split=split,
            parallel_instances=parallel_instances,
            max_guidelines=max_guidelines,
            guideline_dropout=guideline_dropout,
            seed=seed,
            prompt_template=prompt_template,
            ensure_positives_on_train=ensure_positives_on_train,
            sample_only_gold_guidelines=sample_only_gold_guidelines,
            dataset_name=dataset_name,
            scorer=scorer,
            task_definitions=task_definitions,
            task_target=task_target,
            definitions=GUIDELINES,
            examples=EXAMPLES,
            **kwargs,
        )

src/tasks/e3c/data_loader.py
hitz-zentroa/GoLLIE
GUIDELINES = {
    "e3c_disease": {
        "en": [
            "A definite pathologic process with a characteristic set of signs and symptoms. Examples: tumor, tumour,"
            " vomiting, swelling, epistaxis, ascites, headache, fever, cyst, diplopia."
        ]
    },
}

src/tasks/e3c/guidelines.py
hitz-zentroa/GoLLIE
GUIDELINES = {
    "e3c_disease": {"en": ["A definite pathologic process with a characteristic set of signs and symptoms."]},
}

EXAMPLES = {
    "e3c_disease_examples": {
        "en": [
            "tumor",
            "tumour",
            "vomiting",
            "swelling",
            "epistaxis",
            "ascites",
            "headache",
            "fever",
            "cyst",
            "diplopia",
        ]
    }
}

src/tasks/e3c/guidelines_gold.py
hitz-zentroa/GoLLIE
from typing import List

from ..utils_typing import Entity, dataclass


"""Entity definitions

The entity definitions are derived from the official E3C corpus guidelines:
https://github.com/hltfbk/E3C-Corpus/blob/main/documentation/CLINICALENTITY_ANNOTATION_GUIDELINES_v1.1.pdf

"""


@dataclass
class ClinicalEntity(Entity):
    """{e3c_disease}"""

    span: str  # {e3c_disease_examples}


ENTITY_DEFINITIONS: List[Entity] = [ClinicalEntity]

# __all__ = list(map(str, [*ENTITY_DEFINITIONS]))

src/tasks/e3c/prompts.py
hitz-zentroa/GoLLIE
from typing import Dict, List, Type

from src.tasks.e3c.prompts import ENTITY_DEFINITIONS
from src.tasks.utils_scorer import SpanScorer
from src.tasks.utils_typing import Entity


class E3CEntityScorer(SpanScorer):
    """CoNLL03 Entity identification and classification scorer."""

    valid_types: List[Type] = ENTITY_DEFINITIONS

    def __call__(self, reference: List[Entity], predictions: List[Entity]) -> Dict[str, Dict[str, float]]:
        output = super().__call__(reference, predictions)
        return {"entities": output["spans"]}

src/tasks/e3c/scorer.py
hitz-zentroa/GoLLIE
from . import data_loader, prompts, scorer


__all__ = ["data_loader", "prompts", "scorer"]

src/tasks/fabner/__init__.py
hitz-zentroa/GoLLIE
from typing import Dict, List, Tuple, Type, Union

from src.tasks.fabner.guidelines import GUIDELINES
from src.tasks.fabner.guidelines_gold import EXAMPLES
from src.tasks.fabner.prompts import (
    ENTITY_DEFINITIONS,
    Application,
    Biomedical,
    ConceptPrinciples,
    EnablingTechnology,
    EngineeringFeatures,
    MachineEquipment,
    ManufacturingProcess,
    ManufacturingStandards,
    Material,
    MechanicalProperties,
    ProcessCharacterization,
    ProcessParameters,
)
from src.tasks.label_encoding import rewrite_labels

from ..utils_data import DatasetLoader, Sampler
from ..utils_typing import Entity


def get_fabner_hf(
    split: str,
    ENTITY_TO_CLASS_MAPPING: Dict[str, Type[Entity]],
) -> Tuple[List[List[str]], List[List[Entity]]]:
    """
    Get the fabner dataset from the huggingface datasets library
    Args:
        split (str): The path_or_split to load. Can be one of `train`, `validation` or `test`.
    Returns:
        (List[str],List[Entity]): The text and the entities
    """
    from datasets import load_dataset

    dataset = load_dataset("DFKI-SLT/fabner", name="fabner_bio")
    dataset["train"].features["ner_tags"].feature.names

    id2label = dict(enumerate(dataset["train"].features["ner_tags"].feature.names))

    dataset_sentences: List[List[str]] = []
    dataset_entities: List[List[Entity]] = []

    for example in dataset[split]:
        words = example["tokens"]
        # Ensure IOB2 encoding
        labels = rewrite_labels(labels=[id2label[label] for label in example["ner_tags"]], encoding="iob2")

        # Get labeled word spans
        spans = []
        for i, label in enumerate(labels):
            if label == "O":
                continue
            elif label.startswith("B-"):
                spans.append([label[2:], i, i + 1])
            elif label.startswith("I-"):
                spans[-1][2] += 1
            else:
                raise ValueError(f"Found an unexpected label: {label}")

        # Get entities
        entities = []
        for label, start, end in spans:
            entities.append(ENTITY_TO_CLASS_MAPPING[label](span=" ".join(words[start:end])))

        dataset_sentences.append(words)
        dataset_entities.append(entities)

    return dataset_sentences, dataset_entities


class FabNerDatasetLoader(DatasetLoader):
    """
    A `DatasetLoader` for the FabNER dataset.

    Args:
        split (`str`):
            The split to load. Can be one of `train`, `validation` or `test`.

    Raises:
        `ValueError`:
            raised when a not defined value found.
    """

    ENTITY_TO_CLASS_MAPPING = None

    def __init__(self, path_or_split: str, **kwargs) -> None:
        self.ENTITY_TO_CLASS_MAPPING = {
            "MATE": Material,
            "MANP": ManufacturingProcess,
            "MACEQ": MachineEquipment,
            "APPL": Application,
            "FEAT": EngineeringFeatures,
            "PRO": MechanicalProperties,
            "CHAR": ProcessCharacterization,
            "PARA": ProcessParameters,
            "ENAT": EnablingTechnology,
            "CONPRI": ConceptPrinciples,
            "MANS": ManufacturingStandards,
            "BIOP": Biomedical,
        }

        self.elements = {}

        dataset_words, dataset_entities = get_fabner_hf(
            split=path_or_split,
            ENTITY_TO_CLASS_MAPPING=self.ENTITY_TO_CLASS_MAPPING,
        )

        for id, (words, entities) in enumerate(zip(dataset_words, dataset_entities)):
            self.elements[id] = {
                "id": id,
                "doc_id": id,
                "text": " ".join(words),
                "entities": entities,
                "gold": entities,
            }


class FabNerSampler(Sampler):
    """
    A data `Sampler` for the FabNER dataset.

    Args:
        dataset_loader (`FabNerDatasetLoader`):
            The dataset loader that contains the data information.
        task (`str`, optional):
            The task to sample. It must be one of the following: NER, VER, RE, EE.
            Defaults to `None`.
        split (`str`, optional):
            The path_or_split to sample. It must be one of the following: "train", "dev" or
            "test". Depending on the path_or_split the sampling strategy differs. Defaults to
            `"train"`.
        parallel_instances (`Union[int, Tuple[int, int]]`, optional):
            The number of sentences sampled in parallel. Options:

                * **`int`**: The amount of elements that will be sampled in parallel.
                * **`tuple`**: The range of elements that will be sampled in parallel.

            Defaults to 1.
        max_guidelines (`int`, optional):
            The number of guidelines to append to the example at the same time. If `-1`
            is given then all the guidelines are appended. Defaults to `-1`.
        guideline_dropout (`float`, optional):
            The probability to dropout a guideline definition for the given example. This
            is only applied on training. Defaults to `0.0`.
        seed (`float`, optional):
            The seed to sample the examples. Defaults to `0`.
        prompt_template (`str`, optional):
            The path to the prompt template. Defaults to `"templates/prompt.txt"`.
        ensure_positives_on_train (bool, optional):
            Whether to ensure that the guidelines of annotated examples are not removed.
            Defaults to `False`.
        dataset_name (str, optional):
            The name of the dataset. Defaults to `None`.
        scorer (`str`, optional):
           The scorer class import string. Defaults to `None`.
        sample_only_gold_guidelines (`bool`, optional):
            Whether to sample only guidelines of present annotations. Defaults to `False`.
    """

    def __init__(
        self,
        dataset_loader: FabNerDatasetLoader,
        task: str = None,
        split: str = "train",
        parallel_instances: Union[int, Tuple[int, int]] = 1,
        max_guidelines: int = -1,
        guideline_dropout: float = 0.0,
        seed: float = 0,
        prompt_template: str = "templates/prompt.txt",
        ensure_positives_on_train: bool = False,
        dataset_name: str = None,
        scorer: str = None,
        sample_only_gold_guidelines: bool = False,
        **kwargs,
    ) -> None:
        assert task in [
            "NER",
        ], f"FabNER only supports NER task. {task} is not supported."

        task_definitions, task_target = {
            "NER": (ENTITY_DEFINITIONS, "entities"),
        }[task]

        super().__init__(
            dataset_loader=dataset_loader,
            task=task,
            split=split,
            parallel_instances=parallel_instances,
            max_guidelines=max_guidelines,
            guideline_dropout=guideline_dropout,
            seed=seed,
            prompt_template=prompt_template,
            ensure_positives_on_train=ensure_positives_on_train,
            sample_only_gold_guidelines=sample_only_gold_guidelines,
            dataset_name=dataset_name,
            scorer=scorer,
            task_definitions=task_definitions,
            task_target=task_target,
            definitions=GUIDELINES,
            examples=EXAMPLES,
            **kwargs,
        )

src/tasks/fabner/data_loader.py
hitz-zentroa/GoLLIE
GUIDELINES = {
    "fabner_material": {
        "en": [
            "Refers to any substance or compound used in manufacturing processes that has a significant impact on the"
            " process, properties, or characteristics of the end product. Materials can include metals, alloys,"
            " polymers, ceramics, composites, and more."
        ]
    },
    "fabner_manufacturingprocess": {
        "en": ["Refers to a specific technique, method, or procedure used in the production of goods."]
    },
    "fabner_machineequipment": {
        "en": [
            "Any necessary items required for a particular purpose, including assembled parts and set of tools."
            " Equipments are required to set up not only the mechanical machines, but also for setting up any relevant"
            " operation or task."
        ]
    },
    "fabner_application": {
        "en": [
            "Refers to the specific field or industry in which a manufacturing process is utilized or intended to be"
            " utilized. It signifies the purpose, domain, or sector where the manufacturing process finds its use."
        ]
    },
    "fabner_engineeringfeatures": {
        "en": [
            "Encompass distinct terms or phrases that convey specific information related to the structural,"
            " functional, or visual properties of a manufacturing process."
        ]
    },
    "fabner_mechanicalproperties": {
        "en": [
            "Refers to specific attributes or characteristics of materials that describe their behavior under"
            " mechanical forces and conditions. These properties play a crucial role in manufacturing processes and"
            " are often used to evaluate the suitability and performance of materials for different applications."
        ]
    },
    "fabner_processcharacterization": {
        "en": [
            "Refers to the detailed analysis and measurement of various aspects of manufacturing processes to"
            " understand their performance, quality, and characteristics."
        ]
    },
    "fabner_processparameters": {
        "en": [
            "Process parameters refer to specific quantitative values or attributes that play a significant role in"
            " manufacturing processes."
        ]
    },
    "fabner_enablingtechnology": {
        "en": [
            "Specific tools, methods, processes, or technologies that play a pivotal role in facilitating or enhancing"
            " various aspects of manufacturing processes."
        ]
    },
    "fabner_conceptprinciples": {
        "en": ["Refers to key concepts, principles, and elements related to manufacturing processes."]
    },
    "fabner_manufacturingstandards": {
        "en": [
            "Refer to specific guidelines, specifications, benchmarks, and terminology used within the manufacturing"
            " industry to ensure uniformity, quality, and compatibility of products, processes, and systems."
        ]
    },
    "fabner_biomedical": {"en": ["Refers to any biomedical entity involved in the manufacturing process."]},
}

src/tasks/fabner/guidelines.py
hitz-zentroa/GoLLIE
GUIDELINES = {
    "fabner_material": {
        "en": [
            "Refers to any substance or compound used in manufacturing processes that has a significant impact on the"
            " process, properties, or characteristics of the end product. Materials can include metals, alloys,"
            " polymers, ceramics, composites, and more."
        ]
    },
    "fabner_manufacturingprocess": {
        "en": ["Refers to a specific technique, method, or procedure used in the production of goods."]
    },
    "fabner_machineequipment": {
        "en": [
            "Any necessary items required for a particular purpose, including assembled parts and set of tools."
            " Equipments are required to set up not only the mechanical machines, but also for setting up any relevant"
            " operation or task."
        ]
    },
    "fabner_application": {
        "en": [
            "Refers to the specific field or industry in which a manufacturing process is utilized or intended to be"
            " utilized. It signifies the purpose, domain, or sector where the manufacturing process finds its use."
        ]
    },
    "fabner_engineeringfeatures": {
        "en": [
            "Encompass distinct terms or phrases that convey specific information related to the structural,"
            " functional, or visual properties of a manufacturing process."
        ]
    },
    "fabner_mechanicalproperties": {
        "en": [
            "Refers to specific attributes or characteristics of materials that describe their behavior under"
            " mechanical forces and conditions. These properties play a crucial role in manufacturing processes and"
            " are often used to evaluate the suitability and performance of materials for different applications."
        ]
    },
    "fabner_processcharacterization": {
        "en": [
            "Refers to the detailed analysis and measurement of various aspects of manufacturing processes to"
            " understand their performance, quality, and characteristics."
        ]
    },
    "fabner_processparameters": {
        "en": [
            "Process parameters refer to specific quantitative values or attributes that play a significant role in"
            " manufacturing processes."
        ]
    },
    "fabner_enablingtechnology": {
        "en": [
            "Specific tools, methods, processes, or technologies that play a pivotal role in facilitating or enhancing"
            " various aspects of manufacturing processes."
        ]
    },
    "fabner_conceptprinciples": {
        "en": ["Refers to key concepts, principles, and elements related to manufacturing processes."]
    },
    "fabner_manufacturingstandards": {
        "en": [
            "Refer to specific guidelines, specifications, benchmarks, and terminology used within the manufacturing"
            " industry to ensure uniformity, quality, and compatibility of products, processes, and systems."
        ]
    },
    "fabner_biomedical": {"en": ["Refers to any biomedical entity involved in the manufacturing process."]},
}

EXAMPLES = {
    "fabner_material_examples": {
        "en": ["as", "be", "material", "powder", "alloy", "melt pool", "polymer", "filament", "composites", "alloys"]
    },
    "fabner_manufacturingprocess_examples": {
        "en": [
            "AM",
            "additive manufacturing",
            "SLM",
            "3D printing",
            "fabrication",
            "FFF",
            "manufacturing",
            "3D printed",
            "LPBF",
            "additively manufactured",
        ]
    },
    "fabner_machineequipment_examples": {
        "en": [
            "components",
            "component",
            "nozzle",
            "beam",
            "powder bed",
            "machine",
            "tool",
            "printer",
            "core",
            "3D printer",
        ]
    },
    "fabner_application_examples": {
        "en": [
            "mechanical",
            "imaging",
            "coating",
            "industrial",
            "support",
            "aerospace",
            "cell",
            "industry",
            "set",
            "contact",
        ]
    },
    "fabner_engineeringfeatures_examples": {
        "en": [
            "design",
            "designed",
            "texture",
            "lattice structures",
            "designs",
            "dimensions",
            "scaffolds",
            "topology optimization",
            "lattice structure",
            "cellular structures",
        ]
    },
    "fabner_mechanicalproperties_examples": {
        "en": [
            "porosity",
            "strength",
            "density",
            "hardness",
            "residual stress",
            "surface roughness",
            "stress",
            "tensile strength",
            "fatigue",
            "pores",
        ]
    },
    "fabner_processcharacterization_examples": {
        "en": [
            "accuracy",
            "optical",
            "measurement",
            "spatter",
            "SEM",
            "testing",
            "accurate",
            "scanning electron microscopy",
            "X-ray diffraction",
            "energy absorption",
        ]
    },
    "fabner_processparameters_examples": {
        "en": [
            "temperature",
            "layer",
            "build",
            "range",
            "temperatures",
            "laser power",
            "area",
            "resolution",
            "weight",
            "power",
        ]
    },
    "fabner_enablingtechnology_examples": {
        "en": [
            "laser",
            "simulation",
            "modeling",
            "simulations",
            "CT",
            "laser scan",
            "modelling",
            "Laser",
            "rapid prototyping",
            "CAD",
        ]
    },
    "fabner_conceptprinciples_examples": {
        "en": [
            "process",
            "model",
            "mechanical properties",
            "samples",
            "materials",
            "microstructure",
            "properties",
            "surface",
            "parameters",
            "fabricated",
        ]
    },
    "fabner_manufacturingstandards_examples": {
        "en": [
            "ISO",
            "benchmark",
            "ISO 25178-2",
            "file",
            "STL",
            "files",
            "meter",
            "Benchmark",
            "Standard Tessellation Language",
            "ISO standards",
        ]
    },
    "fabner_biomedical_examples": {
        "en": [
            "bone",
            "dendrites",
            "starch",
            "dendrite",
            "bone scaffolds",
            "osteoblast cells",
            "osteoblast",
            "bone defects",
            "subcutaneous",
            "horizontal specimens",
        ]
    },
}

src/tasks/fabner/guidelines_gold.py
hitz-zentroa/GoLLIE
from typing import List

from ..utils_typing import Entity, dataclass


"""Entity definitions

The entity definitions are derived from the official fabner guidelines:
https://par.nsf.gov/servlets/purl/10290810

"""


@dataclass
class Material(Entity):
    """{fabner_material}"""

    span: str  # {fabner_material_examples}


@dataclass
class ManufacturingProcess(Entity):
    """{fabner_manufacturingprocess}"""

    span: str  # {fabner_manufacturingprocess_examples}


@dataclass
class MachineEquipment(Entity):
    """{fabner_machineequipment}"""

    span: str  # {fabner_machineequipment_examples}


@dataclass
class Application(Entity):
    """{fabner_application}"""

    span: str  # {fabner_application_examples}


@dataclass
class EngineeringFeatures(Entity):
    """{fabner_engineeringfeatures}"""

    span: str  # {fabner_engineeringfeatures_examples}


@dataclass
class MechanicalProperties(Entity):
    """{fabner_mechanicalproperties}"""

    span: str  # {fabner_mechanicalproperties_examples}


@dataclass
class ProcessCharacterization(Entity):
    """{fabner_processcharacterization}"""

    span: str  # {fabner_processcharacterization_examples}


@dataclass
class ProcessParameters(Entity):
    """{fabner_processparameters}"""

    span: str  # {fabner_processparameters_examples}


@dataclass
class EnablingTechnology(Entity):
    """{fabner_enablingtechnology}"""

    span: str  # {fabner_enablingtechnology_examples}


@dataclass
class ConceptPrinciples(Entity):
    """{fabner_conceptprinciples}"""

    span: str  # {fabner_conceptprinciples_examples}


@dataclass
class ManufacturingStandards(Entity):
    """{fabner_manufacturingstandards}"""

    span: str  # {fabner_manufacturingstandards_examples}


@dataclass
class Biomedical(Entity):
    """{fabner_biomedical}"""

    span: str  # {fabner_biomedical_examples}


ENTITY_DEFINITIONS: List[Entity] = [
    Material,
    ManufacturingProcess,
    MachineEquipment,
    Application,
    EngineeringFeatures,
    MechanicalProperties,
    ProcessCharacterization,
    ProcessParameters,
    EnablingTechnology,
    ConceptPrinciples,
    ManufacturingStandards,
    Biomedical,
]

# __all__ = list(map(str, [*ENTITY_DEFINITIONS]))

src/tasks/fabner/prompts.py
hitz-zentroa/GoLLIE
from typing import Dict, List, Type

from src.tasks.fabner.prompts import ENTITY_DEFINITIONS
from src.tasks.utils_scorer import SpanScorer
from src.tasks.utils_typing import Entity


class FabNerEntityScorer(SpanScorer):
    """CoNLL03 Entity identification and classification scorer."""

    valid_types: List[Type] = ENTITY_DEFINITIONS

    def __call__(self, reference: List[Entity], predictions: List[Entity]) -> Dict[str, Dict[str, float]]:
        output = super().__call__(reference, predictions)
        return {"entities": output["spans"]}

src/tasks/fabner/scorer.py
hitz-zentroa/GoLLIE
from . import data_loader, prompts, scorer


__all__ = ["data_loader", "prompts", "scorer"]

src/tasks/harveyner/__init__.py
hitz-zentroa/GoLLIE
from typing import Tuple, Union

from src.tasks.conll03.data_loader import load_conll_tsv
from src.tasks.harveyner.guidelines import GUIDELINES
from src.tasks.harveyner.guidelines_gold import EXAMPLES
from src.tasks.harveyner.prompts import (
    ENTITY_DEFINITIONS,
    Area,
    Point,
    River,
    Road,
)

from ..utils_data import DatasetLoader, Sampler


class HarveyNerDatasetLoader(DatasetLoader):
    """
    A `DatasetLoader` for the HarveyNER dataset.

    Args:
        split (`str`):
            The split to load. Can be one of `train`, `validation` or `test`.

    Raises:
        `ValueError`:
            raised when a not defined value found.
    """

    ENTITY_TO_CLASS_MAPPING = None

    def __init__(self, path_or_split: str, **kwargs) -> None:
        self.ENTITY_TO_CLASS_MAPPING = {
            "POINT": Point,
            "AREA": Area,
            "RIVER": River,
            "ROAD": Road,
        }

        self.elements = {}

        dataset_words, dataset_entities = load_conll_tsv(
            path=path_or_split,
            include_misc=True,
            ENTITY_TO_CLASS_MAPPING=self.ENTITY_TO_CLASS_MAPPING,
        )

        for id, (words, entities) in enumerate(zip(dataset_words, dataset_entities)):
            self.elements[id] = {
                "id": id,
                "doc_id": id,
                "text": " ".join(words),
                "entities": entities,
                "gold": entities,
            }


class HarveyNerSampler(Sampler):
    """
    A data `Sampler` for the HarveyNER dataset.

    Args:
        dataset_loader (`HarveyNERDatasetLoader`):
            The dataset loader that contains the data information.
        task (`str`, optional):
            The task to sample. It must be one of the following: NER, VER, RE, EE.
            Defaults to `None`.
        split (`str`, optional):
            The path_or_split to sample. It must be one of the following: "train", "dev" or
            "test". Depending on the path_or_split the sampling strategy differs. Defaults to
            `"train"`.
        parallel_instances (`Union[int, Tuple[int, int]]`, optional):
            The number of sentences sampled in parallel. Options:

                * **`int`**: The amount of elements that will be sampled in parallel.
                * **`tuple`**: The range of elements that will be sampled in parallel.

            Defaults to 1.
        max_guidelines (`int`, optional):
            The number of guidelines to append to the example at the same time. If `-1`
            is given then all the guidelines are appended. Defaults to `-1`.
        guideline_dropout (`float`, optional):
            The probability to dropout a guideline definition for the given example. This
            is only applied on training. Defaults to `0.0`.
        seed (`float`, optional):
            The seed to sample the examples. Defaults to `0`.
        prompt_template (`str`, optional):
            The path to the prompt template. Defaults to `"templates/prompt.txt"`.
        ensure_positives_on_train (bool, optional):
            Whether to ensure that the guidelines of annotated examples are not removed.
            Defaults to `False`.
        dataset_name (str, optional):
            The name of the dataset. Defaults to `None`.
        scorer (`str`, optional):
           The scorer class import string. Defaults to `None`.
        sample_only_gold_guidelines (`bool`, optional):
            Whether to sample only guidelines of present annotations. Defaults to `False`.
    """

    def __init__(
        self,
        dataset_loader: HarveyNerDatasetLoader,
        task: str = None,
        split: str = "train",
        parallel_instances: Union[int, Tuple[int, int]] = 1,
        max_guidelines: int = -1,
        guideline_dropout: float = 0.0,
        seed: float = 0,
        prompt_template: str = "templates/prompt.txt",
        ensure_positives_on_train: bool = False,
        dataset_name: str = None,
        scorer: str = None,
        sample_only_gold_guidelines: bool = False,
        **kwargs,
    ) -> None:
        assert task in [
            "NER",
        ], f"HarveyNER only supports NER task. {task} is not supported."

        task_definitions, task_target = {
            "NER": (ENTITY_DEFINITIONS, "entities"),
        }[task]

        super().__init__(
            dataset_loader=dataset_loader,
            task=task,
            split=split,
            parallel_instances=parallel_instances,
            max_guidelines=max_guidelines,
            guideline_dropout=guideline_dropout,
            seed=seed,
            prompt_template=prompt_template,
            ensure_positives_on_train=ensure_positives_on_train,
            sample_only_gold_guidelines=sample_only_gold_guidelines,
            dataset_name=dataset_name,
            scorer=scorer,
            task_definitions=task_definitions,
            task_target=task_target,
            definitions=GUIDELINES,
            examples=EXAMPLES,
            **kwargs,
        )

src/tasks/harveyner/data_loader.py
hitz-zentroa/GoLLIE
GUIDELINES = {
    "harveyner_point": {
        "en": [
            "Refers to a location that is a building, a landmark, an intersection of two roads, an intersection of a"
            " river with a lake/reservoir/ocean, or a specifc address. Ignore generic company/franchise names unless"
            " it is accompanied with a precise location, for example, HEB at Kirkwood Drive. However, non-franchised"
            " small businesses with only one unique location are considered as a point. Ignore any locations in the"
            " Twitter username, unless @ does not refer to a Twitter account name. For example, I am @ XXX High"
            " School."
        ]
    },
    "harveyner_area": {
        "en": [
            "Refers to all the named entities of cities, neighborhoods, super neighborhoods, geographic divisions etc."
            " The following locations, Lake Houston, Barker Reservoir, and Addick’s Reservoir, are annotated as areas"
            " due to their significant size while all other lakes/reservoirs are not considered as areas."
        ]
    },
    "harveyner_road": {
        "en": [
            "Refers to a road/avenue/street or a section of a road/avenue/street when the tweet does not provide an"
            " exact location on that road."
        ]
    },
    "harveyner_river": {
        "en": [
            "Refers to a river or a section of a river when the tweet does not imply there is an intersection between"
            " the river and other places."
        ]
    },
}

src/tasks/harveyner/guidelines.py
hitz-zentroa/GoLLIE
GUIDELINES = {
    "harveyner_point": {
        "en": [
            "Refers to a location that is a building, a landmark, an intersection of two roads, an intersection of a"
            " river with a lake/reservoir/ocean, or a specifc address. Ignore generic company/franchise names unless"
            " it is accompanied with a precise location, for example, HEB at Kirkwood Drive. However, non-franchised"
            " small businesses with only one unique location are considered as a point. Ignore any locations in the"
            " Twitter username, unless @ does not refer to a Twitter account name. For example, I am @ XXX High"
            " School."
        ]
    },
    "harveyner_area": {
        "en": [
            "Refers to all the named entities of cities, neighborhoods, super neighborhoods, geographic divisions etc."
            " The following locations, Lake Houston, Barker Reservoir, and Addick’s Reservoir, are annotated as areas"
            " due to their significant size while all other lakes/reservoirs are not considered as areas."
        ]
    },
    "harveyner_road": {
        "en": [
            "Refers to a road/avenue/street or a section of a road/avenue/street when the tweet does not provide an"
            " exact location on that road."
        ]
    },
    "harveyner_river": {
        "en": [
            "Refers to a river or a section of a river when the tweet does not imply there is an intersection between"
            " the river and other places."
        ]
    },
}

EXAMPLES = {
    "harveyner_point_examples": {
        "en": [
            "GRB",
            "GEORGE R. BROWN",
            "Lakewood Church",
            "Bayou Oaks",
            "Northgate Subdivision S of toll road",
            "TerryHS",
            "GRB Convention Center",
            "@GRBCC",
            "Addicks",
            "TERRY HIGH SCHOOL",
        ]
    },
    "harveyner_area_examples": {
        "en": [
            "Sienna Plantation",
            "Galveston",
            "Addicks",
            "Barker",
            "Fort Bend",
            "Chambers",
            "Fort Bend County",
            "Pecan Grove",
            "Brazoria",
            "Dickinson",
        ]
    },
    "harveyner_river_examples": {
        "en": [
            "Buffalo Bayou",
            "Brazos River",
            "Brays Bayou",
            "Cypress Creek",
            "Spring Creek",
            "Addicks",
            "Colorado River",
            "San Jacinto",
            "WhiteOakBayou",
            "Brazos",
        ]
    },
    "harveyner_road_examples": {
        "en": ["Barker Cypress", "I-10", "US 290 WB", "SH-71", "I-45", "105", "59", "Fry", "US-59 South", "Summer St"]
    },
}

src/tasks/harveyner/guidelines_gold.py
hitz-zentroa/GoLLIE
from typing import List

from ..utils_typing import Entity, dataclass


"""Entity definitions

The entity definitions are derived from the official harveyner guidelines:
https://aclanthology.org/2022.naacl-main.243.pdf

"""


@dataclass
class Point(Entity):
    """{harveyner_point}"""

    span: str  # {harveyner_point_examples}


class Area(Entity):
    """{harveyner_area}"""

    span: str  # {harveyner_area_examples}


class Road(Entity):
    """{harveyner_road}"""

    span: str  # {harveyner_road_examples}


class River(Entity):
    """{harveyner_river}"""

    span: str  # {harveyner_river_examples}


ENTITY_DEFINITIONS: List[Entity] = [
    Point,
    Area,
    Road,
    River,
]

# __all__ = list(map(str, [*ENTITY_DEFINITIONS]))

src/tasks/harveyner/prompts.py
hitz-zentroa/GoLLIE
from typing import Dict, List, Type

from src.tasks.harveyner.prompts import ENTITY_DEFINITIONS
from src.tasks.utils_scorer import SpanScorer
from src.tasks.utils_typing import Entity


class HarveyNEREntityScorer(SpanScorer):
    """CoNLL03 Entity identification and classification scorer."""

    valid_types: List[Type] = ENTITY_DEFINITIONS

    def __call__(self, reference: List[Entity], predictions: List[Entity]) -> Dict[str, Dict[str, float]]:
        output = super().__call__(reference, predictions)
        return {"entities": output["spans"]}

src/tasks/harveyner/scorer.py
hitz-zentroa/GoLLIE
from typing import List


def to_iob_encoding(tags: List[str]) -> List[str]:
    # From IOB2 or BILOU
    prev_tag_b: str = "O"
    prev_tag_t: str = ""

    for i in range(len(tags)):
        tag = tags[i]
        if tag == "O":
            prev_tag_b = "O"
            prev_tag_t = ""
        else:
            try:
                b, t = tag.split("-")
            except ValueError:
                raise ValueError(f"Error tag {tag}, unable to split the tag in 2 fields.")
            if (b == "B" or b == "U") and prev_tag_b != "O" and prev_tag_t == t:
                tags[i] = f"B-{t}"
            else:
                tags[i] = f"I-{t}"

            prev_tag_b = b
            prev_tag_t = t

    return tags


def to_iob2_encoding(tags: List[str]) -> List[str]:
    # From IOB or BILOU
    prev_tag_b: str = "O"
    prev_tag_t: str = ""
    for i in range(len(tags)):
        tag = tags[i]
        if tag == "O":
            prev_tag_b = "O"
            prev_tag_t = ""
        else:
            try:
                b, t = tag.split("-")
            except ValueError:
                raise ValueError(f"Error tag {tag}, unable to split the tag in 2 fields.")

            if (b == "B" or b == "U") or ((prev_tag_b == "O") or (prev_tag_t != "" and prev_tag_t != t)):
                tags[i] = f"B-{t}"
            else:
                tags[i] = f"I-{t}"

            prev_tag_b = b
            prev_tag_t = t

    return tags


def to_bilou_encoding(tags: List[str]) -> List[str]:
    # From IOB or IOB2

    prev_word_tag_tmp: str = ""
    for i in range(len(tags)):
        tag = tags[i]
        if tag == "O":
            if prev_word_tag_tmp != "":
                try:
                    prev_b, prev_t = prev_word_tag_tmp.split("-")
                except ValueError:
                    raise ValueError(f"Error in tag {prev_word_tag_tmp}, unable to split the tag in 2 fields.")

                if prev_b == "B":
                    tags[i - 1] = f"U-{prev_t}"
                else:
                    tags[i - 1] = f"L-{prev_t}"

            prev_word_tag_tmp: str = ""

        else:
            try:
                b, t = tag.split("-")
            except ValueError:
                raise ValueError(f"Error in tag {prev_word_tag_tmp}, unable to split the tag in 2 fields.")

            if prev_word_tag_tmp == "":
                if b == "U":
                    prev_word_tag_tmp = ""
                else:
                    prev_word_tag_tmp = f"B-{t}"

            else:
                try:
                    prev_b, prev_t = prev_word_tag_tmp.split("-")
                except ValueError:
                    raise ValueError(f"Error in tag {prev_word_tag_tmp}, unable to split the tag in 2 fields.")

                if b == "U":
                    if prev_b == "B":
                        tags[i - 1] = f"U-{prev_t}"
                    else:
                        tags[i - 1] = f"L-{prev_t}"

                    prev_word_tag_tmp = ""

                elif b == "B":
                    if prev_b == "B":
                        tags[i - 1] = f"U-{prev_t}"
                    else:
                        tags[i - 1] = f"L-{prev_t}"

                    prev_word_tag_tmp = f"B-{t}"

                else:
                    if prev_t != t:
                        if prev_b == "B":
                            tags[i - 1] = f"U-{prev_t}"
                        else:
                            tags[i - 1] = f"L-{prev_t}"

                        prev_word_tag_tmp = f"B-{t}"
                    else:
                        if prev_b == "B":
                            tags[i - 1] = f"B-{prev_t}"
                        else:
                            tags[i - 1] = f"I-{prev_t}"

                        prev_word_tag_tmp = f"I-{t}"

    if prev_word_tag_tmp != "":
        try:
            prev_b, prev_t = prev_word_tag_tmp.split("-")
        except ValueError:
            raise ValueError(f"Error in tag {prev_word_tag_tmp}, unable to split the tag in 2 fields.")

        if prev_b == "B":
            tags[-1] = f"U-{prev_t}"
        else:
            tags[-1] = f"L-{prev_t}"

    return tags


def rewrite_labels(labels, encoding: str = "iob2") -> List[str]:
    if encoding.lower() == "iob":
        return to_iob_encoding(labels)
    elif encoding.lower() == "iob2":
        return to_iob2_encoding(labels)
    elif encoding.lower() == "bilou":
        return to_bilou_encoding(labels)
    else:
        raise NotImplementedError(f"Encoding {encoding} not supported. Supported encodings [IOB,IOB2,BILOU]")

src/tasks/label_encoding.py
hitz-zentroa/GoLLIE
from . import data_loader, prompts, scorer


__all__ = ["data_loader", "prompts", "scorer"]

src/tasks/mitmovie/__init__.py
hitz-zentroa/GoLLIE
from typing import Dict, List, Tuple, Type, Union

from src.tasks.label_encoding import rewrite_labels
from src.tasks.mitmovie.guidelines import GUIDELINES
from src.tasks.mitmovie.guidelines_gold import EXAMPLES
from src.tasks.mitmovie.prompts import (
    ENTITY_DEFINITIONS,
    Actor,
    Character,
    Director,
    Entity,
    Genre,
    Plot,
    Rating,
    RatingsAverage,
    Review,
    Song,
    Tittle,
    Trailer,
    Year,
)

from ..utils_data import DatasetLoader, Sampler


def read_tsv(filepath) -> Tuple[List[List[str]], List[List[str]]]:
    """
    READ tsv file in MIT format (Label\tWord)
    Args:
        filepath: Path to the file
    Returns: List of words, List of labels
    """

    dataset_words: List[List[str]] = []
    dataset_labels: List[List[str]] = []
    with open(filepath, "r", encoding="utf-8") as f:
        words = []
        labels = []
        for line in f:
            line = line.strip()
            if line.startswith("-DOCSTART-") or line == "" or line == "\n":
                if words:
                    dataset_words.append(words)
                    dataset_labels.append(labels)
                    words = []
                    labels = []
            else:
                try:
                    label, word = line.split()
                except ValueError:
                    raise ValueError(f"Cannot path_or_split line: {line}")
                if word:
                    words.append(word)
                    labels.append(label)
        if words:
            dataset_words.append(words)
            dataset_labels.append(labels)

    print(f"Read {len(dataset_words)} sentences from {filepath}")

    dataset_labels = [rewrite_labels(labels, encoding="iob2") for labels in dataset_labels]

    return dataset_words, dataset_labels


def load_mit_tsv(
    path: str,
    ENTITY_TO_CLASS_MAPPING: Dict[str, Type[Entity]],
) -> Tuple[List[List[str]], List[List[Entity]]]:
    """
    Load the mit dataset from a tsv file
    Args:
        path (str): The path to the tsv file
        include_misc (bool): Whether to include the MISC entity type. Defaults to `True`.
    Returns:
        (List[str],List[Union[Location,Organization,Person,Miscellaneous]]): The text and the entities
    """
    dataset_sentences: List[List[str]] = []
    dataset_entities: List[List[Entity]] = []

    dataset_words, dataset_labels = read_tsv(path)

    for words, labels in zip(dataset_words, dataset_labels):
        # Some of the CoNLL02-03 datasets are in IOB1 format instead of IOB2,
        # we convert them to IOB2, so we don't have to deal with this later.
        labels = rewrite_labels(labels=labels, encoding="iob2")
        # Get labeled word spans
        spans = []
        for i, label in enumerate(labels):
            if label == "O":
                continue
            elif label.startswith("B-"):
                spans.append([label[2:], i, i + 1])
            elif label.startswith("I-"):
                spans[-1][2] += 1
            else:
                raise ValueError(f"Found an unexpected label: {label}")

        # Get entities
        entities = []
        for label, start, end in spans:
            entities.append(ENTITY_TO_CLASS_MAPPING[label](span=" ".join(words[start:end])))

        dataset_sentences.append(words)
        dataset_entities.append(entities)

    return dataset_sentences, dataset_entities


class MitMovieDatasetLoader(DatasetLoader):
    """
    A `DatasetLoader` for the MIT Movie dataset.

    Args:
        split (`str`):
            The split to load. Can be one of `train`, `validation` or `test`.

    Raises:
        `ValueError`:
            raised when a not defined value found.
    """

    ENTITY_TO_CLASS_MAPPING = None

    def __init__(self, path_or_split: str, **kwargs) -> None:
        self.ENTITY_TO_CLASS_MAPPING = {
            "ACTOR": Actor,
            "CHARACTER": Character,
            "DIRECTOR": Director,
            "GENRE": Genre,
            "PLOT": Plot,
            "RATING": Rating,
            "RATINGS_AVERAGE": RatingsAverage,
            "REVIEW": Review,
            "SONG": Song,
            "TITLE": Tittle,
            "TRAILER": Trailer,
            "YEAR": Year,
        }

        self.elements = {}

        dataset_words, dataset_entities = load_mit_tsv(
            path=path_or_split,
            ENTITY_TO_CLASS_MAPPING=self.ENTITY_TO_CLASS_MAPPING,
        )

        for id, (words, entities) in enumerate(zip(dataset_words, dataset_entities)):
            self.elements[id] = {
                "id": id,
                "doc_id": id,
                "text": " ".join(words),
                "entities": entities,
                "gold": entities,
            }


class MitMovieSampler(Sampler):
    """
    A data `Sampler` for the MIT Movie dataset.

    Args:
        dataset_loader (`MitMovieDatasetLoader`):
            The dataset loader that contains the data information.
        task (`str`, optional):
            The task to sample. It must be one of the following: NER, VER, RE, EE.
            Defaults to `None`.
        split (`str`, optional):
            The path_or_split to sample. It must be one of the following: "train", "dev" or
            "test". Depending on the path_or_split the sampling strategy differs. Defaults to
            `"train"`.
        parallel_instances (`Union[int, Tuple[int, int]]`, optional):
            The number of sentences sampled in parallel. Options:

                * **`int`**: The amount of elements that will be sampled in parallel.
                * **`tuple`**: The range of elements that will be sampled in parallel.

            Defaults to 1.
        max_guidelines (`int`, optional):
            The number of guidelines to append to the example at the same time. If `-1`
            is given then all the guidelines are appended. Defaults to `-1`.
        guideline_dropout (`float`, optional):
            The probability to dropout a guideline definition for the given example. This
            is only applied on training. Defaults to `0.0`.
        seed (`float`, optional):
            The seed to sample the examples. Defaults to `0`.
        prompt_template (`str`, optional):
            The path to the prompt template. Defaults to `"templates/prompt.txt"`.
        ensure_positives_on_train (bool, optional):
            Whether to ensure that the guidelines of annotated examples are not removed.
            Defaults to `False`.
        dataset_name (str, optional):
            The name of the dataset. Defaults to `None`.
        scorer (`str`, optional):
           The scorer class import string. Defaults to `None`.
        sample_only_gold_guidelines (`bool`, optional):
            Whether to sample only guidelines of present annotations. Defaults to `False`.
    """

    def __init__(
        self,
        dataset_loader: MitMovieDatasetLoader,
        task: str = None,
        split: str = "train",
        parallel_instances: Union[int, Tuple[int, int]] = 1,
        max_guidelines: int = -1,
        guideline_dropout: float = 0.0,
        seed: float = 0,
        prompt_template: str = "templates/prompt.txt",
        ensure_positives_on_train: bool = False,
        dataset_name: str = None,
        scorer: str = None,
        sample_only_gold_guidelines: bool = False,
        **kwargs,
    ) -> None:
        assert task in [
            "NER",
        ], f"Mit Movie only supports NER task. {task} is not supported."

        task_definitions, task_target = {
            "NER": (ENTITY_DEFINITIONS, "entities"),
        }[task]

        super().__init__(
            dataset_loader=dataset_loader,
            task=task,
            split=split,
            parallel_instances=parallel_instances,
            max_guidelines=max_guidelines,
            guideline_dropout=guideline_dropout,
            seed=seed,
            prompt_template=prompt_template,
            ensure_positives_on_train=ensure_positives_on_train,
            sample_only_gold_guidelines=sample_only_gold_guidelines,
            dataset_name=dataset_name,
            scorer=scorer,
            task_definitions=task_definitions,
            task_target=task_target,
            definitions=GUIDELINES,
            examples=EXAMPLES,
            **kwargs,
        )

src/tasks/mitmovie/data_loader.py
hitz-zentroa/GoLLIE
GUIDELINES = {
    "mit_actor": {
        "en": [
            "Refers to an individual who plays a role in a movie, including main cast members, supporting actors, and"
            " cameo appearances. If an individual is known for their work in television, theater, or other media but"
            " also appears in movies, they can be considered an Actor in this context."
        ]
    },
    "mit_character": {
        "en": [
            "Refers to a fictional or real individual portrayed in a movie's narrative or storyline. It includes main"
            " characters, supporting characters, antagonists, and any other roles that contribute to the plot or"
            " development of the movie. Characters can be human or non-human, and they represent the individuals with"
            " specific roles and identities within the movie's fictional universe."
        ]
    },
    "mit_director": {
        "en": [
            "Refers to an individual who is responsible for overseeing and guiding the artistic vision and creative"
            " elements of a movie. The director plays a pivotal role in shaping the film's narrative, performances,"
            " and overall execution. The Director entity includes the names of both film directors and television"
            " directors, as long as they are relevant to the movie context."
        ]
    },
    "mit_genre": {
        "en": [
            "Refers to a category or classification that characterizes the style, theme, or tone of a movie. Genres"
            " are used to identify and categorize movies based on their common characteristics and storytelling"
            " elements. Genres can be broad categories (e.g., Action, Comedy, Drama) or more specific sub-genres."
        ]
    },
    "mit_plot": {
        "en": [
            "A plot element or theme is a recurring concept, event, or motif that plays a significant role in the"
            " development of a story, narrative, or work of fiction. For example what happens in a movie (murder, love"
            " at first sight, cars that fly, horse racing), the type of characters in the movie (vampires, zombies,"
            " fugitives, cats), where the movie was recorded (British, American), among others."
        ]
    },
    "mit_rating": {
        "en": [
            "Refers to the official classification or content rating assigned to a movie by a recognized film rating"
            " organization. Ratings provide guidance to audiences and parents about the movie's suitability based on"
            " its content, themes, and age-appropriateness. Ratings can include various classifications such as MPAA"
            " (Motion Picture Association) ratings like R-Rated, PG-13, PG, or other regional and international rating"
            " systems."
        ]
    },
    "mit_ratings_average": {
        "en": [
            "Refers to the collective or average assessment of a movie's quality, popularity, or critical acclaim"
            " based on reviews, ratings, or awards from various sources. This entity includes general descriptions of"
            " the overall reception of the movie, such as 3-star rating, award-winning, critically heralded, box"
            " office hit, or similar expressions that provide an indication of the movie's reception."
        ]
    },
    "mit_review": {
        "en": [
            "Refers to an evaluation or assessment of a movie's quality, performance, or other aspects. It includes"
            " expressions or phrases that indicate personal opinions, critiques, or judgments about a movie's overall"
            " merit, such as greatest, excellent, disappointing, and similar evaluative terms. It also includes"
            " expressions such as five stars rated, lowest rated or recommended, which provide an indication of the"
            " reviewer's assessment of the movie."
        ]
    },
    "mit_song": {
        "en": [
            "Refers to a musical composition or piece of music that is associated with a movie. This includes"
            " individual songs featured in the movie's soundtrack, as well as musical groups and singers who perform"
            " those songs. The Song entity class encompasses both original songs created specifically for the movie"
            " and pre-existing songs used in the movie's soundtrack."
        ]
    },
    "mit_title": {
        "en": [
            "Refers to the name or label that uniquely identifies a specific movie or film production. It includes the"
            " official title of the movie, which is used for identification, promotion, and reference purposes. The"
            " Title entity class encompasses the full names of movies, as well as any abbreviations or alternative"
            " titles that are commonly associated with the movie."
        ]
    },
    "mit_trailer": {
        "en": [
            "Refers to a short promotional video or preview of a movie that is designed to generate interest and"
            " excitement among the audience before the movie's release. This entity also includes mentions to"
            " trailers, such as 'trailer','preview', 'teaser' or 'highlights'."
        ]
    },
    "mit_year": {
        "en": [
            "Refers to a numerical representation of the calendar year in which a specific movie was released or"
            " produced. The Year entity class encompasses four-digit numerical values that indicate the specific year"
            " of a movie's release or production."
        ]
    },
}

src/tasks/mitmovie/guidelines.py
hitz-zentroa/GoLLIE
GUIDELINES = {
    "mit_actor": {
        "en": [
            "Refers to an individual who plays a role in a movie, including main cast members, supporting actors, and"
            " cameo appearances. If an individual is known for their work in television, theater, or other media but"
            " also appears in movies, they can be considered an Actor in this context."
        ]
    },
    "mit_character": {
        "en": [
            "Refers to a fictional or real individual portrayed in a movie's narrative or storyline. It includes main"
            " characters, supporting characters, antagonists, and any other roles that contribute to the plot or"
            " development of the movie. Characters can be human or non-human, and they represent the individuals with"
            " specific roles and identities within the movie's fictional universe."
        ]
    },
    "mit_director": {
        "en": [
            "Refers to an individual who is responsible for overseeing and guiding the artistic vision and creative"
            " elements of a movie. The director plays a pivotal role in shaping the film's narrative, performances,"
            " and overall execution. The Director entity includes the names of both film directors and television"
            " directors, as long as they are relevant to the movie context."
        ]
    },
    "mit_genre": {
        "en": [
            "Refers to a category or classification that characterizes the style, theme, or tone of a movie. Genres"
            " are used to identify and categorize movies based on their common characteristics and storytelling"
            " elements. Genres can be broad categories (e.g., Action, Comedy, Drama) or more specific sub-genres."
        ]
    },
    "mit_plot": {
        "en": [
            "A plot element or theme is a recurring concept, event, or motif that plays a significant role in the"
            " development of a story, narrative, or work of fiction. For example what happens in a movie (murder, love"
            " at first sight, cars that fly, horse racing), the type of characters in the movie (vampires, zombies,"
            " fugitives, cats), where the movie was recorded (British, American), among others."
        ]
    },
    "mit_rating": {
        "en": [
            "Refers to the official classification or content rating assigned to a movie by a recognized film rating"
            " organization. Ratings provide guidance to audiences and parents about the movie's suitability based on"
            " its content, themes, and age-appropriateness. Ratings can include various classifications such as MPAA"
            " (Motion Picture Association) ratings like R-Rated, PG-13, PG, or other regional and international rating"
            " systems."
        ]
    },
    "mit_ratings_average": {
        "en": [
            "Refers to the collective or average assessment of a movie's quality, popularity, or critical acclaim"
            " based on reviews, ratings, or awards from various sources. This entity includes general descriptions of"
            " the overall reception of the movie, such as 3-star rating, award-winning, critically heralded, box"
            " office hit, or similar expressions that provide an indication of the movie's reception."
        ]
    },
    "mit_review": {
        "en": [
            "Refers to an evaluation or assessment of a movie's quality, performance, or other aspects. It includes"
            " expressions or phrases that indicate personal opinions, critiques, or judgments about a movie's overall"
            " merit, such as greatest, excellent, disappointing, and similar evaluative terms. It also includes"
            " expressions such as five stars rated, lowest rated or recommended, which provide an indication of the"
            " reviewer's assessment of the movie."
        ]
    },
    "mit_song": {
        "en": [
            "Refers to a musical composition or piece of music that is associated with a movie. This includes"
            " individual songs featured in the movie's soundtrack, as well as musical groups and singers who perform"
            " those songs. The Song entity class encompasses both original songs created specifically for the movie"
            " and pre-existing songs used in the movie's soundtrack."
        ]
    },
    "mit_title": {
        "en": [
            "Refers to the name or label that uniquely identifies a specific movie or film production. It includes the"
            " official title of the movie, which is used for identification, promotion, and reference purposes. The"
            " Title entity class encompasses the full names of movies, as well as any abbreviations or alternative"
            " titles that are commonly associated with the movie."
        ]
    },
    "mit_trailer": {
        "en": [
            "Refers to a short promotional video or preview of a movie that is designed to generate interest and"
            " excitement among the audience before the movie's release. This entity also includes mentions to"
            " trailers, such as 'trailer','preview', 'teaser' or 'highlights'."
        ]
    },
    "mit_year": {
        "en": [
            "Refers to a numerical representation of the calendar year in which a specific movie was released or"
            " produced. The Year entity class encompasses four-digit numerical values that indicate the specific year"
            " of a movie's release or production."
        ]
    },
}

EXAMPLES = {
    "mit_actor_examples": {
        "en": [
            "johnny depp",
            "brad pitt",
            "tom hanks",
            "tom cruise",
            "clint eastwood",
            "leonardo dicaprio",
            "will smith",
            "harrison ford",
            "bruce willis",
            "julia roberts",
        ]
    },
    "mit_character_examples": {
        "en": [
            "james bond",
            "indiana jones",
            "jason bourne",
            "harry potter",
            "batman",
            "shrek",
            "buzz lightyear",
            "donkey",
            "jason",
            "jack sparrow",
        ]
    },
    "mit_director_examples": {
        "en": [
            "steven spielberg",
            "james cameron",
            "directed",
            "clint eastwood",
            "alfred hitchcock",
            "orson welles",
            "tim burton",
            "stanley kubrick",
            "george lucas",
            "michael bay",
        ]
    },
    "mit_genre_examples": {
        "en": [
            "action",
            "comedy",
            "horror",
            "fantasy",
            "science fiction",
            "drama",
            "thriller",
            "adventure",
            "documentary",
            "war",
        ]
    },
    "mit_plot_examples": {
        "en": [
            "murder",
            "vampires",
            "ghosts",
            "levitation",
            "love at first sight",
            "robot",
            "vampire",
            "war",
            "police",
            "wilderness",
        ]
    },
    "mit_rating_examples": {
        "en": ["r", "pg 13", "pg", "g", "nc 17", "unrated", "r rated", "g rated", "rated r", "rated"]
    },
    "mit_ratings_average_examples": {
        "en": [
            "average",
            "ok",
            "mediocre",
            "highly rated",
            "critically acclaimed",
            "very good",
            "really good",
            "must see",
            "four",
            "good",
        ]
    },
    "mit_review_examples": {
        "en": [
            "good",
            "best",
            "must see",
            "classic",
            "funniest",
            "popular",
            "oscar",
            "most popular",
            "top rated",
            "funny",
        ]
    },
    "mit_song_examples": {
        "en": [
            "soundtrack",
            "my heart will go on",
            "theme song",
            "i will always love you",
            "john williams",
            "somewhere out there",
            "soundtracks",
            "score",
            "my favorite things",
            "linkin park",
        ]
    },
    "mit_title_examples": {
        "en": [
            "inception",
            "the shining",
            "wall e",
            "star wars",
            "titanic",
            "harry potter",
            "shrek",
            "pretty woman",
            "lord of the rings",
            "pirates of the caribbean",
        ]
    },
    "mit_trailer_examples": {
        "en": [
            "trailer",
            "preview",
            "trailers",
            "shortest trailer",
            "star wars",
            "theme song",
            "review",
            "previews",
            "clips",
            "shogun movie trailer",
        ]
    },
    "mit_year_examples": {
        "en": ["1950 s", "1970 s", "2010 s", "1980 s", "1990 s", "1960 s", "1990s", "2010", "2000", "1990"]
    },
}

src/tasks/mitmovie/guidelines_gold.py
hitz-zentroa/GoLLIE
from typing import List

from ..utils_typing import Entity, dataclass


"""Entity definitions

In the absence of public guidelines, the guidelines have been defined by the CoLLIE authors.
Dataset available at: https://groups.csail.mit.edu/sls/downloads/movie/

"""


@dataclass
class Actor(Entity):
    """{mit_actor}"""

    span: str  # {mit_actor_examples}


@dataclass
class Character(Entity):
    """{mit_character}"""

    span: str  # {mit_character_examples}


@dataclass
class Director(Entity):
    """{mit_director}"""

    span: str  # {mit_director_examples}


@dataclass
class Genre(Entity):
    """{mit_genre}"""

    span: str  # {mit_genre_examples}


@dataclass
class Plot(Entity):
    """{mit_plot}"""

    span: str  # {mit_plot_examples}


@dataclass
class Rating(Entity):
    """{mit_rating}"""

    span: str  # {mit_rating_examples}


@dataclass
class RatingsAverage(Entity):
    """{mit_ratings_average}"""

    span: str  # {mit_ratings_average_examples}


@dataclass
class Review(Entity):
    """{mit_review}"""

    span: str  # {mit_review_examples}


@dataclass
class Song(Entity):
    """{mit_song}"""

    span: str  # {mit_song_examples}


@dataclass
class Tittle(Entity):
    """{mit_title}"""

    span: str  # {mit_title_examples}


@dataclass
class Trailer(Entity):
    """{mit_trailer}"""

    span: str  # {mit_trailer_examples}


@dataclass
class Year(Entity):
    """{mit_year}"""

    span: str  # {mit_year_examples}


ENTITY_DEFINITIONS: List[Entity] = [
    Actor,
    Character,
    Director,
    Genre,
    Plot,
    Rating,
    RatingsAverage,
    Review,
    Song,
    Tittle,
    Trailer,
    Year,
]

# __all__ = list(map(str, [*ENTITY_DEFINITIONS]))

src/tasks/mitmovie/prompts.py
hitz-zentroa/GoLLIE
from typing import Dict, List, Type

from src.tasks.mitmovie.prompts import ENTITY_DEFINITIONS
from src.tasks.utils_scorer import SpanScorer
from src.tasks.utils_typing import Entity


class MitMovieEntityScorer(SpanScorer):
    """CoNLL03 Entity identification and classification scorer."""

    valid_types: List[Type] = ENTITY_DEFINITIONS

    def __call__(self, reference: List[Entity], predictions: List[Entity]) -> Dict[str, Dict[str, float]]:
        output = super().__call__(reference, predictions)
        return {"entities": output["spans"]}

src/tasks/mitmovie/scorer.py
hitz-zentroa/GoLLIE
from . import data_loader, prompts, scorer


__all__ = ["data_loader", "prompts", "scorer"]

src/tasks/mitrestaurant/__init__.py
hitz-zentroa/GoLLIE
from typing import Tuple, Union

from src.tasks.mitmovie.data_loader import load_mit_tsv
from src.tasks.mitrestaurant.guidelines import GUIDELINES
from src.tasks.mitrestaurant.guidelines_gold import EXAMPLES
from src.tasks.mitrestaurant.prompts import (
    ENTITY_DEFINITIONS,
    Amenity,
    Cuisine,
    Dish,
    Hours,
    Location,
    Price,
    Rating,
    RestaurantName,
)

from ..utils_data import DatasetLoader, Sampler


class MitRestaurantDatasetLoader(DatasetLoader):
    """
    A `DatasetLoader` for the MIT Restaurant dataset.

    Args:
        split (`str`):
            The split to load. Can be one of `train`, `validation` or `test`.

    Raises:
        `ValueError`:
            raised when a not defined value found.
    """

    ENTITY_TO_CLASS_MAPPING = None

    def __init__(self, path_or_split: str, **kwargs) -> None:
        self.ENTITY_TO_CLASS_MAPPING = {
            "Rating": Rating,
            "Amenity": Amenity,
            "Location": Location,
            "Restaurant_Name": RestaurantName,
            "Price": Price,
            "Hours": Hours,
            "Dish": Dish,
            "Cuisine": Cuisine,
        }

        self.elements = {}

        dataset_words, dataset_entities = load_mit_tsv(
            path=path_or_split,
            ENTITY_TO_CLASS_MAPPING=self.ENTITY_TO_CLASS_MAPPING,
        )

        for id, (words, entities) in enumerate(zip(dataset_words, dataset_entities)):
            self.elements[id] = {
                "id": id,
                "doc_id": id,
                "text": " ".join(words),
                "entities": entities,
                "gold": entities,
            }


class MitRestaurantSampler(Sampler):
    """
    A data `Sampler` for the MIT Restaurant dataset.

    Args:
        dataset_loader (`MitRestaurantDatasetLoader`):
            The dataset loader that contains the data information.
        task (`str`, optional):
            The task to sample. It must be one of the following: NER, VER, RE, EE.
            Defaults to `None`.
        split (`str`, optional):
            The path_or_split to sample. It must be one of the following: "train", "dev" or
            "test". Depending on the path_or_split the sampling strategy differs. Defaults to
            `"train"`.
        parallel_instances (`Union[int, Tuple[int, int]]`, optional):
            The number of sentences sampled in parallel. Options:

                * **`int`**: The amount of elements that will be sampled in parallel.
                * **`tuple`**: The range of elements that will be sampled in parallel.

            Defaults to 1.
        max_guidelines (`int`, optional):
            The number of guidelines to append to the example at the same time. If `-1`
            is given then all the guidelines are appended. Defaults to `-1`.
        guideline_dropout (`float`, optional):
            The probability to dropout a guideline definition for the given example. This
            is only applied on training. Defaults to `0.0`.
        seed (`float`, optional):
            The seed to sample the examples. Defaults to `0`.
        prompt_template (`str`, optional):
            The path to the prompt template. Defaults to `"templates/prompt.txt"`.
        ensure_positives_on_train (bool, optional):
            Whether to ensure that the guidelines of annotated examples are not removed.
            Defaults to `False`.
        dataset_name (str, optional):
            The name of the dataset. Defaults to `None`.
        scorer (`str`, optional):
           The scorer class import string. Defaults to `None`.
        sample_only_gold_guidelines (`bool`, optional):
            Whether to sample only guidelines of present annotations. Defaults to `False`.
    """

    def __init__(
        self,
        dataset_loader: MitRestaurantDatasetLoader,
        task: str = None,
        split: str = "train",
        parallel_instances: Union[int, Tuple[int, int]] = 1,
        max_guidelines: int = -1,
        guideline_dropout: float = 0.0,
        seed: float = 0,
        prompt_template: str = "templates/prompt.txt",
        ensure_positives_on_train: bool = False,
        dataset_name: str = None,
        scorer: str = None,
        sample_only_gold_guidelines: bool = False,
        **kwargs,
    ) -> None:
        assert task in [
            "NER",
        ], f"Mit Restaurant only supports NER task. {task} is not supported."

        task_definitions, task_target = {
            "NER": (ENTITY_DEFINITIONS, "entities"),
        }[task]

        super().__init__(
            dataset_loader=dataset_loader,
            task=task,
            split=split,
            parallel_instances=parallel_instances,
            max_guidelines=max_guidelines,
            guideline_dropout=guideline_dropout,
            seed=seed,
            prompt_template=prompt_template,
            ensure_positives_on_train=ensure_positives_on_train,
            sample_only_gold_guidelines=sample_only_gold_guidelines,
            dataset_name=dataset_name,
            scorer=scorer,
            task_definitions=task_definitions,
            task_target=task_target,
            definitions=GUIDELINES,
            examples=EXAMPLES,
            **kwargs,
        )

src/tasks/mitrestaurant/data_loader.py
hitz-zentroa/GoLLIE
GUIDELINES = {
    "mit_rating": {
        "en": [
            "Refers to any specific qualitative or quantitative evaluation of a restaurant's quality, performance, or"
            " popularity, such as star ratings (5 start, 2-star), adjectives (e.g., good, excellent, best), and"
            " phrases that describe the restaurant's overall rating or customer experience (e.g., busy, top rated)."
        ]
    },
    "mit_amenity": {
        "en": [
            "Refers to any specific feature, service, or facility offered by a restaurant to enhance the customer"
            " experience. It includes various expressions , such as accessibility options (handicapped accessible),"
            " special events (karaoke night), promotions (happy hour), or reservation requirements (do not require"
            " reservation), family-friendly services (kid friendly, kids eat free), smoking areas, parking facilities"
            " (valet parking), offerings (serve alcohol, buffet, takeout, free refills on soda), restaurant type"
            " (romantic, sports bar,..)  and other special or unique features that distinguish the restaurant."
        ]
    },
    "mit_location": {
        "en": [
            "Refers to any specific place, area, city, neighborhood, or geographical location mentioned in the context"
            " of a restaurant. It includes various location expressions, such as city names (e.g., San Diego,Oakland),"
            " general areas (e.g., in town, in the area), relative positions (e.g., closest to my area, near me), or"
            " any other references to a place associated with the restaurant."
        ]
    },
    "mit_restaurantname": {
        "en": [
            "Refers to the specific names or titles of restaurants mentioned in the text. It includes names of"
            " individual restaurants (hampton diner), chains, or franchises (mcdonalds) that uniquely identify a food"
            " establishment."
        ]
    },
    "mit_price": {
        "en": [
            "Refers to any specific information or descriptors related to the cost or pricing of the restaurant's menu"
            " items. It includes various expressions, such as numerical price ranges (e.g., prices no more than"
            " 15.00), qualitative price indicators (e.g., cheap, moderately priced), and any other price-related"
            " descriptions (e.g., affordable, expensive)."
        ]
    },
    "mit_hours": {
        "en": [
            "Refers to any specific time-related information related to the operating hours or availability of the"
            " restaurant. It includes expressions denoting specific times (e.g., 8pm), time intervals (e.g., until"
            " 11pm), and general time descriptions (e.g., late at night, open right now)."
        ]
    },
    "mit_dish": {
        "en": [
            "Refers to any specific food item, meal, or culinary creation offered by the restaurant. It includes"
            " various food dishes, beverages, desserts, condiments, toppings, and other menu items that can be ordered"
            " and served."
        ]
    },
    "mit_cuisine": {
        "en": [
            "Refers to any specific style or type of food associated with a particular culture, region, or culinary"
            " tradition. It includes various types of cuisines, such as chinese, italian, mexican, fast food,"
            " seafood, etc."
        ]
    },
}

src/tasks/mitrestaurant/guidelines.py
hitz-zentroa/GoLLIE
GUIDELINES = {
    "mit_rating": {
        "en": [
            "Refers to any specific qualitative or quantitative evaluation of a restaurant's quality, performance, or"
            " popularity, such as star ratings (5 start, 2-star), adjectives (e.g., good, excellent, best), and"
            " phrases that describe the restaurant's overall rating or customer experience (e.g., busy, top rated)."
        ]
    },
    "mit_amenity": {
        "en": [
            "Refers to any specific feature, service, or facility offered by a restaurant to enhance the customer"
            " experience. It includes various expressions , such as accessibility options (handicapped accessible),"
            " special events (karaoke night), promotions (happy hour), or reservation requirements (do not require"
            " reservation), family-friendly services (kid friendly, kids eat free), smoking areas, parking facilities"
            " (valet parking), offerings (serve alcohol, buffet, takeout, free refills on soda), restaurant type"
            " (romantic, sports bar,..)  and other special or unique features that distinguish the restaurant."
        ]
    },
    "mit_location": {
        "en": [
            "Refers to any specific place, area, city, neighborhood, or geographical location mentioned in the context"
            " of a restaurant. It includes various location expressions, such as city names (e.g., San Diego,Oakland),"
            " general areas (e.g., in town, in the area), relative positions (e.g., closest to my area, near me), or"
            " any other references to a place associated with the restaurant."
        ]
    },
    "mit_restaurantname": {
        "en": [
            "Refers to the specific names or titles of restaurants mentioned in the text. It includes names of"
            " individual restaurants (hampton diner), chains, or franchises (mcdonalds) that uniquely identify a food"
            " establishment."
        ]
    },
    "mit_price": {
        "en": [
            "Refers to any specific information or descriptors related to the cost or pricing of the restaurant's menu"
            " items. It includes various expressions, such as numerical price ranges (e.g., prices no more than"
            " 15.00), qualitative price indicators (e.g., cheap, moderately priced), and any other price-related"
            " descriptions (e.g., affordable, expensive)."
        ]
    },
    "mit_hours": {
        "en": [
            "Refers to any specific time-related information related to the operating hours or availability of the"
            " restaurant. It includes expressions denoting specific times (e.g., 8pm), time intervals (e.g., until"
            " 11pm), and general time descriptions (e.g., late at night, open right now)."
        ]
    },
    "mit_dish": {
        "en": [
            "Refers to any specific food item, meal, or culinary creation offered by the restaurant. It includes"
            " various food dishes, beverages, desserts, condiments, toppings, and other menu items that can be ordered"
            " and served."
        ]
    },
    "mit_cuisine": {
        "en": [
            "Refers to any specific style or type of food associated with a particular culture, region, or culinary"
            " tradition. It includes various types of cuisines, such as chinese, italian, mexican, fast food,"
            " seafood, etc."
        ]
    },
}

EXAMPLES = {
    "mit_rating_examples": {
        "en": [
            "good",
            "best",
            "nice",
            "great",
            "five star",
            "highest rated",
            "highly rated",
            "5 star",
            "good reviews",
            "4 star",
        ]
    },
    "mit_amenity_examples": {
        "en": [
            "kid friendly",
            "carry out",
            "take out",
            "buffet",
            "outdoor dining",
            "romantic",
            "valet parking",
            "parking",
            "reservations",
            "smoking area",
        ]
    },
    "mit_location_examples": {
        "en": [
            "nearby",
            "closest",
            "nearest",
            "around here",
            "in the area",
            "close by",
            "near here",
            "within 5 miles",
            "local",
            "in town",
        ]
    },
    "mit_restaurantname_examples": {
        "en": [
            "mcdonalds",
            "burger king",
            "wendys",
            "taco bell",
            "olive garden",
            "dennys",
            "subway",
            "pizza hut",
            "applebees",
            "red lobster",
        ]
    },
    "mit_price_examples": {
        "en": [
            "cheap",
            "expensive",
            "reasonably",
            "cheapest",
            "affordable",
            "inexpensive",
            "moderately",
            "good",
            "reasonable",
            "great",
        ]
    },
    "mit_hours_examples": {
        "en": [
            "lunch",
            "open",
            "dinner",
            "breakfast",
            "late night",
            "tonight",
            "open late",
            "close",
            "open right now",
            "still open",
        ]
    },
    "mit_dish_examples": {
        "en": ["pizza", "sushi", "steak", "burger", "ice cream", "tacos", "hamburger", "sandwich", "burgers", "beer"]
    },
    "mit_cuisine_examples": {
        "en": [
            "chinese",
            "italian",
            "mexican",
            "fast food",
            "seafood",
            "indian",
            "japanese",
            "pizza",
            "sushi",
            "vegan",
        ]
    },
}

src/tasks/mitrestaurant/guidelines_gold.py
hitz-zentroa/GoLLIE
from typing import List

from ..utils_typing import Entity, dataclass


"""Entity definitions

In the absence of public guidelines, the guidelines have been defined by the CoLLIE authors.
Dataset available at: https://groups.csail.mit.edu/sls/downloads/restaurant/

"""


@dataclass
class Rating(Entity):
    """{mit_rating}"""

    span: str  # {mit_rating_examples}


@dataclass
class Amenity(Entity):
    """{mit_amenity}"""

    span: str  # {mit_amenity_examples}


@dataclass
class Location(Entity):
    """{mit_location}"""

    span: str  # {mit_location_examples}


@dataclass
class RestaurantName(Entity):
    """{mit_restaurantname}"""

    span: str  # {mit_restaurantname_examples}


@dataclass
class Price(Entity):
    """{mit_price}"""

    span: str  # {mit_price_examples}


@dataclass
class Hours(Entity):
    """{mit_hours}"""

    span: str  # {mit_hours_examples}


@dataclass
class Dish(Entity):
    """{mit_dish}"""

    span: str  # {mit_dish_examples}


@dataclass
class Cuisine(Entity):
    """{mit_cuisine}"""

    span: str  # {mit_cuisine_examples}


ENTITY_DEFINITIONS: List[Entity] = [Rating, Amenity, Location, RestaurantName, Price, Hours, Dish, Cuisine]

# __all__ = list(map(str, [*ENTITY_DEFINITIONS]))

src/tasks/mitrestaurant/prompts.py
hitz-zentroa/GoLLIE
from typing import Dict, List, Type

from src.tasks.mitrestaurant.prompts import ENTITY_DEFINITIONS
from src.tasks.utils_scorer import SpanScorer
from src.tasks.utils_typing import Entity


class MitRestaurantEntityScorer(SpanScorer):
    """CoNLL03 Entity identification and classification scorer."""

    valid_types: List[Type] = ENTITY_DEFINITIONS

    def __call__(self, reference: List[Entity], predictions: List[Entity]) -> Dict[str, Dict[str, float]]:
        output = super().__call__(reference, predictions)
        return {"entities": output["spans"]}

src/tasks/mitrestaurant/scorer.py
hitz-zentroa/GoLLIE
from typing import Dict, List, Type

from ..utils_typing import Span


class Person(Span):
    """{multiconer_person}"""

    span: str


class Group(Span):
    """{multiconer_group}"""

    span: str


class Location(Span):
    """{multiconer_location}"""

    span: str


class CreativeWork(Span):
    """{multiconer_creative_work}"""

    span: str


class Product(Span):
    """{multiconer_product}"""

    span: str


class Medical(Span):
    """{multiconer_medical}"""

    span: str


COARSE_ENTITY_DEFINITIONS: List[Type] = [Person, Group, Location, CreativeWork, Product, Medical]


class Scientist(Person):
    """{multiconer_scientist}"""

    span: str


class Artist(Person):
    """{multiconer_artist}"""

    span: str


class Athlete(Person):
    """{multiconer_athlete}"""

    span: str


class Politician(Person):
    """{multiconer_politician}"""

    span: str


class Cleric(Person):
    """{multiconer_cleric}"""

    span: str


class SportsManager(Person):
    """{multiconer_sports_manager}"""

    span: str


class OtherPerson(Person):
    """{multiconer_other_person}"""

    span: str


class MusicalGroup(Group):
    """{multiconer_musical_group}"""

    span: str


class PublicCorp(Group):
    """{multiconer_public_corp}"""

    span: str


class PrivateCorp(Group):
    """{multiconer_private_corp}"""

    span: str


class AerospaceManufacturer(Group):
    """{multiconer_aerospace_manufacturer}"""

    span: str


class SportsGroup(Group):
    """{multiconer_sports_group}"""

    span: str


class OtherGroup(Group):
    """{multiconer_other_group}"""

    span: str


class VisualWork(CreativeWork):
    """{multiconer_visual_work}"""

    span: str


class MusicalWork(CreativeWork):
    """{multiconer_musical_work}"""

    span: str


class WrittenWork(CreativeWork):
    """{multiconer_written_work}"""

    span: str


class ArtWork(CreativeWork):
    """{multiconer_art_work}"""

    span: str


class Software(CreativeWork):
    """{multiconer_software}"""

    span: str


class Facility(Location):
    """{multiconer_facility}"""

    span: str


class HumanSettlement(Location):
    """{multiconer_human_settlement}"""

    span: str


class Station(Location):
    """{multiconer_station}"""

    span: str


class OtherLocation(Location):
    """{multiconer_other_location}"""

    span: str


class Clothing(Product):
    """{multiconer_clothing}"""

    span: str


class Vehicle(Product):
    """{multiconer_vehicle}"""

    span: str


class Food(Product):
    """{multiconer_food}"""

    span: str


class Drink(Product):
    """{multiconer_drink}"""

    span: str


class OtherProduct(Product):
    """{multiconer_other_product}"""

    span: str


class MedicationOrVaccine(Medical):
    """{multiconer_medication_or_vaccine}"""

    span: str


class MedicalProcedure(Medical):
    """{multiconer_medical_procedure}"""

    span: str


class AnatomicalStructure(Medical):
    """{multiconer_anatomical_structure}"""

    span: str


class Symptom(Medical):
    """{multiconer_symtpom}"""

    span: str


class Disease(Medical):
    """{multiconer_disease}"""

    span: str


ENTITY_DEFINITIONS: List[Type] = [
    Scientist,
    Artist,
    Athlete,
    Politician,
    Cleric,
    SportsManager,
    OtherPerson,
    MusicalGroup,
    PublicCorp,
    PrivateCorp,
    AerospaceManufacturer,
    SportsGroup,
    OtherGroup,
    VisualWork,
    MusicalWork,
    WrittenWork,
    ArtWork,
    Software,
    Facility,
    HumanSettlement,
    Station,
    OtherLocation,
    Clothing,
    Vehicle,
    Food,
    Drink,
    OtherProduct,
    MedicationOrVaccine,
    MedicalProcedure,
    AnatomicalStructure,
    Symptom,
    Disease,
]

FINE_TO_COARSE_ENTITIES: Dict[Type, Type] = {_def: _def.__base__ for _def in ENTITY_DEFINITIONS}
COARSE_TO_FINE_ENTITIES: Dict[Type, List[Type]] = {}
for fine, coarse in FINE_TO_COARSE_ENTITIES.items():
    if coarse not in COARSE_TO_FINE_ENTITIES:
        COARSE_TO_FINE_ENTITIES[coarse] = []
    COARSE_TO_FINE_ENTITIES[coarse].append(fine)

src/tasks/multiconer2/prompts.py
hitz-zentroa/GoLLIE
from . import data_loader, prompts, scorer


__all__ = ["data_loader", "prompts", "scorer"]

src/tasks/multinerd/__init__.py
hitz-zentroa/GoLLIE
from typing import Dict, List, Tuple, Type, Union

from src.tasks.label_encoding import rewrite_labels
from src.tasks.multinerd.guidelines import GUIDELINES
from src.tasks.multinerd.guidelines_gold import EXAMPLES
from src.tasks.multinerd.prompts import (
    ENTITY_DEFINITIONS,
    Animal,
    Biological,
    Celestial,
    Disease,
    Event,
    Food,
    Instrument,
    Location,
    Media,
    Mythological,
    Organization,
    Person,
    Plant,
    Time,
    Vehicle,
)

from ..utils_data import DatasetLoader, Sampler
from ..utils_typing import Entity


def get_multinerd_hf(
    split: str,
    language: str,
    ENTITY_TO_CLASS_MAPPING: Dict[str, Type[Entity]],
) -> Tuple[List[List[str]], List[List[Entity]]]:
    """
    Get the Multi-NERD dataset from the huggingface datasets library
    Args:
        split (str): The path_or_split to load. Can be one of `train`, `validation` or `test`.
        language (str): The language to load. Can be one of Chinese (zh), Dutch (nl), English (en), French (fr),
                        German (de), Italian (it), Polish (pl), Portuguese (pt), Russian (ru), Spanish (es).
    Returns:
        (List[str],List[Entity]): The text and the entities
    """
    from datasets import load_dataset

    dataset = load_dataset("Babelscape/multinerd")
    # Dirty fix to prevent errors reading the labels
    label2id = {
        "O": 0,
        "B-PER": 1,
        "I-PER": 2,
        "B-ORG": 3,
        "I-ORG": 4,
        "B-LOC": 5,
        "I-LOC": 6,
        "B-ANIM": 7,
        "I-ANIM": 8,
        "B-BIO": 9,
        "I-BIO": 10,
        "B-CEL": 11,
        "I-CEL": 12,
        "B-DIS": 13,
        "I-DIS": 14,
        "B-EVE": 15,
        "I-EVE": 16,
        "B-FOOD": 17,
        "I-FOOD": 18,
        "B-INST": 19,
        "I-INST": 20,
        "B-MEDIA": 21,
        "I-MEDIA": 22,
        "B-MYTH": 23,
        "I-MYTH": 24,
        "B-PLANT": 25,
        "I-PLANT": 26,
        "B-TIME": 27,
        "I-TIME": 28,
        "B-VEHI": 29,
        "I-VEHI": 30,
    }

    id2label = {v: k for k, v in label2id.items()}

    dataset_sentences: List[List[str]] = []
    dataset_entities: List[List[Entity]] = []

    for example in dataset[split]:
        lang = example["lang"]
        if lang != language:
            continue
        words = example["tokens"]
        # Ensure IOB2 encoding
        labels = rewrite_labels(labels=[id2label[label] for label in example["ner_tags"]], encoding="iob2")

        # Get labeled word spans
        spans = []
        for i, label in enumerate(labels):
            if label == "O":
                continue
            elif label.startswith("B-"):
                spans.append([label[2:], i, i + 1])
            elif label.startswith("I-"):
                spans[-1][2] += 1
            else:
                raise ValueError(f"Found an unexpected label: {label}")

        # Get entities
        entities = []
        for label, start, end in spans:
            entities.append(ENTITY_TO_CLASS_MAPPING[label](span=" ".join(words[start:end])))

        dataset_sentences.append(words)
        dataset_entities.append(entities)

    return dataset_sentences, dataset_entities


class MultinerdDatasetLoader(DatasetLoader):
    """
    A `DatasetLoader` for the Multi-Nerd dataset.

    Args:
        split (`str`):
            The split to load. Can be one of `train`, `validation` or `test`.

    Raises:
        `ValueError`:
            raised when a not defined value found.
    """

    ENTITY_TO_CLASS_MAPPING = None

    def __init__(self, path_or_split: str, **kwargs) -> None:
        self.ENTITY_TO_CLASS_MAPPING = {
            "PER": Person,
            "ORG": Organization,
            "LOC": Location,
            "ANIM": Animal,
            "BIO": Biological,
            "CEL": Celestial,
            "DIS": Disease,
            "EVE": Event,
            "FOOD": Food,
            "INST": Instrument,
            "MEDIA": Media,
            "MYTH": Mythological,
            "PLANT": Plant,
            "TIME": Time,
            "VEHI": Vehicle,
        }

        self.elements = {}

        dataset_words, dataset_entities = get_multinerd_hf(
            split=path_or_split,
            language=kwargs["language"],
            ENTITY_TO_CLASS_MAPPING=self.ENTITY_TO_CLASS_MAPPING,
        )

        for id, (words, entities) in enumerate(zip(dataset_words, dataset_entities)):
            self.elements[id] = {
                "id": id,
                "doc_id": id,
                "text": " ".join(words),
                "entities": entities,
                "gold": entities,
            }


class MultinerdSampler(Sampler):
    """
    A data `Sampler` for the Multi-NERD dataset.

    Args:
        dataset_loader (`MultinerdDatasetLoader`):
            The dataset loader that contains the data information.
        task (`str`, optional):
            The task to sample. It must be one of the following: NER, VER, RE, EE.
            Defaults to `None`.
        split (`str`, optional):
            The path_or_split to sample. It must be one of the following: "train", "dev" or
            "test". Depending on the path_or_split the sampling strategy differs. Defaults to
            `"train"`.
        parallel_instances (`Union[int, Tuple[int, int]]`, optional):
            The number of sentences sampled in parallel. Options:

                * **`int`**: The amount of elements that will be sampled in parallel.
                * **`tuple`**: The range of elements that will be sampled in parallel.

            Defaults to 1.
        max_guidelines (`int`, optional):
            The number of guidelines to append to the example at the same time. If `-1`
            is given then all the guidelines are appended. Defaults to `-1`.
        guideline_dropout (`float`, optional):
            The probability to dropout a guideline definition for the given example. This
            is only applied on training. Defaults to `0.0`.
        seed (`float`, optional):
            The seed to sample the examples. Defaults to `0`.
        prompt_template (`str`, optional):
            The path to the prompt template. Defaults to `"templates/prompt.txt"`.
        ensure_positives_on_train (bool, optional):
            Whether to ensure that the guidelines of annotated examples are not removed.
            Defaults to `False`.
        dataset_name (str, optional):
            The name of the dataset. Defaults to `None`.
        scorer (`str`, optional):
           The scorer class import string. Defaults to `None`.
        sample_only_gold_guidelines (`bool`, optional):
            Whether to sample only guidelines of present annotations. Defaults to `False`.
    """

    def __init__(
        self,
        dataset_loader: MultinerdDatasetLoader,
        task: str = None,
        split: str = "train",
        parallel_instances: Union[int, Tuple[int, int]] = 1,
        max_guidelines: int = -1,
        guideline_dropout: float = 0.0,
        seed: float = 0,
        prompt_template: str = "templates/prompt.txt",
        ensure_positives_on_train: bool = False,
        dataset_name: str = None,
        scorer: str = None,
        sample_only_gold_guidelines: bool = False,
        **kwargs,
    ) -> None:
        assert task in [
            "NER",
        ], f"Multi-Nerd only supports NER task. {task} is not supported."

        task_definitions, task_target = {
            "NER": (ENTITY_DEFINITIONS, "entities"),
        }[task]

        super().__init__(
            dataset_loader=dataset_loader,
            task=task,
            split=split,
            parallel_instances=parallel_instances,
            max_guidelines=max_guidelines,
            guideline_dropout=guideline_dropout,
            seed=seed,
            prompt_template=prompt_template,
            ensure_positives_on_train=ensure_positives_on_train,
            sample_only_gold_guidelines=sample_only_gold_guidelines,
            dataset_name=dataset_name,
            scorer=scorer,
            task_definitions=task_definitions,
            task_target=task_target,
            definitions=GUIDELINES,
            examples=EXAMPLES,
            **kwargs,
        )

src/tasks/multinerd/data_loader.py
hitz-zentroa/GoLLIE
GUIDELINES = {
    "multinerd_person": {"en": ["People."]},
    "multinerd_organization": {
        "en": ["Associations, companies, agencies, institutions, nationalities and religious or political groups."]
    },
    "multinerd_location": {
        "en": [
            "Physical locations (e.g. mountains, bodies of water), geopolitical entities (e.g. cities, states), and"
            " facilities (e.g. bridges, buildings, airports)"
        ]
    },
    "multinerd_animal": {"en": ["Breeds of dogs, cats and other animals, including their scientific names."]},
    "multinerd_biological": {
        "en": ["Genus of fungus, bacteria and protoctists, families of viruses, and other biological entities."]
    },
    "multinerd_celestial": {
        "en": ["Planets, stars, asteroids, comets, nebulae, galaxies and other astronomical objects."]
    },
    "multinerd_disease": {
        "en": [
            "Physical, mental, infectious, non-infectious, deficiency, inherited, degenerative, social and"
            " self-inflicted diseases."
        ]
    },
    "multinerd_event": {"en": ["Sport events, battles, wars and other events."]},
    "multinerd_food": {"en": ["Foods and drinks."]},
    "multinerd_instrument": {
        "en": ["Technological instruments, mechanical instruments, musical instruments, and other tools."]
    },
    "multinerd_media": {
        "en": ["Titles of films, books, magazines, songs and albums, fictional characters and languages."]
    },
    "multinerd_plant": {"en": ["Types of trees, flowers, and other plants, including their scientific names."]},
    "multinerd_mythological": {"en": ["Mythological and religious entities."]},
    "multinerd_time": {
        "en": [
            "Specific and well-defined time intervals, such as eras, historical periods, centuries, years and"
            " important days. No months and days of the week."
        ]
    },
    "multinerd_vehicle": {"en": ["Cars, motorcycles and other vehicles."]},
}

src/tasks/multinerd/guidelines.py
hitz-zentroa/GoLLIE
GUIDELINES = {
    "multinerd_person": {"en": ["People."]},
    "multinerd_organization": {
        "en": ["Associations, companies, agencies, institutions, nationalities and religious or political groups."]
    },
    "multinerd_location": {
        "en": [
            "Physical locations (e.g. mountains, bodies of water), geopolitical entities (e.g. cities, states), and"
            " facilities (e.g. bridges, buildings, airports)"
        ]
    },
    "multinerd_animal": {"en": ["Breeds of dogs, cats and other animals, including their scientific names."]},
    "multinerd_biological": {
        "en": ["Genus of fungus, bacteria and protoctists, families of viruses, and other biological entities."]
    },
    "multinerd_celestial": {
        "en": ["Planets, stars, asteroids, comets, nebulae, galaxies and other astronomical objects."]
    },
    "multinerd_disease": {
        "en": [
            "Physical, mental, infectious, non-infectious, deficiency, inherited, degenerative, social and"
            " self-inflicted diseases."
        ]
    },
    "multinerd_event": {"en": ["Sport events, battles, wars and other events."]},
    "multinerd_food": {"en": ["Foods and drinks."]},
    "multinerd_instrument": {
        "en": ["Technological instruments, mechanical instruments, musical instruments, and other tools."]
    },
    "multinerd_media": {
        "en": ["Titles of films, books, magazines, songs and albums, fictional characters and languages."]
    },
    "multinerd_plant": {"en": ["Types of trees, flowers, and other plants, including their scientific names."]},
    "multinerd_mythological": {"en": ["Mythological and religious entities."]},
    "multinerd_time": {
        "en": [
            "Specific and well-defined time intervals, such as eras, historical periods, centuries, years and"
            " important days. No months and days of the week."
        ]
    },
    "multinerd_vehicle": {"en": ["Cars, motorcycles and other vehicles."]},
}

EXAMPLES = {
    "multinerd_person_examples": {
        "en": [
            "George Hampson",
            "Edward Meyrick",
            "Breuning",
            "Francis Walker",
            "Aristotle",
            "Roger Ebert",
            "Plato",
            "Charles Darwin",
            "Herbert Druce",
            "Barack Obama",
        ]
    },
    "multinerd_organization_examples": {
        "en": [
            "European Union",
            "BBC",
            "Boston Red Sox",
            "New York Yankees",
            "Pittsburgh Pirates",
            "Los Angeles Dodgers",
            "Chicago Cubs",
            "Chicago White Sox",
            "Philadelphia Phillies",
            "Detroit Tigers",
        ]
    },
    "multinerd_location_examples": {
        "en": [
            "United States",
            "Australia",
            "Brazil",
            "China",
            "United Kingdom",
            "India",
            "Mexico",
            "Ecuador",
            "Japan",
            "Europe",
        ]
    },
    "multinerd_animal_examples": {
        "en": ["moth", "sea snail", "Erebidae", "Noctuidae", "mine", "fish", "ant", "livestock", "bird", "cattle"]
    },
    "multinerd_biological_examples": {
        "en": [
            "Escherichia coli",
            "E. coli",
            "Plasmodium",
            "Aspergillus",
            "HIV",
            "Streptococcus pneumoniae",
            "Agrobacterium",
            "Salmonella",
            "Streptomyces",
            "Rhizobium",
        ]
    },
    "multinerd_celestial_examples": {
        "en": ["Earth", "Sun", "Jupiter", "Moon", "Milky Way", "Mars", "Venus", "white dwarf", "Earth 's", "sun"]
    },
    "multinerd_disease_examples": {
        "en": [
            "cancer",
            "anxiety",
            "schizophrenia",
            "tuberculosis",
            "Alzheimer 's disease",
            "Parkinson 's disease",
            "asthma",
            "obesity",
            "pneumonia",
            "stroke",
        ]
    },
    "multinerd_event_examples": {
        "en": [
            "World War II",
            "World War I",
            "American Civil War",
            "FA Cup",
            "Vietnam War",
            "French Revolution",
            "UEFA Champions League",
            "World Series",
            "Second World War",
            "Ligue 1",
        ]
    },
    "multinerd_food_examples": {
        "en": ["sugar", "salt", "beef", "fruit", "pork", "honey", "milk", "vinegar", "coffee", "meat"]
    },
    "multinerd_instrument_examples": {
        "en": [
            "LaserDisc",
            "Android",
            "Voyager 2",
            "Voyager 1",
            "Laserdisc",
            "Apple II",
            "Super Nintendo Entertainment System",
            "AIM-9 Sidewinder",
            "Nasdaq",
            "MESSENGER",
        ]
    },
    "multinerd_media_examples": {
        "en": [
            "Rotten Tomatoes",
            "Oxford English Dictionary",
            "TV Guide",
            "Entertainment Weekly",
            "Rolling Stone",
            "American Idol",
            "Chicago Sun-Times",
            "Diagnostic and Statistical Manual of Mental Disorders",
            "Computer Gaming World",
            "GamePro",
        ]
    },
    "multinerd_mythological_examples": {
        "en": [
            "Thomas Aquinas",
            "Zeus",
            "Satan",
            "Ra",
            "Aphrodite",
            "John the Baptist",
            "Hermes",
            "Hercules",
            "Artemis",
            "Osiris",
        ]
    },
    "multinerd_plant_examples": {
        "en": ["plant", "wheat", "shrub", "rice", "deciduous", "oak", "legume", "flowering plant", "garlic", "maize"]
    },
    "multinerd_time_examples": {
        "en": [
            "Middle Ages",
            "Great Depression",
            "Renaissance",
            "Bronze Age",
            "half-life",
            "Christmas",
            "Neolithic",
            "puberty",
            "Enlightenment",
            "Age of Enlightenment",
        ]
    },
    "multinerd_vehicle_examples": {
        "en": [
            "Airbus A380",
            "Titanic",
            "Tupolev Tu-144",
            "Mayflower",
            "Dornier Do X",
            "Boeing 747",
            "Cessna 152",
            "GE Dash 8 Series",
            "Porsche 911",
            "Ford Model T",
        ]
    },
}

src/tasks/multinerd/guidelines_gold.py
hitz-zentroa/GoLLIE
from typing import List

from ..utils_typing import Entity, dataclass


"""Entity definitions

The entity definitions are derived from the official Multi-NERD guidelines:
https://aclanthology.org/2022.findings-naacl.60.pdf

"""


@dataclass
class Person(Entity):
    """{multinerd_person}"""

    span: str  # {multinerd_person_examples}


@dataclass
class Organization(Entity):
    """{multinerd_organization}"""

    span: str  # {multinerd_organization_examples}


@dataclass
class Location(Entity):
    """{multinerd_location}"""

    span: str  # {multinerd_location_examples}


@dataclass
class Animal(Entity):
    """{multinerd_animal}"""

    span: str  # {multinerd_animal_examples}


@dataclass
class Biological(Entity):
    """{multinerd_biological}"""

    span: str  # {multinerd_biological_examples}


@dataclass
class Celestial(Entity):
    """{multinerd_celestial}"""

    span: str  # {multinerd_celestial_examples}


@dataclass
class Disease(Entity):
    """{multinerd_disease}"""

    span: str  # {multinerd_disease_examples}


@dataclass
class Event(Entity):
    """{multinerd_event}"""

    span: str  # {multinerd_event_examples}


@dataclass
class Food(Entity):
    """{multinerd_food}"""

    span: str  # {multinerd_food_examples}


@dataclass
class Instrument(Entity):
    """{multinerd_instrument}"""

    span: str  # {multinerd_instrument_examples}


@dataclass
class Media(Entity):
    """{multinerd_media}"""

    span: str  # {multinerd_media_examples}


@dataclass
class Plant(Entity):
    """{multinerd_plant}"""

    span: str  # {multinerd_plant_examples}


@dataclass
class Mythological(Entity):
    """{multinerd_mythological}"""

    span: str  # {multinerd_mythological_examples}


@dataclass
class Time(Entity):
    """{multinerd_time}"""

    span: str  # {multinerd_time_examples}


@dataclass
class Vehicle(Entity):
    """{multinerd_vehicle}"""

    span: str  # {multinerd_vehicle_examples}


ENTITY_DEFINITIONS: List[Entity] = [
    Person,
    Location,
    Organization,
    Animal,
    Biological,
    Celestial,
    Disease,
    Event,
    Food,
    Instrument,
    Media,
    Plant,
    Mythological,
    Time,
    Vehicle,
]

# __all__ = list(map(str, [*ENTITY_DEFINITIONS]))

src/tasks/multinerd/prompts.py
hitz-zentroa/GoLLIE
from typing import Dict, List, Type

from src.tasks.multinerd.prompts import ENTITY_DEFINITIONS
from src.tasks.utils_scorer import SpanScorer
from src.tasks.utils_typing import Entity


class MultinerdEntityScorer(SpanScorer):
    """CoNLL03 Entity identification and classification scorer."""

    valid_types: List[Type] = ENTITY_DEFINITIONS

    def __call__(self, reference: List[Entity], predictions: List[Entity]) -> Dict[str, Dict[str, float]]:
        output = super().__call__(reference, predictions)
        return {"entities": output["spans"]}

src/tasks/multinerd/scorer.py
hitz-zentroa/GoLLIE
from . import data_loader, prompts, scorer


__all__ = ["data_loader", "prompts", "scorer"]

src/tasks/ncbidisease/__init__.py
hitz-zentroa/GoLLIE
from typing import Dict, List, Tuple, Type, Union

from src.tasks.label_encoding import rewrite_labels
from src.tasks.ncbidisease.guidelines import GUIDELINES
from src.tasks.ncbidisease.guidelines_gold import EXAMPLES
from src.tasks.ncbidisease.prompts import (
    ENTITY_DEFINITIONS,
    Disease,
)

from ..utils_data import DatasetLoader, Sampler
from ..utils_typing import Entity


def get_ncbidisease_hf(
    split: str,
    ENTITY_TO_CLASS_MAPPING: Dict[str, Type[Entity]],
) -> Tuple[List[List[str]], List[List[Entity]]]:
    """
    Get the NCBI-disease dataset from the huggingface datasets library
    Args:
        split (str): The path_or_split to load. Can be one of `train`, `validation` or `test`.
    Returns:
        (List[str],List[Union[Location,Organization,Person,Miscellaneous]]): The text and the entities
    """
    from datasets import load_dataset

    dataset = load_dataset("ncbi_disease")
    id2label = dict(enumerate(dataset["train"].features["ner_tags"].feature.names))
    dataset_sentences: List[List[str]] = []
    dataset_entities: List[List[Entity]] = []

    for example in dataset[split]:
        words = example["tokens"]
        # Ensure IOB2 encoding
        labels = rewrite_labels(labels=[id2label[label] for label in example["ner_tags"]], encoding="iob2")

        # Get labeled word spans
        spans = []
        for i, label in enumerate(labels):
            if label == "O":
                continue
            elif label.startswith("B-"):
                spans.append([label[2:], i, i + 1])
            elif label.startswith("I-"):
                spans[-1][2] += 1
            else:
                raise ValueError(f"Found an unexpected label: {label}")

        # Get entities
        entities = []
        for label, start, end in spans:
            entities.append(ENTITY_TO_CLASS_MAPPING[label](span=" ".join(words[start:end])))

        dataset_sentences.append(words)
        dataset_entities.append(entities)

    return dataset_sentences, dataset_entities


class NcbiDiseaseDatasetLoader(DatasetLoader):
    """
    A `DatasetLoader` for the ncbi-disease dataset.

    Args:
        split (`str`):
            The split to load. Can be one of `train`, `validation` or `test`.

    Raises:
        `ValueError`:
            raised when a not defined value found.
    """

    ENTITY_TO_CLASS_MAPPING = None

    def __init__(self, path_or_split: str, **kwargs) -> None:
        self.ENTITY_TO_CLASS_MAPPING = {
            "Disease": Disease,
        }

        self.elements = {}

        dataset_words, dataset_entities = get_ncbidisease_hf(
            split=path_or_split,
            ENTITY_TO_CLASS_MAPPING=self.ENTITY_TO_CLASS_MAPPING,
        )

        for id, (words, entities) in enumerate(zip(dataset_words, dataset_entities)):
            self.elements[id] = {
                "id": id,
                "doc_id": id,
                "text": " ".join(words),
                "entities": entities,
                "gold": entities,
            }


class NcbiDiseaseSampler(Sampler):
    """
    A data `Sampler` for the NCBI Disease dataset.

    Args:
        dataset_loader (`NcbiDiseaseDatasetLoader`):
            The dataset loader that contains the data information.
        task (`str`, optional):
            The task to sample. It must be one of the following: NER, VER, RE, EE.
            Defaults to `None`.
        split (`str`, optional):
            The path_or_split to sample. It must be one of the following: "train", "dev" or
            "test". Depending on the path_or_split the sampling strategy differs. Defaults to
            `"train"`.
        parallel_instances (`Union[int, Tuple[int, int]]`, optional):
            The number of sentences sampled in parallel. Options:

                * **`int`**: The amount of elements that will be sampled in parallel.
                * **`tuple`**: The range of elements that will be sampled in parallel.

            Defaults to 1.
        max_guidelines (`int`, optional):
            The number of guidelines to append to the example at the same time. If `-1`
            is given then all the guidelines are appended. Defaults to `-1`.
        guideline_dropout (`float`, optional):
            The probability to dropout a guideline definition for the given example. This
            is only applied on training. Defaults to `0.0`.
        seed (`float`, optional):
            The seed to sample the examples. Defaults to `0`.
        prompt_template (`str`, optional):
            The path to the prompt template. Defaults to `"templates/prompt.txt"`.
        ensure_positives_on_train (bool, optional):
            Whether to ensure that the guidelines of annotated examples are not removed.
            Defaults to `False`.
        dataset_name (str, optional):
            The name of the dataset. Defaults to `None`.
        scorer (`str`, optional):
           The scorer class import string. Defaults to `None`.
        sample_only_gold_guidelines (`bool`, optional):
            Whether to sample only guidelines of present annotations. Defaults to `False`.
    """

    def __init__(
        self,
        dataset_loader: NcbiDiseaseDatasetLoader,
        task: str = None,
        split: str = "train",
        parallel_instances: Union[int, Tuple[int, int]] = 1,
        max_guidelines: int = -1,
        guideline_dropout: float = 0.0,
        seed: float = 0,
        prompt_template: str = "templates/prompt.txt",
        ensure_positives_on_train: bool = False,
        dataset_name: str = None,
        scorer: str = None,
        sample_only_gold_guidelines: bool = False,
        **kwargs,
    ) -> None:
        assert task in [
            "NER",
        ], f"Ncbi Disease only supports NER task. {task} is not supported."

        task_definitions, task_target = {
            "NER": (ENTITY_DEFINITIONS, "entities"),
        }[task]

        super().__init__(
            dataset_loader=dataset_loader,
            task=task,
            split=split,
            parallel_instances=parallel_instances,
            max_guidelines=max_guidelines,
            guideline_dropout=guideline_dropout,
            seed=seed,
            prompt_template=prompt_template,
            ensure_positives_on_train=ensure_positives_on_train,
            sample_only_gold_guidelines=sample_only_gold_guidelines,
            dataset_name=dataset_name,
            scorer=scorer,
            task_definitions=task_definitions,
            task_target=task_target,
            definitions=GUIDELINES,
            examples=EXAMPLES,
            **kwargs,
        )

src/tasks/ncbidisease/data_loader.py
hitz-zentroa/GoLLIE
GUIDELINES = {
    "ncbi_disease": {
        "en": [
            (
                "A disease is a particular abnormal condition that negatively affects the structure or function of all"
                " or part of an \norganism, and that is not immediately due to any external injury. Disease mentions,"
                " include Specific Disease (i.e., \nclear-cell renal cell carcinoma), Disease Class (i.e., cystic"
                " kidney diseases), Composite mentions (i.e., prostatic, \npancreas, skin, and lung cancer), and"
                " Modifier (i.e., hereditary breast cancer families)."
            ),
            (
                "A disease refers to an unusual physiological state that detrimentally impacts the structure or"
                " function of all or \npart of an organism, and which isn't immediately caused by any external injury."
                " Disease-related mentions encompass \nSpecific Disease (for example, clear-cell renal cell"
                " carcinoma), Disease Class (such as cystic kidney diseases), \nComposite mentions (like prostatic,"
                " pancreas, skin, and lung cancer), and Modifier (for instance, hereditary breast cancer \nfamilies)."
            ),
            (
                "A disease refers to a distinct anomalous state that detrimentally impacts the structure or function"
                " of an entire \norganism, not resulting from any immediate external injury. Disease mentions"
                " encompass Specific Disease (for example, \nclear-cell renal cell carcinoma), Disease Class (such as"
                " cystic kidney diseases), Composite mentions (like prostatic, \npancreas, skin, and lung cancer), and"
                " Modifier (for instance, hereditary breast cancer families)."
            ),
            (
                "A disease refers to a distinct irregular condition that unfavorably impacts the structure or function"
                " of all or part \nof an organism, and is not immediately attributable to any external trauma. Disease"
                " mentions encompass Specific \nDisease (for example, clear-cell renal cell carcinoma), Disease Class"
                " (for instance, cystic kidney diseases), Composite \nmentions (such as prostatic, pancreas, skin, and"
                " lung cancer), and Modifier (for example, hereditary breast cancer \nfamilies)."
            ),
            (
                "A disease refers to a distinct irregularity that unfavorably impacts the structure or function of an"
                " entire \norganism or a portion of it, and is not immediately attributable to any external injury."
                " Disease mentions encompass Specific \nDisease (for example, clear-cell renal cell carcinoma),"
                " Disease Class (such as cystic kidney diseases), Composite \nmentions (like prostatic, pancreas,"
                " skin, and lung cancer), and Modifier (like hereditary breast cancer families)."
            ),
        ]
    }
}

src/tasks/ncbidisease/guidelines.py
hitz-zentroa/GoLLIE
GUIDELINES = {
    "ncbi_disease": {
        "en": [
            "A disease is a particular abnormal condition that negatively affects the structure or function of all or"
            " part of an organism, and that is not immediately due to any external injury. Disease mentions, include"
            " Specific Disease (i.e., clear-cell renal cell carcinoma), Disease Class (i.e., cystic kidney diseases),"
            " Composite mentions (i.e., prostatic, pancreas, skin, and lung cancer), and Modifier (i.e., hereditary"
            " breast cancer families). "
        ]
    },
}

EXAMPLES = {
    "ncbi_disease_examples": {
        "en": ["DM", "DMD", "APC", "ALD", "PWS", "WAS", "myotonic dystrophy", "G6PD deficiency", "HD", "PKU"]
    }
}

src/tasks/ncbidisease/guidelines_gold.py
hitz-zentroa/GoLLIE
from typing import List

from ..utils_typing import Entity, dataclass


"""Entity definitions

The entity definitions are derived from the official NCBI-Disease guidelines:
https://www.ncbi.nlm.nih.gov/pmc/articles/PMC3951655/
Disease definition from: https://en.wikipedia.org/wiki/Disease which is based on the definitions from:
 [1] "Disease" at Dorland's Medical Dictionary
 [2] White, Tim (19 December 2014). "What is the Difference Between an "Injury" and "Disease" for Comcare Commonwealth
 Compensation Claims?". Tindall Gask Bentley. Archived from the original on 27 October 2017. Retrieved 6 November 2017.


"""


@dataclass
class Disease(Entity):
    """{ncbi_disease}"""

    span: str  # {ncbi_disease_examples}


ENTITY_DEFINITIONS: List[Entity] = [Disease]

# __all__ = list(map(str, [*ENTITY_DEFINITIONS]))

src/tasks/ncbidisease/prompts.py
hitz-zentroa/GoLLIE
from typing import Dict, List, Type

from src.tasks.ncbidisease.prompts import ENTITY_DEFINITIONS
from src.tasks.utils_scorer import SpanScorer
from src.tasks.utils_typing import Entity


class NcbiDiseaseEntityScorer(SpanScorer):
    """CoNLL03 Entity identification and classification scorer."""

    valid_types: List[Type] = ENTITY_DEFINITIONS

    def __call__(self, reference: List[Entity], predictions: List[Entity]) -> Dict[str, Dict[str, float]]:
        output = super().__call__(reference, predictions)
        return {"entities": output["spans"]}

src/tasks/ncbidisease/scorer.py
hitz-zentroa/GoLLIE
from . import data_loader, prompts, scorer


__all__ = ["data_loader", "prompts", "scorer"]

src/tasks/ontonotes/__init__.py
hitz-zentroa/GoLLIE
from typing import Dict, List, Tuple, Type, Union

from src.tasks.label_encoding import rewrite_labels
from src.tasks.ontonotes.guidelines import GUIDELINES
from src.tasks.ontonotes.guidelines_gold import EXAMPLES
from src.tasks.ontonotes.prompts import (
    ENTITY_DEFINITIONS,
    GPE,
    NORP,
    Cardinal,
    Date,
    Event,
    Facility,
    Language,
    Law,
    Location,
    Money,
    Ordinal,
    Organization,
    Percentage,
    Person,
    Product,
    Quantity,
    Time,
    WorkOfArt,
)

from ..utils_data import DatasetLoader, Sampler
from ..utils_typing import Entity


def get_ontonotes_hf(
    split: str,
    ENTITY_TO_CLASS_MAPPING: Dict[str, Type[Entity]],
) -> Tuple[List[List[str]], List[List[Entity]]]:
    """
    Get the OntoNotes dataset from the huggingface datasets library
    Args:
        split (str): The path_or_split to load. Can be one of `train`, `validation` or `test`.
    Returns:
        (List[str],List[Union[Location,Organization,Person,Miscellaneous]]): The text and the entities
    """
    from datasets import load_dataset

    dataset = load_dataset("conll2012_ontonotesv5", "english_v12")
    id2label = dict(enumerate(dataset["train"].features["sentences"][0]["named_entities"].feature.names))
    dataset_sentences: List[List[str]] = []
    dataset_entities: List[List[Entity]] = []

    for example in dataset[split]:
        for sentence in example["sentences"]:
            words = sentence["words"]
            # Ensure IOB2 encoding
            labels = rewrite_labels(labels=[id2label[label] for label in sentence["named_entities"]], encoding="iob2")

            # Get labeled word spans
            spans = []
            for i, label in enumerate(labels):
                if label == "O":
                    continue
                elif label.startswith("B-"):
                    spans.append([label[2:], i, i + 1])
                elif label.startswith("I-"):
                    spans[-1][2] += 1
                else:
                    raise ValueError(f"Found an unexpected label: {label}")

            # Get entities
            entities = []
            for label, start, end in spans:
                entities.append(ENTITY_TO_CLASS_MAPPING[label](span=" ".join(words[start:end])))

            dataset_sentences.append(words)
            dataset_entities.append(entities)

    return dataset_sentences, dataset_entities


class OntoNotesDatasetLoader(DatasetLoader):
    """
    A `DatasetLoader` for the OntoNotesV5 dataset.

    Args:
        split (`str`):
            The split to load. Can be one of `train`, `validation` or `test`.

    Raises:
        `ValueError`:
            raised when a not defined value found.
    """

    ENTITY_TO_CLASS_MAPPING = None

    def __init__(self, path_or_split: str, **kwargs) -> None:
        self.ENTITY_TO_CLASS_MAPPING = {
            "PERSON": Person,
            "NORP": NORP,
            "FAC": Facility,
            "ORG": Organization,
            "GPE": GPE,
            "LOC": Location,
            "PRODUCT": Product,
            "DATE": Date,
            "TIME": Time,
            "PERCENT": Percentage,
            "MONEY": Money,
            "QUANTITY": Quantity,
            "ORDINAL": Ordinal,
            "CARDINAL": Cardinal,
            "EVENT": Event,
            "WORK_OF_ART": WorkOfArt,
            "LAW": Law,
            "LANGUAGE": Language,
        }

        self.elements = {}

        dataset_words, dataset_entities = get_ontonotes_hf(
            split=path_or_split,
            ENTITY_TO_CLASS_MAPPING=self.ENTITY_TO_CLASS_MAPPING,
        )

        for id, (words, entities) in enumerate(zip(dataset_words, dataset_entities)):
            self.elements[id] = {
                "id": id,
                "doc_id": id,
                "text": " ".join(words),
                "entities": entities,
                "gold": entities,
            }


class OntoNotesSampler(Sampler):
    """
    A data `Sampler` for the OntoNotesv5 dataset.

    Args:
        dataset_loader (`OntoNotesDatasetLoader`):
            The dataset loader that contains the data information.
        task (`str`, optional):
            The task to sample. It must be one of the following: NER, VER, RE, EE.
            Defaults to `None`.
        split (`str`, optional):
            The path_or_split to sample. It must be one of the following: "train", "dev" or
            "test". Depending on the path_or_split the sampling strategy differs. Defaults to
            `"train"`.
        parallel_instances (`Union[int, Tuple[int, int]]`, optional):
            The number of sentences sampled in parallel. Options:

                * **`int`**: The amount of elements that will be sampled in parallel.
                * **`tuple`**: The range of elements that will be sampled in parallel.

            Defaults to 1.
        max_guidelines (`int`, optional):
            The number of guidelines to append to the example at the same time. If `-1`
            is given then all the guidelines are appended. Defaults to `-1`.
        guideline_dropout (`float`, optional):
            The probability to dropout a guideline definition for the given example. This
            is only applied on training. Defaults to `0.0`.
        seed (`float`, optional):
            The seed to sample the examples. Defaults to `0`.
        prompt_template (`str`, optional):
            The path to the prompt template. Defaults to `"templates/prompt.txt"`.
        ensure_positives_on_train (bool, optional):
            Whether to ensure that the guidelines of annotated examples are not removed.
            Defaults to `False`.
        dataset_name (str, optional):
            The name of the dataset. Defaults to `None`.
        scorer (`str`, optional):
           The scorer class import string. Defaults to `None`.
        sample_only_gold_guidelines (`bool`, optional):
            Whether to sample only guidelines of present annotations. Defaults to `False`.
    """

    def __init__(
        self,
        dataset_loader: OntoNotesDatasetLoader,
        task: str = None,
        split: str = "train",
        parallel_instances: Union[int, Tuple[int, int]] = 1,
        max_guidelines: int = -1,
        guideline_dropout: float = 0.0,
        seed: float = 0,
        prompt_template: str = "templates/prompt.txt",
        ensure_positives_on_train: bool = False,
        dataset_name: str = None,
        scorer: str = None,
        sample_only_gold_guidelines: bool = False,
        **kwargs,
    ) -> None:
        assert task in [
            "NER",
        ], f"OntoNotes only supports NER task. {task} is not supported."

        task_definitions, task_target = {
            "NER": (ENTITY_DEFINITIONS, "entities"),
        }[task]

        super().__init__(
            dataset_loader=dataset_loader,
            task=task,
            split=split,
            parallel_instances=parallel_instances,
            max_guidelines=max_guidelines,
            guideline_dropout=guideline_dropout,
            seed=seed,
            prompt_template=prompt_template,
            ensure_positives_on_train=ensure_positives_on_train,
            sample_only_gold_guidelines=sample_only_gold_guidelines,
            dataset_name=dataset_name,
            scorer=scorer,
            task_definitions=task_definitions,
            task_target=task_target,
            definitions=GUIDELINES,
            examples=EXAMPLES,
            **kwargs,
        )

src/tasks/ontonotes/data_loader.py
hitz-zentroa/GoLLIE
GUIDELINES = {
    "ontonotes_person": {"en": ["People, including fictional."]},
    "ontonotes_norp": {"en": ["Nationalities or religious or political groups."]},
    "ontonotes_facility": {"en": ["Buildings, airports, highways, bridges, etc."]},
    "ontonotes_organization": {"en": ["Companies, agencies, institutions, etc."]},
    "ontonotes_gpe": {"en": ["Countries, cities, states."]},
    "ontonotes_location": {"en": ["Non-GPE locations, mountain ranges, bodies of water."]},
    "ontonotes_product": {"en": ["Objects, vehicles, foods, etc. (Not services)."]},
    "ontonotes_event": {"en": ["Named hurricanes, battles, wars, sports events, etc."]},
    "ontonotes_work_of_art": {"en": ["Titles of books, songs, etc."]},
    "ontonotes_law": {"en": ["Named documents made into laws."]},
    "ontonotes_language": {"en": ["Any named language."]},
    "ontonotes_date": {"en": ["Absolute or relative dates or periods."]},
    "ontonotes_time": {"en": ["Times smaller than a day."]},
    "ontonotes_percent": {"en": ["Percentage, including ”%“."]},
    "ontonotes_money": {"en": ["Monetary values, including unit."]},
    "ontonotes_quantity": {"en": ["Measurements, as of weight or distance."]},
    "ontonotes_ordinal": {"en": ["first, second, third, First, fourth, fifth, Second, seventh, eighth, sixth."]},
    "ontonotes_cardinal": {"en": ["two, one, three, One, four, five, six, seven, Two, half."]},
}

src/tasks/ontonotes/guidelines.py
hitz-zentroa/GoLLIE
GUIDELINES = {
    "ontonotes_person": {"en": ["People, including fictional."]},
    "ontonotes_norp": {"en": ["Nationalities or religious or political groups."]},
    "ontonotes_facility": {"en": ["Buildings, airports, highways, bridges, etc."]},
    "ontonotes_organization": {"en": ["Companies, agencies, institutions, etc."]},
    "ontonotes_gpe": {"en": ["Countries, cities, states."]},
    "ontonotes_location": {"en": ["Non-GPE locations, mountain ranges, bodies of water."]},
    "ontonotes_product": {"en": ["Objects, vehicles, foods, etc. (Not services)."]},
    "ontonotes_event": {"en": ["Named hurricanes, battles, wars, sports events, etc."]},
    "ontonotes_work_of_art": {"en": ["Titles of books, songs, etc."]},
    "ontonotes_law": {"en": ["Named documents made into laws."]},
    "ontonotes_language": {"en": ["Any named language."]},
    "ontonotes_date": {"en": ["Absolute or relative dates or periods."]},
    "ontonotes_time": {"en": ["Times smaller than a day."]},
    "ontonotes_percent": {"en": ["Percentage, including ”%“."]},
    "ontonotes_money": {"en": ["Monetary values, including unit."]},
    "ontonotes_quantity": {"en": ["Measurements, as of weight or distance."]},
    "ontonotes_ordinal": {"en": ["first, second, third, First, fourth, fifth, Second, seventh, eighth, sixth."]},
    "ontonotes_cardinal": {"en": ["two, one, three, One, four, five, six, seven, Two, half."]},
}

EXAMPLES = {
    "ontonotes_person_examples": {
        "en": ["Bush", "Clinton", "Barak", "Noriega", "Putin", "Milosevic", "Ehud Barak", "Chen", "Reagan", "Gore"]
    },
    "ontonotes_norp_examples": {
        "en": [
            "Chinese",
            "Japanese",
            "American",
            "Israeli",
            "Palestinian",
            "Iraqi",
            "British",
            "Palestinians",
            "Taiwanese",
            "Americans",
        ]
    },
    "ontonotes_facility_examples": {
        "en": [
            "the White House",
            "Pentagon",
            "the Bay Bridge",
            "Guantanamo",
            "Mir",
            "the United Nations",
            "rue Saint - Denis",
            "the Fourth Nuclear Power Plant",
            "the Zhengtai Railway",
            "Candlestick Park",
        ]
    },
    "ontonotes_organization_examples": {
        "en": ["Congress", "Senate", "Treasury", "Xinhua News Agency", "House", "Fed", "UN", "Pentagon", "CNN", "IBM"]
    },
    "ontonotes_gpe_examples": {
        "en": ["U.S.", "China", "Taiwan", "Japan", "US", "New York", "Iraq", "Washington", "Israel", "Hong Kong"]
    },
    "ontonotes_location_examples": {
        "en": [
            "Europe",
            "the Middle East",
            "Asia",
            "Africa",
            "West",
            "Middle East",
            "Earth",
            "Eastern Europe",
            "Gulf",
            "Mideast",
        ]
    },
    "ontonotes_product_examples": {
        "en": [
            "Cole",
            "Galileo",
            "Kursk",
            "the `` USS Cole ''",
            "the `` USS Cole",
            "the USS Cole",
            "Discovery",
            "Atlantis",
            "Enigma",
            "C3CRM",
        ]
    },
    "ontonotes_date_examples": {
        "en": [
            "today",
            "yesterday",
            "last year",
            "this year",
            "Friday",
            "September",
            "annual",
            "Monday",
            "August",
            "1988",
        ]
    },
    "ontonotes_time_examples": {
        "en": [
            "tonight",
            "morning",
            "night",
            "this morning",
            "last night",
            "afternoon",
            "evening",
            "Tonight",
            "hours",
            "this afternoon",
        ]
    },
    "ontonotes_percent_examples": {
        "en": ["10 %", "8 %", "50 %", "15 %", "9 %", "5 %", "20 %", "40 %", "60 %", "25 %"]
    },
    "ontonotes_money_examples": {
        "en": [
            "$ 1 million",
            "$ 1 billion",
            "1,000",
            "10,000",
            "$ 200 million",
            "$ 150 million",
            "$ 100 million",
            "$ 50 million",
            "$ 2 billion",
            "$ 500 million",
        ]
    },
    "ontonotes_quantity_examples": {
        "en": ["two", "three", "four", "ten", "one", "six", "gallon", "1 degree below zero", "two miles", "five"]
    },
    "ontonotes_ordinal_examples": {
        "en": ["first", "second", "third", "First", "fourth", "fifth", "Second", "seventh", "eighth", "sixth"]
    },
    "ontonotes_cardinal_examples": {
        "en": ["two", "one", "three", "One", "four", "five", "six", "seven", "Two", "half"]
    },
    "ontonotes_event_examples": {
        "en": [
            "New Year",
            "Hurricane Hugo",
            "World War II",
            "Hugo",
            "the World Series",
            "Series",
            "the Hundred Regiments Offensive",
            "New Year 's Eve",
            "Watergate",
            "Iran - Contra",
        ]
    },
    "ontonotes_work_of_art_examples": {
        "en": [
            "Hardball",
            "Late Edition",
            "Yasukuni Shrine",
            "D.T.",
            "Newsnight",
            "Between You and Me",
            "Letting Them Go",
            "Batibot",
            "Three Sixty",
            "Headliners",
        ]
    },
    "ontonotes_law_examples": {
        "en": [
            "Chapter 11",
            "RICO",
            "Constitution",
            "Gramm - Rudman",
            "Superfund",
            "Section 89",
            "the Voting Rights Act",
            "constitution",
            "the First Amendment",
            "Bankruptcy Code",
        ]
    },
    "ontonotes_language_examples": {
        "en": [
            "English",
            "Taiwanese",
            "Chinese",
            "Filipino",
            "French",
            "Portuguese",
            "Russian",
            "Spanish",
            "Mandarin",
            "Arabic",
        ]
    },
}

src/tasks/ontonotes/guidelines_gold.py
hitz-zentroa/GoLLIE
from typing import List

from ..utils_typing import Entity, dataclass


"""Entity definitions

The entity definitions are derived from the official OntoNotes 5.0 guidelines:
https://catalog.ldc.upenn.edu/docs/LDC2013T19/OntoNotes-Release-5.0.pdf
"""


@dataclass
class Person(Entity):
    """{ontonotes_person}"""

    span: str  # {ontonotes_person_examples}


@dataclass
class NORP(Entity):
    """{ontonotes_norp}"""

    span: str  # {ontonotes_norp_examples}


@dataclass
class Facility(Entity):
    """{ontonotes_facility}"""

    span: str  # {ontonotes_facility_examples}


@dataclass
class Organization(Entity):
    """{ontonotes_organization}"""

    span: str  # {ontonotes_organization_examples}


@dataclass
class GPE(Entity):
    """{ontonotes_gpe}"""

    span: str  # {ontonotes_gpe_examples}


@dataclass
class Location(Entity):
    """{ontonotes_location}"""

    span: str  # {ontonotes_location_examples}


@dataclass
class Product(Entity):
    """{ontonotes_product}"""

    span: str  # {ontonotes_product_examples}


@dataclass
class Event(Entity):
    """{ontonotes_event}"""

    span: str  # {ontonotes_event_examples}


@dataclass
class WorkOfArt(Entity):
    """{ontonotes_work_of_art}"""

    span: str  # {ontonotes_work_of_art_examples}


@dataclass
class Law(Entity):
    """{ontonotes_law}"""

    span: str  # {ontonotes_law_examples}


@dataclass
class Language(Entity):
    """{ontonotes_language}"""

    span: str  # {ontonotes_language_examples}


@dataclass
class Date(Entity):
    """{ontonotes_date}"""

    span: str  # {ontonotes_date_examples}


@dataclass
class Time(Entity):
    """{ontonotes_time}"""

    span: str  # {ontonotes_time_examples}


@dataclass
class Percentage(Entity):
    """{ontonotes_percent}"""

    span: str  # {ontonotes_percent_examples}


@dataclass
class Money(Entity):
    """{ontonotes_money}"""

    span: str  # {ontonotes_money_examples}


@dataclass
class Quantity(Entity):
    """{ontonotes_quantity}"""

    span: str  # {ontonotes_quantity_examples}


@dataclass
class Ordinal(Entity):
    """{ontonotes_ordinal}"""

    span: str  # {ontonotes_ordinal_examples}


@dataclass
class Cardinal(Entity):
    """{ontonotes_cardinal}"""

    span: str  # {ontonotes_cardinal_examples}


ENTITY_DEFINITIONS: List[Entity] = [
    Person,
    Organization,
    Location,
    NORP,
    Facility,
    GPE,
    Product,
    Event,
    WorkOfArt,
    Law,
    Language,
    Date,
    Time,
    Percentage,
    Money,
    Quantity,
    Ordinal,
    Cardinal,
]

# __all__ = list(map(str, [*ENTITY_DEFINITIONS]))

src/tasks/ontonotes/prompts.py
hitz-zentroa/GoLLIE
from typing import Dict, List, Type

from src.tasks.ontonotes.prompts import ENTITY_DEFINITIONS
from src.tasks.utils_scorer import SpanScorer
from src.tasks.utils_typing import Entity


class OntoNotesEntityScorer(SpanScorer):
    """CoNLL03 Entity identification and classification scorer."""

    valid_types: List[Type] = ENTITY_DEFINITIONS

    def __call__(self, reference: List[Entity], predictions: List[Entity]) -> Dict[str, Dict[str, float]]:
        output = super().__call__(reference, predictions)
        return {"entities": output["spans"]}

src/tasks/ontonotes/scorer.py
hitz-zentroa/GoLLIE
from . import data_loader, prompts, scorer


__all__ = ["data_loader", "prompts", "scorer"]

src/tasks/rams/__init__.py
hitz-zentroa/GoLLIE
import inspect
import json
from typing import Tuple, Union

from src.tasks.rams.guidelines import GUIDELINES
from src.tasks.rams.prompts import (
    EVENT_DEFINITIONS,
    AccidentCrash,
    Agreement,
    ArrestJailDetain,
    ArtifactFailure,
    Attack,
    Collaborate,
    CommandOrder,
    CommitmentPromiseExpressIntent,
    Convene,
    Coup,
    DamageDestroy,
    Demonstrate,
    Die,
    Discussion,
    DiseaseOutbreak,
    Elect,
    EndPossition,
    FireExplosion,
    Formation,
    FuneralVigil,
    GenericCrime,
    InitiateJudicialProcess,
    Injure,
    InvestigateCrime,
    JudicialConsequences,
    Legislate,
    ManufactureArtifact,
    MediaStatement,
    MedicalIntervention,
    Negotiate,
    Prevarication,
    PublicStatementInPerson,
    RequestAdvice,
    SensoryObserve,
    Shortage,
    Spy,
    StartPossition,
    TargetAimAt,
    ThreatenCoerce,
    Transaction,
    TransferMoney,
    TransferOwnership,
    TransportArtifact,
    TransportPerson,
    Vote,
    Yield,
)

from ..utils_data import DatasetLoader, Sampler


class RAMSDatasetLoader(DatasetLoader):
    """
    A `DatasetLoader` for the RAMS dataset.

    Args:
        path (`str`):
            The location of the dataset directory.

    Raises:
        `ValueError`:
            raised when a not defined value found.
    """

    _EVENT_CONSTANTS_MAPPING = {"n/a": None, "place": "place"}
    EVENT_TO_CLASS_MAPPING = {
        "artifactexistence.artifactfailure": {
            "class": ArtifactFailure,
            "subtypes": {"mechanicalfailure": "MechanicalFailure"},
            "mechanicalartifact": "artifact",
            "instrument": "instrument",
        },
        "artifactexistence.damagedestroy": {
            "class": DamageDestroy,
            "subtypes": {"damage": "Damage", "destroy": "Destroy"},
            "damagerdestroyer": "agent",
            "damager": "agent",
            "destroyer": "agent",
            "artifact": "artifact",
            "instrument": "instrument",
        },
        "artifactexistence.shortage": {
            "class": Shortage,
            "subtypes": {"shortage": "Shortage"},
            "experiencer": "experiencer",
            "supply": "supply",
        },
        "conflict.attack": {
            "class": Attack,
            "subtypes": {
                "airstrikemissilestrike": "AirStrikeMissileStrike",
                "biologicalchemicalpoisonattack": "BiologicalChemicalPoisonAttack",
                "bombing": "Bombing",
                "firearmattack": "FireArmAttack",
                "hanging": "Hanging",
                "invade": "Invade",
                "selfdirectedbattle": "SelfDirectedBattle",
                "setfire": "SetFire",
                "stabbing": "Stabbing",
                "stealrobhijack": "StealRobHijack",
                "strangling": "Strangling",
            },
            "attacker": "attacker",
            "target": "target",
            "instrument": "instrument",
            "artifact": "artifact",
        },
        "conflict.coup": {
            "class": Coup,
            "subtypes": {"coup": "Coup"},
            "deposedentity": "deposed_entity",
            "deposingentity": "deposing_entity",
        },
        "conflict.demonstrate": {
            "class": Demonstrate,
            "subtypes": {"marchprotestpoliticalgathering": "MarchProtestPoliticalGathering"},
            "demonstrator": "demonstrator",
        },
        "conflict.yield": {
            "class": Yield,
            "subtypes": {"retreat": "Retreat", "surrender": "Surrender"},
            "yielder": "agent",
            "retreater": "agent",
            "surrenderer": "agent",
            "recipient": "recipient",
            "origin": "origin",
            "destination": "destination",
        },
        "contact.collaborate": {
            "class": Collaborate,
            "subtypes": {"correspondence": "Correspondence", "meet": "Meet"},
            "participant": "participants",
        },
        "contact.commandorder": {
            "class": CommandOrder,
            "subtypes": {
                "broadcast": "Broadcast",
                "correspondence": "Correspondence",
                "meet": "Meet",
            },
            "communicator": "communicator",
            "recipient": "recipient",
            "topic": "topic",
        },
        "contact.commitmentpromiseexpressintent": {
            "class": CommitmentPromiseExpressIntent,
            "subtypes": {
                "broadcast": "Broadcast",
                "correspondence": "Correspondence",
                "meet": "Meet",
            },
            "communicator": "communicator",
            "recipient": "recipient",
            "topic": "topic",
        },
        "contact.discussion": {
            "class": Discussion,
            "subtypes": {"correspondence": "Correspondence", "meet": "Meet"},
            "participant": "participants",
        },
        "contact.funeralvigil": {
            "class": FuneralVigil,
            "subtypes": {"meet": "Meet"},
            "participant": "participants",
            "deceased": "deceased",
        },
        "contact.mediastatement": {
            "class": MediaStatement,
            "subtypes": {"broadcast": "Broadcast"},
            "communicator": "communicator",
            "recipient": "recipient",
        },
        "contact.negotiate": {
            "class": Negotiate,
            "subtypes": {"correspondence": "Correspondence", "meet": "Meet"},
            "participant": "participants",
            "topic": "topic",
        },
        "contact.prevarication": {
            "class": Prevarication,
            "subtypes": {
                "broadcast": "Broadcast",
                "correspondence": "Correspondence",
                "meet": "Meet",
            },
            "communicator": "communicator",
            "recipient": "recipient",
            "topic": "topic",
        },
        "contact.publicstatementinperson": {
            "class": PublicStatementInPerson,
            "subtypes": {"broadcast": "Broadcast"},
            "communicator": "communicator",
            "recipient": "recipient",
        },
        "contact.requestadvise": {
            "class": RequestAdvice,
            "subtypes": {
                "broadcast": "Broadcast",
                "correspondence": "Correspondence",
                "meet": "Meet",
            },
            "communicator": "communicator",
            "recipient": "recipient",
            "topic": "topic",
        },
        "contact.threatencoerce": {
            "class": ThreatenCoerce,
            "subtypes": {
                "broadcast": "Broadcast",
                "correspondence": "Correspondence",
                "meet": "Meet",
            },
            "communicator": "communicator",
            "recipient": "recipient",
            "topic": "topic",
        },
        "disaster.accidentcrash": {
            "class": AccidentCrash,
            "subtypes": {"accidentcrash": "AccidentCrash"},
            "driverpassenger": "driver_or_passengers",
            "vehicle": "vehicle",
            "crashobject": "crash_object",
        },
        "disaster.diseaseoutbreak": {
            "class": DiseaseOutbreak,
            "subtypes": {"diseaseoutbreak": "DiseaseOutbreak"},
            "disease": "disease",
            "victim": "victim",
        },
        "disaster.fireexplosion": {
            "class": FireExplosion,
            "subtypes": {"fireexplosion": "FireExplosion"},
            "fireexplosionobject": "fire_explosion_object",
            "instrument": "instrument",
        },
        "genericcrime.genericcrime": {
            "class": GenericCrime,
            "subtypes": {"genericcrime": "GenericCrime"},
            "perpetrator": "perpetrator",
            "victim": "victim",
        },
        "government.agreements": {
            "class": Agreement,
            "subtypes": {
                "acceptagreementcontractceasefire": "Accept",
                "rejectnullifyagreementcontractceasefire": "RejectNullify",
                "violateagreement": "Violate",
            },
            "participant": "participants",
            "rejecternullifier": "rejecter_nullifier",
            "violator": "violator",
            "otherparticipant": "other_participant",
        },
        "government.convene": {
            "class": Convene,
            "subtypes": {"convene": "Convene"},
            "convener": "convener",
            "convenedthing": "convened_entity",
        },
        "government.formation": {
            "class": Formation,
            "subtypes": {"mergegpe": "Merge", "startgpe": "Start"},
            "gpe": "gpe",
            "founder": "founder",
            "participant": "participants",
        },
        "government.legislate": {
            "class": Legislate,
            "subtypes": {"legislate": "Legislate"},
            "governmentbody": "government_body",
            "law": "law",
        },
        "government.spy": {
            "class": Spy,
            "subtypes": {"spy": "Spy"},
            "spy": "spy",
            "observedentity": "observed_entity",
            "beneficiary": "beneficiary",
        },
        "government.vote": {
            "class": Vote,
            "subtypes": {"castvote": "CastVote", "violationspreventvote": "PreventVote"},
            "voter": "voter",
            "candidate": "candidate",
            "ballot": "ballot",
            "result": "result",
            "preventer": "preventer",
        },
        "inspection.sensoryobserve": {
            "class": SensoryObserve,
            "subtypes": {
                "inspectpeopleorganization": "InspectPeopleOrganization",
                "monitorelection": "MonitorElection",
                "physicalinvestigateinspect": "PhysicalInvestigateInspect",
            },
            "observer": "observer",
            "observedentity": "observed_entity",
            "monitor": "observer",
            "monitoredentity": "observed_entity",
            "inspector": "observer",
            "inspectedentity": "observed_entity",
        },
        "inspection.targetaimat": {
            "class": TargetAimAt,
            "subtypes": {"targetaimat": "TargetAimAt"},
            "targeter": "targeter",
            "target": "target",
            "instrument": "instrument",
        },
        "justice.arrestjaildetain": {
            "class": ArrestJailDetain,
            "subtypes": {"arrestjaildetain": "ArrestJailDetain"},
            "jailer": "jailer",
            "detainee": "detainee",
            "crime": "crime",
        },
        "justice.initiatejudicialprocess": {
            "class": InitiateJudicialProcess,
            "subtypes": {"chargeindict": "ChargeIndict", "trialhearing": "TrialHearing"},
            "prosecutor": "prosecutor",
            "defendant": "defendant",
            "judgecourt": "judge_court",
            "crime": "crime",
        },
        "justice.investigate": {
            "class": InvestigateCrime,
            "subtypes": {"investigatecrime": "InvestigateCrime"},
            "investigator": "investigator",
            "defendant": "defendant",
            "crime": "crime",
        },
        "justice.judicialconsequences": {
            "class": JudicialConsequences,
            "subtypes": {
                "convict": "Convict",
                "execute": "Execute",
                "extradite": "Extradite",
            },
            "judgecourt": "judge_court",
            "defendant": "defendant",
            "crime": "crime",
            "executioner": "judge_court",
            "extraditer": "judge_court",
            "origin": "origin",
            "destination": "destination",
        },
        "life.die": {
            "class": Die,
            "subtypes": {
                "deathcausedbyviolentevents": "Violent",
                "nonviolentdeath": "NonViolent",
            },
            "victim": "victim",
            "killer": "killer",
            "medicalissue": "medical_issue",
            "instrument": "instrument",
        },
        "life.injure": {
            "class": Injure,
            "subtypes": {
                "illnessdegradationhungerthirst": "HungerThirst",
                "illnessdegradationphysical": "Physical",
                "illnessdegredationsickness": "Sickness",
                "injurycausedbyviolentevents": "Violent",
            },
            "victim": "victim",
            "injurer": "injurer",
            "medicalissue": "medical_issue",
            "disease": "disease",
            "instrument": "instrument",
        },
        "manufacture.artifact": {
            "class": ManufactureArtifact,
            "subtypes": {
                "build": "Build",
                "createintellectualproperty": "IntellectualProperty",
                "createmanufacture": "CreateManufacture",
            },
            "manufacturer": "manufacturer",
            "artifact": "artifact",
            "instrument": "instrument",
        },
        "medical.intervention": {
            "class": MedicalIntervention,
            "subtypes": {"intervention": "Intervention"},
            "treater": "treater",
            "patient": "patient",
            "medicalissue": "medical_issue",
            "instrument": "instrument",
        },
        "movement.transportartifact": {
            "class": TransportArtifact,
            "subtypes": {
                "bringcarryunload": "BringCarryUnload",
                "disperseseparate": "DisperseSeparate",
                "fall": "Fall",
                "grantentry": "GrantEntry",
                "hide": "Hide",
                "lossofcontrol": "LostOfControl",
                "nonviolentthrowlaunch": "NonViolentThrowLaunch",
                "prevententry": "PreventEntry",
                "preventexit": "PreventExit",
                "receiveimport": "ReceiveImport",
                "sendsupplyexport": "SendSupplyExport",
                "smuggleextract": "SmuggleExtract",
            },
            "transporter": "transporter",
            "artifact": "artifact",
            "vehicle": "vehicle",
            "origin": "origin",
            "destination": "destination",
            "hidingplace": "hidding_place",
            "controller": "transporter",
            "controlledthing": "artifact",
            "preventer": "preventer",
        },
        "movement.transportperson": {
            "class": TransportPerson,
            "subtypes": {
                "bringcarryunload": "BringCarryUnload",
                "disperseseparate": "DisperseSeparate",
                "evacuationrescue": "EvacuationRescue",
                "fall": "Fall",
                "grantentryasylum": "GrantedAsylum",
                "hide": "Hide",
                "prevententry": "PreventEntry",
                "preventexit": "PreventExit",
                "selfmotion": "SelfMotion",
                "smuggleextract": "SmuggleExtract",
            },
            "transporter": "transporter",
            "passenger": "passenger",
            "vehicle": "vehicle",
            "origin": "origin",
            "destination": "destination",
            "granter": "granter",
            "hidingplace": "hidding_place",
            "preventer": "preventer",
        },
        "personnel.elect": {
            "class": Elect,
            "subtypes": {"winelection": "WinElection"},
            "voter": "voter",
            "candidate": "candidate",
        },
        "personnel.endposition": {
            "class": EndPossition,
            "subtypes": {"firinglayoff": "FiringLayOff", "quitretire": "QuitRetire"},
            "employee": "employee",
            "placeofemployment": "organization",
        },
        "personnel.startposition": {
            "class": StartPossition,
            "subtypes": {"hiring": "Hiring"},
            "employee": "employee",
            "placeofemployment": "organization",
        },
        "transaction.transaction": {
            "class": Transaction,
            "subtypes": {
                "embargosanction": "EmbargoSanction",
                "giftgrantprovideaid": "GiftGrantProvideAid",
                "transfercontrol": "TransferControl",
            },
            "participant": "participants",
            "beneficiary": "beneficiary",
            "preventer": "preventer",
            "giver": "giver",
            "recipient": "recipient",
            "artifactmoney": "artifact",
            "territoryorfacility": "artifact",
        },
        "transaction.transfermoney": {
            "class": TransferMoney,
            "subtypes": {
                "borrowlend": "BorrowLend",
                "embargosanction": "EmbargoSanction",
                "giftgrantprovideaid": "GiftGrantProvideAid",
                "payforservice": "PayForService",
                "purchase": "Purchase",
            },
            "giver": "giver",
            "recipient": "recipient",
            "beneficiary": "beneficiary",
            "money": "money",
            "preventer": "preventer",
        },
        "transaction.transferownership": {
            "class": TransferOwnership,
            "subtypes": {
                "borrowlend": "BorrowLend",
                "embargosanction": "EmbargoSanction",
                "giftgrantprovideaid": "GiftGrantProvideAid",
                "purchase": "Purchase",
            },
            "giver": "giver",
            "recipient": "recipient",
            "beneficiary": "beneficiary",
            "artifact": "artifact",
            "preventer": "preventer",
        },
    }

    def __init__(self, path: str, **kwargs) -> None:
        self.elements = {}

        with open(path, "rt") as in_f:
            for line in in_f:
                line = json.loads(line.strip())
                key = line["doc_key"]
                tokens = [token for sentence in line["sentences"] for token in sentence]
                text = " ".join(tokens)
                events = []
                for event in line["evt_triggers"]:
                    trigger_start, trigger_end = event[:2]
                    event_type = event[-1][0][0].split(".")
                    _event_full_type = ".".join(event_type[:2])

                    info = self.EVENT_TO_CLASS_MAPPING[_event_full_type]
                    _inst = {param: [] for param in inspect.signature(info["class"]).parameters.keys()}
                    _inst["mention"] = " ".join(tokens[trigger_start : trigger_end + 1])

                    # Find the correct event subtype
                    if event_type[-1] in info["subtypes"]:
                        value = info["subtypes"][event_type[-1]]
                    elif event_type[-1] in self._EVENT_CONSTANTS_MAPPING:
                        value = self._EVENT_CONSTANTS_MAPPING[event_type[-1]]
                    else:
                        raise ValueError(f"Event subtype {event_type[-1]} not found!")
                    _inst["subtype"] = value

                    # Find the arguments of the event
                    for argument in line["gold_evt_links"]:
                        if argument[0] != [trigger_start, trigger_end]:
                            continue

                        arg_text = " ".join(tokens[argument[1][0] : argument[1][1] + 1])
                        role = argument[-1][11:]
                        if role in info:
                            value = info[role]
                        elif role in self._EVENT_CONSTANTS_MAPPING:
                            value = self._EVENT_CONSTANTS_MAPPING[role]
                        else:
                            raise ValueError(f"Argument {event_type}:{role} not found!")

                        try:
                            _inst[value].append(arg_text)
                        except KeyError:
                            raise KeyError(f"{event_type}-{value}")

                    events.append(info["class"](**_inst))

                self.elements[key] = {"id": key, "doc_id": key, "text": text, "labels": events, "gold": events}


class RAMSSampler(Sampler):
    """
    A data `Sampler` for the RAMS dataset.

    Args:
        dataset_loader (`RAMSDatasetLoader`):
            The dataset loader that contains the data information.
        task (`str`, optional):
            The task to sample. It can only be EAE. Defaults to `"EAE"`.
        split (`str`, optional):
            The split to sample. It must be one of the following: "train", "dev" or
            "test". Depending on the split the sampling strategy differs. Defaults to
            `"train"`.
        parallel_instances (`Union[int, Tuple[int, int]]`, optional):
            The number of sentences sampled in parallel. Options:

                * **`int`**: The amount of elements that will be sampled in parallel.
                * **`tuple`**: The range of elements that will be sampled in parallel.

            Defaults to 1.
        max_guidelines (`int`, optional):
            The number of guidelines to append to the example at the same time. If `-1`
            is given then all the guidelines are appended. Defaults to `-1`.
        guideline_dropout (`float`, optional):
            The probability to dropout a guideline definition for the given example. This
            is only applied on training. Defaults to `0.0`.
        seed (`float`, optional):
            The seed to sample the examples. Defaults to `0`.
        prompt_template (`str`, optional):
            The path to the prompt template. Defaults to `"templates/prompt_eae.txt"`.
        ensure_positives_on_train (bool, optional):
            Whether to ensure that the guidelines of annotated examples are not removed.
            Defaults to `True`.
        dataset_name (str, optional):
            The name of the dataset. Defaults to `None`.
        scorer (`str`, optional):
           The scorer class import string. Defaults to `None`.
        sample_only_gold_guidelines (`bool`, optional):
            Whether to sample only guidelines of present annotations. Defaults to `False`.
    """

    def __init__(
        self,
        dataset_loader: RAMSDatasetLoader,
        task: str = "EAE",
        split: str = "train",
        parallel_instances: Union[int, Tuple[int, int]] = 1,
        max_guidelines: int = -1,
        guideline_dropout: float = 0.0,
        seed: float = 0,
        prompt_template: str = "templates/prompt_eae.txt",
        ensure_positives_on_train: bool = True,
        dataset_name: str = None,
        scorer: str = None,
        sample_only_gold_guidelines: bool = True,
        **kwargs,
    ) -> None:
        assert task in [
            "EAE",
        ], f"{task} must be 'EAE'."

        super().__init__(
            dataset_loader=dataset_loader,
            task=task,
            split=split,
            parallel_instances=parallel_instances,
            max_guidelines=max_guidelines,
            guideline_dropout=guideline_dropout,
            seed=seed,
            prompt_template=prompt_template,
            ensure_positives_on_train=ensure_positives_on_train,
            sample_only_gold_guidelines=sample_only_gold_guidelines,
            dataset_name=dataset_name,
            scorer=scorer,
            task_definitions=EVENT_DEFINITIONS,
            task_target="labels",
            definitions=GUIDELINES,
            **kwargs,
        )

src/tasks/rams/data_loader.py
hitz-zentroa/GoLLIE
GUIDELINES = {
    "rams_accidentcrash": {
        "en": [
            (
                "An AccidentCrash (disaster) Event occurs when a Driver/Passenger in a Vehicle crashes into"
                ' CrashObject at some \nplace. The only possible event subtype is: "AccidentCrash".'
            ),
            (
                "An Unfortunate Incident (calamity) takes place when an Individual in a Motor Vehicle collides with an"
                ' Obstacle at a \ncertain location. The sole event subtype is designated as: "UnfortunateIncident".'
            ),
            (
                "A DisasterEvent takes place when an individual driving or riding in a Vehicle collides with an Object"
                ' at a specific \nlocation. The sole event subtype within this category is labeled as "AccidentCrash".'
            ),
            (
                "An UnintentionalCollision (calamity) incident takes place when an Operator/Passenger within a"
                " \nTransportationDevice impacts a CollisionObject at a particular location. The sole feasible event"
                ' subtype is: \n"UnintentionalCollision".'
            ),
            (
                "An Unfortunate Incident takes place when an Individual within a Vehicle collides with an Object at a"
                ' particular \nlocation. The sole event subtype encompassed in this category is: "Unfortunate'
                ' Incident".'
            ),
        ]
    },
    "rams_agreement": {
        "en": [
            (
                "An Agreement (government) Event occurs when a Participant signed, rejected, nullified or violated an"
                ' agreement at \nsome Place. The possible Event subtypes are: "Accept", "RejectNullify", "Violate" or'
                " None."
            ),
            (
                "An Agreement (government) Event transpires when a Participant enters into, refuses, annuls, or"
                ' breaches an \nagreement at a particular Location. The various Event subtypes include: "Accept",'
                ' "RejectNullify", "Violate", or None.'
            ),
            (
                "An Agreement (government) Event transpires when a Participant executes, refuses, annuls, or breaches"
                ' an \nagreement at a specific Location. The possible Event subtypes encompass: "Accept",'
                ' "RejectNullify", "Violate", or None.'
            ),
            (
                "An accord (government) occurrence transpires when a participant signs, declines, nullifies, or"
                ' breaches an \nagreement at a particular location. The feasible event subtypes encompass: "Accept",'
                ' "RejectNullify", "Violate", or \nabsence of a subtype.'
            ),
            (
                "An Agreement (government) Event transpires when a Participant enters into, refuses, annuls, or"
                ' breaches an \nagreement at a specific Location. The various Event subtypes encompass: "Accept",'
                ' "RejectNullify", "Violate", or None.'
            ),
        ]
    },
    "rams_arrestjaildetain": {
        "en": [
            (
                "An ArrestJailDetain (justice) Event occurs when a Jailer arrested or jailed a Detainee for a Crime at"
                ' some Place. The \nonly possible event subtype is: "ArrestJailDetain".'
            ),
            (
                'A "ArrestJailDetain (justice) Event" takes place when a Jailer apprehends and imprisons a Detainee'
                ' for a Crime at a \nspecific Location. The sole event subtype that can occur is "ArrestJailDetain.'
            ),
            (
                "A detainment event takes place when a jailer apprehends or imprisons a detainee due to a criminal"
                " offense at a \nspecific location. The sole event subtype of this nature is denoted as"
                ' "ArrestJailDetain.'
            ),
            (
                "A detainment event takes place when a jailer apprehends or imprisons a detainee for a criminal"
                " offense at a specific \nlocation. The sole event subtype of this nature is designated as"
                ' "ArrestJailDetain.'
            ),
            (
                "A detainment event takes place when a jailer apprehends or imprisons a detainee for a criminal"
                ' offense at a specific \nlocation. The sole event subtype of this kind is labeled "ArrestJailDetain".'
            ),
        ]
    },
    "rams_artifactfailure": {
        "en": [
            (
                "An ArtifactFailure (artifact) Event occurs whenever a (mechanical) Artifact failed due to some"
                ' Instrument at some \nPlace. The possible Event subtypes are: "MechanicalFailure"or None.'
            ),
            (
                "An ArtifactFailure (artifact) Event takes place whenever a Mechanical Artifact encounters a failure"
                " due to a \nspecific Instrument in a particular Place. The potential subtypes of this Event are"
                ' "MechanicalFailure" or "None.'
            ),
            (
                "An ArtifactFailure (artifact) Event takes place whenever a Mechanical Artifact experiences a failure"
                " as a result \nof an Instrument in a specific Place. The possible subtypes of this Event are either"
                ' "MechanicalFailure" or None.'
            ),
            (
                "An ArtifactFailure (artifact) Incident takes place whenever a Mechanical Artifact encounters a"
                " failure because \nof a specific Instrument situated in a particular Place. The potential subtypes of"
                ' this Event are \n"MechanicalFailure" or "None.'
            ),
            (
                "An ArtifactFailure (artifact) Event takes place whenever a Mechanical Artifact experiences a failure"
                " due to the \nuse of a specific Instrument in a particular Place. The possible subtypes of this Event"
                ' are either "MechanicalFailure" \nor "None.'
            ),
        ]
    },
    "rams_attack": {
        "en": [
            (
                "An Attack (conflict) Event occurs when an Attacker attacks a Target with some Instrument at some"
                ' Place. The possible \nEvent subtypes are: "AirStrikeMissileStrike",'
                ' "BiologicalChemicalPoisonAttack", "Bombing", "FireArmAttack", \n"Hanging", "Invade",'
                ' "SelfDirectedBattle", "SetFire", "Stabbing", "StealRobHijack", "Strangling"or None.'
            ),
            (
                "An Assault (conflict) Incident transpires when an Assailant launches an offensive against a Target"
                " utilizing a \nspecific Tool at a particular Location. The feasible Event subtypes encompass:"
                ' "AirStrikeMissileStrike", \n"BiologicalChemicalPoisonAttack", "Bombing", "FireArmAssault",'
                ' "Hanging", "Invasion", "Self-InflictedCombat", "SettingFire", "Stabbing", \n"TheftRobberyHijacking",'
                ' "Strangulation" or None.'
            ),
            (
                "An Assault (conflict) Incident transpires when an Aggressor launches an attack on a Target using a"
                " particular \nInstrument at a specific Location. The possible Subtypes of this Event are:"
                ' "AirStrikeMissileStrike", \n"BiologicalChemicalPoisonAttack", "Bombing", "FireArmAttack", "Hanging",'
                ' "Invasion", "Self-InflictedCombat", "SetFire", "Stabbing", \n"TheftRobberyHijacking", or'
                ' "Strangling", or it may be classified as None.'
            ),
            (
                "An Assault (conflict) Incident transpires when an Aggressor launches an Offensive against a Target"
                " utilizing a \nSpecific Weapon at a Particular Location. The conceivable Event subtypes encompass:"
                ' "AirStrikeMissileStrike", \n"BiologicalChemicalPoisonAttack", "Bombing", "FireArmAttack", "Hanging",'
                ' "Invade", "SelfDirectedBattle", "SetFire", "Stabbing", \n"StealRobHijack", "Strangling" or Absence.'
            ),
            (
                "An Assault (conflict) Incident transpires when an Aggressor targets a Victim with a specific Weapon"
                ' at a particular \nLocation. The possible subtypes of this Event include: "AirStrikeMissileStrike",'
                ' "BiologicalChemicalPoisonAttack", \n"Bombing", "FireArmAttack", "Hanging", "Invasion",'
                ' "Self-InflictedCombat", "SetFire", "Stabbing", \n"TheftRobberyHijacking", "Strangling" or "None".'
            ),
        ]
    },
    "rams_collaborate": {
        "en": [
            (
                "A Collaborate (contact) Event occurs when some Participants communicates (remotely or face-to-face)"
                ' at some \nPlace. The possible Event subtypes are: "Correspondence", "Meet"or None.'
            ),
            (
                "A collaborative event takes place when a group of participants engages in communication, either in"
                " person or \nremotely, at a specific location. The various subtypes of this event can include"
                ' "Correspondence," "Meet," or none.'
            ),
            (
                "A collaborative event takes place when a group of participants engages in communication, either in"
                " person or \nremotely, at a specific location. The various subtypes of this event include:"
                ' "Correspondence," "Meet," or none.'
            ),
            (
                "A collaborative event takes place when a group of participants engages in communication, either in"
                " person or \nremotely, at a specific location. The various subtypes of this event can include:"
                ' "Correspondence," "Meet," or none.'
            ),
            (
                "A collaborative event takes place when a group of participants engages in communication, either in"
                " person or \nremotely, at a specific location. The possible subtypes of this event are:"
                ' "Correspondence", "Meet", or none.'
            ),
        ]
    },
    "rams_commandorder": {
        "en": [
            (
                "A CommandOrder (contact) Event occurs when a Communicator communicates with/to Recipient about a"
                ' Topic at some \nPlace. The possible Event subtypes are: "Broadcast", "Correspondence", "Meet"or'
                " None."
            ),
            (
                "A CommandOrder event takes place when a Communicator engages in communication with/to a Recipient"
                " regarding a \nspecific Topic at a particular Place. The possible subtypes of this event are:"
                ' "Broadcast", "Correspondence", "Meet", or \nnone.'
            ),
            (
                "A CommandOrder event takes place when a Communicator engages in communication with/to a Recipient"
                " concerning a \nspecific Topic at a particular Place. The possible subtypes of this event are:"
                ' "Broadcast", "Correspondence", "Meet", or \nnone.'
            ),
            (
                "A CommandOrder (contact) Event takes place when a Communicator interacts with/to a Recipient"
                ' concerning a Topic at \na particular Place. The potential Event subtypes include: "Broadcast",'
                ' "Correspondence", "Meet", or "None".'
            ),
            (
                "A CommandOrder event takes place when a communicator engages with a recipient to discuss a specific"
                " topic at a \nparticular location. The various event subtypes that can arise from this interaction"
                ' include: "Broadcast," \n"Correspondence," "Meet," or no subtype specified.'
            ),
        ]
    },
    "rams_commitmentpromiseexpressintent": {
        "en": [
            (
                "A CommitmentPromiseExpressIntent (contact) Event occurs when a Communicator commits, promises,"
                " expresses an \nintent to Recipient about a Topic at some Place. The possible Event subtypes are:"
                ' "Broadcast", "Correspondence", \n"Meet"or None.'
            ),
            (
                "A CommitmentPromiseExpressIntent event takes place when a Communicator makes a commitment, promise,"
                " or \nexpresses their intent to a Recipient regarding a specific Topic at a particular Place. The"
                ' various Event subtypes include: \n"Broadcast," "Correspondence," "Meet," or "None.'
            ),
            (
                "A CommitmentPromiseExpressIntent event takes place when a Communicator makes a commitment, pledges,"
                " or \nexpresses their intent to a Recipient regarding a specific Topic at a particular Place. The"
                ' possible subtypes of this event are \n"Broadcast," "Correspondence," "Meet," or none.'
            ),
            (
                "A CommitmentPromiseExpressIntent event takes place when a Communicator makes a commitment, pledges,"
                " or \nexpresses their intent to a Recipient regarding a specific Topic at a particular Place. The"
                ' various Event subtypes include: \n"Broadcast," "Correspondence," "Meet," or "None.'
            ),
            (
                "A CommitmentPromiseExpressIntent event takes place when a Communicator makes a commitment, expresses"
                " an \nintention, or promises to a Recipient regarding a particular Topic at a specific Place. The"
                ' various Event subtypes include: \n"Broadcast," "Correspondence," "Meet," or none.'
            ),
        ]
    },
    "rams_convene": {
        "en": [
            (
                "A Convene (government) Event occurs when a Convener convened a ConvenedEntity at some Place. The only"
                ' possible \nevent subtype is: "Convene".'
            ),
            (
                "A government event known as a Convene occurs when a Convener brings together a ConvenedEntity at a"
                ' specific \nlocation. The sole event subtype of this kind is labeled "Convene.'
            ),
            (
                'A government event called "Convene" takes place when a convener brings together a ConvenedEntity at a'
                ' specific \nlocation. The sole event subtype of this kind is "Convene.'
            ),
            (
                "A government gathering, referred to as a Convene Event, takes place when a Convener brings together a"
                ' \nConvenedEntity at a specific location. The sole event subtype of this kind is labeled as "Convene.'
            ),
            (
                "A government event known as a Convene occurs when a Convener brings together a ConvenedEntity at a"
                ' specific \nlocation. The sole event subtype associated with this situation is "Convene".'
            ),
        ]
    },
    "rams_coup": {
        "en": [
            (
                "A Coup (conflict) Event occurs when a DeposedEntity was desposed by a DeposingEntity at some Place."
                ' The only \npossible event subtype is: "Coup".'
            ),
            (
                'A "Coup (conflict)" event takes place when a "DeposedEntity" is deposed by a "DeposingEntity" at a'
                ' specific \n"Place." The sole event subtype associated with this scenario is "Coup.'
            ),
            (
                'A "Coup (conflict)" event takes place when a "DeposedEntity" is deposed by a "DeposingEntity" at a'
                ' specific \n"Place." The sole possible subtype of this event is "Coup.'
            ),
            (
                'A "Coup (conflict)" event takes place when a "DeposedEntity" is deposed by a "DeposingEntity" at a'
                ' specific \n"Place". The sole possible subtype of this event is "Coup".'
            ),
            (
                'A "Coup (conflict)" event takes place when a "DeposedEntity" is deposed by a "DeposingEntity" at a'
                ' specific \n"Place". The sole event subtype that can occur is "Coup".'
            ),
        ]
    },
    "rams_damagedestroy": {
        "en": [
            (
                "A DamageDestroy (artifact) Event occurs when an Artifact is damaged or destroyed by some Agent"
                " (damager or \ndestroyer) using an Instrument at some Place. The possible Event subtypes are:"
                ' "Damage", "Destroy" or None.'
            ),
            (
                "A DamageDestroy (artifact) Event transpires when an Artifact endures harm or is annihilated by an"
                " Agent (the \ninflictor or demolisher) utilizing an Instrument at a particular Place. The attainable"
                ' Event subtypes include: "Damage," \n"Destroy," or None.'
            ),
            (
                "A DamageDestroy (artifact) Event transpires when an Artifact endures harm or is annihilated by an"
                " Agent (the one \ninflicting damage or destruction) employing an Instrument at a particular Place."
                ' The potential Event subtypes encompass: \n"Damage," "Destroy," or None.'
            ),
            (
                "A DamageDestroy (artifact) Event transpires when an Artifact endures harm or is annihilated by an"
                " Agent (the one \ninflicting damage or destruction) employing an Instrument at some Location. The"
                ' potential Event subtypes encompass: \n"Damage," "Destroy," or Absence of either.'
            ),
            (
                "A DamageDestroy (artifact) Event transpires when an Artifact endures harm or is annihilated by an"
                " Agent (the one \ninflicting damage or causing destruction) employing an Instrument at a particular"
                ' Place. The feasible Event subtypes \nencompass: "Damage," "Destroy," or Absent.'
            ),
        ]
    },
    "rams_demonstrate": {
        "en": [
            (
                "A Demonstrate (conflict) Event occurs when a Demonstrator(s) protest at some Place. The possible"
                ' Event subtypes \nare: "MarchProtestPoliticalGathering"or None.'
            ),
            (
                "A Conflict Event takes place when a group of Demonstrators gather at a specific Location to protest."
                ' The possible \nsubtypes of this event are either "MarchProtestPoliticalGathering" or no specific'
                " subtype."
            ),
            (
                "A Conflict Event takes place when Demonstrators gather in protest at a specific Location. The"
                ' possible subtypes of \nthis event are either "MarchProtestPoliticalGathering" or none.'
            ),
            (
                "A Conflict Event takes place when a group of Demonstrators gather at a specific Location to protest."
                ' The possible \nsubtypes of this event are: "MarchProtestPoliticalGathering" or none.'
            ),
            (
                "A Conflict Event takes place when a group of Demonstrators gather at a specific Location to protest."
                ' The possible \nsubtypes of this event are either "MarchProtestPoliticalGathering" or none.'
            ),
        ]
    },
    "rams_die": {
        "en": [
            (
                "A Die (life) Event occurs when a Victim died in some Place. The reason of the death can be: killed by"
                " a Killer or by a \nMedicalIssue, with or without an Instrument. The possible Event subtypes are:"
                ' "Violent", "NonViolent"or None.'
            ),
            (
                "A Die (life) Event takes place when an individual, the Victim, passes away at a specific location."
                " The cause of their \ndemise can be attributed to either a Killer or a MedicalIssue, with or without"
                ' the use of an Instrument. The various Event \nsubtypes include: "Violent", "NonViolent", or none.'
            ),
            (
                "A Die (life) Event transpires when a Victim expires at a specific Location. The cause of their demise"
                " can be \nattributed to either a Killer or a Medical Issue, with or without the use of an Instrument."
                ' The various Event subtypes encompass: \n"Violent", "NonViolent", or None.'
            ),
            (
                "A Life Event takes place when an individual passes away at a specific Location. The cause of their"
                " demise may be due to \nbeing killed by an assailant, an unfortunate medical issue, or a combination"
                " of both, with or without the use of a tool. The \npossible subtypes of such events include:"
                ' "Violent", "Non-Violent", or none.'
            ),
            (
                "A Die (life) Event transpires when a Victim experiences death at a particular Place. The cause of"
                " death can be \nattributed to either a Killer or a MedicalIssue, with or without the involvement of"
                ' an Instrument. The various Event subtypes \ninclude: "Violent", "NonViolent", or none.'
            ),
        ]
    },
    "rams_discussion": {
        "en": [
            (
                "A Discussion (contact) Event occurs when some Participants discuss at some place. The possible Event"
                ' subtypes are: \n"Correspondence", "Meet"or None.'
            ),
            (
                "A conversation event, referred to as a Discussion (contact) Event, takes place when a group of"
                " individuals engage in \ndialogue at a specific location. There are three possible subtypes of this"
                ' event: "Correspondence", "Meet", or none.'
            ),
            (
                "A communication event, referred to as a Discussion (contact) Event, takes place when a group of"
                " individuals engage \nin a conversation at a specific location. The three possible subtypes of this"
                ' event are: "Correspondence", "Meet", or \nno specified subtype.'
            ),
            (
                "A communication (contact) event takes place when a number of individuals engage in conversation at a"
                ' particular \nlocation. The possible subtypes of this event are: "Correspondence", "Meet", or no'
                " specific subtype."
            ),
            (
                "A conversation (contact) event takes place when a few individuals engage in dialogue at a specific"
                ' location. The \npotential subtypes of this event are: "Correspondence", "Meet", or "None".'
            ),
        ]
    },
    "rams_diseaseoutbreak": {
        "en": [
            (
                "A DiseaseOutbreak (disaster) Event occurs when a Disease broke out among some Victims at some Place."
                ' The only \npossible event subtype is: "DiseaseOutbreak".'
            ),
            (
                "A DiseaseOutbreak (disaster) incident takes place when a particular illness spreads amongst a group"
                " of \nindividuals situated in a specific location. The sole event subtype within this category is"
                ' labeled as "DiseaseOutbreak".'
            ),
            (
                "A Disease Outbreak (disaster) takes place when an illness spreads among a group of individuals within"
                " a specific \nlocation. The sole event subtype that can arise from this situation is"
                ' "DiseaseOutbreak.'
            ),
            (
                "A Disease Outbreak event takes place when an illness spreads among a group of individuals within a"
                " specific \nlocation. The sole event subtype associated with this scenario is labeled as"
                ' "DiseaseOutbreak.'
            ),
            (
                "A Disease Outbreak (disaster) takes place when an illness spreads among a group of individuals in a"
                ' specific \nlocation. The sole event subtype within this category is labeled "DiseaseOutbreak.'
            ),
        ]
    },
    "rams_elect": {
        "en": [
            (
                "An Elect (personell) Event occurs when a Voter elects a Candidate in some Place. The possible Event"
                ' subtypes are: \n"WinElection"or None.'
            ),
            (
                "An Elect (personnel) Event takes place when a Voter selects a Candidate at a specific Location. The"
                ' possible \nsubtypes of this Event are either "WinElection" or none.'
            ),
            (
                "An Elect (personnel) Event transpires when a Voter selects a Candidate within a specific Location."
                ' The possible \nEvent subtypes encompass either "WinElection" or none.'
            ),
            (
                "An Elect (personnel) Event transpires when a Voter selects a Candidate within a specific Place. The"
                ' potential Event \nsubtypes encompass "WinElection" or none.'
            ),
            (
                "An Elect (personnel) Event takes place when a Voter chooses a Candidate within a specific Location."
                ' The possible \nsubtypes of this Event are either "WinElection" or none.'
            ),
        ]
    },
    "rams_endpossition": {
        "en": [
            (
                "An EndPossition (personell) Event occurs when an Employee stops working on a Organization in some"
                ' Place. The \npossible Event subtypes are: "FiringLayOff", "QuitRetire"or None.'
            ),
            (
                "A EndPossition (personnel) Event takes place when a Staff member concludes their work with an"
                ' Organization in a \nspecific Location. The potential Event subtypes include: "FiringLayOff",'
                ' "QuitRetire", or None.'
            ),
            (
                "An EndPosition (personnel) Event transpires when a Staff member concludes their work with an Entity"
                ' in a certain \nLocation. The possible Event subtypes include "FiringLayOff", "QuitRetire", or None.'
            ),
            (
                "A EndPossition (personnel) Event transpires when a Staff member concludes their work at a Business"
                ' within a \nLocation. The potential Event subcategories include: "FiringLayOff", "QuitRetire", or'
                " None."
            ),
            (
                "A EndPossition (personnel) Event takes place when a Staff member concludes their work with an Entity"
                ' in a Location. \nThe possible Event subtypes include: "FiringLayOff", "QuitRetire", or None.'
            ),
        ]
    },
    "rams_fireexplosion": {
        "en": [
            (
                "A FireExplosion (disaster) Event occurs when a FireExplosionObject caught fire or exploded from an"
                ' Instrument at \nsome Place. The only possible event subtype is: "FireExplosion".'
            ),
            (
                "A FireExplosion catastrophe takes place when a FireExplosionObject undergoes combustion or detonation"
                " as a \nconsequence of a Device at a certain Location. The sole event subtype possibility is:"
                ' "FireExplosion".'
            ),
            (
                "A FireExplosionCatastrophe takes place when a FireExplosionArtifact experiences a blaze or detonation"
                " \noriginating from a Tool at a certain Location. The sole feasible event subtype is:"
                ' "FireExplosion".'
            ),
            (
                "A FireExplosion catastrophe takes place when a FireExplosionObject experiences a blaze or detonation"
                " as a \nconsequence of being utilized at a particular Location. The sole feasible event subtype is:"
                ' "FireExplosion".'
            ),
            (
                "A FireExplosion catastrophe takes place when a FireExplosionObject experiences a blaze or detonation"
                " as a \nconsequence of being utilized at a specific Location. The sole feasible event subtype is:"
                ' "FireExplosion".'
            ),
        ]
    },
    "rams_formation": {
        "en": [
            (
                "A Formation (government) Event occurs when some Participant (GPEs) are merged or when a GPE is"
                ' formed/started by a \nFounder at some Place. The possible Event subtypes are: "Merge", "Start",'
                " None."
            ),
            (
                "A government formation event takes place when one or more participating entities (GPEs) undergo a"
                " process of merger \nor when a new GPE is established by a founder at a specific location. The"
                ' various subtypes of such events include \n"Merge," "Start," and "None.'
            ),
            (
                "A government formation event takes place when one or more participating political entities (GPEs)"
                " undergo a \nprocess of merging or when a new GPE is established by an initiator, known as the"
                " founder, at a specific location. The various \ntypes of event subtypes that can occur during this"
                ' process include "Merge," "Start," and "None.'
            ),
            (
                "A government reorganization event takes place when one or more participating entities (GPEs) undergo"
                " a process of \nconsolidation or when a new GPE is established by an initiator (Founder) at a"
                ' specific location. The various subtypes of this event \ninclude "Merge," "Start," and "None.'
            ),
            (
                "A government reorganization event takes place when one or more participating entities (GPEs) undergo"
                " a process of \nmerger or when a new GPE is established by an initiator (Founder) at a specific"
                ' location. The various subtypes of such events \ninclude "Merge," "Start," and "None.'
            ),
        ]
    },
    "rams_funeralvigil": {
        "en": [
            (
                "A FuneralVigil (contact) Event occurs when some Participants communicate during a funeral or vigil"
                ' for Deceased at \nsome Place. The possible Event subtypes are: "Meet"or None.'
            ),
            (
                "A FuneralVigil (contact) Event takes place when a group of Participants engage in communication"
                " during a funeral or \nvigil for Deceased at a specific Place. The possible subtypes of this Event"
                ' are either "Meet" or "None.'
            ),
            (
                "A FuneralVigil (contact) Event takes place when specific Participants interact with each other during"
                ' a funeral or \nvigil for Deceased at a particular Place. The possible Event subtypes include "Meet"'
                ' or "None.'
            ),
            (
                "A FuneralVigil (contact) Event transpires when specific Participants interact during a commemoration"
                ' for \nDeceased at a particular Place. The potential Event subtypes include "Meet" or none.'
            ),
            (
                "A FuneralVigil (contact) Event takes place when specific Participants interact during a funeral or"
                ' vigil for the \nDeceased at a particular Place. The possible Event subtypes include "Meet" or none.'
            ),
        ]
    },
    "rams_genericcrime": {
        "en": [
            (
                "A GenericCrime (crime) Event occurs when a Perpetrator commits a crime against a Victim at some"
                ' Place. The only \npossible event subtype is: "GenericCrime".'
            ),
            (
                "A GenericCrime (crime) incident takes place when a Perpetrator carries out an illegal act against a"
                ' Victim at a \nspecific Location. The sole event subtype is classified as "GenericCrime".'
            ),
            (
                "A GenericCrime (crime) Incident transpires when a Perpetrator carries out a crime against a Victim at"
                ' a specific \nLocation. The sole event subtype is: "GenericCrime".'
            ),
            (
                "A GenericCrime (crime) incident takes place when a Perpetrator carries out a criminal act against a"
                ' Victim at a \nspecific Location. The sole event subtype is labeled "GenericCrime".'
            ),
            (
                "A GenericCrime (crime) incident takes place when a Perpetrator carries out an offense against a"
                ' Victim at a specific \nLocation. The sole event subtype is labeled as "GenericCrime".'
            ),
        ]
    },
    "rams_initiatejudicialprocess": {
        "en": [
            (
                "An InitiateJudicialProcess (justice) Event occurs when a Prosecutor charged, indicted, tried or"
                " initiated a \njudicial process pertaining to a Defendant before a JudgeCourt for a Crime in some"
                ' Place. The possible Event subtypes are: \n"ChargeIndict", "TrialHearing"or None.'
            ),
            (
                "A JudicialProcessInitiation (justice) event takes place when a Prosecutor institutes legal"
                " proceedings against \na Defendant in the presence of a Judge in a Court for an alleged Crime at a"
                ' specific Location. The possible subtypes of \nthis event are: "ChargeIndict", "TrialHearing", or'
                " none."
            ),
            (
                "An InitiateJudicialProcess (justice) Event takes place when a Prosecutor commences a legal procedure"
                " involving a \nDefendant in the presence of a Judge within a Court setting because of a Crime that"
                ' has occurred at a specific Location. The \npossible Event subtypes include: "ChargeIndict,"'
                ' "TrialHearing," or none.'
            ),
            (
                "A JudicialProcessInitiation (justice) event takes place when a Prosecutor instigates, charges, tries,"
                " or \ncommences a legal proceeding involving a Defendant in a Courtroom before a Judge for a Crime at"
                ' a specific Location. The \npossible Event subtypes include: "ChargeIndict", "TrialHearing", or none.'
            ),
            (
                "A JudicialProcessInitiation (justice) event takes place when a Prosecutor instigates, indicts, tries,"
                " or \ncommences a legal procedure involving a Defendant in the presence of a JudgeCourt for an"
                " offense that occurs in a specific \nLocation. The potential subtypes of this event are"
                ' "ChargeIndict," "TrialHearing," or none.'
            ),
        ]
    },
    "rams_injure": {
        "en": [
            (
                "An Injure (life) Event occurs when a Victim is injured in some Place. The reasons can be: have"
                " extreme hunger or \nthirst, some physical degradation and sickness or illness by some MedicalIssue"
                " either infected by some Disease or imposed \nby some Injurer. The injures can be caused by some"
                ' Instrument. The possible Event subtypes are: "HungerThirst", \n"Physical", "Sickness", "Violent"or'
                " None"
            ),
            (
                "An Injure (life) Event takes place when a Victim suffers an injury at a particular Location. The"
                " causes may include: \nexperiencing extreme hunger or thirst, undergoing physical deterioration or"
                " illness due to a MedicalIssue, which could be the \nresult of an infection from a Disease or"
                " inflicted by an Injurer. The injuries may be inflicted by a specific Instrument. The \npossible"
                ' Event subtypes include: "HungerThirst", "Physical", "Sickness", "Violent", or "None".'
            ),
            (
                "An Injury (life) Event takes place when a Victim suffers an injury at a specific Location. The"
                " possible causes may \ninclude: experiencing extreme hunger or thirst, undergoing physical"
                " deterioration or sickness due to a MedicalIssue, \nwhich could be caused by an infection from a"
                " Disease or inflicted by an Injurer. The injuries may be inflicted by a \nparticular Instrument. The"
                ' various Event subtypes include: "HungerThirst", "Physical", "Sickness", "Violent", or "None".'
            ),
            (
                "An Injury (life) Event takes place when a Victim suffers an injury at a specific Location. These"
                " events can be \nattributed to several causes, such as severe hunger or thirst, physical"
                " deterioration, illness resulting from a \nMedicalIssue, which could be due to an infection from a"
                " Disease or inflicted by an Injurer. The injuries may be inflicted by a \nparticular Instrument. The"
                ' various subtypes of these events include: "HungerThirst", "Physical", "Sickness", "Violent", \nor'
                ' "None".'
            ),
            (
                "An Injury (life) Event transpires when a Victim suffers an injury at a specific Location. The"
                " potential causes can \nencompass: experiencing extreme hunger or thirst, enduring physical"
                " deterioration or sickness, or being afflicted by a \nMedical Issue, either infected by a Disease or"
                " subjected to harm by an Injurer. The injuries may be inflicted by a particular \nInstrument. The"
                ' possible Event subtypes include: "HungerThirst", "Physical", "Sickness", "Violent", or "None".'
            ),
        ]
    },
    "rams_investigatecrime": {
        "en": [
            (
                "An InvestigateCrime (justice) Event occurs when a Investigator investigated a Defendant for a Crime"
                ' in some Place. \nThe possible Event subtypes are: "InvestigateCrime"or None.'
            ),
            (
                "An Event of InvestigateCrime (related to justice) takes place when an Investigator looks into the"
                " actions of a \nDefendant concerning a Crime at a specific Location. The possible subtypes of this"
                ' Event are "InvestigateCrime" or none.'
            ),
            (
                "An InvestigateCrime (justice) Event transpires when an Investigator explores a Defendant's"
                " involvement in a \nCrime at a specific Location. The potential Event subtypes encompass"
                ' "InvestigateCrime" or none.'
            ),
            (
                "An Event of InvestigateCrime (related to justice) takes place when an Investigator explores the"
                " actions of a \nDefendant pertaining to a Crime at a specific Location. The potential subtypes of"
                ' this Event are "InvestigateCrime" or none.'
            ),
            (
                "A InvestigateCrime (justice) Event transpires when an Investigator examines a Defendant regarding a"
                ' Crime in a \nspecific Location. The possible Event subtypes include "InvestigateCrime" or none.'
            ),
        ]
    },
    "rams_judicialconsequences": {
        "en": [
            (
                "A JudicialConsequence (justice) Event occurs when a JudgeCourt decided the consequences, convicted,"
                " executed or \nextradited a Defendant for a Crime in some Place. The possible Event subtypes are:"
                ' "Convict", "Execute", "Extradite"or None.'
            ),
            (
                "A LegalOutcome event transpires when a CourtOfficer determines the consequences, convicts, sentences,"
                " or \nextradites an AccusedIndividual for a Crime at a particular Location. The potential event"
                ' subtypes encompass: "Convict," \n"Sentence," "Extradite," or None.'
            ),
            (
                "A JudicialConsequence event takes place when a JudgeCourt makes a determination of the consequences,"
                " finds \nguilty, imposes a penalty, or extradites a Defendant for a Crime in a specific Location. The"
                ' various Event subtypes include: \n"Convict", "Execute", "Extradite", or "None".'
            ),
            (
                "A JudicialConsequence event takes place when a JudgeCourt reaches a decision, imposing penalties such"
                " as \nconviction, execution, or extradition on a Defendant for a Crime that occurred at a specific"
                ' Place. The various subtypes of this \nevent could be: "Convict", "Execute", "Extradite", or none.'
            ),
            (
                "A JudicialConsequence event takes place when a JudgeCourt reaches a decision, imposing penalties such"
                " as \nconviction, execution, or extradition on a Defendant for a crime that occurred in a specific"
                ' location. The various subtypes of \nthis event can be categorized as: "Convict", "Execute",'
                ' "Extradite", or none.'
            ),
        ]
    },
    "rams_legislate": {
        "en": [
            (
                "A Legislate (government) Event occurs when a Government legislature enacted a Law in some Place. The"
                ' only possible \nevent subtype is: "Legislate".'
            ),
            (
                'A "Legislate (government) Event" takes place when a government legislature passes a law in a specific'
                ' location. The \nsole event subtype associated with this occurrence is called "Legislate.'
            ),
            (
                'A "Legislate (government) Event" takes place when a government legislature institutes a law in a'
                ' specific \nlocation. The sole event subtype is labeled "Legislate.'
            ),
            (
                "A Legislative Event transpires when a governing body's legislature promulgates a statute in a"
                ' specific locale. The \nsole event subtype is designated as "Legislate.'
            ),
            (
                "A Legislative Event transpires when a governing body's legislative branch enacts a law within a"
                ' specific location. \nThe sole event subtype is designated as "Legislate.'
            ),
        ]
    },
    "rams_manufactureartifact": {
        "en": [
            (
                "A ManufactureArtifact (manufacture) Event occurs when a Manufacturer manufactured, created or"
                ' produced an \nArtifact using an Instrument at some Place. The possible Event subtypes are: "Build",'
                ' "IntellectualProperty", \n"CreateManufacture"or None.'
            ),
            (
                "A ManufactureArtifact event takes place when a Manufacturer engages in the production of an Artifact"
                " through the \nuse of an Instrument at a specific Location. The various Event subtypes that can be"
                ' encountered include: "Build", \n"IntellectualProperty", "CreateManufacture", or none.'
            ),
            (
                "A ManufactureArtifact (manufacture) event transpires when a Manufacturer engages in the production of"
                ' an \nArtifact via an Instrument at a specific Location. The event\'s subtypes may include: "Build",'
                ' "IntellectualProperty", \n"CreateManufacture", or none.'
            ),
            (
                "A ManufactureArtifact event takes place when a Manufacturer engages in the production, creation, or"
                " \nmanufacturing of an Artifact utilizing a specific Instrument at a particular Place. The various"
                ' Event subtypes that can arise from \nthis process include: "Build", "IntellectualProperty",'
                ' "CreateManufacture", or none.'
            ),
            (
                "A ManufactureArtifact (manufacture) event takes place when a Manufacturer engages in the production"
                " of an \nArtifact through the use of an Instrument at a specific Place. The various Event subtypes"
                ' that can arise from this process \ninclude: "Build", "IntellectualProperty", "CreateManufacture", or'
                " none."
            ),
        ]
    },
    "rams_mediastatement": {
        "en": [
            (
                "A MediaStatement (contact) Event occurs when a Communicator communicates something on media to some"
                ' Recipients at \nsome Place. The possible Event subtypes are: "Broadcast"or None.'
            ),
            (
                "A MediaStatement event takes place when a Communicator conveys information through media to a group"
                ' of Recipients \nat a specific Location. The possible subtypes of this event are either "Broadcast"'
                " or none."
            ),
            (
                "A MediaStatement event takes place when a communicator conveys information through media to a group"
                ' of recipients \nat a specific location. The possible subtypes of this event are either "Broadcast"'
                " or none."
            ),
            (
                "A MediaStatement (contact) Event transpires when a Communicator conveys information via media to a"
                ' collection of \nRecipients at a particular Place. The potential Event subtypes encompass "Broadcast"'
                " or none."
            ),
            (
                "A MediaStatement (contact) Event takes place when a Communicator conveys information through media to"
                ' a group of \nRecipients at a specific Location. The potential subtypes of this Event are "Broadcast"'
                " or none."
            ),
        ]
    },
    "rams_medicalintervention": {
        "en": [
            (
                "A MedicalIntervention (medical) Event occurs when a Treater treated a Patient for a MedicalIssue by"
                ' means of some \nInstrument at some Place. The only possible event subtype is: "Intervention".'
            ),
            (
                "A Medical Intervention (medical) incident takes place when a Treater addresses a Patient's Medical"
                " Issue through \nthe use of a specific Instrument at a particular Place. The sole event subtype"
                ' within this category is labeled \n"Intervention.'
            ),
            (
                "A Medical Intervention (medical) incident takes place when a Treater handles a Patient's Medical"
                " Issue using a \nspecific Tool at a particular Location. The sole event subtype that can arise from"
                ' this situation is: "Intervention.'
            ),
            (
                "A MedicalIntervention (medical) Incident takes place when a Treater manages a Patient's MedicalIssue"
                " utilizing a \nparticular Instrument at a specific Location. The sole possible event subtype is:"
                ' "Intervention.'
            ),
            (
                "A MedicalIntervention event takes place when a healthcare provider, referred to as the Treater,"
                " addresses a \nPatient's MedicalIssue through the use of a specific tool or method, known as the"
                " Instrument, at a particular location called \nthe Place. The sole event subtype within this context"
                ' is labeled "Intervention.'
            ),
        ]
    },
    "rams_negotiate": {
        "en": [
            (
                "A Negotiate (contact) Event occurs when some Participants participate on a negotiation about some"
                ' Topic at some \nPlace. The possible Event subtypes are: "Correspondence", "Meet"or None.'
            ),
            (
                "A Negotiation Participation Event takes place when specific individuals engage in discussions about a"
                " particular \nsubject at a certain location. The possible subtypes of this event include:"
                ' "Correspondence", "Meeting", or none.'
            ),
            (
                "A Negotiate (contact) Event takes place when specific Participants engage in negotiations concerning"
                " a \nparticular Topic at a particular Place. The possible subtypes of this Event include"
                ' "Correspondence," "Meet," or none.'
            ),
            (
                "A Negotiate (contact) Event takes place when specific Participants engage in negotiations concerning"
                " a \nparticular Topic at a particular Place. The possible subtypes of this Event are"
                ' "Correspondence", "Meet", or none.'
            ),
            (
                "A Negotiate (interaction) Event takes place when specific individuals engage in discussions about a"
                " particular \nsubject at a particular location. The possible subtypes of this event are"
                ' "Correspondence", "Meet", or none.'
            ),
        ]
    },
    "rams_prevarication": {
        "en": [
            (
                "A Prevarication (contact) Event occurs when a Communicator prevaricate about some Topic to a"
                ' Recipient at some \nPlace. The possible Event subtypes are: "Broadcast", "Correspondence", "Meet"or'
                " None."
            ),
            (
                "A Deception (contact) Incident transpires when a Communicator engages in deceptive communication"
                " about a \nparticular subject to a Recipient in a particular setting. The various subtypes of this"
                ' event could be: "Broadcast", \n"Correspondence", "Meet" or absence of any.'
            ),
            (
                "A Deception (communication) Incident transpires when an Informer distorts the truth about a"
                " particular subject to \na Receiver at a specific Location. The possible subcategories of this event"
                ' include: "Broadcast", \n"Correspondence", "Gathering", or "None".'
            ),
            (
                "A Misrepresentation Incident transpires when an Informer misleads another party, the Receiver,"
                " regarding a \nparticular subject matter at a specific location. The various subtypes of this event"
                ' include: "Broadcast," \n"Correspondence," "Meet," or none.'
            ),
            (
                "A Deception (communication) Incident transpires when an Informer engages in deceptive discourse about"
                " a \nparticular subject to a Receiver at a specific location. The various subtypes of this event"
                ' could be: "Broadcast", \n"Correspondence", "Gathering", or "None".'
            ),
        ]
    },
    "rams_publicstatementinperson": {
        "en": [
            (
                "A PublicStatementInPerson (contact) Event occurs when a Communicator gives a public statement to a"
                ' Recipient at \nsome Place. The possible Event subtypes are: "Broadcast"or None.'
            ),
            (
                "A PublicStatementInPerson event takes place when a communicator delivers a public statement to a"
                ' recipient at a \nspecific location. The possible subtypes of this event are "Broadcast" or none.'
            ),
            (
                "A PublicStatementInPerson event transpires when a Communicator delivers a public statement to a"
                ' Recipient at a \nspecific Location. The possible subtypes of this event are "Broadcast" or absence'
                " of any subtype."
            ),
            (
                "A PublicStatementInPerson event takes place when a communicator delivers a public statement to a"
                ' recipient at a \nspecific location. The possible subtypes of this event are "Broadcast" or none.'
            ),
            (
                "A PublicStatementInPerson event takes place when a communicator delivers a public statement to a"
                ' recipient at a \nspecific location. The possible subtypes of this event are "Broadcast" or none.'
            ),
        ]
    },
    "rams_requestadvice": {
        "en": [
            (
                "A RequestAdvice (contact) Event occurs when a Communicator requests something or gives advice about a"
                ' Topic to a \nRecipient at some Place. The possible Event subtypes are: "Broadcast",'
                ' "Correspondence", "Meet"or None.'
            ),
            (
                "A RequestAdvice event takes place when a Communicator engages in a discussion about a specific"
                " subject with a \nRecipient at a particular location. The various subtypes of this event could be:"
                ' "Broadcast", "Correspondence", "Meet", or \nnone.'
            ),
            (
                "A RequestAdvice event takes place when a Communicator engages in a conversation, asking for guidance"
                " or providing \nadvice on a specific subject to a Recipient at a particular location. The various"
                ' subtypes of this event could be: \n"Broadcast", "Correspondence", "Meet", or none.'
            ),
            (
                "A RequestAdvice event takes place when a Communicator engages in a conversation, seeking information"
                " or offering \nadvice on a specific topic with a Recipient at a particular location. The various"
                ' subtypes of this event could be: \n"Broadcast", "Correspondence", "Meet", or none.'
            ),
            (
                "A RequestAdvice event takes place when a Communicator seeks information or offers advice on a"
                " specific topic and \ndirects it towards a Recipient at a particular location. The various subtypes"
                ' of this event could be: "Broadcast", \n"Correspondence", "Meet", or none.'
            ),
        ]
    },
    "rams_sensoryobserve": {
        "en": [
            (
                "A SensoryObserve (inspection) Event occurs when a Observer observed, inspected or monitored a"
                ' ObservedEntity in \nsome Place. The possible Event subtypes are: "InspectPeopleOrganization'
                ' "MonitorElection", \n"PhysicalInvestigateInspect"or None.'
            ),
            (
                "A SensoryObserve (inspection) event takes place when an Observer examines, scrutinizes, or keeps"
                " track of a \nObservedEntity within a specific Place. The various Event subtypes that can arise from"
                ' this include: \n"InspectPeopleOrganization", "MonitorElection", "PhysicalInvestigateInspect", or'
                " none."
            ),
            (
                "A SensoryObserve (inspection) event takes place when an Observer examines, observes, or keeps an eye"
                " on a \nObservedEntity within a specific Place. The different event subtypes include:"
                ' "InspectPeopleOrganization", \n"MonitorElection", "PhysicalInvestigateInspect", or none.'
            ),
            (
                "A SensoryObserve (inspection) event takes place when an Observer examines, observes, or keeps track"
                " of a \nObservedEntity within a specific Place. The various Event subtypes include:"
                ' "InspectPeopleOrganization", "MonitorElection", \n"PhysicalInvestigateInspect", or none.'
            ),
            (
                "A SensoryObserve (inspection) event takes place when an Observer examines, scrutinizes, or keeps an"
                " eye on a \nObservedEntity within a specific Place. The different event subtypes include:"
                ' "InspectPeopleOrganization", \n"MonitorElection", "PhysicalInvestigateInspect", or none.'
            ),
        ]
    },
    "rams_shortage": {
        "en": [
            (
                "A Shortage (artifact) Event occurs when a Experiencer experienced a shortage of Supply in some Place."
                ' The only \npossible event subtype is: "Shortage".'
            ),
            (
                'An Artifact Event called "Shortage" takes place when an Experiencer faces a scarcity of Supply in a'
                ' specific Place. \nThe sole event subtype associated with this situation is simply titled "Shortage".'
            ),
            (
                "An Artifact Shortage Incident takes place when an Experiencer faces a scarcity of Supplies in a"
                ' specific Location. \nThe sole event subtype that can arise is labeled "Shortage.'
            ),
            (
                "An Artifact Shortage Incident takes place when an Experiencer faces a scarcity of Supply in a"
                ' specific Location. The \nsole event subtype that can arise is labeled "Shortage.'
            ),
            (
                "An Artifact Shortage Incident takes place when an Experiencer faces a scarcity of Supplies in a"
                ' specific Location. \nThe sole event subtype that can occur is labeled "Shortage.'
            ),
        ]
    },
    "rams_spy": {
        "en": [
            (
                "An Spy (government) Event occurs when a Spy spied on a ObservedEntity to the benefit of a Beneficiary"
                ' in some Place. \nThe only possible event subtype is: "Spy".'
            ),
            (
                "An Intelligence Gathering incident takes place when a spy secretly collects information about an"
                " ObservedEntity \nfor the advantage of a Beneficiary in a specific Location. The sole event subtype"
                ' within this category is labeled as \n"Spy".'
            ),
            (
                "An Intelligence Gathering Incident takes place when a Spy secretly collects information about an"
                " ObservedEntity \nfor the advantage of a Beneficiary in a specific Location. The sole event subtype"
                ' that can arise is labeled as "Spy".'
            ),
            (
                "An Intelligence Gathering incident takes place when a spy secretly collects information about an"
                " ObservedEntity \nfor the advantage of a Beneficiary at a specific Location. The sole event subtype"
                ' within this category is labeled as \n"Spy.'
            ),
            (
                "An Intelligence Gathering Incident takes place when a Spy secretly collects information about an"
                " ObservedEntity \nfor the advantage of a Beneficiary in a specific Location. The sole event subtype"
                ' is labeled as "Spy".'
            ),
        ]
    },
    "rams_startpossition": {
        "en": [
            (
                "An StartPossition (personell) Event occurs when an Employee starts working on a Organization in some"
                ' Place. The \npossible Event subtypes are: "Hiring"or None.'
            ),
            (
                "A StartPossition (personnel) Event takes place when an Employee commences work for an Organization at"
                ' a specific \nLocation. The potential Event subtypes encompass "Hiring" or none.'
            ),
            (
                "A StartPossition (personnel) Event takes place when an Employee commences work at an Organization in"
                ' a specific \nLocation. The potential Event subtypes include "Hiring" or none.'
            ),
            (
                "A StartPossition (personnel) Event takes place when an Employee commences work for an Organization at"
                ' a specific \nLocation. The potential Event subtypes encompass "Hiring" or none.'
            ),
            (
                "A StartPossition (personnel) Event transpires when an Employee commences work for an Organization at"
                ' a specific \nLocation. The potential Event subtypes encompass "Hiring" or none.'
            ),
        ]
    },
    "rams_targetaimat": {
        "en": [
            (
                "A TargetAimAt (inspection) Event occurs when a Targeter physically targeted a Target with a"
                ' Instrument at some \nPlace. The only possible event subtype is: "TargetAimAt".'
            ),
            (
                "A TargetAimAt event takes place when a Targeter engages in the action of targeting an object,"
                " utilizing a specific \nTool at a certain Location. The sole event subtype associated with this"
                ' occurrence is labeled "TargetAimAt.'
            ),
            (
                "A TargetAimAt event takes place when a Targeter engages in the act of targeting an object, utilizing"
                " a specific tool \nor device at a particular location. The sole event subtype associated with this"
                ' occurrence is labeled "TargetAimAt.'
            ),
            (
                "A TargetAimAt event takes place when a Targeter engages in the action of targeting a Target with a"
                " specific \nInstrument at a particular Place. The sole event subtype associated with this occurrence"
                ' is labeled "TargetAimAt".'
            ),
            (
                "A TargetAimAt event takes place when a Targeter specifically directs a Target with a Tool at a"
                ' certain Location. The \nsole event subtype that can occur is labeled "TargetAimAt".'
            ),
        ]
    },
    "rams_threatencoerce": {
        "en": [
            (
                "A ThreatenCoerce (contact) Event occurs when a Communicator threats or coerces a Recipient about a"
                ' Topic at some \nPlace. The possible Event subtypes are: "Broadcast", "Correspondence", "Meet"or'
                " None."
            ),
            (
                "A ThreatenCoerce event takes place when a Communicator engages in threats or coercion against a"
                " Recipient \nregarding a specific Topic at a particular Place. The various subtypes of this event"
                ' include: "Broadcast," \n"Correspondence," "Meet," or none.'
            ),
            (
                "A ThreatenCoerce event takes place when a Communicator applies threats or coercion to a Recipient"
                " regarding a \nspecific Topic at a particular Place. The various Event subtypes that can be"
                ' encountered include: "Broadcast", \n"Correspondence", "Meet", or "None".'
            ),
            (
                "A ThreatenCoerce event transpires when a Communicator employs threats or coercion against a Recipient"
                ' regarding a \nTopic at a particular Place. The various Event subtypes encompass: "Broadcast",'
                ' "Correspondence", "Meet", or "None".'
            ),
            (
                "A ThreatenCoerce event transpires when a Communicator imposes threats or coercion on a Recipient"
                " concerning a \nTopic within a specific Place. The potential subtypes of this event include:"
                ' "Broadcast", "Correspondence", "Meet", or \nnone.'
            ),
        ]
    },
    "rams_transaction": {
        "en": [
            (
                "A Transaction (transaction) Event occurs when a transaction of some Artifact occurred between some"
                " Participants \nfor the benefit of a Beneficiary at some Place. The Participants can be characterized"
                " into Giver and Recipient. There \ncan be a preventer that prevents the transaction to occur. The"
                ' possible Event subtypes are: "EmbargoSanction", \n"GiftGrantProvideAid", "TransferControl"or None.'
            ),
            (
                "A Transaction Event takes place when an exchange of an Artifact transpires between specific"
                " Participants, \nresulting in a benefit for a Beneficiary at a particular Place. The Participants can"
                " be divided into two roles: Giver and \nRecipient. Additionally, there may be a preventer who"
                ' prevents the transaction from happening. The possible Event subtypes \ninclude: "EmbargoSanction",'
                ' "GiftGrantProvideAid", "TransferControl", or none.'
            ),
            (
                "A Transaction Event takes place when an exchange of an Artifact happens between particular"
                " Participants for the \nadvantage of a Beneficiary at a specific Location. The Participants can be"
                " divided into Giver and Recipient. Additionally, \nthere could be a preventer who stops the"
                ' transaction from happening. The possible Event subtypes include: \n"EmbargoSanction",'
                ' "GiftGrantProvideAid", "TransferControl", or none.'
            ),
            (
                "A Transaction Event takes place whenever an exchange of an Artifact happens between specific"
                " Participants for the \nadvantage of a Beneficiary at a particular Place. The Participants can be"
                " categorized into Giver and Recipient. \nAdditionally, there may be a preventer who prevents the"
                ' transaction from happening. The possible Event subtypes include: \n"EmbargoSanction",'
                ' "GiftGrantProvideAid", "TransferControl", or none.'
            ),
            (
                "A Transaction Event transpires when an exchange of an Artifact takes place between specific"
                " Participants, \nultimately providing a benefit to a Beneficiary at a particular Place. The"
                " Participants can be categorized into two roles: \nGiver and Recipient. Furthermore, there may be a"
                " preventer who hinders the transaction from happening. The possible \nEvent subtypes include:"
                ' "EmbargoSanction", "GiftGrantProvideAid", "TransferControl", or none.'
            ),
        ]
    },
    "rams_transfermoney": {
        "en": [
            (
                "A TransferMoney (transaction) Event occurs when a Giver gives Money to a Recipient for the benefit of"
                " Beneficiary in \nsome Place. There can be a preventer that prevents the transaction to occur. The"
                ' possible Event subtypes are: \n"BorrowLend", "EmbargoSanction", "GiftGrantProvideAid",'
                ' "PayForService", "Purchase"or None.'
            ),
            (
                "A TransferMoney event transpires when a Giver provides funds to a Recipient, ultimately benefiting a"
                " Beneficiary \nat a specific Location. There may be an obstacle, preventing the event from taking"
                ' place. The potential subtypes of \nthis event encompass: "BorrowLend", "EmbargoSanction",'
                ' "GiftGrantProvideAid", "PayForService", "Purchase", or \nnone.'
            ),
            (
                "A TransferMoney event transpires when a Giver provides funds to a Recipient, ultimately benefiting a"
                " Beneficiary \nin a specific Location. There may be a hindrance, termed a preventer, which impedes"
                ' the event from taking place. The \nvarious Event subtypes encompass: "BorrowLend",'
                ' "EmbargoSanction", "GiftGrantProvideAid", "PayForService", \n"Purchase", or none.'
            ),
            (
                "A TransferMoney event takes place when a Giver provides money to a Recipient for the advantage of"
                " Beneficiary at a \nspecific Location. There may be an obstacle that hinders the event from"
                ' occurring. The possible subtypes of this event \ninclude: "BorrowLend", "EmbargoSanction",'
                ' "GiftGrantProvideAid", "PayForService", "Purchase", or none.'
            ),
            (
                "A TransferMoney event takes place when a Giver provides money to a Recipient, ultimately benefiting"
                " the \nBeneficiary in a specific Location. There may be an obstacle, or preventer, that hinders the"
                ' event from occurring. The various \nEvent subtypes include: "BorrowLend", "EmbargoSanction",'
                ' "GiftGrantProvideAid", "PayForService", "Purchase", or \nnone.'
            ),
        ]
    },
    "rams_transferownership": {
        "en": [
            (
                "A TransferOwnership (transaction) Event occurs when a Giver gives some Artifact to a Recipient for"
                " the benefit of \nBeneficiary in some Place. There can be a preventer that prevents the transaction"
                ' to occur. The possible Event subtypes are: \n"BorrowLend", "EmbargoSanction", "GiftGrantProvideAid",'
                ' "Purchase"or None.'
            ),
            (
                "A TransferOwnership event takes place when a giver transfers an artifact to a recipient for the"
                " advantage of a \nbeneficiary in a specific location. There may be a preventer that hinders the event"
                ' from happening. The possible event subtypes \ninclude: "BorrowLend", "EmbargoSanction",'
                ' "GiftGrantProvideAid", "Purchase", or none.'
            ),
            (
                "A TransferOwnership event takes place when a Giver transfers an Artifact to a Recipient, ultimately"
                " benefiting the \nBeneficiary in a specific Place. There may be a preventer who hinders the event"
                ' from occurring. The possible Event subtypes \ninclude: "BorrowLend", "EmbargoSanction",'
                ' "GiftGrantProvideAid", "Purchase", or "None".'
            ),
            (
                "A TransferOwnership event takes place when a giver transfers an artifact to a recipient for the"
                " advantage of a \nbeneficiary within a specific location. There may be a preventer who hinders the"
                ' event from happening. The possible subtypes of \nthis event are: "BorrowLend", "EmbargoSanction",'
                ' "GiftGrantProvideAid", "Purchase", or none.'
            ),
            (
                "A TransferOwnership event takes place when a Giver transfers an Artifact to a Recipient, ultimately"
                " benefiting the \nBeneficiary in a specific Place. There may be a preventer who hinders the event"
                ' from occurring. The possible Event subtypes \ninclude: "BorrowLend", "EmbargoSanction",'
                ' "GiftGrantProvideAid", "Purchase", or "None".'
            ),
        ]
    },
    "rams_transportartifact": {
        "en": [
            (
                "A TransportArtifact (transport) Event occurs when a Transporter transports an Artifact from the"
                " Origin to the \nDestination in some Vehicle. For some cases a Preventer entity can prevent the"
                ' Transporter to transport the Artifact. The \npossible Event subtypes are: "BringCarryUnload",'
                ' "DisperseSeparate", "Fall", "GrantEntry", "Hide", "LostOfControl", \n"NonViolentThrowLaunch",'
                ' "PreventEntry", "PreventExit", "ReceiveImport", "SendSupplyExport", "SmuggleExtract"or None.'
            ),
            (
                "A TransportArtifact (transport) Event takes place when a Transporter moves an Artifact from the"
                " Source to the \nDestination using a Vehicle. In certain situations, a Preventer may hinder the"
                " Transporter's ability to transport the \nArtifact. The possible Event subtypes include:"
                ' "BringCarryUnload", "DisperseSeparate", "Fall", "GrantEntry", "Hide", \n"LostOfControl",'
                ' "NonViolentThrowLaunch", "PreventEntry", "PreventExit", "ReceiveImport", "SendSupplyExport",'
                ' \n"SmuggleExtract", or none.'
            ),
            (
                "A TransportArtifact (transport) Event transpires when a Transporter moves an Artifact from the Origin"
                " to the \nDestination utilizing a Vehicle. In certain instances, a Preventer entity can impede the"
                " Transporter's ability to transport \nthe Artifact. The following Event subtypes may occur:"
                ' "BringCarryUnload", "DisperseSeparate", "Fall", \n"GrantEntry", "Hide", "LostOfControl",'
                ' "NonViolentThrowLaunch", "PreventEntry", "PreventExit", "ReceiveImport", \n"SendSupplyExport",'
                ' "SmuggleExtract", or None.'
            ),
            (
                "A TransportArtifact (transport) Event takes place when a Transporter moves an Artifact from the"
                " Origin to the \nDestination using a Vehicle. In certain situations, a Preventer entity can hinder"
                " the Transporter's ability to transport the \nArtifact. The possible Event subtypes include:"
                ' "BringCarryUnload", "DisperseSeparate", "Fall", "GrantEntry", "Hide", \n"LostOfControl",'
                ' "NonViolentThrowLaunch", "PreventEntry", "PreventExit", "ReceiveImport", "SendSupplyExport",'
                ' \n"SmuggleExtract" or None.'
            ),
            (
                "A TransportArtifact (transport) Event takes place when a Transporter moves an Artifact from the"
                " Origin to the \nDestination using a Vehicle. In certain situations, a Preventer may hinder the"
                " Transporter's ability to transport the \nArtifact. The possible Event subtypes include:"
                ' "BringCarryUnload", "DisperseSeparate", "Fall", "GrantEntry", "Hide", \n"LostOfControl",'
                ' "NonViolentThrowLaunch", "PreventEntry", "PreventExit", "ReceiveImport", "SendSupplyExport",'
                ' \n"SmuggleExtract", or None.'
            ),
        ]
    },
    "rams_transportperson": {
        "en": [
            (
                "A TransportPerson (transport) Event occurs when a Transporter transports itself or a Passanger from"
                " the Origin to \nthe Destination in some Vehicle. For some cases a Preventer entity can prevent the"
                " Transporter to transport itself or \nthe Passenger. The possible Event subtypes are:"
                ' "BringCarryUnload", "DisperseSeparate", "EvacuationRescue", \n"Fall", "GrantedAsylum", "Hide",'
                ' "PreventEntry", "PreventExit", "SelfMotion", "SmuggleExtract"or None.'
            ),
            (
                "A TransportPerson (transport) Event occurs when a Transporter moves itself or a Passenger from the"
                " Origin to the \nDestination using a Vehicle. In certain situations, a Preventer may hinder the"
                " Transporter's ability to transport itself or the \nPassenger. The possible Event subtypes include:"
                ' "BringCarryUnload", "DisperseSeparate", "EvacuationRescue", "Fall", \n"GrantedAsylum", "Hide",'
                ' "PreventEntry", "PreventExit", "SelfMotion", "SmuggleExtract", or None.'
            ),
            (
                "A TransportPerson (transport) Event takes place when a Transporter moves itself or a Passenger from"
                " the Origin to \nthe Destination using a Vehicle. In specific situations, a Preventer may hinder the"
                " Transporter's ability to \ntransport itself or the Passenger. The various Event subtypes include:"
                ' "BringCarryUnload", "DisperseSeparate", \n"EvacuationRescue", "Fall", "GrantedAsylum", "Hide",'
                ' "PreventEntry", "PreventExit", "SelfMotion", "SmuggleExtract", or None.'
            ),
            (
                "A TransportPerson (transport) Event takes place when a Transporter moves itself or a Passenger from"
                " the Origin to \nthe Destination using a Vehicle. In certain instances, a Preventer may hinder the"
                " Transporter's ability to transport \nitself or the Passenger. The various Event subtypes include:"
                ' "BringCarryUnload", "DisperseSeparate", \n"EvacuationRescue", "Fall", "GrantedAsylum", "Hide",'
                ' "PreventEntry", "PreventExit", "SelfMotion", "SmuggleExtract", or None.'
            ),
            (
                "A TransportPerson (transport) Event takes place when a Transporter moves itself or a Passenger from"
                " the Origin to \nthe Destination using a Vehicle. In certain situations, a Preventer may hinder the"
                " Transporter's ability to transport \nitself or the Passenger. The possible Event subtypes include:"
                ' "BringCarryUnload", "DisperseSeparate", \n"EvacuationRescue", "Fall", "GrantedAsylum", "Hide",'
                ' "PreventEntry", "PreventExit", "SelfMotion", "SmuggleExtract", or None.'
            ),
        ]
    },
    "rams_vote": {
        "en": [
            (
                "A Vote (government) Event occurs when a Voter votes for a Candidate on a Ballot with a Result in some"
                " Place. This event \nalso handles the situations when a Preventer prevents a Voter for voting. The"
                ' possible Event subtypes are: "CastVote", \n"PreventVote", None.'
            ),
            (
                "A Government Voting Event takes place when a Voter casts a vote for a Candidate on an Election"
                " Ballot, leading to a \nspecific Result in a particular Location. This event also addresses"
                " situations in which a Preventer obstructs a Voter from \nvoting. The various subtypes of this event"
                ' are: "CastVote", "PreventVote", and "None".'
            ),
            (
                "A Governmental Voting Event takes place when a Voter casts a vote for a Candidate on an Election"
                " Ballot, leading to a \nspecific Outcome in a particular Location. This event also addresses"
                " instances where a Preventer impedes a Voter's ability \nto vote. The various subtypes of this event"
                ' are: "CastVote", "PreventVote", and "None".'
            ),
            (
                "A Government Voting Event takes place when a Voter selects a Candidate on a Ballot, leading to a"
                " Result in a specific \nLocation. This event category also addresses instances when a Preventer stops"
                ' a Voter from casting their vote. The various \nsubtypes of these events are: "CastVote",'
                ' "PreventVote", and "None".'
            ),
            (
                "A Governmental Voting Event takes place when a Voter casts a vote for a Candidate on an Election"
                " Ballot, leading to a \nspecific Outcome in a particular Location. This event also addresses"
                " circumstances in which a Preventer prevents a Voter \nfrom voting. The various subtypes of this"
                ' event are: "CastVote", "PreventVote", and "None".'
            ),
        ]
    },
    "rams_yield": {
        "en": [
            (
                "A Yield (conflict) Event occurs when an Agent (yielder or surrender) yields or surrenders to a"
                " Recipient at some \nPlace. But also, when an Agent (retreater) retreats from Origin to some"
                ' Destination. The possible Event subtypes are: \n"Surrender", "Retreat"or None'
            ),
            "A Yield (conflict) Event transpires when an Agent, either a yield",
            "A Yield (conflict) Event transpires when an Agent, either a yield",
            (
                "A Yield (conflict) Event transpires when an Agent, either a yielder or a surrenderer, yields or"
                " surrenders to a \nRecipient at a particular Place. Furthermore, it occurs when an Agent, a"
                " retreater, retreats from an Origin to a Destination. \nThe potential Event subtypes include:"
                ' "Surrender," "Retreat," or None.'
            ),
            (
                "A Yield (conflict) Event transpires when an Agent, either a yielder or surrenderer, yields or"
                " surrenders to a \nRecipient at a particular Place. Furthermore, it occurs when an Agent, designated"
                " as a retreater, retreats from an Origin to a \nDestination. The potential Event subtypes encompass"
                ' "Surrender," "Retreat," or none of these.'
            ),
        ]
    },
}

src/tasks/rams/guidelines.py
hitz-zentroa/GoLLIE
GUIDELINES = {
    "rams_accidentcrash": {
        "en": [
            "An AccidentCrash (disaster) Event occurs when a Driver/Passenger in a Vehicle crashes into CrashObject\n"
            "at some place.\n\n"
            'The only possible event subtype is: "AccidentCrash".\n'
        ]
    },
    "rams_agreement": {
        "en": [
            "An Agreement (government) Event occurs when a Participant signed, rejected, nullified or violated an\n"
            "agreement at some Place.\n\n"
            'The possible Event subtypes are: "Accept", "RejectNullify", "Violate" or None.\n'
        ]
    },
    "rams_arrestjaildetain": {
        "en": [
            "An ArrestJailDetain (justice) Event occurs when a Jailer arrested or jailed a Detainee for a Crime at\n"
            "some Place.\n\n"
            'The only possible event subtype is: "ArrestJailDetain".\n'
        ]
    },
    "rams_artifactfailure": {
        "en": [
            "An ArtifactFailure (artifact) Event occurs whenever a (mechanical) Artifact failed due to some\n"
            "Instrument at some Place.\n\n"
            'The possible Event subtypes are: "MechanicalFailure"or None.\n'
        ]
    },
    "rams_attack": {
        "en": [
            "An Attack (conflict) Event occurs when an Attacker attacks a Target with some Instrument at some"
            ' Place.\n\nThe possible Event subtypes are: "AirStrikeMissileStrike",\n'
            '"BiologicalChemicalPoisonAttack", "Bombing", "FireArmAttack", "Hanging", "Invade",\n'
            '"SelfDirectedBattle", "SetFire", "Stabbing", "StealRobHijack", "Strangling"or None.\n'
        ]
    },
    "rams_collaborate": {
        "en": [
            "A Collaborate (contact) Event occurs when some Participants communicates (remotely or face-to-face) at\n"
            "some Place.\n\n"
            'The possible Event subtypes are: "Correspondence", "Meet"or None.\n'
        ]
    },
    "rams_commandorder": {
        "en": [
            "A CommandOrder (contact) Event occurs when a Communicator communicates with/to Recipient about a Topic\n"
            "at some Place.\n\n"
            'The possible Event subtypes are: "Broadcast", "Correspondence", "Meet"or None.\n'
        ]
    },
    "rams_commitmentpromiseexpressintent": {
        "en": [
            "A CommitmentPromiseExpressIntent (contact) Event occurs when a Communicator commits, promises,\n"
            "expresses an intent to Recipient about a Topic at some Place.\n\n"
            'The possible Event subtypes are: "Broadcast", "Correspondence", "Meet"or None.\n'
        ]
    },
    "rams_convene": {
        "en": [
            "A Convene (government) Event occurs when a Convener convened a ConvenedEntity at some Place. \n\n"
            'The only possible event subtype is: "Convene".\n'
        ]
    },
    "rams_coup": {
        "en": [
            "A Coup (conflict) Event occurs when a DeposedEntity was desposed by a DeposingEntity at some Place. \n\n"
            'The only possible event subtype is: "Coup".\n'
        ]
    },
    "rams_damagedestroy": {
        "en": [
            "A DamageDestroy (artifact) Event occurs when an Artifact is damaged or destroyed by some Agent\n"
            "(damager or destroyer) using an Instrument at some Place.\n\n"
            'The possible Event subtypes are: "Damage", "Destroy" or None.\n'
        ]
    },
    "rams_demonstrate": {
        "en": [
            "A Demonstrate (conflict) Event occurs when a Demonstrator(s) protest at some Place. \n\n"
            'The possible Event subtypes are: "MarchProtestPoliticalGathering"or None.\n'
        ]
    },
    "rams_die": {
        "en": [
            "A Die (life) Event occurs when a Victim died in some Place. The reason of the death can be: killed by a\n"
            "Killer or by a MedicalIssue, with or without an Instrument. \n\n"
            'The possible Event subtypes are: "Violent", "NonViolent"or None.\n'
        ]
    },
    "rams_discussion": {
        "en": [
            "A Discussion (contact) Event occurs when some Participants discuss at some place.\n\n"
            'The possible Event subtypes are: "Correspondence", "Meet"or None.\n'
        ]
    },
    "rams_diseaseoutbreak": {
        "en": [
            "A DiseaseOutbreak (disaster) Event occurs when a Disease broke out among some Victims at some Place.\n\n"
            'The only possible event subtype is: "DiseaseOutbreak".\n'
        ]
    },
    "rams_elect": {
        "en": [
            "An Elect (personell) Event occurs when a Voter elects a Candidate in some Place.\n\n"
            'The possible Event subtypes are: "WinElection"or None.\n'
        ]
    },
    "rams_endpossition": {
        "en": [
            "An EndPossition (personell) Event occurs when an Employee stops working on a Organization in some\n"
            "Place.\n\n"
            'The possible Event subtypes are: "FiringLayOff", "QuitRetire"or None.\n'
        ]
    },
    "rams_fireexplosion": {
        "en": [
            "A FireExplosion (disaster) Event occurs when a FireExplosionObject caught fire or exploded from an\n"
            "Instrument at some Place.\n\n"
            'The only possible event subtype is: "FireExplosion".\n'
        ]
    },
    "rams_formation": {
        "en": [
            "A Formation (government) Event occurs when some Participant (GPEs) are merged or when a GPE is\n"
            "formed/started by a Founder at some Place.\n\n"
            'The possible Event subtypes are: "Merge", "Start", None.\n'
        ]
    },
    "rams_funeralvigil": {
        "en": [
            "A FuneralVigil (contact) Event occurs when some Participants communicate during a funeral or vigil for\n"
            "Deceased at some Place.\n\n"
            'The possible Event subtypes are: "Meet"or None.\n'
        ]
    },
    "rams_genericcrime": {
        "en": [
            "A GenericCrime (crime) Event occurs when a Perpetrator commits a crime against a Victim at some"
            ' Place.\n\nThe only possible event subtype is: "GenericCrime".\n'
        ]
    },
    "rams_initiatejudicialprocess": {
        "en": [
            "An InitiateJudicialProcess (justice) Event occurs when a Prosecutor charged, indicted, tried or\n"
            "initiated a judicial process pertaining to a Defendant before a JudgeCourt for a Crime in some Place.\n\n"
            'The possible Event subtypes are: "ChargeIndict", "TrialHearing"or None.\n'
        ]
    },
    "rams_injure": {
        "en": [
            "An Injure (life) Event occurs when a Victim is injured in some Place. The reasons can be: have extreme\n"
            "hunger or thirst, some physical degradation and sickness or illness by some MedicalIssue either\n"
            "infected by some Disease or imposed by some Injurer. The injures can be caused by some Instrument.\n\n"
            'The possible Event subtypes are: "HungerThirst", "Physical", "Sickness", "Violent"or None\n'
        ]
    },
    "rams_investigatecrime": {
        "en": [
            "An InvestigateCrime (justice) Event occurs when a Investigator investigated a Defendant for a Crime in\n"
            "some Place.\n\n"
            'The possible Event subtypes are: "InvestigateCrime"or None.\n'
        ]
    },
    "rams_judicialconsequences": {
        "en": [
            "A JudicialConsequence (justice) Event occurs when a JudgeCourt decided the consequences, convicted,\n"
            "executed or extradited a Defendant for a Crime in some Place.\n\n"
            'The possible Event subtypes are: "Convict", "Execute", "Extradite"or None.\n'
        ]
    },
    "rams_legislate": {
        "en": [
            "A Legislate (government) Event occurs when a Government legislature enacted a Law in some Place.\n\n"
            'The only possible event subtype is: "Legislate".\n'
        ]
    },
    "rams_manufactureartifact": {
        "en": [
            "A ManufactureArtifact (manufacture) Event occurs when a Manufacturer manufactured, created or produced\n"
            "an Artifact using an Instrument at some Place.\n\n"
            'The possible Event subtypes are: "Build", "IntellectualProperty", "CreateManufacture"or None.\n'
        ]
    },
    "rams_mediastatement": {
        "en": [
            "A MediaStatement (contact) Event occurs when a Communicator communicates something on media to some\n"
            "Recipients at some Place.\n\n"
            'The possible Event subtypes are: "Broadcast"or None.\n'
        ]
    },
    "rams_medicalintervention": {
        "en": [
            "A MedicalIntervention (medical) Event occurs when a Treater treated a Patient for a MedicalIssue by\n"
            "means of some Instrument at some Place.\n\n"
            'The only possible event subtype is: "Intervention".\n'
        ]
    },
    "rams_negotiate": {
        "en": [
            "A Negotiate (contact) Event occurs when some Participants participate on a negotiation about some Topic\n"
            "at some Place.\n\n"
            'The possible Event subtypes are: "Correspondence", "Meet"or None.\n'
        ]
    },
    "rams_prevarication": {
        "en": [
            "A Prevarication (contact) Event occurs when a Communicator prevaricate about some Topic to a Recipient\n"
            "at some Place.\n\n"
            'The possible Event subtypes are: "Broadcast", "Correspondence", "Meet"or None.\n'
        ]
    },
    "rams_publicstatementinperson": {
        "en": [
            "A PublicStatementInPerson (contact) Event occurs when a Communicator gives a public statement to a\n"
            "Recipient at some Place.\n\n"
            'The possible Event subtypes are: "Broadcast"or None.\n'
        ]
    },
    "rams_requestadvice": {
        "en": [
            "A RequestAdvice (contact) Event occurs when a Communicator requests something or gives advice about a\n"
            "Topic to a Recipient at some Place.\n\n"
            'The possible Event subtypes are: "Broadcast", "Correspondence", "Meet"or None.\n'
        ]
    },
    "rams_sensoryobserve": {
        "en": [
            "A SensoryObserve (inspection) Event occurs when a Observer observed, inspected or monitored\n"
            "a ObservedEntity in some Place.\n\n"
            'The possible Event subtypes are: "InspectPeopleOrganization "MonitorElection", \n'
            '"PhysicalInvestigateInspect"or None.\n'
        ]
    },
    "rams_shortage": {
        "en": [
            "A Shortage (artifact) Event occurs when a Experiencer experienced a shortage of Supply in some Place.\n\n"
            'The only possible event subtype is: "Shortage".\n'
        ]
    },
    "rams_spy": {
        "en": [
            "An Spy (government) Event occurs when a Spy spied on a ObservedEntity to the benefit of a Beneficiary\n"
            "in some Place.\n\n"
            'The only possible event subtype is: "Spy".\n'
        ]
    },
    "rams_startpossition": {
        "en": [
            "An StartPossition (personell) Event occurs when an Employee starts working on a Organization in some\n"
            "Place.\n\n"
            'The possible Event subtypes are: "Hiring"or None.\n'
        ]
    },
    "rams_targetaimat": {
        "en": [
            "A TargetAimAt (inspection) Event occurs when a Targeter physically targeted a Target with a Instrument\n"
            "at some Place.\n\n"
            'The only possible event subtype is: "TargetAimAt".\n'
        ]
    },
    "rams_threatencoerce": {
        "en": [
            "A ThreatenCoerce (contact) Event occurs when a Communicator threats or coerces a Recipient about a\n"
            "Topic at some Place.\n\n"
            'The possible Event subtypes are: "Broadcast", "Correspondence", "Meet"or None.\n'
        ]
    },
    "rams_transaction": {
        "en": [
            "A Transaction (transaction) Event occurs when a transaction of some Artifact occurred between some\n"
            "Participants for the benefit of a Beneficiary at some Place. The Participants can be characterized\n"
            "into Giver and Recipient. There can be a preventer that prevents the transaction to occur.\n\n"
            'The possible Event subtypes are: "EmbargoSanction", "GiftGrantProvideAid", "TransferControl"or None.\n'
        ]
    },
    "rams_transfermoney": {
        "en": [
            "A TransferMoney (transaction) Event occurs when a Giver gives Money to a Recipient for the benefit of\n"
            "Beneficiary in some Place. There can be a preventer that prevents the transaction to occur.\n\n"
            'The possible Event subtypes are: "BorrowLend", "EmbargoSanction", "GiftGrantProvideAid", \n'
            '"PayForService", "Purchase"or None.\n'
        ]
    },
    "rams_transferownership": {
        "en": [
            "A TransferOwnership (transaction) Event occurs when a Giver gives some Artifact to a Recipient for\n"
            "the benefit of Beneficiary in some Place. There can be a preventer that prevents the transaction to\n"
            "occur.\n\n"
            'The possible Event subtypes are: "BorrowLend", "EmbargoSanction", "GiftGrantProvideAid",\n'
            '"Purchase"or\nNone.\n'
        ]
    },
    "rams_transportartifact": {
        "en": [
            "A TransportArtifact (transport) Event occurs when a Transporter transports an Artifact from the Origin\n"
            "to the Destination in some Vehicle. For some cases a Preventer entity can prevent the Transporter to\n"
            "transport the Artifact.\n\n"
            'The possible Event subtypes are: "BringCarryUnload", "DisperseSeparate",\n'
            '"Fall", "GrantEntry", "Hide", "LostOfControl", "NonViolentThrowLaunch", "PreventEntry",\n'
            '"PreventExit", "ReceiveImport", "SendSupplyExport", "SmuggleExtract"or None.\n'
        ]
    },
    "rams_transportperson": {
        "en": [
            "A TransportPerson (transport) Event occurs when a Transporter transports itself or a Passanger from the\n"
            "Origin to the Destination in some Vehicle. For some cases a Preventer entity can prevent the\n"
            "Transporter to transport itself or the Passenger.\n\n"
            'The possible Event subtypes are: "BringCarryUnload", "DisperseSeparate", "EvacuationRescue", "Fall",\n'
            '"GrantedAsylum", "Hide", "PreventEntry", "PreventExit", "SelfMotion", "SmuggleExtract"or None.\n'
        ]
    },
    "rams_vote": {
        "en": [
            "A Vote (government) Event occurs when a Voter votes for a Candidate on a Ballot with a Result in some\n"
            "Place. This event also handles the situations when a Preventer prevents a Voter for voting.\n\n"
            'The possible Event subtypes are: "CastVote", "PreventVote", None.\n'
        ]
    },
    "rams_yield": {
        "en": [
            "A Yield (conflict) Event occurs when an Agent (yielder or surrender) yields or surrenders to a\n"
            "Recipient at some Place. But also, when an Agent (retreater) retreats from Origin to some Destination.\n"
            '\nThe possible Event subtypes are: "Surrender", "Retreat"or None\n'
        ]
    },
}

src/tasks/rams/guidelines_gold.py
hitz-zentroa/GoLLIE
from typing import List, Union

from ..utils_typing import Event, dataclass


"""Event definitions

The official guidelines are not public. All event definitions are taken from:
https://github.com/raspberryice/gen-arg/blob/main/aida_ontology_cleaned.csv
"""


@dataclass
class ArtifactFailure(Event):
    """{rams_artifactfailure}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    subtype: Union[str, None]  # Possible event subtype.
    artifact: List[str]  # The artifact involved
    instrument: List[str]  # The reason (instrument) for the artifact fail
    place: List[str]  # Where the fail occurred


@dataclass
class DamageDestroy(Event):
    """{rams_damagedestroy}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    subtype: Union[str, None]  # Possible event subtype.
    agent: List[str]  # The damager or destroyer
    artifact: List[str]  # The damaged or destroyed artifact
    instrument: List[str]  # The instrument used for damaging or destroying
    place: List[str]  # Where the event takes place


@dataclass
class Shortage(Event):
    """{rams_shortage}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    subtype: Union[str, None]  # Possible event subtype.
    experiencer: List[str]  # The entity experiencing the shortage of supply
    supply: List[str]  # The entity being supplied
    place: List[str]  # Where the event takes place


@dataclass
class Attack(Event):
    """{rams_attack}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    subtype: Union[str, None]  # Possible event subtype
    attacker: List[str]  # The attacker
    target: List[str]  # The target of the attack
    instrument: List[str]  # The instrument used in the attack
    place: List[str]  # Where the event takes place
    artifact: List[str]  # Only in "StealRobHijack". The artifact being stolen, ...


@dataclass
class Coup(Event):
    """{rams_coup}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    subtype: Union[str, None]  # Possible event subtype
    deposed_entity: List[str]  # The entity being deposed
    deposing_entity: List[str]  # The deposing entity
    place: List[str]  # Where the event takes place


@dataclass
class Demonstrate(Event):
    """{rams_demonstrate}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    subtype: Union[str, None]  # Possible event subtype
    demonstrator: List[str]  # The Demonstrators that take part in the protest
    place: List[str]  # Where the protest takes place


@dataclass
class Yield(Event):
    """{rams_yield}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    subtype: Union[str, None]  # Possible event subtype
    agent: List[str]  # The yielder, surrender or retreater agent
    recipient: List[str]  # To who/where the agent yields or surrenders
    place: List[str]  # Where the event takes place
    origin: List[str]  # Only in "Retreat". From where the agent is retreated.
    destination: List[str]  # Only in "Retreat". To where the agent is retreated.


@dataclass
class Collaborate(Event):
    """{rams_collaborate}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    subtype: Union[str, None]  # Possible event subtype
    participants: List[str]  # The participants of the Event
    place: List[str]  # Where the Event takes place


@dataclass
class CommandOrder(Event):
    """{rams_commandorder}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    subtype: Union[str, None]  # Possible event subtype
    communicator: List[str]  # The Communicator agent
    recipient: List[str]  # The recipient of the communication event
    topic: List[str]  # The topic of the communication
    place: List[str]  # Where the Event takes place


@dataclass
class CommitmentPromiseExpressIntent(Event):
    """{rams_commitmentpromiseexpressintent}"""

    mention: str  # # The text span that most clearly expresses (triggers) the event
    subtype: Union[str, None]  # Possible event subtype
    communicator: List[str]  # The Communicator agent
    recipient: List[str]  # The recipient of the communication event
    topic: List[str]  # The topic of the communication
    place: List[str]  # Where the event takes place


@dataclass
class Discussion(Event):
    """{rams_discussion}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    subtype: Union[str, None]  # Possible event subtype
    participants: List[str]  # The participants of the discussion
    place: List[str]  # Where the event takes place


@dataclass
class FuneralVigil(Event):
    """{rams_funeralvigil}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    subtype: Union[str, None]  # Possible event subtype
    participants: List[str]  # The participants of the communicative event
    deceased: List[str]  # The deceased person(s)
    place: List[str]  # Where the event takes place


@dataclass
class MediaStatement(Event):
    """{rams_mediastatement}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    subtype: Union[str, None]  # Possible event subtype
    communicator: List[str]  # The communication agent
    recipient: List[str]  # The recipient of the communication event
    place: List[str]  # Where the event takes place


@dataclass
class Negotiate(Event):
    """{rams_negotiate}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    subtype: Union[str, None]  # Possible event subtype
    participants: List[str]  # The participants of the negotiation
    topic: List[str]  # The topic of the negotiation
    place: List[str]  # Where the event takes place


@dataclass
class Prevarication(Event):
    """{rams_prevarication}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    subtype: Union[str, None]  # Possible event subtype
    communicator: List[str]  # The communicator of the prevarication
    recipient: List[str]  # The recipient of the communication event
    topic: List[str]  # The topic of the prevarication
    place: List[str]  # Where the event takes place


@dataclass
class PublicStatementInPerson(Event):
    """{rams_publicstatementinperson}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    subtype: Union[str, None]  # Possible event subtype
    communicator: List[str]  # The communicator of the public statement
    recipient: List[str]  # The recipient of the statement
    place: List[str]  # Where the event takes place


@dataclass
class RequestAdvice(Event):
    """{rams_requestadvice}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    subtype: Union[str, None]  # Possible event subtype
    communicator: List[str]  # The communicator of the request or advice
    recipient: List[str]  # The recipient of the request or advice
    topic: List[str]  # The topic requested or adviced
    place: List[str]  # Where the event takes place


@dataclass
class ThreatenCoerce(Event):
    """{rams_threatencoerce}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    subtype: Union[str, None]  # Possible event subtype
    communicator: List[str]  # The threatening or coercing agent
    recipient: List[str]  # The threatened or coerced entity
    topic: List[str]  # The topic of the threat or coerce
    place: List[str]  # Where the event takes place


@dataclass
class AccidentCrash(Event):
    """{rams_accidentcrash}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    subtype: Union[str, None]  # Possible event subtype
    driver_or_passengers: List[str]  # The driven/passengers of the vehicle
    vehicle: List[str]  # The crashed vehicle
    crash_object: List[str]  # The object to which the vehicle crashed
    place: List[str]  # Where the event takes place


@dataclass
class DiseaseOutbreak(Event):
    """{rams_diseaseoutbreak}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    subtype: Union[str, None]  # Possible event subtype
    disease: List[str]  # The disease of the outbreak
    victim: List[str]  # The victims of the outbreak
    place: List[str]  # Where the event takes place


@dataclass
class FireExplosion(Event):
    """{rams_fireexplosion}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    subtype: Union[str, None]  # Possible event subtype
    fire_explosion_object: List[str]  # The object that caught fire or exploded
    instrument: List[str]  # The instrument used to generate the explosion
    place: List[str]  # Where the event takes place


@dataclass
class GenericCrime(Event):
    """{rams_genericcrime}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    subtype: Union[str, None]  # Possible event subtype
    perpetrator: List[str]  # The person commiting the crime
    victim: List[str]  # The victim of the crime
    place: List[str]  # Where the event takes place


@dataclass
class Agreement(Event):
    """{rams_agreement}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    subtype: Union[str, None]  # Possible event subtype
    participants: List[str]  # Only in "Accept" or None. List of participants
    rejecter_nullifier: List[str]  # Only in "RejectNullify". The rejecter or nullifier.
    violator: List[str]  # Only in "Violate". The agreement violator.
    other_participant: List[str]  # Only in "RejectNullify" or "Violate". The rest.
    place: List[str]  # Where the event takes place


@dataclass
class Convene(Event):
    """{rams_convene}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    subtype: Union[str, None]  # Possible event subtype
    convener: List[str]  # The convener
    convened_entity: List[str]  # The convened entity
    place: List[str]  # Where the event takes place


@dataclass
class Formation(Event):
    """{rams_formation}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    subtype: Union[str, None]  # Possible event subtype
    gpe: List[str]  # Only in "Start" or None. The founded Geo-Political entity
    participants: List[str]  # Only in "Merge". The merged participants.
    founder: List[str]  # Only in "Start" or None. The founder of the GPE.
    place: List[str]  # Where the event takes place


@dataclass
class Legislate(Event):
    """{rams_legislate}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    subtype: Union[str, None]  # Possible event subtype
    government_body: List[str]  # The Government body that enacted the law
    law: List[str]  # The law enacted by the government
    place: List[str]  # Where the event takes place


@dataclass
class Spy(Event):
    """{rams_spy}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    subtype: Union[str, None]  # Possible event subtype
    spy: List[str]  # The spying entity
    observed_entity: List[str]  # The entity being observed (spied)
    beneficiary: List[str]  # The entity that benefits from the information
    place: List[str]  # Where the event takes place


@dataclass
class Vote(Event):
    """{rams_vote}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    subtype: Union[str, None]  # Possible event subtype
    voter: List[str]  # The person or entity who votes.
    candidate: List[str]  # The candidate, entity being voted.
    ballot: List[str]  # The ballot
    result: List[str]  # The result of the ballot
    preventer: List[str]  # Only in "PreventVote". The vote preventer.
    place: List[str]  # Where the event takes place


@dataclass
class SensoryObserve(Event):
    """{rams_sensoryobserve}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    subtype: Union[str, None]  # Possible event subtype
    observer: List[str]  # The observer entity
    observed_entity: List[str]  # The observed entity
    place: List[str]  # Where the event takes place


@dataclass
class TargetAimAt(Event):
    """{rams_targetaimat}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    subtype: Union[str, None]  # Possible event subtype
    targeter: List[str]  # The agent of the event
    target: List[str]  # The entity being physically targeted
    instrument: List[str]  # The instrument used
    place: List[str]  # Where the event takes place


@dataclass
class ArrestJailDetain(Event):
    """{rams_arrestjaildetain}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    subtype: Union[str, None]  # Possible event subtype
    jailer: List[str]  # The person who arrested or jailed the detainee
    detainee: List[str]  # The person being arrested or jailed
    crime: List[str]  # The reason (crime) for the arresting or jailing
    place: List[str]  # Where the event takes place


@dataclass
class InitiateJudicialProcess(Event):
    """{rams_initiatejudicialprocess}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    subtype: Union[str, None]  # Possible event subtype
    prosecutor: List[str]  # The prosecutor entity
    defendant: List[str]  # The defendant of the process
    judge_court: List[str]  # The judge or court in charge
    crime: List[str]  # The crime
    place: List[str]  # Where the event takes place


@dataclass
class InvestigateCrime(Event):
    """{rams_investigatecrime}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    subtype: Union[str, None]  # Possible event subtype
    investigator: List[str]  # The investigator of the crime
    defendant: List[str]  # The person investigated
    crime: List[str]  # The crime investigated
    place: List[str]  # Where the event takes place


@dataclass
class JudicialConsequences(Event):
    """{rams_judicialconsequences}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    subtype: Union[str, None]  # Possible event subtype
    judge_court: List[str]  # The judge, court, extraditer or executioner
    defendant: List[str]  # The defendant of the judicial process
    crime: List[str]  # The crime
    origin: List[str]  # Only in "Extradite". The location of origin
    destination: List[str]  # Only in "Extradite". The location of destination.
    place: List[str]  # Where the event takes place


@dataclass
class Die(Event):
    """{rams_die}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    subtype: Union[str, None]  # Possible event subtype
    victim: List[str]  # The person who died
    killer: List[str]  # If any, the killer
    medical_issue: List[str]  # If died by medicall issues, the issue
    instrument: List[str]  # If killed with an instrument, the instrument
    place: List[str]  # Where the event takes place


@dataclass
class Injure(Event):
    """{rams_injure}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    subtype: Union[str, None]  # Possible event subtype
    victim: List[str]  # The injured person
    medical_issue: List[str]  # If any, the medical issue
    disease: List[str]  # If any, the disease causing the medical issue
    injurer: List[str]  # If any, the agent causing the injure
    instrument: List[str]  # The instrument used to injure
    place: List[str]  # Where the event takes place


@dataclass
class ManufactureArtifact(Event):
    """{rams_manufactureartifact}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    subtype: Union[str, None]  # Possible event subtype
    manufacturer: List[str]  # The entity that created the artifact
    artifact: List[str]  # The artifact being created
    instrument: List[str]  # The instrument used to create the artifact
    place: List[str]  # Where the event takes place


@dataclass
class MedicalIntervention(Event):
    """{rams_medicalintervention}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    subtype: Union[str, None]  # Possible event subtype
    treater: List[str]  # The entity in charge of the intervention
    patient: List[str]  # The treated entity
    medical_issue: List[str]  # The reason for the intervention
    instrument: List[str]  # The instrument used on the intervention
    place: List[str]  # Where the event takes place


@dataclass
class TransportArtifact(Event):
    """{rams_transportartifact}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    subtype: Union[str, None]  # Possible event subtype
    transporter: List[str]  # The entity carrying the transport event
    artifact: List[str]  # The entity being transported
    vehicle: List[str]  # The vehicle used to transport the entity
    origin: List[str]  # The origin place
    destination: List[str]  # The destination place
    hidding_place: List[str]  # Only in "Hide". The place to hide the artifact.
    preventer: List[str]  # Only in "PreventEntry" and "PreventExit". The preventer.


@dataclass
class TransportPerson(Event):
    """{rams_transportperson}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    subtype: Union[str, None]  # Possible event subtype
    transporter: List[str]  # The entity carrying the transport event
    passenger: List[str]  # If any, the passenger transported
    vehicle: List[str]  # The vehicle used to transport the entity
    origin: List[str]  # The origin place
    destination: List[str]  # The destination place
    granter: List[str]  # Only in "GrantedAsylum". The granter of asylum permission
    hidding_place: List[str]  # Only in "Hide". The place to hide the passenger.
    preventer: List[str]  # Only in "PreventEntry" and "PreventExit". The preventer.


@dataclass
class Elect(Event):
    """{rams_elect}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    subtype: Union[str, None]  # Possible event subtype
    voter: List[str]  # The entity who elects the candidate
    candidate: List[str]  # The candidate being elected
    place: List[str]  # Where the event takes place


@dataclass
class EndPossition(Event):
    """{rams_endpossition}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    subtype: Union[str, None]  # Possible event subtype
    employee: List[str]  # The entity that has stopped working for the organization
    organization: List[str]  # The organization to which the employee worked for
    place: List[str]  # Where the event takes place


@dataclass
class StartPossition(Event):
    """{rams_startpossition}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    subtype: Union[str, None]  # Possible event subtype
    employee: List[str]  # The entity that has started working on the organization
    organization: List[str]  # The organization to which the employee works
    place: List[str]  # Where the event takes place


@dataclass
class Transaction(Event):
    """{rams_transaction}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    subtype: Union[str, None]  # Possible event subtype
    participants: List[str]  # Only if subtype is None. The participants.
    giver: List[str]  # The entity giving the artifact
    recipient: List[str]  # The entity receiving the artifact
    beneficiary: List[str]  # The entity that benefits from the transaction (other than recipient)
    artifact: List[str]  # The entity being transferred (artifact, money or territory)
    preventer: List[str]  # Only in "EmbargoSanction". The entity that prevents the transaction.
    place: List[str]  # Where the event takes place


@dataclass
class TransferMoney(Event):
    """{rams_transfermoney}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    subtype: Union[str, None]  # Possible event subtype
    giver: List[str]  # The entity giving the money
    recipient: List[str]  # The entity receiving the money
    beneficiary: List[str]  # The entity that benefits from the transaction (other than recipient)
    money: List[str]  # The money amount
    preventer: List[str]  # Only in "EmbargoSanction". The entity that prevents the transaction.
    place: List[str]  # Where the event takes place


@dataclass
class TransferOwnership(Event):
    """{rams_transferownership}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    subtype: Union[str, None]  # Possible event subtype
    giver: List[str]  # The entity giving the artifact
    recipient: List[str]  # The entity receiving the artifact
    beneficiary: List[str]  # The entity that benefits from the transaction (other than recipient)
    artifact: List[str]  # The artifact being transferred
    preventer: List[str]  # Only in "EmbargoSanction". The entity that prevents the transaction.
    place: List[str]  # Where the event takes place


EVENT_DEFINITIONS: List[Event] = [
    ArtifactFailure,
    DamageDestroy,
    Shortage,
    Attack,
    Coup,
    Demonstrate,
    Yield,
    Collaborate,
    CommandOrder,
    CommitmentPromiseExpressIntent,
    Discussion,
    FuneralVigil,
    MediaStatement,
    Negotiate,
    Prevarication,
    PublicStatementInPerson,
    RequestAdvice,
    ThreatenCoerce,
    AccidentCrash,
    DiseaseOutbreak,
    FireExplosion,
    GenericCrime,
    Agreement,
    Convene,
    Formation,
    Legislate,
    Spy,
    Vote,
    SensoryObserve,
    TargetAimAt,
    ArrestJailDetain,
    InitiateJudicialProcess,
    InvestigateCrime,
    JudicialConsequences,
    Die,
    Injure,
    ManufactureArtifact,
    MedicalIntervention,
    TransportArtifact,
    TransportPerson,
    Elect,
    EndPossition,
    StartPossition,
    Transaction,
    TransferMoney,
    TransferOwnership,
]

# __all__ = list(map(str, [*EVENT_DEFINITIONS]))

src/tasks/rams/prompts.py
hitz-zentroa/GoLLIE
from typing import List, Type

from src.tasks.rams.prompts import EVENT_DEFINITIONS
from src.tasks.utils_scorer import EventScorer


class RAMSEventScorer(EventScorer):
    """RAMS Argument classification scorer."""

    valid_types: List[Type] = EVENT_DEFINITIONS

src/tasks/rams/scorer.py
hitz-zentroa/GoLLIE
from . import data_loader, prompts, scorer


__all__ = ["data_loader", "prompts", "scorer"]

src/tasks/tacred/__init__.py
hitz-zentroa/GoLLIE
import json
from collections import defaultdict
from typing import Tuple, Union

import rich

from src.tasks.tacred.guidelines import GUIDELINES
from src.tasks.tacred.prompts import TEMPLATE_DEFINITIONS, OrganizationTemplate, PersonTemplate

from ..utils_data import DatasetLoader, Sampler
from ..utils_typing import Relation, Template, dataclass


@dataclass
class NoneRelation(Relation):
    arg1: str
    arg2: str


class TACREDDatasetLoader(DatasetLoader):
    """
    A `DatasetLoader` for the TACRED dataset.

    Args:
        path (`str`):
            The location of the dataset directory.

    Raises:
        `ValueError`:
            raised when a not defined value found.
    """

    TEMPLATE_TO_CLASS_MAPPING = {
        "PERSON": {
            "class": PersonTemplate,
            "per:alternate_names": "alternate_names",
            "per:date_of_birth": "date_of_birth",
            "per:age": "age",
            "per:country_of_birth": "country_of_birth",
            "per:stateorprovince_of_birth": "state_or_province_of_birth",
            "per:city_of_birth": "city_of_birth",
            "per:origin": "origin",
            "per:date_of_death": "date_of_death",
            "per:country_of_death": "country_of_death",
            "per:stateorprovince_of_death": "state_or_province_of_death",
            "per:city_of_death": "city_of_death",
            "per:cause_of_death": "cause_of_death",
            "per:countries_of_residence": "countries_of_residence",
            "per:stateorprovinces_of_residence": "states_or_provinces_of_residence",
            "per:cities_of_residence": "cities_of_residence",
            "per:schools_attended": "schools_attended",
            "per:title": "title",
            "per:employee_of": "employee_or_member_of",
            "per:religion": "religion",
            "per:spouse": "spouse",
            "per:children": "children",
            "per:parents": "parents",
            "per:siblings": "siblings",
            "per:other_family": "other_family",
            "per:charges": "charges",
        },
        "ORGANIZATION": {
            "class": OrganizationTemplate,
            "org:alternate_names": "alternate_names",
            "org:political/religious_affiliation": "political_or_religious_affiliation",
            "org:top_members/employees": "top_members_employees",
            "org:number_of_employees/members": "number_of_employees_members",
            "org:members": "members",
            "org:member_of": "member_of",
            "org:subsidiaries": "subsidiaries",
            "org:parents": "parents",
            "org:founded_by": "founded_by",
            "org:founded": "date_founded",
            "org:dissolved": "date_dissolved",
            "org:country_of_headquarters": "country_of_headquarters",
            "org:stateorprovince_of_headquarters": "state_or_province_of_headquarters",
            "org:city_of_headquarters": "city_of_headquarters",
            "per:cities_of_residence": "city_of_headquarters",  # Bug on data
            "org:shareholders": "shareholders",
            "org:website": "website",
        },
    }

    def __init__(self, path: str, **kwargs) -> None:
        self.elements = {}
        with open(path, "rt", encoding="utf-8") as in_f:
            data = json.load(in_f)

        data_by_sentence = defaultdict(list)
        for line in data:
            text = " ".join(line["token"])
            data_by_sentence[text].append(line)

        for text, instances in data_by_sentence.items():
            text = text.replace("-LRB-", "(").replace("-RRB-", ")").replace("-LSB-", "[").replace("-RSB-", "]")
            templates = {}
            for inst in instances:
                if inst["relation"] == "no_relation":
                    continue

                subj = (
                    " ".join(inst["token"][inst["subj_start"] : inst["subj_end"] + 1])
                    .replace("-LRB-", "(")
                    .replace("-RRB-", ")")
                    .replace("-LSB-", "[")
                    .replace("-RSB-", "]")
                )
                if f"{subj}-{inst['subj_type']}" not in templates:
                    _info = self.TEMPLATE_TO_CLASS_MAPPING[inst["subj_type"]]
                    templates[f"{subj}-{inst['subj_type']}"] = {"query": subj, "_info": _info}
                _info = templates[f"{subj}-{inst['subj_type']}"]["_info"]

                obj = (
                    " ".join(inst["token"][inst["obj_start"] : inst["obj_end"] + 1])
                    .replace("-LRB-", "(")
                    .replace("-RRB-", ")")
                    .replace("-LSB-", "[")
                    .replace("-RSB-", "]")
                )
                slot_name = _info[inst["relation"]]
                if slot_name not in templates[f"{subj}-{inst['subj_type']}"]:
                    templates[f"{subj}-{inst['subj_type']}"][slot_name] = []

                templates[f"{subj}-{inst['subj_type']}"][slot_name].append(obj)

            key = str(hash(text))
            labels, gold = [], []
            for temp in templates.values():
                _info = temp.pop("_info")
                try:
                    template: Template = _info["class"](**temp)
                except Exception as e:
                    rich.print(_info)
                    rich.print(temp)
                    raise e
                template.assert_typing_constraints()
                labels.append(template)
                gold.append(template.query)

            self.elements[key] = {
                "id": key,
                "doc_id": instances[0]["docid"],
                "text": text,
                "labels": labels,
                "gold": labels,
            }

        assert len(self.elements) == len(data_by_sentence), "The hash function failed."


class TACREDSampler(Sampler):
    """
    A data `Sampler` for the TACRED dataset.

    Args:
        dataset_loader (`TACREDDatasetLoader`):
            The dataset loader that contains the data information.
        task (`str`, optional):
            The task to sample. It must be one of the following: NER, VER, RE, EE.
            Defaults to `None`.
        split (`str`, optional):
            The split to sample. It must be one of the following: "train", "dev" or
            "test". Depending on the split the sampling strategy differs. Defaults to
            `"train"`.
        parallel_instances (`Union[int, Tuple[int, int]]`, optional):
            The number of sentences sampled in parallel. Options:

                * **`int`**: The amount of elements that will be sampled in parallel.
                * **`tuple`**: The range of elements that will be sampled in parallel.

            Defaults to 1.
        max_guidelines (`int`, optional):
            The number of guidelines to append to the example at the same time. If `-1`
            is given then all the guidelines are appended. Defaults to `-1`.
        guideline_dropout (`float`, optional):
            The probability to dropout a guideline definition for the given example. This
            is only applied on training. Defaults to `0.0`.
        seed (`float`, optional):
            The seed to sample the examples. Defaults to `0`.
        prompt_template (`str`, optional):
            The path to the prompt template. Defaults to `"templates/prompt.txt"`.
        ensure_positives_on_train (bool, optional):
            Whether to ensure that the guidelines of annotated examples are not removed.
            Defaults to `False`.
        dataset_name (str, optional):
            The name of the dataset. Defaults to `None`.
        scorer (`str`, optional):
           The scorer class import string. Defaults to `None`.
        sample_only_gold_guidelines (`bool`, optional):
            Whether to sample only guidelines of present annotations. Defaults to `False`.
    """

    def __init__(
        self,
        dataset_loader: TACREDDatasetLoader,
        task: str = None,
        split: str = "train",
        parallel_instances: Union[int, Tuple[int, int]] = 1,
        max_guidelines: int = -1,
        guideline_dropout: float = 0.0,
        seed: float = 0,
        prompt_template: str = "templates/prompt_tacred.txt",
        ensure_positives_on_train: bool = True,
        dataset_name: str = None,
        scorer: str = None,
        sample_only_gold_guidelines: bool = False,
        **kwargs,
    ) -> None:
        assert task in [
            "SF",
        ], f"{task} can only be 'SF'."

        super().__init__(
            dataset_loader=dataset_loader,
            task=task,
            split=split,
            parallel_instances=parallel_instances,
            max_guidelines=max_guidelines,
            guideline_dropout=guideline_dropout,
            seed=seed,
            prompt_template=prompt_template,
            ensure_positives_on_train=ensure_positives_on_train,
            sample_only_gold_guidelines=sample_only_gold_guidelines,
            dataset_name=dataset_name,
            scorer=scorer,
            task_definitions=TEMPLATE_DEFINITIONS,
            task_target="labels",
            definitions=GUIDELINES,
            **kwargs,
        )

src/tasks/tacred/data_loader.py
hitz-zentroa/GoLLIE
GUIDELINES = {
    "tacred_person": {
        "en": [
            "Person templates encodes the information about the given query Person entity.",
            "Person templates store the data related to the specific Person query entity.",
            "Person templates store the data related to the specific Person entity in question.",
            "Person templates store the data relevant to the specified Person entity in question.",
            "Person templates store the data relevant to the specified Person entity in question.",
        ]
    },
    "tacred_person_alternate_names": {
        "en": [
            (
                "Names used to refer to the query person that are distinct from the 'official' name. Including:"
                " aliases, stage names, \nalternate transliterations, abbreviations, ..."
            ),
            (
                "Text: \"Names employed to identify the person being queried that diverge from the 'official'"
                " appellation. \nIncorporating: aliases, stage names, alternative transliterations, abbreviations, ..."
            ),
            (
                "Text: \"Names employed to identify the inquiry individual that are separate from the 'official' name."
                " \nIncorporating: aliases, stage names, alternative transliterations, abbreviations, ..."
            ),
            (
                "Alternate names employed for identifying the subject, separate from their formal designation."
                " Including: \nassumed names, professional pseudonyms, diverse script renditions, abbreviations, and"
                " other comparable appellations."
            ),
            (
                "Alternative names employed to identify the individual subject to the query, separate from their"
                " formal or \nrecognized name. Inclusive of: pseudonyms, stage personas, alternative transcriptions,"
                " abbreviations, and so forth."
            ),
        ]
    },
    "tacred_person_date_of_birth": {
        "en": [
            "The date on which the query person was born.",
            "The birthdate of the individual in question.",
            "The specific day when the individual in question came into existence.",
            "The specific day when the individual being investigated was born.",
            "The specific day when the individual being investigated was born.",
        ]
    },
    "tacred_person_age": {
        "en": [
            "A reported age of the query person.",
            "An indicated age for the individual in question.",
            "An indicated age for the individual being searched.",
            "An indicated age for the individual in question.",
            "An stated age for the individual in question.",
        ]
    },
    "tacred_person_country_of_birth": {
        "en": [
            "The country in which the query person was born.",
            "The nation in which the individual being queried was born.",
            "The nation where the individual in question was born.",
            "The nation where the individual in question was born.",
            "The nation where the individual in question was born.",
        ]
    },
    "tacred_person_state_or_province_of_birth": {
        "en": [
            "The geopolitical entity at state or province level in which the assigned person was born.",
            (
                "The political subdivision, either at the state or provincial level, where the specified individual"
                " was born."
            ),
            "The political subdivision, either at the state or provincial level, where the individual was born.",
            (
                "The political subdivision within a country, either at the state or province level, where the"
                " specified individual \nwas born."
            ),
            (
                "The political subdivision, either at the state or provincial level, where the specified individual"
                " was born."
            ),
        ]
    },
    "tacred_person_city_of_birth": {
        "en": [
            (
                "The geopolitical entity at the municipality level (city, town, or village) in which the query person"
                " was born"
            ),
            (
                "The political jurisdiction at the local level (such as a city, town, or village) where the individual"
                " being searched \nwas born."
            ),
            (
                "The political jurisdiction within the local community, such as a city, town, or village, where the"
                " individual being \nresearched was born."
            ),
            (
                "The political subdivision, specifically at the local level (such as a city, town, or village) where"
                " the person of \ninterest was born."
            ),
            (
                "The political subdivision, specifically at the municipal level (including cities, towns, and"
                " villages), where \nthe individual being researched was born."
            ),
        ]
    },
    "tacred_person_origin": {
        "en": [
            "The nationality and/or ethnicity of the assigned person.",
            "The individual's nationality and/or ethnic background, as determined by the assigned person.",
            "The specific cultural or ethnic background of the individual in question.",
            "The individual's citizenship and/or cultural background, as determined by the assigned care provider.",
            "The individual's citizenship and/or cultural background assigned to them.",
        ]
    },
    "tacred_person_date_of_death": {
        "en": [
            "The date of the query person's death.",
            "The specific date when the individual in question passed away.",
            "The specific day when the individual in question passed away.",
            "The specific date when the individual in question passed away.",
            "The specific date when the individual in question passed away.",
        ]
    },
    "tacred_person_country_of_death": {
        "en": [
            "The country in which the assigned person died.",
            "The nation where the specified individual passed away.",
            "The nation where the specified individual passed away.",
            "The nation where the specified individual passed away.",
            "The nation where the specified individual passed away.",
        ]
    },
    "tacred_person_state_or_province_of_death": {
        "en": [
            "The geopolitical entity at state or province level in which the assigned person died.",
            (
                "The political subdivision, either at the state or provincial level, where the specified individual"
                " passed away."
            ),
            (
                "The political subdivision, either at the state or provincial level, where the specified individual"
                " passed away."
            ),
            "The political subdivision, either a state or province, in which the specified individual passed away.",
            "The political subdivision, either a state or province, in which the specified individual passed away.",
        ]
    },
    "tacred_person_city_of_death": {
        "en": [
            "The geopolitical entity at the level of city, town, village in which the assigned person died.",
            (
                "The political subdivision, specifically within the context of a city, town, or village, in which the"
                " individual \npassed away."
            ),
            (
                "The political subdivision, encompassing a city, town, or village, in which the specified individual"
                " passed away."
            ),
            (
                "The political subdivision, within the context of a city, town, or village, in which the specified"
                " individual passed \naway."
            ),
            "The political subdivision, specifically at the municipal level, in which the individual passed away.",
        ]
    },
    "tacred_person_cause_of_death": {
        "en": [
            "The explicit cause of death for the query person.",
            "The specific reason for the query individual's passing.",
            "The specific reason for the query individual's death.",
            "The specific reason for the query individual's passing.",
            "The specific reason for the query individual's passing.",
        ]
    },
    "tacred_person_countries_of_residence": {
        "en": [
            "All countries in which the query person has lived.",
            "All nations where the subject has resided.",
            'Alternate text: "The complete list of nations where the subject has resided.',
            "A comprehensive list of all nations where the subject has resided.",
            'Text paraphrase: "The complete list of nations where the subject has resided.',
        ]
    },
    "tacred_person_states_or_provinces_of_residence": {
        "en": [
            "Geopolitical entities at the state or province level in which the assigned person has lived.",
            "Political divisions, such as states or provinces, where the specified individual has resided.",
            "Political divisions, such as states or provinces, where the specified individual has resided.",
            "Political divisions, such as states or provinces, where the specified individual has resided.",
            "Political divisions, such as states or provinces, where the specified individual has resided.",
        ]
    },
    "tacred_person_cities_of_residence": {
        "en": [
            "Geopolitical entities at the level of city, town, or village in which the query person has lived.",
            (
                "Political divisions within local communities, such as cities, towns, and villages, where the searched"
                " individual \nhas resided."
            ),
            (
                "Political divisions on a local level, such as cities, towns, and villages, where the person of"
                " interest has resided."
            ),
            (
                "Political divisions on a local level, such as cities, towns, and villages where the searched"
                " individual has \nresided."
            ),
            (
                "Political divisions on a local level, such as cities, towns, and villages where the search person has"
                " resided."
            ),
        ]
    },
    "tacred_person_schools_attended": {
        "en": [
            "Any school (college, high school, university, etc.) that the query person has attended.",
            (
                "All educational institutions, including colleges, high schools, and universities, that the individual"
                " in \nquestion has attended."
            ),
            (
                "All educational institutions, including colleges, high schools, and universities, that the individual"
                " in \nquestion has attended."
            ),
            (
                "All educational institutions, including colleges, high schools, and universities, that the individual"
                " in \nquestion has attended."
            ),
            (
                "All educational institutions, including colleges, high schools, and universities, that the individual"
                " in \nquestion has attended."
            ),
        ]
    },
    "tacred_person_title": {
        "en": [
            (
                "Official or unofficial name(s) of the employment or membership positions that have been held by the"
                " query person."
            ),
            (
                "Specific title(s) or unofficial designations of the job roles or organizational memberships that the"
                " individual \nof interest has occupied."
            ),
            "Names of positions, either official or unofficial, that the individual of interest has occupied.",
            (
                "Title(s) of the job roles or memberships that the individual in question has occupied, whether they"
                " were formal or \ninformal positions."
            ),
            (
                "Specific titles or informal designations of the jobs or memberships that the individual in question"
                " has held."
            ),
        ]
    },
    "tacred_person_employee_or_member_of": {
        "en": [
            (
                "The organizations or geopolitical entities (governments) of which the query person has been an"
                " employee or member."
            ),
            (
                "The organizations or political entities, such as governments, for which the individual in question"
                " has been \nemployed or held membership."
            ),
            (
                "The organizations and geopolitical entities, such as governments, for which the individual in"
                " question has been \nemployed or held membership."
            ),
            (
                "The organizations or political entities, such as governments, for which the individual in question"
                " has been \nemployed or held membership."
            ),
            (
                "The entities, either organizational or geopolitical in nature, for which the individual in question"
                " has served as \nan employee or member."
            ),
        ]
    },
    "tacred_person_religion": {
        "en": [
            "The religion to which the query person has belonged.",
            "The religious affiliation that the individual in question has adhered to.",
            "The faith tradition that the individual in question has embraced.",
            "The religious affiliation that the individual in question has embraced.",
            "The faith tradition that the individual in question has embraced.",
        ]
    },
    "tacred_person_spouse": {
        "en": [
            "The spouse(s) of the query person.",
            "The spouse(s) associated with the individual in question.",
            "The husband or wife of the individual in question.",
            "The husband or wife of the individual in question.",
            "The husband or wife of the individual in question.",
        ]
    },
    "tacred_person_children": {
        "en": [
            "The children of the query person, including adopted and step-children.",
            "The offspring of the inquiry individual, comprising adopted and step-children.",
            "The offspring of the inquiry individual, encompassing adopted and step-children.",
            "The offspring of the inquiry individual, encompassing embraced and step-kids.",
            "The offspring of the inquiry individual, incorporating embraced and step-kids.",
        ]
    },
    "tacred_person_parents": {
        "en": [
            (
                "The parents of the query person. In addition to biological parents, step-parents and adoptive parents"
                " are also \nacceptable answers."
            ),
            (
                "The individuals who are the biological parents of the person in question may include not only their"
                " natural parents \nbut also step-parents and those who have adopted them."
            ),
            (
                "The individuals who are the biological parents, as well as any step-parents or adoptive parents, of"
                " the person in \nquestion."
            ),
            (
                "The individuals who are the birth parents of the individual in question. Besides the biological"
                " parents, any \nguardians who have legally assumed the role of a parent, such as step-parents or"
                " those who have adopted the person, are also \nconsidered valid responses."
            ),
            (
                "The individuals who are the biological parents of the person in question. Furthermore, step-parents"
                " and adoptive \nparents are also considered valid responses."
            ),
        ]
    },
    "tacred_person_siblings": {
        "en": [
            "The brothers and sisters of the query person.",
            "The siblings of the individual in question.",
            "The siblings of the individual in question.",
            "The siblings of the individual in question.",
            "The siblings of the individual in question.",
        ]
    },
    "tacred_person_other_family": {
        "en": [
            (
                "Family other than siblings, parents, children, and spouse (or former spouse). Including: brothers and"
                " and sisters \nin law, grandparents, grandchildren, cousins, aunts, uncles, etc."
            ),
            (
                "Family members beyond siblings, parents, children, and spouses (or former spouses), encompassing:"
                " brothers and \nsisters-in-law, grandparents, grandchildren, cousins, aunts, uncles, and other"
                " extended relatives."
            ),
            (
                "Non-sibling family members, excluding parents, children, and spouses (current or former). Including:"
                " brothers- \nand sisters-in-law, grandparents, grandchildren, cousins, aunts, uncles, and other"
                " extended relatives."
            ),
            (
                "Family members beyond siblings, parents, children, and spouses (or former spouses). Incorporating:"
                " brothers and \nsisters-in-law, grandparents, grandchildren, cousins, aunts, uncles, and other"
                " extended relatives."
            ),
            (
                "Family members beyond siblings, parents, children, and spouses (or former spouses). Inclusive of:"
                " brothers and \nsisters-in-law, grandparents, grandchildren, cousins, aunts, uncles, and more."
            ),
        ]
    },
    "tacred_person_charges": {
        "en": [
            "The charges or crimes (alleged or convicted) of the query person.",
            "The alleged or proven offenses committed by the individual in question.",
            "The alleged or proven offenses committed by the individual in question.",
            "The alleged or confirmed offenses committed by the individual in question.",
            "The alleged or proven offenses committed by the individual in question.",
        ]
    },
    "tacred_organization_template": {
        "en": [
            "Organization templates encodes the information about the given query Organization entity.",
            "Organization templates store data related to the specified Organization entity.",
            "Organization templates store data related to the specific Organization entity being queried.",
            "Organization templates store the data related to the specified Organization entity in a coded format.",
            "Organization templates store data related to the specific Organization entity in question.",
        ]
    },
    "tacred_organization_alternate_names": {
        "en": [
            (
                "Any name used to refer the query organization that is distinct from the 'official' name. Including:"
                " former names, \naliases, alternate spellings, acronyms, abbreviations, ..."
            ),
            (
                "Any name employed to address the entity subject to the query, which is separate from its 'official'"
                " designation. \nIncorporating: previous appellations, nicknames, alternate orthographies,"
                " abbreviations, acronyms, and other comparable \nforms."
            ),
            (
                "Any name employed to address the inquiry organization that differs from its 'official' designation."
                " \nIncorporating: previous names, nicknames, multiple spellings, abbreviations, acronyms, and other"
                " alternative forms."
            ),
            (
                "A phrase used to identify the inquiry organization, which is separate from its formal title."
                " Including: previous \nnames, assumed names, alternative spellings, abbreviations, acronyms, and so"
                " forth."
            ),
            (
                "A description of the term referring to any name, other than the official one, by which the entity"
                " subject to the query \nis known. This may encompass previous designations, alternate appellations,"
                " diverse orthographic renderings, \nabbreviations, and acronyms."
            ),
        ]
    },
    "tacred_organnization_political_or_religious_affiliation": {
        "en": [
            "Ideological groups with which the organization is associated.",
            'Alternative phrasing: "Political factions connected to the institution.',
            "Political factions connected to the aforementioned organization.",
            "Political factions that have connections with the company.",
            "Political factions connected to the aforementioned entity.",
        ]
    },
    "tacred_organization_top_members_employees": {
        "en": [
            (
                "The persons in high-level, leading positions at the query organization. Top member position imply a"
                " level of \ndecision-making authority over the entire organization."
            ),
            (
                "The individuals holding prominent, top-level roles within the target organization possess a"
                " considerable degree \nof decision-making power that extends across the entire entity."
            ),
            (
                "The individuals holding prominent, top-level roles within the target organization possess a"
                " considerable degree \nof decision-making power that extends across the entire entity."
            ),
            (
                "The individuals holding prominent, top-level roles within the target organization possess a"
                " significant degree \nof decision-making power that extends across the entire entity."
            ),
            (
                "The individuals holding prominent, top-level roles within the targeted organization possess a"
                " considerable \ndegree of decision-making power that extends across the entire entity."
            ),
        ]
    },
    "tacred_organization_number_of_employees_members": {
        "en": [
            "The total number of people who are employed by or have membership in an organization.",
            (
                "The aggregate count of individuals who hold either employment or membership status within a given"
                " organization."
            ),
            "The aggregate count of individuals who hold employment or membership within a particular organization.",
            (
                "The aggregate count of individuals who hold either employment or membership status within a"
                " particular \norganization."
            ),
            "The aggregate count of individuals who hold employment or membership within a particular organization.",
        ]
    },
    "tacred_organization_members": {
        "en": [
            (
                "Organization or Geopolitical entities that are members of the query organization. While similar to"
                " subsidiaries, \nmembers is different because correct fillers are distinct entities capable of"
                " autonomously ending their membership with \nthe query organization."
            ),
            (
                "Members of the query organization represent geopolitical entities or distinct organizations that are"
                " part of the \nlarger group. Although these members share some similarities with subsidiaries, it's"
                " essential to note that they are \nseparate entities with the ability to independently terminate"
                " their affiliation with the main organization."
            ),
            (
                "Organization or Geopolitical entities that form part of the query organization. Although reminiscent"
                " of \nsubsidiaries, the distinction lies in these being separate and distinct entities with the"
                " capacity to independently terminate \ntheir affiliation with the query organization."
            ),
            (
                "Organization or Geopolitical entities that are part of the query organization. Though similar to"
                " subsidiaries, \nmembers differ in that they are separate entities with the ability to independently"
                " terminate their membership with the \nquery organization."
            ),
            (
                "Members of the query organization represent geopolitical entities or separate organizations that are"
                " part of the \nlarger group. Although similar to subsidiaries, the distinction lies in the fact that"
                " these members are independent \nentities that can freely terminate their association with the main"
                " organization."
            ),
        ]
    },
    "tacred_organization_member_of": {
        "en": [
            "Organizations or geopolitical entities of which the query organization is a member itself.",
            "Organizations and geopolitical entities that the queried organization is part of.",
            "Organizations and geopolitical entities that the queried organization belongs to as a member.",
            "Organizations and geopolitical entities that the queried organization belongs to as a member.",
            (
                "Organizations and geopolitical entities that are part of the same larger group or political entity as"
                " the queried \norganization."
            ),
        ]
    },
    "tacred_organization_subsidiaries": {
        "en": [
            "Organizations that are subsidiaries of the query organization.",
            "Companies that are part of the main organization being queried.",
            "Companies that operate as subsidiaries of the entity in question.",
            "Companies that are part of the main organization being queried.",
            "Companies that are subsidiaries of the main organization being queried.",
        ]
    },
    "tacred_organization_parents": {
        "en": [
            "Organizations or geopolitical entities of which the query organization is a subsidiary.",
            (
                "Organizations or political entities that are part of a larger structure, with the query organization"
                " being a \nsubordinate component."
            ),
            "Organizations or political entities that have the query organization as a subordinate component.",
            (
                "Organizations or political entities that are part of a larger structure, with the query organization"
                " being a \nsubordinate component."
            ),
            (
                "Organizations and geopolitical entities in which the queried organization serves as a subordinate"
                " component."
            ),
        ]
    },
    "tacred_organization_founded_by": {
        "en": [
            "The person, organization, or geopolitical entity that founded the query organization.",
            "The individual, group, or political entity that established the inquiry organization.",
            "The individual, group, or political entity responsible for establishing the queried organization.",
            "The individual, group, or political entity responsible for establishing the inquiry organization.",
            "The individual, group, or political entity responsible for establishing the subject organization.",
        ]
    },
    "tacred_organization_date_founded": {
        "en": [
            "The date on which the query organization was founded.",
            "The year in which the inquiry organization came into existence.",
            "The establishment year of the inquiry organization.",
            "The year in which the requested company was established.",
            "The establishment year of the inquiry organization.",
        ]
    },
    "tacred_organization_date_dissolved": {
        "en": [
            (
                "The date on which the query organization was dissolved. When companies merge to form a new entity,"
                " the original \ncompanies should be considered as dissolved."
            ),
            (
                "The dissolution date of the merged organization. When businesses combine to create a new entity, the"
                " original \ncompanies are considered to have ceased operations."
            ),
            (
                "The effective date of the merger, whereby the participating organizations ceased to exist as separate"
                " entities and \nbecame part of the newly formed organization."
            ),
            (
                "The date when the inquiry organization underwent dissolution. Notably, when businesses combine to"
                " create a new \nentity, the initial companies ought to be viewed as having dissolved."
            ),
            "The effective date of the merger, marking the dissolution of the constituent organizations.",
        ]
    },
    "tacred_organization_country_of_headquarters": {
        "en": [
            "Countries on which the headquarters of the query organization are located.",
            "The specific countries where the main offices of the organization being inquired about are situated.",
            "The specific countries where the main offices of the target organization can be found.",
            "The specific nations where the main offices of the requested organization can be found.",
            "The specific nations where the main offices of the target organization can be found.",
        ]
    },
    "tacred_organization_state_or_province_of_headquarters": {
        "en": [
            "Location of the headquarters of the query organization at the state or province level.",
            (
                "The requested information pertains to the specific location of the main office of the subject"
                " organization within \nthe context of a state or province."
            ),
            (
                "The requested information pertains to the specific state or province where the subject organization's"
                " main \nheadquarters is situated."
            ),
            (
                'Paraphrase: "The specific location of the main office of the organization in question, within a'
                " particular state or \nprovince."
            ),
            (
                "The requested information pertains to the specific location of the main office of the organization in"
                " question, \ndetailing its placement at the state or provincial tier within the given country."
            ),
        ]
    },
    "tacred_organization_city_of_headquarters": {
        "en": [
            "Location of the headquarters of the assigned organization at the city, town, or village level.",
            (
                'The paraphrased text is: "Specify the precise location of the main office of the selected'
                " organization, detailing \nits position at the city, town, or village level."
            ),
            (
                "Specifying the precise location of the main offices of the assigned organization, within the context"
                " of a city, \ntown, or village."
            ),
            (
                "The specific location of the main office of the aforementioned organization within a city, town, or"
                " village."
            ),
            (
                "The specific location of the main offices of the designated organization within a city, town, or"
                " village."
            ),
        ]
    },
    "tacred_organization_shareholders": {
        "en": [
            "Any organization, person, or geopolitical entity that holds shares of the organization.",
            (
                "Any entity, whether an individual, corporation, or political faction, that possesses shares in the"
                " organization."
            ),
            "Any group, individual, or political entity that owns shares in the company.",
            "Any individual, group, or political entity that owns a portion of the organization.",
            "Any individual, group, or political entity that owns a portion of the organization.",
        ]
    },
    "tacred_organization_website": {
        "en": [
            "An official top level URL for the organization's website.",
            "A definitive, highest-level web address for the company's official online presence.",
            "A definitive, highest-level web address for the organization's official website.",
            "A definitive, highest-level web address for the organization's official website.",
            "A definitive, highest-level web address for the organization's official website.",
        ]
    },
}

src/tasks/tacred/guidelines.py
hitz-zentroa/GoLLIE
GUIDELINES = {
    "tacred_person": {"en": ["Person templates encodes the information about the given query Person entity."]},
    "tacred_person_alternate_names": {
        "en": [
            "Names used to refer to the query person that are distinct from the 'official' name.\n"
            "Including: aliases, stage names, alternate transliterations, abbreviations, ..."
        ]
    },
    "tacred_person_date_of_birth": {"en": ["The date on which the query person was born."]},
    "tacred_person_age": {"en": ["A reported age of the query person."]},
    "tacred_person_country_of_birth": {"en": ["The country in which the query person was born."]},
    "tacred_person_state_or_province_of_birth": {
        "en": ["The geopolitical entity at state or province level in which the assigned person was born."]
    },
    "tacred_person_city_of_birth": {
        "en": [
            "The geopolitical entity at the municipality level (city, town, or village) in which the query person\n"
            "was born"
        ]
    },
    "tacred_person_origin": {"en": ["The nationality and/or ethnicity of the assigned person."]},
    "tacred_person_date_of_death": {"en": ["The date of the query person's death."]},
    "tacred_person_country_of_death": {"en": ["The country in which the assigned person died."]},
    "tacred_person_state_or_province_of_death": {
        "en": ["The geopolitical entity at state or province level in which the assigned person died."]
    },
    "tacred_person_city_of_death": {
        "en": ["The geopolitical entity at the level of city, town, village in which the assigned person died."]
    },
    "tacred_person_cause_of_death": {"en": ["The explicit cause of death for the query person."]},
    "tacred_person_countries_of_residence": {"en": ["All countries in which the query person has lived."]},
    "tacred_person_states_or_provinces_of_residence": {
        "en": ["Geopolitical entities at the state or province level in which the assigned person has lived."]
    },
    "tacred_person_cities_of_residence": {
        "en": ["Geopolitical entities at the level of city, town, or village in which the query person has lived."]
    },
    "tacred_person_schools_attended": {
        "en": ["Any school (college, high school, university, etc.) that the query person has attended."]
    },
    "tacred_person_title": {
        "en": [
            "Official or unofficial name(s) of the employment or membership positions that have been held by the\n"
            "query person."
        ]
    },
    "tacred_person_employee_or_member_of": {
        "en": [
            "The organizations or geopolitical entities (governments) of which the query person has been an employee\n"
            "or member."
        ]
    },
    "tacred_person_religion": {"en": ["The religion to which the query person has belonged."]},
    "tacred_person_spouse": {"en": ["The spouse(s) of the query person."]},
    "tacred_person_children": {"en": ["The children of the query person, including adopted and step-children."]},
    "tacred_person_parents": {
        "en": [
            "The parents of the query person. In addition to biological parents, step-parents and adoptive parents\n"
            "are also acceptable answers."
        ]
    },
    "tacred_person_siblings": {"en": ["The brothers and sisters of the query person."]},
    "tacred_person_other_family": {
        "en": [
            "Family other than siblings, parents, children, and spouse (or former spouse). Including: brothers and\n"
            "and sisters in law, grandparents, grandchildren, cousins, aunts, uncles, etc."
        ]
    },
    "tacred_person_charges": {"en": ["The charges or crimes (alleged or convicted) of the query person."]},
    "tacred_organization_template": {
        "en": ["Organization templates encodes the information about the given query Organization entity."]
    },
    "tacred_organization_alternate_names": {
        "en": [
            "Any name used to refer the query organization that is distinct from the 'official' name.\n"
            "Including: former names, aliases, alternate spellings, acronyms, abbreviations, ..."
        ]
    },
    "tacred_organnization_political_or_religious_affiliation": {
        "en": ["Ideological groups with which the organization is associated."]
    },
    "tacred_organization_top_members_employees": {
        "en": [
            "The persons in high-level, leading positions at the query organization. Top member position imply\n"
            "a level of decision-making authority over the entire organization."
        ]
    },
    "tacred_organization_number_of_employees_members": {
        "en": ["The total number of people who are employed by or have membership in an organization."]
    },
    "tacred_organization_members": {
        "en": [
            "Organization or Geopolitical entities that are members of the query organization. While similar to\n"
            "subsidiaries, members is different because correct fillers are distinct entities capable of\n"
            "autonomously ending their membership with the query organization."
        ]
    },
    "tacred_organization_member_of": {
        "en": ["Organizations or geopolitical entities of which the query organization is a member itself."]
    },
    "tacred_organization_subsidiaries": {"en": ["Organizations that are subsidiaries of the query organization."]},
    "tacred_organization_parents": {
        "en": ["Organizations or geopolitical entities of which the query organization is a subsidiary."]
    },
    "tacred_organization_founded_by": {
        "en": ["The person, organization, or geopolitical entity that founded the query organization."]
    },
    "tacred_organization_date_founded": {"en": ["The date on which the query organization was founded."]},
    "tacred_organization_date_dissolved": {
        "en": [
            "The date on which the query organization was dissolved. When companies merge to form a new entity,\n"
            "the original companies should be considered as dissolved."
        ]
    },
    "tacred_organization_country_of_headquarters": {
        "en": ["Countries on which the headquarters of the query organization are located."]
    },
    "tacred_organization_state_or_province_of_headquarters": {
        "en": ["Location of the headquarters of the query organization at the state or province level."]
    },
    "tacred_organization_city_of_headquarters": {
        "en": ["Location of the headquarters of the assigned organization at the city, town, or village level."]
    },
    "tacred_organization_shareholders": {
        "en": ["Any organization, person, or geopolitical entity that holds shares of the organization."]
    },
    "tacred_organization_website": {"en": ["An official top level URL for the organization's website."]},
}

src/tasks/tacred/guidelines_gold.py
hitz-zentroa/GoLLIE
from typing import List, Optional

from ..utils_typing import Name, String, Template, Value, dataclass


"""Relation definitions

The relation definitions are derived from the official TACRED guidelines:
https://tac.nist.gov/2014/KBP/ColdStart/guidelines/TAC_KBP_2014_Slot_Descriptions_V1.4.pdf
"""


@dataclass
class PersonTemplate(Template):
    """{tacred_person}"""

    query: str  # The Person entity query
    alternate_names: Optional[List[Name]] = None
    """{tacred_person_alternate_names}"""
    date_of_birth: Optional[Value] = None
    """{tacred_person_date_of_birth}"""
    age: Optional[Value] = None
    """{tacred_person_age}"""
    country_of_birth: Optional[Name] = None
    """{tacred_person_country_of_birth}"""
    state_or_province_of_birth: Optional[Name] = None
    """{tacred_person_state_or_province_of_birth}"""
    city_of_birth: Optional[Name] = None
    """{tacred_person_city_of_birth}"""
    origin: Optional[List[Name]] = None
    """{tacred_person_origin}"""
    date_of_death: Optional[Value] = None
    """{tacred_person_date_of_death}"""
    country_of_death: Optional[Name] = None
    """{tacred_person_country_of_death}"""
    state_or_province_of_death: Optional[Name] = None
    """{tacred_person_state_or_province_of_death}"""
    city_of_death: Optional[Name] = None
    """{tacred_person_city_of_death}"""
    cause_of_death: Optional[String] = None
    """{tacred_person_cause_of_death}"""
    countries_of_residence: Optional[List[Name]] = None
    """{tacred_person_countries_of_residence}"""
    states_or_provinces_of_residence: Optional[List[Name]] = None
    """{tacred_person_states_or_provinces_of_residence}"""
    cities_of_residence: Optional[List[Name]] = None
    """{tacred_person_cities_of_residence}"""
    schools_attended: Optional[List[Name]] = None
    """{tacred_person_schools_attended}"""
    title: Optional[List[String]] = None
    """{tacred_person_title}"""
    employee_or_member_of: Optional[List[Name]] = None
    """{tacred_person_employee_or_member_of}"""
    religion: Optional[String] = None
    """{tacred_person_religion}"""
    spouse: Optional[List[Name]] = None
    """{tacred_person_spouse}"""
    children: Optional[List[Name]] = None
    """{tacred_person_children}"""
    parents: Optional[List[Name]] = None
    """{tacred_person_parents}"""
    siblings: Optional[List[Name]] = None
    """{tacred_person_siblings}"""
    other_family: Optional[List[Name]] = None
    """{tacred_person_other_family}"""
    charges: Optional[List[String]] = None
    """{tacred_person_charges}"""


@dataclass
class OrganizationTemplate(Template):
    """{tacred_organization_template}"""

    query: str  # The Organization entity query
    alternate_names: Optional[List[Name]] = None
    """{tacred_organization_alternate_names}"""
    political_or_religious_affiliation: Optional[List[Name]] = None
    """{tacred_organnization_political_or_religious_affiliation}"""
    top_members_employees: Optional[List[Name]] = None
    """{tacred_organization_top_members_employees}"""
    number_of_employees_members: Optional[Value] = None
    """{tacred_organization_number_of_employees_members}"""
    members: Optional[List[Name]] = None
    """{tacred_organization_members}"""
    member_of: Optional[List[Name]] = None
    """{tacred_organization_member_of}"""
    subsidiaries: Optional[List[Name]] = None
    """{tacred_organization_subsidiaries}"""
    parents: Optional[List[Name]] = None
    """{tacred_organization_parents}"""
    founded_by: Optional[List[Name]] = None
    """{tacred_organization_founded_by}"""
    date_founded: Optional[Value] = None
    """{tacred_organization_date_founded}"""
    date_dissolved: Optional[Value] = None
    """{tacred_organization_date_dissolved}"""
    country_of_headquarters: Optional[Name] = None
    """{tacred_organization_country_of_headquarters}"""
    state_or_province_of_headquarters: Optional[Name] = None
    """{tacred_organization_state_or_province_of_headquarters}"""
    city_of_headquarters: Optional[Name] = None
    """{tacred_organization_city_of_headquarters}"""
    shareholders: Optional[List[Name]] = None
    """{tacred_organization_shareholders}"""
    website: Optional[String] = None
    """{tacred_organization_website}"""


TEMPLATE_DEFINITIONS = [PersonTemplate, OrganizationTemplate]

src/tasks/tacred/prompts.py
hitz-zentroa/GoLLIE
from typing import List, Type

from src.tasks.tacred.prompts import TEMPLATE_DEFINITIONS
from src.tasks.utils_scorer import TemplateScorer


class TACREDTemplateScorer(TemplateScorer):
    """TACRED Template scorer."""

    valid_types: List[Type] = TEMPLATE_DEFINITIONS

src/tasks/tacred/scorer.py
hitz-zentroa/GoLLIE
import inspect
import logging
import math
import random
import re
from string import Formatter
from typing import Any, Dict, List, Set, Tuple, Type, Union

import black
import numpy as np
from jinja2 import Template

from src.tasks.utils_typing import cast_to


class DatasetLoader:
    """An abstract class for dataset loaders."""

    def __iter__(self):
        for elem in self.elements.values():
            yield elem

    def __len__(self):
        return len(self.elements)

    def __getitem__(self, idx: Union[int, str]) -> Dict[str, Any]:
        if isinstance(idx, int) or isinstance(idx, slice):
            return list(self.elements.values())[idx]  # Not very efficient
        else:
            return self.elements[idx]


class Sampler:
    """
    A generic data `Sampler` class.

    Args:
        dataset_loader (`DatasetLoader`):
            The dataset loader that contains the data information.
        task (`str`, optional):
            The task to sample. Defaults to `None`.
        split (`str`, optional):
            The split to sample. It must be one of the following: "train", "dev" or
            "test". Depending on the split the sampling strategy differs. Defaults to
            `"train"`.
        parallel_instances (`Union[int, Tuple[int, int]]`, optional):
            The number of sentences sampled in parallel. Options:

                * **`int`**: The amount of elements that will be sampled in parallel.
                * **`tuple`**: The range of elements that will be sampled in parallel.

            Defaults to 1.
        max_guidelines (`int`, optional):
            The number of guidelines to append to the example at the same time. If `-1`
            is given then all the guidelines are appended. Defaults to `-1`.
        sample_total_guidelines (`int`, optional):
            The total number of guidelines to sample. If `-1` is given then all the
            guidelines are sampled. Defaults to `-1`.
        guideline_dropout (`float`, optional):
            The probability to dropout a guideline definition for the given example. This
            is only applied on training. Defaults to `0.0`.
        seed (`float`, optional):
            The seed to sample the examples. Defaults to `0`.
        prompt_template (`str`, optional):
            The path to the prompt template. Defaults to `"templates/prompt_eae.txt"`.
        ensure_positives_on_train (bool, optional):
            Whether to ensure that the guidelines of annotated examples are not removed.
            Defaults to `True`.
        dataset_name (str, optional):
            The name of the dataset. Defaults to `None`.
        scorer (`str`, optional):
           The scorer class import string. Defaults to `None`.
        sample_only_gold_guidelines (`bool`, optional):
            Whether to sample only guidelines of present annotations. Defaults to `False`.
        task_definitions (`List[Type]`, optional):
            The task definitions or guidelines. Defaults to `None`.
        task_target (`str`, optional):
            The key of the target task annotations in the dict outputed by the
            `DatasetLoader`. This is useful when the `DataLoader` returns annotations for
            different tasks. Defaults to "labels".
        remove_guidelines (`bool`, optional):
            Whether or not to remove guideline information. This is usefull for building the
            baseline. Defaults to `False`.
        is_coarse_to_fine (`bool`, optional):
            Whether or not the task is coarse_to_fine classification. Defaults to `False`.
        coarse_to_fine (`Dict[Type, List[Type]]`, optional):
            If `is_coarse_to_fine` this argument contains the information to map from coarse
            labels to fine labels. Defaults to `None`.
        fine_to_coarse (`Dict[Type, Type]`, optional):
            If `is_coarse_to_fine` this argument contains the information to map from fine
            labels to coarse labels. Defaults to `None`.
        lang (`str`, optional):
            Language of the guidelines to sample. Defaults to `"en"`.
        definitions (`Dict[str, Any]`, optional):
            Dictionary from where to sample the guideline definitions. Defaults to None.
        include_examples_prob (float, optional):
            Whether or not include examples in the guidelines. Defaults to `0.0`.
        examples (`Dict[str, Any]`, optional):
            Dictionary from where to sample the examples. Defaults to None.
        label_noise_prob (`float`, optional):
            The probability to hide the label names. Defaults to `0.0`.

    Raises:
        ValueError:
            raised when no task definitions are given.
    """

    def __init__(
        self,
        dataset_loader: DatasetLoader,
        task: str = None,
        split: str = "train",
        parallel_instances: Union[int, Tuple[int, int]] = 1,
        max_guidelines: int = -1,
        sample_total_guidelines: int = -1,
        guideline_dropout: float = 0.0,
        seed: float = 0,
        prompt_template: str = "templates/prompt.txt",
        ensure_positives_on_train: bool = False,
        sample_only_gold_guidelines: bool = False,
        dataset_name: str = None,
        scorer: str = None,
        task_definitions: List[Type] = None,
        task_target: str = "labels",
        remove_guidelines: bool = False,
        is_coarse_to_fine: bool = False,
        coarse_to_fine: Dict[Type, List[Type]] = None,
        fine_to_coarse: Dict[Type, Type] = None,
        lang: str = "en",
        definitions: Dict[str, Any] = None,
        include_examples_prob: float = 0.0,
        examples: Dict[str, Any] = None,
        label_noise_prob: float = 0.0,
        coarse_dropout: float = 0.0,
        **kwargs,
    ) -> None:
        self.loader = dataset_loader
        self.task = task
        assert split in [
            "train",
            "dev",
            "test",
        ], f"{split} must be either 'train', 'dev' or 'test'."
        self.split = split
        if isinstance(parallel_instances, int):
            parallel_instances = (1, parallel_instances)
        self.parallel_instances = tuple(parallel_instances)
        self.guideline_dropout = guideline_dropout
        self.coarse_dropout = coarse_dropout
        self.seed = seed
        if not task_definitions or not len(task_definitions):
            raise ValueError("task_definitions argument must not be None or empty")
        self.task_definitions = task_definitions
        self.task_target = task_target

        if max_guidelines < 0 or max_guidelines > len(self.task_definitions):
            self.max_guidelines = len(self.task_definitions)
        else:
            self.max_guidelines = max_guidelines
        if sample_total_guidelines < 0 or sample_total_guidelines > len(self.task_definitions):
            self.sample_total_guidelines = len(self.task_definitions)
        else:
            self.sample_total_guidelines = sample_total_guidelines
        self.ensure_positives_on_train = ensure_positives_on_train
        self.sample_only_gold_guidelines = sample_only_gold_guidelines

        with open(prompt_template, "rt") as f:
            self.template = Template(f.read())

        self.dataset_name = dataset_name
        self.scorer_cls = scorer

        # Maping information for coarse --> fine tasks such as EAE or RC
        self.is_coarse_to_fine = is_coarse_to_fine
        self._coarse_to_fine = coarse_to_fine
        self._fine_to_coarse = fine_to_coarse

        self._black_mode = black.Mode()
        self.remove_guidelines = remove_guidelines
        # self._remove_guidelines_re = re.compile(r'"""(.+\n?)*"""')
        self._remove_guidelines_re = re.compile(r'"""[^"]+"""')
        self._remove_guidelines_fn = lambda x: self._remove_guidelines_re.sub("", x).replace("\n    \n", "\n")

        self._remove_comments_re = re.compile(r"#.+?\n")
        self._remove_comments_fn = lambda x: self._remove_comments_re.sub("\n", x)

        self._remove_empty_comments_re = re.compile(r"#()*\n")
        self._remove_empty_comments_fn = lambda x: self._remove_empty_comments_re.sub("\n", x)
        self._formatter = Formatter()

        self.lang = lang
        self.definitions = definitions
        if not self.definitions:
            raise ValueError("You must provide definitions for your guidelines!")
        self.include_examples_prob = include_examples_prob
        # Make 1.0 prob on example sampling in evaluation for reproducibility
        if self.include_examples_prob > 0 and self.split != "train":
            self.include_examples_prob = 1.0
        self.examples = examples
        if include_examples_prob > 0 and not self.examples:
            logging.warn(
                "`include_examples_prob` is > 0 but `examples` is None. If you want to include examples, you must"
                " provide examples. `include_examples_prob` has been changed to 0.0"
            )
            self.include_examples_prob = 0

        self.label_noise_prob = label_noise_prob
        self._class_label_re = re.compile(r"class (\w+)")

    def _sample(self, instances):
        # _gold refers to specifc gold information that is used in the template (depends on the task)
        _gold: List[Any] = [gold for inst in instances for gold in inst["gold"]]
        # positive_guidelines referst just to the guidelines definitions of the labels in the example
        positive_guidelines: Set[Type] = {type(ann) for inst in instances for ann in inst[self.task_target]}
        if self.is_coarse_to_fine:
            coarse_guidelines: Set[Type] = {self._fine_to_coarse[_def] for _def in positive_guidelines}

        guidelines: List[Type] = [*self.task_definitions]

        # The variable all_guidelines makes compatible the coarse-to-fine with normal tasks
        all_guidelines = [guidelines] if not self.is_coarse_to_fine else coarse_guidelines
        for guidelines in all_guidelines:
            if self.is_coarse_to_fine:
                if self.coarse_dropout and random.random() < self.coarse_dropout:
                    continue
                # In case of `is_coarse_to_fine` the guidelines variable is a single type
                coarse_type = guidelines
                guidelines = self._coarse_to_fine[coarse_type]

            if self.sample_only_gold_guidelines:
                # This may defer with `positive_guidelines` because we can apply this after coarse-to-fine conversion
                guidelines = [
                    definition
                    for definition in guidelines
                    if any(isinstance(ann, definition) for inst in instances for ann in inst[self.task_target])
                ]

            # Reduce the ammount of labels by sampling. We can make sure positive guidelines are sampled using `ensure_positives_on_train`
            if self.sample_total_guidelines < len(guidelines) and not self.sample_only_gold_guidelines:
                p = np.asarray(
                    [
                        (100.0 if _def in positive_guidelines and self.ensure_positives_on_train else 0.0)
                        for _def in guidelines
                    ]
                )
                p += 1.0 / p.shape[0]
                p /= p.sum()
                guidelines = np.random.choice(
                    np.asarray(guidelines),
                    size=(self.sample_total_guidelines,),
                    replace=False,
                    p=p,
                ).tolist()

            # Shuffle the guidelines
            random.shuffle(guidelines)
            # Split the guidelines according to `max_guidelines`
            splits = math.ceil(len(guidelines) / self.max_guidelines)
            for i in range(splits):
                _guidelines = guidelines[i * self.max_guidelines : (i + 1) * self.max_guidelines]
                # Apply guideline dropout
                if self.split == "train":
                    _guidelines_dropout = [
                        _def
                        for _def in _guidelines
                        if random.random() > self.guideline_dropout
                        or (_def in positive_guidelines and self.ensure_positives_on_train)
                    ]

                    if len(_guidelines_dropout) == 0 and len(_guidelines) > 0:
                        # Ensure at least one guideline is used
                        _guidelines_dropout.append(random.choice(_guidelines))
                    _guidelines = _guidelines_dropout

                _ann = [ann for inst in instances for ann in inst[self.task_target] if type(ann) in _guidelines]
                _text = " ".join([inst["text"] for inst in instances]).strip()
                # This makes the gold information useful for coarse to fine tasks
                if self.is_coarse_to_fine:
                    _gold = [cast_to(ann, coarse_type) for ann in _ann]

                # Remove the chances for hallucination because the task is classification
                if self.is_coarse_to_fine and not len(_ann):
                    continue

                _guidelines = [inspect.getsource(definition) for definition in _guidelines]
                # Apply definition paraphrases if train
                _definitions = {
                    key: random.choice(value[self.lang]) if self.split == "train" else value[self.lang][0]
                    for key, value in self.definitions.items()
                }
                # Sample few-shot examples if train (add epsilon for not sampling a 0.0)
                if min(random.random() + 1e-6, 1.0) <= self.include_examples_prob:
                    _examples = {
                        key: (
                            f"""Such as: "{'", "'.join(random.sample(value[self.lang], k=min(5,len(value[self.lang]))))}" """
                            if self.split == "train"
                            else f"""Such as: "{'", "'.join(value[self.lang][:5])}" """
                        )
                        for key, value in self.examples.items()
                    }
                else:
                    # _examples = {key: "" for key in self.examples.keys()}
                    _examples = {
                        key[1]: ""
                        for definition in _guidelines
                        for key in self._formatter.parse(definition)
                        if key[1] is not None and "example" in key[1]
                    }
                _repl = {**_examples, **_definitions}
                _guidelines = [definition.format(**_repl) for definition in _guidelines]
                # If no examples are provide, empty comments are created, the following line removes them
                _guidelines = {self._remove_empty_comments_fn(definition) for definition in _guidelines}

                # Remove definitions for baseline
                if self.remove_guidelines:
                    _guidelines = [self._remove_guidelines_fn(definition) for definition in _guidelines]
                    _guidelines = [self._remove_comments_fn(definition) for definition in _guidelines]

                text = self.template.render(guidelines=_guidelines, text=_text, annotations=_ann, gold=_gold)
                # Apply label noise (but keem them if we are removing the guidelines)
                if self.split == "train" and self.label_noise_prob > 0.0 and not self.remove_guidelines:
                    pretext_idx = text.index("\ntext =")
                    results_idx = text.index("\nresult =")
                    _pretext = text[:pretext_idx]
                    _intext = text[pretext_idx:results_idx]
                    _postext = text[results_idx:]
                    class_names = self._class_label_re.findall(_pretext)
                    random.shuffle(class_names)
                    i = 1
                    for name in class_names:
                        if random.random() <= self.label_noise_prob:
                            _pretext = _pretext.replace(f"{name}", f"LABEL_{i}")
                            _postext = _postext.replace(f"{name}(", f"LABEL_{i}(")
                            i += 1
                    text = _pretext + _intext + _postext

                yield {
                    "ids": [inst["id"] for inst in instances],
                    "task_id": f"{self.dataset_name}_{self.task}",
                    "scorer_cls": self.scorer_cls,
                    "labels": black.format_str(_ann.__repr__(), mode=self._black_mode),
                    "text": black.format_str(text, mode=self._black_mode),
                    "unlabelled_sentence": _text,
                }

    def __iter__(self):
        random.seed(self.seed)
        np.random.seed(self.seed)
        instances = []
        total_inst = random.randint(*self.parallel_instances)
        prev_id = None
        for elem in self.loader:
            # Prevent mixing sentences from different documents. TODO: generalize
            if (len(instances) == total_inst) or (prev_id is not None and elem["doc_id"] != prev_id):
                for samp in self._sample(instances):
                    yield samp
                instances = []
                total_inst = random.randint(*self.parallel_instances)

            instances.append(elem)
            prev_id = elem["doc_id"]

        if len(instances):
            for samp in self._sample(instances):
                yield samp

src/tasks/utils_data.py
hitz-zentroa/GoLLIE
from typing import Any, Dict, List, Type, Union

from typing_extensions import override

from .utils_typing import Entity, Event, Relation, Template, Value


class Scorer:
    """An abstract class for scorers."""

    def __call__(self, reference: List[Any], predictions: List[Any]) -> Dict[str, Dict[str, float]]:
        """
        Computes the scores for the given reference and predictions.

        Args:
            reference (`List[Any]`):
                The list of reference or gold annotations.
            predictions (`List[Any]`):
                The list of predicted annotations.

        Returns:
            Dict[str, Dict[str, float]]:
                A `dict` object that contains the scores (usually Precision, Recall and
                F1-score) for the evaluated tasks.
        """
        raise NotImplementedError("This method must be implemented.")

    def _filter_valid_types(self, elems: List[Any]) -> List[Any]:
        """
        Filters the non-valid annotations for a given task.

        Args:
            elems (`List[Any]`):
                List of annotations to filter based on the predefined valid types.

        Returns:
            `List[Union[Entity, Value]]`:
                List of the filtered annotations to keep only the ones defined by the
                valid types.
        """
        return [elem for elem in elems if any(isinstance(elem, _type) for _type in self.valid_types)]


class SpanScorer(Scorer):
    """A general scorer implementation for span identification and classification
    tasks.
    """

    valid_types: List[Type] = [Entity, Value]

    @override
    def __call__(
        self,
        reference: List[Union[Entity, Value]],
        predictions: List[Union[Entity, Value]],
    ) -> Dict[str, Dict[str, float]]:
        if not len(reference) or (len(reference) and not isinstance(reference[0], list)):
            reference = [reference]
        if not len(predictions) or (len(predictions) and not isinstance(predictions[0], list)):
            predictions = [predictions]

        assert len(reference) == len(
            predictions
        ), f"Reference ({len(reference)}) and prediction ({len(predictions)}) amount must be equal."

        tp = total_pos = total_pre = 0

        class_scores = {}

        for ref, pre in zip(reference, predictions):
            ref = self._filter_valid_types(ref)
            pre = self._filter_valid_types(pre)

            ref2 = ref.copy()
            pre2 = pre.copy()

            total_pos += len(ref)
            total_pre += len(pre)
            for entity in pre:
                if entity in ref:
                    tp += 1
                    ref.pop(ref.index(entity))

            # Calculate class scores
            for entity in ref2:
                label = type(entity).__name__
                if label not in class_scores:
                    class_scores[label] = {"tp": 0, "total_pos": 0, "total_pre": 0}
                class_scores[label]["total_pos"] += 1

            for entity in pre2:
                label = type(entity).__name__
                if label not in class_scores:
                    class_scores[label] = {"tp": 0, "total_pos": 0, "total_pre": 0}
                class_scores[label]["total_pre"] += 1
                if entity in ref2:
                    class_scores[label]["tp"] += 1
                    ref2.pop(ref2.index(entity))

        precision = tp / total_pre if total_pre > 0.0 else 0.0
        recall = tp / total_pos if total_pos > 0.0 else 0.0
        f1_score = 2 * precision * recall / (precision + recall) if (precision + recall) > 0.0 else 0.0
        # Calculate class scores
        for label in class_scores:
            class_scores[label]["precision"] = (
                class_scores[label]["tp"] / class_scores[label]["total_pre"]
                if class_scores[label]["total_pre"] > 0.0
                else 0.0
            )
            class_scores[label]["recall"] = (
                class_scores[label]["tp"] / class_scores[label]["total_pos"]
                if class_scores[label]["total_pos"] > 0.0
                else 0.0
            )
            class_scores[label]["f1-score"] = (
                2
                * class_scores[label]["precision"]
                * class_scores[label]["recall"]
                / (class_scores[label]["precision"] + class_scores[label]["recall"])
                if (class_scores[label]["precision"] + class_scores[label]["recall"]) > 0.0
                else 0.0
            )

        return {
            "spans": {"precision": precision, "recall": recall, "f1-score": f1_score, "class_scores": class_scores},
        }


class RelationScorer(SpanScorer):
    """A general scorer implementation for relation identification and
    classification tasks
    """

    valid_types: List[Type] = [Relation]

    @override
    def __call__(self, reference: List[Relation], predictions: List[Relation]) -> Dict[str, Dict[str, float]]:
        output = super().__call__(reference, predictions)
        return {"relations": output["spans"]}


class EventScorer(Scorer):
    """A general scorer implementation for event and argument extraction."""

    valid_types: List[Type] = [Event]

    @override
    def __call__(self, reference: Any, predictions: Any) -> Dict[str, Dict[str, float]]:
        if not len(reference) or (len(reference) and not isinstance(reference[0], list)):
            reference = [reference]
        if not len(predictions) or (len(predictions) and not isinstance(predictions[0], list)):
            predictions = [predictions]

        assert len(reference) == len(
            predictions
        ), f"Reference ({len(reference)}) and prediction ({len(predictions)}) amount must be equal."
        e_tp = e_total_pos = e_total_pre = 0
        a_tp = a_total_pos = a_total_pre = 0

        event_scores = {}

        for ref, pre in zip(reference, predictions):
            ref = self._filter_valid_types(ref)
            pre = self._filter_valid_types(pre)
            ref2 = ref.copy()
            pre2 = pre.copy()

            e_total_pos += len(ref)
            for event in ref:
                a_total_pos += len(event)

            e_total_pre += len(pre)
            for pre_event in pre:
                a_total_pre += len(pre_event)
                if pre_event in ref:
                    ref_event = ref.pop(ref.index(pre_event))
                    e_tp += 1
                    a_tp += len(ref_event & pre_event)

            # Calculate argument scores
            for event in ref2:
                label = type(event).__name__
                if label not in event_scores:
                    event_scores[label] = {"tp": 0, "total_pos": 0, "total_pre": 0}
                event_scores[label]["total_pos"] += 1

            for event in pre2:
                label = type(event).__name__
                if label not in event_scores:
                    event_scores[label] = {"tp": 0, "total_pos": 0, "total_pre": 0}
                event_scores[label]["total_pre"] += 1
                if event in ref2:
                    event_scores[label]["tp"] += 1
                    ref2.pop(ref2.index(event))

        e_precision = e_tp / e_total_pre if e_total_pre > 0.0 else 0.0
        a_precision = a_tp / a_total_pre if a_total_pre > 0.0 else 0.0
        e_recall = e_tp / e_total_pos if e_total_pos > 0.0 else 0.0
        a_recall = a_tp / a_total_pos if a_total_pos > 0.0 else 0.0
        e_f1_score = 2 * e_precision * e_recall / (e_precision + e_recall) if (e_precision + e_recall) > 0.0 else 0.0
        a_f1_score = 2 * a_precision * a_recall / (a_precision + a_recall) if (a_precision + a_recall) > 0.0 else 0.0

        # Calculate event scores
        for label in event_scores:
            event_scores[label]["precision"] = (
                event_scores[label]["tp"] / event_scores[label]["total_pre"]
                if event_scores[label]["total_pre"] > 0.0
                else 0.0
            )
            event_scores[label]["recall"] = (
                event_scores[label]["tp"] / event_scores[label]["total_pos"]
                if event_scores[label]["total_pos"] > 0.0
                else 0.0
            )
            event_scores[label]["f1-score"] = (
                2
                * event_scores[label]["precision"]
                * event_scores[label]["recall"]
                / (event_scores[label]["precision"] + event_scores[label]["recall"])
                if (event_scores[label]["precision"] + event_scores[label]["recall"]) > 0.0
                else 0.0
            )

        return {
            "events": {
                "precision": e_precision,
                "recall": e_recall,
                "f1-score": e_f1_score,
                "class_scores": event_scores,
            },
            "arguments": {
                "precision": a_precision,
                "recall": a_recall,
                "f1-score": a_f1_score,
            },
        }


class TemplateScorer(Scorer):
    """A general scorer implementation for template extraction."""

    valid_types: List[Type] = [Template]

    @override
    def __call__(self, reference: Any, predictions: Any) -> Dict[str, Dict[str, float]]:
        if not len(reference) or (len(reference) and not isinstance(reference[0], list)):
            reference = [reference]
        if not len(predictions) or (len(predictions) and not isinstance(predictions[0], list)):
            predictions = [predictions]

        assert len(reference) == len(
            predictions
        ), f"Reference ({len(reference)}) and prediction ({len(predictions)}) amount must be equal."
        t_tp = t_total_pos = t_total_pre = 0
        s_tp = s_total_pos = s_total_pre = 0

        template_scores = {}

        for ref, pre in zip(reference, predictions):
            ref = self._filter_valid_types(ref)
            pre = self._filter_valid_types(pre)
            ref2 = ref.copy()
            pre2 = pre.copy()

            t_total_pos += len(ref)
            for template in ref:
                s_total_pos += len(template)

            t_total_pre += len(pre)
            for pre_temp in pre:
                s_total_pre += len(pre_temp)
                if pre_temp in ref:
                    ref_temp = ref.pop(ref.index(pre_temp))
                    t_tp += 1
                    s_tp += len(ref_temp & pre_temp)

            # Calculate slot scores
            for template in ref2:
                label = type(template).__name__
                if label not in template_scores:
                    template_scores[label] = {"tp": 0, "total_pos": 0, "total_pre": 0}
                template_scores[label]["total_pos"] += 1

            for template in pre2:
                label = type(template).__name__
                if label not in template_scores:
                    template_scores[label] = {"tp": 0, "total_pos": 0, "total_pre": 0}
                template_scores[label]["total_pre"] += 1
                if template in ref2:
                    template_scores[label]["tp"] += 1
                    ref2.pop(ref2.index(template))

        t_precision = t_tp / t_total_pre if t_total_pre > 0.0 else 0.0
        s_precision = s_tp / s_total_pre if s_total_pre > 0.0 else 0.0
        t_recall = t_tp / t_total_pos if t_total_pos > 0.0 else 0.0
        s_recall = s_tp / s_total_pos if s_total_pos > 0.0 else 0.0
        t_f1_score = 2 * t_precision * t_recall / (t_precision + t_recall) if (t_precision + t_recall) > 0.0 else 0.0
        s_f1_score = 2 * s_precision * s_recall / (s_precision + s_recall) if (s_precision + s_recall) > 0.0 else 0.0

        # Calculate template scores
        for label in template_scores:
            template_scores[label]["precision"] = (
                template_scores[label]["tp"] / template_scores[label]["total_pre"]
                if template_scores[label]["total_pre"] > 0.0
                else 0.0
            )
            template_scores[label]["recall"] = (
                template_scores[label]["tp"] / template_scores[label]["total_pos"]
                if template_scores[label]["total_pos"] > 0.0
                else 0.0
            )
            template_scores[label]["f1-score"] = (
                2
                * template_scores[label]["precision"]
                * template_scores[label]["recall"]
                / (template_scores[label]["precision"] + template_scores[label]["recall"])
                if (template_scores[label]["precision"] + template_scores[label]["recall"]) > 0.0
                else 0.0
            )

        return {
            "templates": {
                "precision": t_precision,
                "recall": t_recall,
                "f1-score": t_f1_score,
                "class_scores": template_scores,
            },
            "slots": {
                "precision": s_precision,
                "recall": s_recall,
                "f1-score": s_f1_score,
            },
        }

src/tasks/utils_scorer.py
hitz-zentroa/GoLLIE
from __future__ import annotations

import dataclasses
import importlib
import inspect
import logging
import re
from collections import defaultdict
from copy import deepcopy
from dataclasses import dataclass as org_dataclass
from typing import Any, Dict, List, Tuple, Type, TypeVar, Union


def dataclass(
    cls=None,
    /,
    *,
    init=True,
    repr=False,
    eq=False,
    order=False,
    unsafe_hash=False,
    frozen=False,
):
    return org_dataclass(
        cls,
        init=init,
        repr=repr,
        eq=eq,
        order=order,
        unsafe_hash=unsafe_hash,
        frozen=frozen,
    )


def cast_to(obj: Any, dtype: Type) -> Any:
    if not isinstance(obj, dtype):
        raise TypeError(f"Type {dtype} must be a parent class of object {obj}.")

    _inst = {param: getattr(obj, param) for param in inspect.signature(dtype).parameters.keys() if hasattr(obj, param)}
    return dtype(**_inst)


class HallucinatedType:
    def __init__(self, *args, **kwargs) -> None:
        pass

    def exists_in(self, text: str) -> bool:
        """
        Checks whether the annotation exists on a given text. This function is used to
        identify model alucinations.

        Args:
            text (`str`):
                The text used to check whether the annotation is an alucionation or not.

        Returns:
            `bool`:
                Whether the annotation exists on the input text or not.
        """
        return True


class Name(str):
    def __repr__(self):
        return f"Name({super().__repr__()})"


class Value(str):
    def __repr__(self) -> str:
        return f"Value({super().__repr__()})"


class String(str):
    def __repr__(self) -> str:
        return f"String({super().__repr__()})"


@dataclass
class Entity:
    """A general class to represent entities."""

    span: str

    def __post_init__(self: Entity) -> None:
        self._allow_partial_match: bool = False

    def __eq__(self: Entity, other: Entity) -> bool:
        self_span = self.span.lower().strip()
        other_span = other.span.lower().strip()
        if self._allow_partial_match:
            return type(self) == type(other) and (self.span in other_span or other_span in self.span)

        return type(self) == type(other) and self_span == other_span

    def __repr__(self) -> str:
        """Returns a string containing only the non-default field values."""
        s = ", ".join(
            f"{field.name}={getattr(self, field.name)!r}"
            for field in dataclasses.fields(self)
            if getattr(self, field.name) != field.default
        )
        return f"{type(self).__name__}({s})"

    def key(self) -> Union[str, None]:
        """
        Return the key span.

        Returns:
            Union[str, None]:
                The span that represents the annotation.
        """
        return self.span.lower()

    def exists_in(self, text: str) -> bool:
        """
        Checks whether the annotation exists on a given text. This function is used to
        identify model alucinations.

        Args:
            text (`str`):
                The text used to check whether the annotation is an alucionation or not.

        Returns:
            `bool`:
                Whether the annotation exists on the input text or not.
        """

        # Check if the span is not None
        if not isinstance(self.span, str):
            return False

        return self.span.lower().strip() in text.lower()

    def index(self, text: str) -> int:
        """
        Returns the first position of the span given the text.

        Args:
            text (`str`):
                The text to search the span on.

        Raises:
            IndexError:
                Raised when the span does not exist in the text

        Returns:
            `Tuple[int, int]`:
                The position of the span in the text.
        """
        if not self.exists_in(text):
            raise IndexError("The span is not in text.")

        pos = text.lower().index(self.span.lower().strip())
        return (pos, pos + len(self.span))


@dataclass
class Relation:
    """A general class to represent relations."""

    arg1: str
    arg2: str

    def __post_init__(self: Relation) -> None:
        self._allow_partial_match: bool = False

    def __eq__(self: Value, other: Value) -> bool:
        if self._allow_partial_match:
            return (
                type(self) == type(other)
                and (self.arg1 in other.arg1 or other.arg1 in self.arg1)
                and (self.arg2 in other.arg2 or other.arg2 in self.arg2)
            )
        return type(self) == type(other) and self.arg1 == other.arg1 and self.arg2 == other.arg2

    def __repr__(self) -> str:
        """Returns a string containing only the non-default field values."""
        s = ", ".join(
            f"{field.name}={getattr(self, field.name)!r}"
            for field in dataclasses.fields(self)
            if getattr(self, field.name) != field.default
        )
        return f"{type(self).__name__}({s})"

    def key(self) -> Union[str, None]:
        """
        Return the key span.

        Returns:
            Union[str, None]:
                The span that represents the annotation.
        """
        return None

    def exists_in(self, text: str) -> bool:
        """
        Checks whether the annotation exists on a given text. This function is used to
        identify model alucinations.

        Args:
            text (`str`):
                The text used to check whether the annotation is an alucionation or not.

        Returns:
            `bool`:
                Whether the annotation exists on the input text or not.
        """
        if not isinstance(self.arg1, str) or not isinstance(self.arg2, str):
            return False

        return self.arg1.lower() in text.lower() and self.arg2.lower() in text.lower()

    def index(self, text: str) -> Tuple[int, int]:
        """
        Returns the first position of the span given the text.

        Args:
            text (`str`):
                The text to search the span on.

        Raises:
            IndexError:
                Raised when the span does not exist in the text

        Returns:
            `Tuple[int, int]`:
                The position of the span in the text.
        """
        if not self.exists_in(text):
            raise IndexError("The span is not in text.")

        pos = text.lower().index(self.arg1.lower().strip())
        return (pos, pos + len(self.arg1))


@dataclass
class Event:
    """A general class to represent events."""

    mention: str

    def __post_init__(self: Event) -> None:
        self._allow_partial_match: bool = False

    def __eq__(self: Event, other: Event) -> bool:
        if self._allow_partial_match:
            # On events we consider partial match if just the type is predicted
            # Some datasets has different guidelines when annotating event triggers
            return type(self) == type(other)
        return type(self) == type(other) and self.key() == other.key()

    def __and__(self: Event, other: Event) -> Event:
        attrs = {
            attr: []
            for attr, values in inspect.getmembers(self)
            if not (attr.startswith("_") or attr in ["mention", "subtype"] or inspect.ismethod(values))
        }
        if self == other:
            for attr in attrs.keys():
                self_values = getattr(self, attr)
                other_values = deepcopy(getattr(other, attr))
                for value in self_values:
                    if self._allow_partial_match:
                        for i, _value in enumerate(other_values):
                            if value in _value or _value in value:
                                attrs[attr].append(value)
                                other_values.pop(i)
                                break
                    else:
                        if value in other_values:
                            attrs[attr].append(value)
                            other_values.pop(other_values.index(value))

        pos_args = []
        if hasattr(self, "mention"):
            pos_args.append(self.mention)
        if hasattr(self, "subtype"):
            pos_args.append(self.subtype)

        return type(self)(*pos_args, **attrs)

    def __len__(self: Event) -> int:
        attrs = {
            attr: values
            for attr, values in inspect.getmembers(self)
            if not (attr.startswith("_") or attr in ["mention", "subtype"] or inspect.ismethod(values))
        }
        _len = 0
        for values in attrs.values():
            _len += len(values)

        return _len

    def __repr__(self) -> str:
        """Returns a string containing only the non-default field values."""
        s = ", ".join(
            f"{field.name}={getattr(self, field.name)!r}"
            for field in dataclasses.fields(self)
            if getattr(self, field.name) != field.default
        )
        return f"{type(self).__name__}({s})"

    def key(self) -> Union[str, None]:
        """
        Return the key span.

        Returns:
            Union[str, None]:
                The span that represents the annotation.
        """
        if self.mention:
            return self.mention.lower()

        return None

    def exists_in(self, text: str) -> Event:
        """
        Checks whether the annotation exists on a given text. This function is used to
        identify model alucinations.

        Args:
            text (`str`):
                The text used to check whether the annotation is an alucionation or not.

        Returns:
            `bool`:
                Whether the annotation exists on the input text or not.
        """
        text = text.lower()
        # Only return None if the mention is defined and is not in the text
        if self.mention and self.mention.lower() not in text:
            return None

        attrs = {
            attr: []
            for attr, values in inspect.getmembers(self)
            if not (attr.startswith("_") or attr in ["mention", "subtype"] or inspect.ismethod(values))
        }
        for attr in attrs.keys():
            self_values = getattr(self, attr)
            if not isinstance(self_values, list):
                import rich

                rich.print(self)
                raise TypeError(f"{attr}:{self_values} is not iterable")
            for value in self_values:
                # Avoid calling lower to a list when the model hallucinates
                if hasattr(value, "lower") and value.lower() in text:
                    attrs[attr].append(value)

        pos_args = []
        if hasattr(self, "mention"):
            pos_args.append(self.mention)
        if hasattr(self, "subtype"):
            pos_args.append(self.subtype)

        return type(self)(*pos_args, **attrs)

    def index(self, text: str) -> int:
        """
        Returns the first position of the span given the text.

        Args:
            text (`str`):
                The text to search the span on.

        Raises:
            IndexError:
                Raised when the span does not exist in the text

        Returns:
            `int`:
                The position of the span in the text.
        """
        # Return len(text) if there is no mention defined
        if not self.mention:
            return (len(text), len(text))

        if not self.exists_in(text):
            raise IndexError("The span is not in text.")

        pos = text.lower().index(self.mention.lower().strip())
        return (pos, pos + len(self.mention))

    def assert_typing_constraints(self) -> bool:
        import typing

        def check_types(var: Any, _type: TypeVar) -> Tuple[bool, Any]:
            if _type is type(None):
                return (var is None, var)
            elif isinstance(_type, type):
                is_correct = isinstance(var, _type)
                if not is_correct:
                    if isinstance(var, list):
                        var = var[0] if len(var) else _type()
                    var = _type(var)
                return (is_correct, var)
            # elif _type is type(None):
            #     return (var is None, var)
            elif isinstance(_type, typing._GenericAlias):
                origin = _type.__origin__
                if origin is typing.Union:
                    is_correct = False
                    _var = None
                    for _t in _type.__args__:
                        _is_correct, tmp_var = check_types(var, _t)
                        if _is_correct:
                            return (_is_correct, tmp_var)
                        if _t is not type(None):
                            _var = tmp_var
                        is_correct |= _is_correct
                    # return any(check_types(var, _t) for _t in _type.__args__)
                    return (is_correct, _var)
                elif origin is list:
                    if not isinstance(var, list):
                        var = [var]
                    for _t in _type.__args__:
                        _var = []
                        is_correct = True
                        for v in var:
                            _is_correct, v = check_types(v, _t)
                            is_correct &= _is_correct
                            _var.append(v)
                        if is_correct:
                            return (is_correct, _var)  # Small difference here to allow empty lists

                    return (is_correct, _var)
                    # return any(all(check_types(v, _t) for v in var) if isinstance(var, list) else False for _t in _type.__args__)
                elif any(issubclass(origin, _t) for _t in [str, int, float, bool]):
                    return check_types(var, origin)
                else:
                    raise ValueError(f"Unsupported type: {origin}")
            else:
                raise ValueError("Only native types or typing module types are supported.")

        correct = True
        for field in dataclasses.fields(self):
            _correct, value = check_types(getattr(self, field.name), field.type)
            correct &= _correct
            setattr(self, field.name, value)

        return correct


@dataclass
class Template:
    """A general class to represent templates."""

    query: str

    def __eq__(self: Template, other: Template) -> bool:
        return type(self) == type(other) and self.key() == other.key()

    def __repr__(self) -> str:
        """Returns a string containing only the non-default field values."""
        s = ", ".join(
            f"{field.name}={getattr(self, field.name)!r}"
            for field in dataclasses.fields(self)
            if getattr(self, field.name) != field.default
        )
        return f"{type(self).__name__}({s})"

    def _get_attributes(self, ignore: bool = False) -> Dict[str, Any]:
        """Returns the non-positional attributes of a Template instance."""
        attrs = {}
        for attr, values in inspect.getmembers(self):
            if attr.startswith("_") or attr in ["query"] or inspect.ismethod(values):
                continue
            if ignore and values:
                attrs[attr] = type(values)()
            else:
                attrs[attr] = values
        return attrs

    def _get_pos_attributes(self: Template) -> List[Any]:
        """Return the positional attributes of a Template instance."""
        return [getattr(self, attr) for attr in ["query"]]

    def __and__(self: Template, other: Template) -> Template:
        attrs = self._get_attributes(ignore=True)
        if self == other:
            for attr in attrs.keys():
                self_values = getattr(self, attr)
                other_values = deepcopy(getattr(other, attr))
                if self_values is None:
                    continue
                if isinstance(self_values, list):
                    if not isinstance(other_values, list):
                        other_values = [other_values]
                    for value in self_values:
                        if value in other_values:
                            attrs[attr].append(value)
                            other_values.pop(other_values.index(value))
                else:
                    if isinstance(other_values, list):
                        attrs[attr] = self_values if self_values in other_values else None
                    else:
                        attrs[attr] = self_values if self_values == other_values else None

        pos_args = self._get_pos_attributes()

        return type(self)(*pos_args, **attrs)

    def __len__(self: Template) -> int:
        attrs = self._get_attributes()
        _len = 0
        for values in attrs.values():
            if not values:
                continue
            _len += len(values) if isinstance(values, list) else 1

        return _len

    def key(self: Template) -> Union[str, None]:
        """
        Return the key span.

        Returns:
            Union[str, None]:
                The span that represents the annotation.
        """
        if self.query:
            return self.query.lower()

        return None

    def exists_in(self: Template, text: str) -> Template:
        """
        Checks whether the annotation exists on a given text. This function is used to
        identify model alucinations.

        Args:
            text (`str`):
                The text used to check whether the annotation is an alucionation or not.

        Returns:
            `bool`:
                Whether the annotation exists on the input text or not.
        """
        text = text.lower()
        # Only return None if the key is defined and is not in the text
        if self.key() and self.key() not in text:
            return None

        attrs = self._get_attributes(ignore=True)
        for attr in attrs.keys():
            self_values = getattr(self, attr)
            if self_values:
                if isinstance(self_values, list):
                    for value in self_values:
                        if value.lower() in text:
                            attrs[attr].append(value)
                else:
                    if self_values.lower() in text:
                        attrs[attr] = self_values

        pos_args = self._get_pos_attributes()

        return type(self)(*pos_args, **attrs)

    def index(self, text: str) -> int:
        """
        Returns the first position of the span given the text.

        Args:
            text (`str`):
                The text to search the span on.

        Raises:
            IndexError:
                Raised when the span does not exist in the text

        Returns:
            `int`:
                The position of the span in the text.
        """
        # Return len(text) if there is no mention defined
        if not self.key():
            return (len(text), len(text))

        if not self.exists_in(text):
            raise IndexError("The span is not in text.")

        pos = text.lower().index(self.key().strip())
        return (pos, pos + len(self.key()))

    def assert_typing_constraints(self) -> bool:
        import typing

        def check_types(var: Any, _type: TypeVar) -> Tuple[bool, Any]:
            if _type is type(None):
                return (var is None, var)
            elif isinstance(_type, type):
                is_correct = isinstance(var, _type)
                if not is_correct:
                    if isinstance(var, list):
                        var = var[0] if len(var) else _type()
                    var = _type(var)
                return (is_correct, var)
            # elif _type is type(None):
            #     return (var is None, var)
            elif isinstance(_type, typing._GenericAlias):
                origin = _type.__origin__
                if origin is typing.Union:
                    is_correct = False
                    _var = None
                    for _t in _type.__args__:
                        _is_correct, tmp_var = check_types(var, _t)
                        if _is_correct:
                            return (_is_correct, tmp_var)
                        if _t is not type(None):
                            _var = tmp_var
                        is_correct |= _is_correct
                    # return any(check_types(var, _t) for _t in _type.__args__)
                    return (is_correct, _var)
                elif origin is list:
                    if not isinstance(var, list):
                        var = [var]
                    for _t in _type.__args__:
                        _var = []
                        is_correct = True
                        for v in var:
                            _is_correct, v = check_types(v, _t)
                            is_correct &= _is_correct
                            _var.append(v)
                        if is_correct:
                            return (is_correct, _var if len(_var) else None)

                    return (is_correct, _var)
                    # return any(all(check_types(v, _t) for v in var) if isinstance(var, list) else False for _t in _type.__args__)
                elif any(issubclass(origin, _t) for _t in [str, int, float, bool]):
                    return check_types(var, origin)
                else:
                    raise ValueError(f"Unsupported type: {origin}")
            else:
                raise ValueError("Only native types or typing module types are supported.")

        correct = True
        for field in dataclasses.fields(self):
            _correct, value = check_types(getattr(self, field.name), field.type)
            correct &= _correct
            setattr(self, field.name, value)

        return correct


@dataclass
class Generic:
    """A general class to represent Generics."""

    def __eq__(self, other) -> bool:
        return type(self) == type(other)

    def __repr__(self) -> str:
        """Returns a string containing only the non-default field values."""
        s = ", ".join(
            f"{field.name}={getattr(self, field.name)!r}"
            for field in dataclasses.fields(self)
            if getattr(self, field.name) != field.default
        )
        return f"{type(self).__name__}({s})"

    def _get_attributes(self, ignore: bool = False) -> Dict[str, Any]:
        """Returns the non-positional attributes of a Generic instance."""
        attrs = {}
        for attr, values in inspect.getmembers(self):
            if attr.startswith("_") or attr in ["query"] or inspect.ismethod(values):
                continue
            if ignore and values:
                attrs[attr] = type(values)()
            else:
                attrs[attr] = values
        return attrs

    def _get_pos_attributes(self) -> List[Any]:
        """Return the positional attributes of a Generic instance."""
        return []

    def __and__(self, other):
        attrs = self._get_attributes(ignore=True)
        if self == other:
            for attr in attrs.keys():
                self_values = getattr(self, attr)
                other_values = deepcopy(getattr(other, attr))
                if self_values is None:
                    continue
                if isinstance(self_values, list):
                    if not isinstance(other_values, list):
                        other_values = [other_values]
                    for value in self_values:
                        if value in other_values:
                            attrs[attr].append(value)
                            other_values.pop(other_values.index(value))
                else:
                    if isinstance(other_values, list):
                        attrs[attr] = self_values if self_values in other_values else None
                    else:
                        attrs[attr] = self_values if self_values == other_values else None

        pos_args = self._get_pos_attributes()

        return type(self)(*pos_args, **attrs)

    def __len__(self) -> int:
        attrs = self._get_attributes()
        _len = 0
        for values in attrs.values():
            if not values:
                continue
            _len += len(values) if isinstance(values, list) else 1

        return _len

    def exists_in(self, text: str):
        """
        Checks whether the annotation exists on a given text. This function is used to
        identify model alucinations.

        Args:
            text (`str`):
                The text used to check whether the annotation is an alucionation or not.

        Returns:
            `bool`:
                Whether the annotation exists on the input text or not.
        """
        text = text.lower()

        attrs = self._get_attributes(ignore=True)
        for attr in attrs.keys():
            self_values = getattr(self, attr)
            if self_values:
                if isinstance(self_values, list):
                    for value in self_values:
                        if value.lower() in text:
                            attrs[attr].append(value)
                else:
                    if self_values.lower() in text:
                        attrs[attr] = self_values

        pos_args = self._get_pos_attributes()

        return type(self)(*pos_args, **attrs)

    def index(self, text: str) -> int:
        """
        Returns the first position of the span given the text.

        Args:
            text (`str`):
                The text to search the span on.

        Raises:
            IndexError:
                Raised when the span does not exist in the text

        Returns:
            `int`:
                The position of the span in the text.
        """

        raise IndexError("Generic template does not support this function yet")

    def assert_typing_constraints(self) -> bool:
        import typing

        def check_types(var: Any, _type: TypeVar) -> Tuple[bool, Any]:
            if _type is type(None):
                return (var is None, var)
            elif isinstance(_type, type):
                is_correct = isinstance(var, _type)
                if not is_correct:
                    if isinstance(var, list):
                        var = var[0] if len(var) else _type()
                    var = _type(var)
                return (is_correct, var)
            # elif _type is type(None):
            #     return (var is None, var)
            elif isinstance(_type, typing._GenericAlias):
                origin = _type.__origin__
                if origin is typing.Union:
                    is_correct = False
                    _var = None
                    for _t in _type.__args__:
                        _is_correct, tmp_var = check_types(var, _t)
                        if _is_correct:
                            return (_is_correct, tmp_var)
                        if _t is not type(None):
                            _var = tmp_var
                        is_correct |= _is_correct
                    # return any(check_types(var, _t) for _t in _type.__args__)
                    return (is_correct, _var)
                elif origin is list:
                    if not isinstance(var, list):
                        var = [var]
                    for _t in _type.__args__:
                        _var = []
                        is_correct = True
                        for v in var:
                            _is_correct, v = check_types(v, _t)
                            is_correct &= _is_correct
                            _var.append(v)
                        if is_correct:
                            return (is_correct, _var if len(_var) else None)

                    return (is_correct, _var)
                    # return any(all(check_types(v, _t) for v in var) if isinstance(var, list) else False for _t in _type.__args__)
                elif any(issubclass(origin, _t) for _t in [str, int, float, bool]):
                    return check_types(var, origin)
                else:
                    raise ValueError(f"Unsupported type: {origin}")
            else:
                raise ValueError("Only native types or typing module types are supported.")

        correct = True
        for field in dataclasses.fields(self):
            _correct, value = check_types(getattr(self, field.name), field.type)
            correct &= _correct
            setattr(self, field.name, value)

        return correct


class AnnotationList(list):
    """
    A class that handles the list of system outputs.

    Args:
        list (List[Any]):
            List of annotations generated by the system.
    """

    SIMPLE_TYPES = [Entity, Value, Relation]
    COMPLEX_TYPES = [Event, Template]

    def __init__(self, elems: List[Any], hallucinated_no: int = 0, parse_error: bool = False):
        self._hallucinated_no = hallucinated_no
        self._parse_error = parse_error
        super().__init__(elems)

    @property
    def hallucinated_no(self) -> int:
        return self._hallucinated_no

    @property
    def parse_error(self) -> bool:
        return self._parse_error

    @staticmethod
    def _load_guidelines(task_module: str) -> Dict[str, TypeVar]:
        mdl = importlib.import_module(task_module)
        names = {x: y for x, y in mdl.__dict__.items() if not x.startswith("_")}

        return names

    def filter_hallucinations(self, text: str) -> AnnotationList:
        _elems = []
        _counts = defaultdict(int)
        for elem in self:
            if any(isinstance(elem, _type) for _type in self.SIMPLE_TYPES):
                elem = elem if elem.exists_in(text) else None
            elif any(isinstance(elem, _type) for _type in self.COMPLEX_TYPES):
                elem = elem.exists_in(text)
            elif isinstance(elem, HallucinatedType) or isinstance(elem, type):
                continue
            else:
                print(elem)
                if hasattr(elem, "exist_in"):
                    elem = elem if elem.exists_in(text) else None
                else:
                    continue

            if elem is not None:
                key = elem.key()
                # Check if key is not None
                if key:
                    _counts[key] += 1
                    text_counts = text.lower().count(key)
                    # Skip element if the maximum allowed elements already exists
                    if text_counts < _counts[key]:
                        continue

                _elems.append(elem)

        return type(self)(_elems, hallucinated_no=len(self) - len(_elems), parse_error=self.parse_error)

    def assert_typing_constraints(self) -> None:
        for elem in self:
            if hasattr(elem, "assert_typing_constraints"):
                elem.assert_typing_constraints()

    @classmethod
    def from_output(
        cls,
        ann: str,
        task_module: str,
        text: str = None,
        filter_hallucinations: bool = False,
        assert_typing_constraints: bool = True,
    ) -> AnnotationList:
        # Import guidelines
        guidelines = cls._load_guidelines(task_module)
        locals().update(guidelines)

        reading = True
        parse_error = False
        while reading:
            try:
                _elems = eval(ann)
                reading = False
            # Handle hallucinations
            except NameError as e:
                name = re.search(r"'\w+'", e.args[0]).group(0).strip("'")
                logging.warning(f"An hallucinated predicted guideline found: {name}")
                locals().update({name: HallucinatedType})
            except Exception:
                _elems = []
                reading = False
                parse_error = True

        self = cls(_elems, parse_error=parse_error)

        if filter_hallucinations:
            if text is None:
                raise ValueError("To filter the hallucinations the text argument must not be None.")
            self = self.filter_hallucinations(text)

        if assert_typing_constraints:
            self.assert_typing_constraints()

        return self

    @classmethod
    def from_gold(
        cls,
        ann: str,
        task_module: str,
        text: str = None,
        filter_hallucinations: bool = False,
        assert_typing_constraints: bool = True,
    ) -> AnnotationList:
        # Import guidelines
        guidelines = cls._load_guidelines(task_module)
        locals().update(guidelines)

        # Remove malformed elements
        _elems = []
        pase_error = False
        for elem in eval(ann):
            try:
                elem = eval(elem)
                _elems.append(elem)
            except SyntaxError:
                logging.warning(f"Found an incorrectly formatted gold: {elem}")
                pase_error = True

        self = cls(_elems, parse_error=pase_error)

        if filter_hallucinations:
            if text is None:
                raise ValueError("To filter the hallucinations the text argument must not be None.")
            self = self.filter_hallucinations(text)

        if assert_typing_constraints:
            self.assert_typing_constraints()

        return self

    def to_string(self) -> str:
        return str(self)

src/tasks/utils_typing.py
hitz-zentroa/GoLLIE
from . import data_loader, prompts, scorer


__all__ = ["data_loader", "prompts", "scorer"]

src/tasks/wikievents/__init__.py
hitz-zentroa/GoLLIE
import inspect
import json
from typing import Tuple, Union

from src.tasks.wikievents.guidelines import GUIDELINES
from src.tasks.wikievents.guidelines_gold import EXAMPLES
from src.tasks.wikievents.prompts import (
    COARSE_EVENT_DEFINITIONS,
    COARSE_TO_FINE_EVENTS,
    ENTITY_DEFINITIONS,
    EVENT_DEFINITIONS,
    FINE_TO_COARSE_EVENTS,
    GPE,
    Abstract,
    Acquit,
    ArrestJailDetain,
    ArtifactExistanceEvent,
    Attack,
    BodyPart,
    ChargeIndict,
    CognitiveEvent,
    CommercialProduct,
    ConflictEvent,
    Contact,
    ContactEvent,
    ControlEvent,
    Convict,
    Crash,
    DamageDestroy,
    Defeat,
    Demonstrate,
    Die,
    DisableDefuse,
    DisasterEvent,
    DiseaseOutbreak,
    Dismantle,
    Donation,
    EndPosition,
    Evacuation,
    ExchangeBuySell,
    Facility,
    GenericCrime,
    GenericCrimeEvent,
    IdentifyCategorize,
    IllegalTransportation,
    ImpedeInterfereWith,
    Infect,
    Information,
    Injure,
    Inspection,
    Intervention,
    InvestigateCrime,
    JobTitle,
    JusticeEvent,
    LifeEvent,
    Location,
    ManufactureAssemble,
    MedicalEvent,
    MedicalHealthIssue,
    Money,
    MovementTransportEvent,
    Numeric,
    Organization,
    Person,
    PersonnelEvent,
    PreventPassage,
    ReleaseParole,
    RequestCommand,
    Research,
    Sentence,
    SideOfConflict,
    StartPosition,
    TeachingTrainingLearning,
    ThreatenCoerce,
    TransactionEvent,
    Transportation,
    TrialHearing,
    Vehicle,
    Weapon,
)

from ..utils_data import DatasetLoader, Sampler


class WikiEventsDatasetLoader(DatasetLoader):
    ENTITY_TO_CLASS_MAPPING = {
        "ABS": Abstract,
        "BOD": BodyPart,
        "COM": CommercialProduct,
        "FAC": Facility,
        "GPE": GPE,
        "INF": Information,
        "LOC": Location,
        "MHI": MedicalHealthIssue,
        "MON": Money,
        "ORG": Organization,
        "PER": Person,
        "SID": SideOfConflict,
        "TTL": JobTitle,
        "VAL": Numeric,
        "VEH": Vehicle,
        "WEA": Weapon,
    }

    _EVENT_CONSTANTS_MAPPING = {
        "trigger": "mention",
        "Place": "place",
    }

    EVENT_TO_CLASS_MAPPING = {
        "ArtifactExistence.DamageDestroyDisableDismantle.Damage": {
            "coarse": ArtifactExistanceEvent,
            "class": DamageDestroy,
            "Artifact": "artifact",
            "Instrument": "instrument",
            "Damager": "agent",
        },
        "ArtifactExistence.DamageDestroyDisableDismantle.Destroy": {
            "coarse": ArtifactExistanceEvent,
            "class": DamageDestroy,
            "Artifact": "artifact",
            "Destroyer": "agent",
            "Instrument": "instrument",
        },
        "ArtifactExistence.DamageDestroyDisableDismantle.DisableDefuse": {
            "coarse": ArtifactExistanceEvent,
            "class": DisableDefuse,
            "Artifact": "artifact",
            "Disabler": "agent",
            "Instrument": "instrument",
        },
        "ArtifactExistence.DamageDestroyDisableDismantle.Dismantle": {
            "coarse": ArtifactExistanceEvent,
            "class": Dismantle,
            "Artifact": "artifact",
            "Components": "components",
            "Dismantler": "agent",
            "Instrument": "instrument",
        },
        "ArtifactExistence.DamageDestroyDisableDismantle.Unspecified": {
            "coarse": ArtifactExistanceEvent,
            "class": DamageDestroy,
            "Artifact": "artifact",
            "Instrument": "instrument",
            "DamagerDestroyer": "agent",
        },
        "ArtifactExistence.ManufactureAssemble.Unspecified": {
            "coarse": ArtifactExistanceEvent,
            "class": ManufactureAssemble,
            "Artifact": "artifact",
            "Components": "components",
            "ManufacturerAssembler": "agent",
        },
        "Cognitive.IdentifyCategorize.Unspecified": {
            "coarse": CognitiveEvent,
            "class": IdentifyCategorize,
            "IdentifiedObject": "object",
            "IdentifiedRole": "role",
            "Identifier": "agent",
        },
        "Cognitive.Inspection.SensoryObserve": {
            "coarse": CognitiveEvent,
            "class": Inspection,
            "Instrument": "instrument",
            "ObservedEntity": "entity",
            "Observer": "agent",
        },
        "Cognitive.Research.Unspecified": {
            "coarse": CognitiveEvent,
            "class": Research,
            "Researcher": "agent",
            "Subject": "subject",
        },
        "Cognitive.TeachingTrainingLearning.Unspecified": {
            "coarse": CognitiveEvent,
            "class": TeachingTrainingLearning,
            "Learner": "patient",
            "TeacherTrainer": "agent",
        },
        "Conflict.Attack.DetonateExplode": {
            "coarse": ConflictEvent,
            "class": Attack,
            "Attacker": "agent",
            "ExplosiveDevice": "explosive_device",
            "Instrument": "instrument",
            "Target": "target",
        },
        "Conflict.Attack.Unspecified": {
            "coarse": ConflictEvent,
            "class": Attack,
            "Attacker": "agent",
            "Instrument": "instrument",
            "Target": "target",
        },
        "Conflict.Defeat.Unspecified": {
            "coarse": ConflictEvent,
            "class": Defeat,
            "Defeated": "defeated",
            "Victor": "victor",
        },
        "Conflict.Demonstrate.DemonstrateWithViolence": {
            "coarse": ConflictEvent,
            "class": Demonstrate,
            "Demonstrator": "demonstrator",
            "Regulator": "regulator",
        },
        "Conflict.Demonstrate.Unspecified": {
            "coarse": ConflictEvent,
            "class": Demonstrate,
            "Demonstrator": "demonstrator",
            "Target": "target",
            "Topic": "topic",
            "Regulator": "regulator",
        },
        "Contact.Contact.Broadcast": {
            "coarse": ContactEvent,
            "class": Contact,
            "Communicator": "participants",
            "Instrument": "instrument",
            "Recipient": "participants",
            "Topic": "topic",
        },
        "Contact.Contact.Correspondence": {
            "coarse": ContactEvent,
            "class": Contact,
            "Participant": "participants",
            "Topic": "topic",
        },
        "Contact.Contact.Meet": {
            "coarse": ContactEvent,
            "class": Contact,
            "Participant": "participants",
            "Topic": "topic",
        },
        "Contact.Contact.Unspecified": {
            "coarse": ContactEvent,
            "class": Contact,
            "Participant": "participants",
            "Topic": "topic",
        },
        "Contact.RequestCommand.Broadcast": {
            "coarse": ContactEvent,
            "class": RequestCommand,
            "Communicator": "communicator",
            "Recipient": "receptor",
            "Topic": "topic",
        },
        "Contact.RequestCommand.Correspondence": {
            "coarse": ContactEvent,
            "class": RequestCommand,
            "Communicator": "communicator",
            "Recipient": "receptor",
            "Topic": "topic",
        },
        "Contact.RequestCommand.Meet": {
            "coarse": ContactEvent,
            "class": RequestCommand,
            "Communicator": "communicator",
            "Recipient": "receptor",
            "Topic": "topic",
        },
        "Contact.RequestCommand.Unspecified": {
            "coarse": ContactEvent,
            "class": RequestCommand,
            "Communicator": "communicator",
            "Recipient": "receptor",
            "Topic": "topic",
        },
        "Contact.ThreatenCoerce.Broadcast": {
            "coarse": ContactEvent,
            "class": ThreatenCoerce,
            "Communicator": "communicator",
            "Recipient": "recipient",
        },
        "Contact.ThreatenCoerce.Correspondence": {
            "coarse": ContactEvent,
            "class": ThreatenCoerce,
            "Communicator": "communicator",
            "Recipient": "recipient",
        },
        "Contact.ThreatenCoerce.Unspecified": {
            "coarse": ContactEvent,
            "class": ThreatenCoerce,
            "Communicator": "communicator",
            "Recipient": "recipient",
        },
        "Control.ImpedeInterfereWith.Unspecified": {
            "coarse": ControlEvent,
            "class": ImpedeInterfereWith,
            "Impeder": "impeder",
        },
        "Disaster.Crash.Unspecified": {
            "coarse": DisasterEvent,
            "class": Crash,
            "CrashObject": "object",
            "Vehicle": "vehicle",
        },
        "Disaster.DiseaseOutbreak.Unspecified": {
            "coarse": DisasterEvent,
            "class": DiseaseOutbreak,
            "Disease": "disease",
            "Victim": "victim",
        },
        "GenericCrime.GenericCrime.GenericCrime": {
            "coarse": GenericCrimeEvent,
            "class": GenericCrime,
            "Perpetrator": "perpetrator",
            "Victim": "victim",
        },
        "Justice.Acquit.Unspecified": {"coarse": JusticeEvent, "class": Acquit, "Defendant": "defendant"},
        "Justice.ArrestJailDetain.Unspecified": {
            "coarse": JusticeEvent,
            "class": ArrestJailDetain,
            "Detainee": "detainee",
            "Jailer": "jailer",
        },
        "Justice.ChargeIndict.Unspecified": {
            "coarse": JusticeEvent,
            "class": ChargeIndict,
            "Defendant": "defendant",
            "JudgeCourt": "judge_court",
            "Prosecutor": "prosecutor",
        },
        "Justice.Convict.Unspecified": {
            "coarse": JusticeEvent,
            "class": Convict,
            "Defendant": "defendant",
            "JudgeCourt": "judge_court",
        },
        "Justice.InvestigateCrime.Unspecified": {
            "coarse": JusticeEvent,
            "class": InvestigateCrime,
            "Defendant": "defendant",
            "Investigator": "investigator",
            "ObservedEntity": "crime",
            "Observer": "investigator",
        },
        "Justice.ReleaseParole.Unspecified": {
            "coarse": JusticeEvent,
            "class": ReleaseParole,
            "Defendant": "defendant",
            "JudgeCourt": "judge_court",
        },
        "Justice.Sentence.Unspecified": {
            "coarse": JusticeEvent,
            "class": Sentence,
            "Defendant": "defendant",
            "JudgeCourt": "judge_court",
        },
        "Justice.TrialHearing.Unspecified": {
            "coarse": JusticeEvent,
            "class": TrialHearing,
            "Defendant": "defendant",
            "JudgeCourt": "judge_court",
            "Prosecutor": "prosecutor",
        },
        "Life.Die.Unspecified": {
            "coarse": LifeEvent,
            "class": Die,
            "Killer": "killer",
            "Victim": "victim",
        },
        "Life.Infect.Unspecified": {"coarse": LifeEvent, "class": Infect, "Victim": "victim"},
        "Life.Injure.Unspecified": {
            "coarse": LifeEvent,
            "class": Injure,
            "BodyPart": "body_part",
            "Injurer": "injurer",
            "Instrument": "instrument",
            "Victim": "victim",
        },
        "Medical.Intervention.Unspecified": {
            "coarse": MedicalEvent,
            "class": Intervention,
            "Patient": "patient",
            "Treater": "treater",
        },
        "Movement.Transportation.Evacuation": {
            "coarse": MovementTransportEvent,
            "class": Evacuation,
            "Destination": "destination",
            "Origin": "origin",
            "PassengerArtifact": "passenger_artifact",
            "Transporter": "transporter",
        },
        "Movement.Transportation.IllegalTransportation": {
            "coarse": MovementTransportEvent,
            "class": IllegalTransportation,
            "Destination": "destination",
            "PassengerArtifact": "passenger_artifact",
            "Transporter": "transporter",
            "Vehicle": "vehicle",
        },
        "Movement.Transportation.PreventPassage": {
            "coarse": MovementTransportEvent,
            "class": PreventPassage,
            "Destination": "destination",
            "Origin": "origin",
            "PassengerArtifact": "passenger_artifact",
            "Preventer": "preventer",
            "Transporter": "transporter",
            "Vehicle": "vehicle",
        },
        "Movement.Transportation.Unspecified": {
            "coarse": MovementTransportEvent,
            "class": Transportation,
            "Destination": "destination",
            "Origin": "origin",
            "PassengerArtifact": "passenger_artifact",
            "Transporter": "transporter",
            "Vehicle": "vehicle",
        },
        "Personnel.EndPosition.Unspecified": {
            "coarse": PersonnelEvent,
            "class": EndPosition,
            "Employee": "employee",
            "PlaceOfEmployment": "organization",
        },
        "Personnel.StartPosition.Unspecified": {
            "coarse": PersonnelEvent,
            "class": StartPosition,
            "Employee": "employee",
            "PlaceOfEmployment": "organization",
            "Position": "position",
        },
        "Transaction.Donation.Unspecified": {
            "coarse": TransactionEvent,
            "class": Donation,
            "ArtifactMoney": "artifact_money",
            "Giver": "giver",
            "Recipient": "recipient",
        },
        "Transaction.ExchangeBuySell.Unspecified": {
            "coarse": TransactionEvent,
            "class": ExchangeBuySell,
            "AcquiredEntity": "adquired_entity",
            "Giver": "giver",
            "PaymentBarter": "payment_barter",
            "Recipient": "recipient",
        },
    }

    def __init__(self, path: str, group_by: str = "sentence", **kwargs) -> None:
        assert group_by in [
            "sentence",
            "document",
        ], "`group_by` must be either 'sentence' or 'document'."

        self.elements = {}

        with open(path, "rt", encoding="utf-8") as in_f:
            for line in in_f:
                line = json.loads(line.strip())

                sent_id = line["doc_id"]
                doc_id = "_".join(sent_id.split("_")[:-1])
                key = doc_id if group_by == "document" else sent_id
                if key not in self.elements:
                    self.elements[key] = {
                        "id": key,
                        "doc_id": doc_id,
                        "text": "",
                        "entities": [],
                        "events": [],
                        "arguments": [],
                        "gold": [],
                    }

                entities = [
                    self.ENTITY_TO_CLASS_MAPPING[entity["entity_type"]](span=entity["text"])
                    for entity in line["entity_mentions"]
                    if entity["entity_type"] in self.ENTITY_TO_CLASS_MAPPING
                ]

                events, arguments = [], []
                for event in line["event_mentions"]:
                    if event["event_type"] not in self.EVENT_TO_CLASS_MAPPING:
                        continue
                    info = self.EVENT_TO_CLASS_MAPPING[event["event_type"]]
                    _inst = {param: [] for param in inspect.signature(info["class"]).parameters.keys()}
                    _inst["mention"] = event["trigger"]["text"]
                    for argument in event["arguments"]:
                        if "OOR" in argument["role"]:
                            continue
                        elif argument["role"] in info:
                            name = info[argument["role"]]
                            _inst[name].append(argument["text"])
                        elif argument["role"] in self._EVENT_CONSTANTS_MAPPING:
                            name = self._EVENT_CONSTANTS_MAPPING[argument["role"]]
                            if name not in _inst:
                                continue
                            _inst[name].append(argument["text"])
                        else:
                            raise ValueError(f"Argument {event['event_type']}:{argument['role']} not found!")

                    events.append(info["coarse"](mention=_inst["mention"]))
                    arguments.append(info["class"](**_inst))

                self.elements[key]["text"] += " " + line["text"].strip()
                self.elements[key]["entities"] += entities
                self.elements[key]["events"] += events
                self.elements[key]["arguments"] += arguments
                self.elements[key]["gold"] += entities  # Is not used anyway


class WikiEventsSampler(Sampler):
    """
    A data `Sampler` for the WikiEvents dataset.

    Args:
        dataset_loader (`WikiEventsDatasetLoader`):
            The dataset loader that contains the data information.
        task (`str`, optional):
            The task to sample. It must be one of the following: NER, EE, EAE.
            Defaults to `None`.
        split (`str`, optional):
            The split to sample. It must be one of the following: "train", "dev" or
            "test". Depending on the split the sampling strategy differs. Defaults to
            `"train"`.
        parallel_instances (`Union[int, Tuple[int, int]]`, optional):
            The number of sentences sampled in parallel. Options:

                * **`int`**: The amount of elements that will be sampled in parallel.
                * **`tuple`**: The range of elements that will be sampled in parallel.

            Defaults to 1.
        max_guidelines (`int`, optional):
            The number of guidelines to append to the example at the same time. If `-1`
            is given then all the guidelines are appended. Defaults to `-1`.
        guideline_dropout (`float`, optional):
            The probability to dropout a guideline definition for the given example. This
            is only applied on training. Defaults to `0.0`.
        seed (`float`, optional):
            The seed to sample the examples. Defaults to `0`.
        prompt_template (`str`, optional):
            The path to the prompt template. Defaults to `"templates/prompt.txt"`.
        ensure_positives_on_train (bool, optional):
            Whether to ensure that the guidelines of annotated examples are not removed.
            Defaults to `False`.
        dataset_name (str, optional):
            The name of the dataset. Defaults to `None`.
        scorer (`str`, optional):
           The scorer class import string. Defaults to `None`.
        sample_only_gold_guidelines (`bool`, optional):
            Whether to sample only guidelines of present annotations. Defaults to `False`.
    """

    def __init__(
        self,
        dataset_loader: WikiEventsDatasetLoader,
        task: str = None,
        split: str = "train",
        parallel_instances: Union[int, Tuple[int, int]] = 1,
        max_guidelines: int = -1,
        guideline_dropout: float = 0.0,
        seed: float = 0,
        ensure_positives_on_train: bool = False,
        dataset_name: str = None,
        scorer: str = None,
        sample_only_gold_guidelines: bool = False,
        **kwargs,
    ) -> None:
        assert task in [
            "NER",
            "EE",
            "EAE",
        ], f"{task} must be either 'NER', 'EE', 'EAE'."

        task_definitions, task_target, task_template = {
            "NER": (ENTITY_DEFINITIONS, "entities", "templates/prompt.txt"),
            "EE": (COARSE_EVENT_DEFINITIONS, "events", "templates/prompt.txt"),
            "EAE": (EVENT_DEFINITIONS, "arguments", "templates/prompt_ace_eae.txt"),
        }[task]

        if task in ["EAE"]:
            is_coarse_to_fine: bool = True
            COARSE_TO_FINE = COARSE_TO_FINE_EVENTS
            FINE_TO_COARSE = FINE_TO_COARSE_EVENTS
        else:
            is_coarse_to_fine = False
            COARSE_TO_FINE = None
            FINE_TO_COARSE = None

        kwargs.pop("prompt_template")

        super().__init__(
            dataset_loader=dataset_loader,
            task=task,
            split=split,
            parallel_instances=parallel_instances,
            max_guidelines=max_guidelines,
            guideline_dropout=guideline_dropout,
            seed=seed,
            prompt_template=task_template,
            ensure_positives_on_train=ensure_positives_on_train,
            sample_only_gold_guidelines=sample_only_gold_guidelines,
            dataset_name=dataset_name,
            scorer=scorer,
            task_definitions=task_definitions,
            task_target=task_target,
            is_coarse_to_fine=is_coarse_to_fine,
            coarse_to_fine=COARSE_TO_FINE,
            fine_to_coarse=FINE_TO_COARSE,
            definitions=GUIDELINES,
            examples=EXAMPLES,
            **kwargs,
        )

src/tasks/wikievents/data_loader.py
hitz-zentroa/GoLLIE
GUIDELINES = {
    "wikievents_ner_abstract": {
        "en": [
            (
                "Abstract, non-tangible artifacts such as software (e.g., programs, tool kits, apps, e-mail),"
                " measureable \nintellectual property, contracts, etc. (nb: does not include laws, which are LAW type)"
            ),
            (
                'Paraphrased text: "Intangible, abstract items like software (for example, programs, toolkits, apps,'
                " email), \nquantifiable intellectual property, contracts, and so on (excluding laws, which are a"
                " separate category)."
            ),
            (
                'Paraphrased text: "Intangible, abstract items like software (for example, programs, toolkits, apps,'
                " email), \nquantifiable intellectual property, contracts, and so on (excluding laws, which are part"
                " of the LEGAL type)."
            ),
            (
                'Paraphrased text: "Intangible, abstract items like software (for example, programs, toolkits, apps,'
                " email), \nintellectual property that can be quantified, agreements, and so on (exclusion: laws,"
                ' which are part of the LEGAL type)".'
            ),
            (
                "Abstract, intangible items like software (e.g., programs, toolkits, apps, email) and measurable"
                " intellectual \nproperty, such as contracts, etc. (excluding laws, which are a separate category)"
            ),
        ]
    },
    "wikievents_ner_bodypart": {
        "en": [
            "An identifiable, living part of a human's or animal's body, such as a eye, ear, neck, leg, etc.",
            "A recognizable, living component of a human or animal's body, like an eye, ear, neck, leg, and so on.",
            (
                "A distinguishable, living component of a human or animal body, such as an eye, ear, neck, leg, and"
                " so on."
            ),
            "A distinguishable, living component of a human or animal body, like an eye, ear, neck, leg, and so on.",
            "A distinct, living component of a human or animal's body, like an eye, ear, neck, or leg, for example.",
        ]
    },
    "wikievents_ner_commercialproduct": {
        "en": [
            (
                "A tangible product or article of trade for which someone pays or barters, or more generally, an"
                " artifact or a thing."
            ),
            (
                "A concrete item that is purchased or exchanged in a trade, or more broadly, a material object or"
                " article."
            ),
            "A concrete item that is bought or exchanged by someone, or more broadly, a material object or article.",
            (
                "A concrete item that is purchased or exchanged in a trade, or more broadly, a material object or"
                " article."
            ),
            (
                "A concrete item that is exchanged for money or another item of value, or more broadly, a material"
                " object or article."
            ),
        ]
    },
    "wikievents_ner_facility": {
        "en": [
            (
                "A functional, primarily man-made structure. Facilities are artifacts falling under the domains of"
                " architecture \nand civil engineering, including more temporary human constructs, such as police"
                " lines and checkpoints."
            ),
            (
                "A functional, predominantly human-created construction. Facilities encompass items classified under"
                " the \nfields of architecture and civil engineering, encompassing even temporary human creations,"
                " like police barriers and \ncheckpoints."
            ),
            (
                "A functional, predominantly human-created construction. Facilities are examples of architecture and"
                " civil \nengineering, encompassing even temporary human-made structures like police barricades and"
                " checkpoints."
            ),
            (
                "A functional, predominantly human-created construction. Facilities encompass items classified under"
                " the \nfields of architecture and civil engineering, as well as transient human-built structures,"
                " like police barriers and \ncheckpoints."
            ),
            (
                "An operational, predominantly human-created object. Facilities encompass items within the scope of"
                " \narchitecture and civil engineering, as well as transitory human creations, like police barriers"
                " and inspection stations."
            ),
        ]
    },
    "wikievents_ner_gpe": {
        "en": [
            (
                "Geopolitical entities such as countries, provinces, states, cities, towns, etc. GPEs are composite"
                " entities, \nconsisting of a physical location, a government, and a population. All three of these"
                " elements must be present for an entity to be \ntagged as a GPE. A GPE entity may be a single"
                " geopolitical entity or a group."
            ),
            (
                "Geopolitical entities, including nations, regions, cities, and towns, are complex entities that"
                " combine a \nspecific location, a governing body, and a population. For an entity to be categorized"
                " as a geopolitical entity, all three of \nthese components must be present. A geopolitical entity can"
                " be either a single, individual entity or a group."
            ),
            (
                "Geopolitical entities, including nations, provinces, states, cities, and towns, are composite"
                " structures \ncomprising a physical location, a governing body, and an inhabited population. For an"
                " entity to be classified as a geopolitical \nentity, all three of these components must be present. A"
                " geopolitical entity can exist as a single distinct entity or a group \nof related entities."
            ),
            (
                "Geopolitical entities, like nations, regions, states, cities, and towns, are composite structures,"
                " comprising a \nspecific location, a governing body, and an inhabiting population. In order to be"
                " categorized as a geopolitical entity, all \nthree of these components must be present. A"
                " geopolitical entity can exist as a single unit or as a group."
            ),
            (
                "Geopolitical entities, including nations, provinces, states, cities, and towns, are multifaceted"
                " entities that \ncomprise a physical location, a governing body, and an inhabiting population. For an"
                " entity to be classified as a \ngeopolitical entity, all three of these components must be present. A"
                " geopolitical entity can exist as a single unit or as a group."
            ),
        ]
    },
    "wikievents_ner_information": {
        "en": [
            (
                "An information object such as a field of study or topic of communication, including thoughts,"
                " opinions, etc."
            ),
            (
                "A subject matter, be it an academic discipline or a subject of discussion, encompassing ideas,"
                " perspectives, and so \non."
            ),
            (
                "A subject matter, encompassing academic disciplines or discussion content, encompassing thoughts and"
                " opinions \nas well."
            ),
            (
                "A subject matter, be it an academic discipline or a theme in communication, encompassing ideas,"
                " beliefs, and \nperspectives."
            ),
            (
                "A subject matter, be it an area of academic focus or a subject of communication encompassing thoughts"
                " and opinions."
            ),
        ]
    },
    "wikievents_ner_location": {
        "en": [
            "Geographical entities such as geographical areas and landmasses, bodies of water.",
            "Geographical features, including regions, landmasses, and bodies of water.",
            "Geographical features, including regions and landmasses, as well as water bodies.",
            "Geographical features, including regions, landmasses, and bodies of water.",
            "Geographical features, including regions, landmasses, and bodies of water.",
        ]
    },
    "wikievents_ner_medicalhealthissue": {
        "en": [
            (
                "Any medical condition or health issue, to include everything from disease to broken bones to fever to"
                " general ill \nhealth, medical errors, even natural causes"
            ),
            (
                "Any and all medical circumstances, encompassing everything from illnesses to fractures, fevers, and"
                " general \nunwellness, along with medical blunders, and even natural causes"
            ),
            (
                "Any and all medical circumstances or health problems, encompassing everything from illnesses to"
                " fractures, \nfevers, and general unwellness, medical blunders, even natural causes"
            ),
            (
                "Any and all medical situations or health problems, encompassing everything from illnesses to broken"
                " bones, \nfevers, and general unwellness, as well as medical blunders, and even natural causes"
            ),
            (
                "Any and all medical circumstances, encompassing everything from illnesses to fractures, fever, and"
                " overall poor \nhealth, as well as medical blunders, and even natural causes"
            ),
        ]
    },
    "wikievents_ner_money": {
        "en": [
            (
                "A monetary payment. The extent of a Money mention includes modifying quantifiers, the amount, and the"
                " currency \nunit, all of which can be optional."
            ),
            (
                "A financial compensation. The scope of a Money reference encompasses adjustable quantifiers, the"
                " payment amount, \nand the currency type, all of which can be discretionary."
            ),
            (
                "A financial compensation. The scope of a Money specification encompasses the adjustment of"
                " quantifiers, the \npayment amount, and the currency type, all of which can be optional elements."
            ),
            (
                "A financial compensation. The scope of a Money reference encompasses modifying determiners, the sum,"
                " and the \ncurrency denomination, all of which can be discretionary."
            ),
            (
                "A financial compensation. The scope of a Money reference encompasses modifying qualifiers, the sum,"
                " and the \ncurrency denomination, all of which can be optional."
            ),
        ]
    },
    "wikievents_ner_organization": {
        "en": [
            (
                "Corporations, agencies, and other groups of people defined by an established organizational"
                " structure. An ORG \nentity may be a single organization or a group. A key feature of an ORG is that"
                " it can change members without changing \nidentity."
            ),
            (
                "Organizations, such as corporations and government agencies, are classified as ORG entities, which"
                " can encompass \na wide range of groups characterized by a specific organizational framework. A"
                " notable attribute of an ORG entity is \nits ability to modify its membership without compromising"
                " its core identity."
            ),
            (
                "Organizations, bureaus, and other collectives characterized by a defined hierarchical arrangement. An"
                " ORG \nentity can encompass a single organization or a group. A crucial aspect of an ORG is its"
                " ability to modify its members without \naltering its core identity."
            ),
            (
                "Organizations, such as corporations and government agencies, are classified as ORG entities. These"
                " entities may \nconsist of a single organization or a group of people, and a defining characteristic"
                " is their ability to modify their \nmembership without compromising their core identity."
            ),
            (
                "Organizations, such as businesses and government agencies, are composed of individuals united under a"
                " specific \nstructure. An ORG entity can represent a single organization or a group of people. A"
                " defining characteristic of an ORG is its \nability to modify its membership without compromising its"
                " core identity."
            ),
        ]
    },
    "wikievents_ner_person": {
        "en": [
            "Person entities are limited to humans. A Person entity may be a single person or a group.",
            (
                'Paraphrase: "Person entities are exclusively reserved for humans. A Person entity can either'
                " represent an \nindividual or a group."
            ),
            (
                "Individual entities are exclusively reserved for humans. A Person entity can represent a single"
                " person or a group of \npeople."
            ),
            (
                "Paraphrase: Person entities are exclusively reserved for human individuals; a Person entity can"
                " either represent \na single person or a group of people."
            ),
            (
                'Paraphrase: "Person entities exclusively encompass human individuals; a Person entity can either'
                " represent a \nsingle person or a group."
            ),
        ]
    },
    "wikievents_ner_side_of_conflict": {
        "en": [
            (
                "The different sides of a conflict, such as philosophical, cultural, ideological, religious,"
                " political, guiding \nphilosophical movement or group orientation. This will encompass sides of the"
                " battle/conflict, sports fans when salient, and \nother such affiliations, in addition to religions,"
                " political parties, and other philosophies."
            ),
            (
                "Various aspects of a disagreement, including philosophical, cultural, ideological, religious,"
                " political, and \nguiding principles of movements or groups, are covered under this definition. It"
                " also encompasses opposing sides in a \nbattle or conflict, enthusiastic supporters of sports teams"
                " when their affiliation is significant, and any other similar \nallegiances, in addition to various"
                " faiths, political parties, and other belief systems."
            ),
            (
                "Various aspects of a disagreement, including philosophical, cultural, ideological, religious,"
                " political, and \nfactional perspectives, are covered under this definition. It encompasses opposing"
                " sides in battles or conflicts, as well as \nloyalty-driven associations like sports fanaticism when"
                " relevant, and other similar allegiances. Moreover, this encompasses \nvarious belief systems like"
                " religions, political parties, and other doctrines."
            ),
            (
                "Various aspects of a disagreement, including philosophical, cultural, ideological, religious,"
                " political, and \nprominent guiding principles or group inclinations, cover the opposing sides of a"
                " conflict, sports fanaticism when \nsignificant, and additional affiliations, as well as religious"
                " beliefs, political parties, and other philosophies."
            ),
            (
                "Various aspects of a disagreement, including philosophical, cultural, ideological, religious,"
                " political, and \nguiding belief systems or group inclinations, are covered under this definition. It"
                " encompasses opposing sides in a \nconflict, sports fan allegiances when significant, and other"
                " relevant affiliations, as well as different religious sects, \npolitical parties, and other"
                " doctrines."
            ),
        ]
    },
    "wikievents_ner_job_title": {
        "en": [
            "A person's title or job role.",
            "The specific position or occupation held by an individual.",
            "The position or occupation of an individual.",
            "A person's position or occupation.",
            "The specific position or occupation held by an individual.",
        ]
    },
    "wikievents_ner_numeric": {
        "en": [
            "A numerical value or non-numerical value such as an informational property such as color or make or URL.",
            (
                "A quantitative or non-quantitative attribute, including characteristics like color, type, or website"
                " address."
            ),
            "A quantitative or non-quantitative value, including properties like color, make, or web address.",
            "A quantitative or non-quantitative value, including properties like color, make, or website address.",
            (
                "A quantitative or non-quantitative value, including characteristics like color, make, or website"
                " address."
            ),
        ]
    },
    "wikievents_ner_vehicle": {
        "en": [
            (
                "A physical device primarily designed to move an object from one location to another, by (for example)"
                " carrying, \nflying, pulling, or pushing the transported object. Vehicle entities may or may not have"
                " their own power source."
            ),
            (
                "A tangible apparatus primarily intended to relocate an object from one place to another, such as by"
                " carrying, \nflying, pulling, or pushing the transported item. Vehicle entities may or may not"
                " possess their own energy source."
            ),
            (
                "A tangible apparatus principally intended to transfer an object from one place to another, such as by"
                " carrying, \nflying, pulling, or pushing the conveyed item. Vehicle entities can either possess or"
                " lack their own power source."
            ),
            (
                "A tangible apparatus principally intended to relocate an object from one place to another, such as by"
                " carrying, \nflying, pulling, or pushing the transported item. Vehicle entities may or may not"
                " possess their own power source."
            ),
            (
                "A tangible apparatus primarily intended to relocate an object from one place to another, such as by"
                " carrying, \nflying, pulling, or pushing the transported item. Vehicle entities can either possess or"
                " lack their own power source."
            ),
        ]
    },
    "wikievents_ner_weapon": {
        "en": [
            "A physical device that is primarily used as an instrument for physically harming or destroying entities.",
            (
                "A tangible apparatus primarily designed for the purpose of inflicting bodily harm or destruction on"
                " entities."
            ),
            (
                "A tangible object primarily designed for the purpose of inflicting physical injury or destruction"
                " upon beings."
            ),
            "A tangible appliance predominantly designed to inflict bodily injury or annihilation upon beings.",
            "A tangible object primarily designed to inflict or cause destruction upon living beings or objects.",
        ]
    },
    "wikievents_ee_artifact_existance": {
        "en": [
            (
                "Any event related to manufacturing, assembling, damagging, destroying, disabling, defusing, or"
                " dismantling an \nartifact or a thing."
            ),
            (
                "Any occurrence associated with the production, assembly, damage, destruction, disabling, disarming,"
                " or taking \napart of an artifact or an object."
            ),
            (
                "Any occurrence associated with the production, assembly, damage, destruction, disabling,"
                " deactivation, or \ndisassembly of an item or object."
            ),
            (
                "Any occurrence associated with the production, assembly, damage, destruction, disabling,"
                " deactivation, or \ndisassembly of an object or item."
            ),
            (
                "Any occurrence associated with the production, assembly, damage, destruction, disabling,"
                " deactivation, or \ndisassembly of an object or item."
            ),
        ]
    },
    "wikievents_ee_cognitive": {
        "en": [
            (
                "The act of establishing the identity of an entity or event, or establishing the relevant category"
                " that an entity or \nevent belongs to, such as being a suspect or a target or an attack. An"
                " observation or inspection event, with any target of \ninspection. Explicit mention of researching"
                " the answer to a question by analyzing literature or testing hypotheses through \nexperiments or"
                " reviewing opinions from experts in a field. The action of teaching a person or animal a particular"
                " skill or type of \nbehavior, where the learner does not have complete mastery of the skill until the"
                " end."
            ),
            (
                "The process of determining the origin or cause of an occurrence, categorizing it, and identifying its"
                " nature, such \nas distinguishing between a suspect, a target, or an attack. A systematic examination"
                " or assessment that involves \nstudying or evaluating any subject or object. The practice of"
                " investigating the solution to a question by studying various \nsources, conducting experiments, or"
                " seeking advice from experts in a given field, ultimately leading to the acquisition of a \nskill or"
                " behavior by a person or animal, even if mastery is not fully achieved."
            ),
            (
                "The process of determining the origin or nature of a person, occurrence, or object, as well as the"
                " classification it \nfalls under, such as being a person of interest, a goal, or the focus of an"
                " attack. A situation involving the examination or \nassessment of any subject or object. The act of"
                " investigating to reach a conclusion by examining various sources, such as \nliterature, experiments,"
                " or expert opinions within a specific domain. The practice of instructing an individual or animal in"
                " a \nparticular ability or behavior, where mastery is not fully achieved until the training is"
                " complete."
            ),
            (
                "The process of determining the origin or nature of a person, object, or event, and the classification"
                " of said entity \ninto a specific category, such as a suspect, target, or attack. A systematic"
                " examination or assessment that may involve \nany subject under scrutiny. The practice of"
                " investigating responses to inquiries by examining literary sources, \nconducting experimental"
                " investigations, or consulting expert perspectives within a given domain. The act of imparting"
                " \nknowledge or skill to an individual, whether human or animal, with the understanding that full"
                " mastery may not be achieved \nuntil the learning process is complete."
            ),
            (
                "The process of determining the origin or cause of a circumstance, categorizing it into appropriate"
                " groups such as \nsuspect, target, or assailant, and engaging in investigative actions or"
                " examinations with any relevant subject. \nAdditionally, the practice of investigating answers to"
                " inquiries by examining written works, testing hypotheses through \nexperimentation, or consulting"
                " expert opinions within a given field. Furthermore, the act of imparting knowledge or skills to an"
                " \nindividual or animal, where the learner's mastery of the skill is not yet complete by the end of"
                " the process."
            ),
        ]
    },
    "wikievents_ee_conflict": {
        "en": [
            (
                "An attack event, a violent physical act causing harm or damage. Defeat in a conflict or an election,"
                " but not a \ngame-style competition. A demonstration, march, protest, or political gathering."
            ),
            (
                "An assault occasion, a forceful physical act leading to injury or destruction. A loss in a contest or"
                " an election, \nthough not in a recreational competition. A display, march, protest, or political"
                " assembly."
            ),
            (
                "An act of aggression, a forceful physical action resulting in injury or destruction. A loss in a"
                " struggle or an \nelection, though not in a recreational contest. A public display, procession,"
                " protest, or political assembly."
            ),
            (
                "An assault occasion, a forceful physical act leading to injury or destruction. A loss in a"
                " disagreement or a voting \nprocess, though not in a recreational contest. A display, march, protest,"
                " or political assembly."
            ),
            (
                "An act of aggression, a forceful physical act that inflicts harm or damage. A loss in a struggle or"
                " an election, but not \nin a recreational contest. A public display of opposition, a march, a"
                " protest, or a political assembly."
            ),
        ]
    },
    "wikievents_ee_contact": {
        "en": [
            (
                "An event where two or more participants communicate over any medium, where the context does not"
                " distinguish the \nnature of the contact or whether the communication is in person or remote. Any"
                " communication over any medium where the \ncommunicator intentionally deceives the addressee or"
                " continue an ongoing deception, either by knowingly conveying false \ninformation, or conveying"
                " information with the intent of the addressee drawing the wrong conclusion, or omitting information."
                " A \nstatement making a request, appeal, command, or order, where the context does not distinguish"
                " whether the communication is \none-way or two-way, or whether it is in person or remote. A statement"
                " offering threat or coercion, where the context does not \ndistinguish whether the communication is"
                " one-way or two-way, or whether it is in person or remote."
            ),
            (
                "An occasion during which two or more individuals engage in communication through any means, without"
                " regard to \nwhether the interaction is face-to-face or remote, or whether the context distinguishes"
                " the type of contact. Any \ninteraction involving deceptive communication, intentionally misleading"
                " the recipient, either by knowingly providing \nfalse information, leading the recipient to the wrong"
                " conclusion, or withholding information. A statement that \nexpresses a request, appeal, command, or"
                " order, where the context does not differentiate between one-way or two-way \ncommunication, or"
                " whether the interaction is in person or remote. A statement that delivers a threat or coercion,"
                " without \ndifferentiating between one-way or two-way communication, or whether the interaction is in"
                " person or remote."
            ),
            (
                "An interaction involving two or more individuals utilizing any means of communication, where the"
                " circumstances do \nnot differentiate between face-to-face or remote communication or the type of"
                " medium used. Any communication process \nin which the sender intentionally misleads the recipient,"
                " either by deliberately providing false information, \nconveying information with the intention of"
                " leading the recipient astray, or withholding relevant information. A statement \nthat expresses a"
                " request, appeal, command, or order, where the context does not distinguish between one-way or"
                " two-way \ncommunication, or whether it occurs in person or at a distance. A statement that delivers"
                " a threat or coercion, where the context does \nnot differentiate between one-way or two-way"
                " communication, or whether it takes place in person or remotely."
            ),
            (
                "An interaction in which multiple participants engage in communication through any means, without"
                " regard to \nwhether the contact is in-person or remote. Any exchange of information through any"
                " medium that involves the intentional \ndeception of the recipient, either by knowingly providing"
                " false information, leading the recipient to the wrong conclusion, \nor withholding information. A"
                " statement that expresses a request, appeal, command, or order, without \ndifferentiating between"
                " one-way or two-way communication, or whether it occurs in-person or remotely. A statement that"
                " issues a \nthreat or employs coercion, without distinguishing between one-way or two-way"
                " communication, or whether it takes place \nin-person or remotely."
            ),
            (
                "An interaction in which two or more individuals engage in communication through any means, where the"
                " context fails \nto differentiate between face-to-face and remote communication or the nature of the"
                " contact. Any communication \nprocess, utilizing any medium, in which the communicator intentionally"
                " deceives the recipient, either by purposefully \nproviding false information, leading the recipient"
                " to the wrong conclusion, or withholding relevant information. A \nstatement that makes a request,"
                " appeal, command, or order, where the context does not specify whether the communication is"
                " \nunidirectional or bidirectional, or whether it occurs in person or at a distance. A statement that"
                " issues a threat or coercion, where \nthe context does not determine whether the communication is"
                " unidirectional or bidirectional, or whether it takes \nplace in person or remotely."
            ),
        ]
    },
    "wikievents_ee_control": {
        "en": [
            (
                "Explicit mention of an entity knowingly placing obstacles, not necessarily physical, that raise"
                " difficulties \nwith respect to the occurrence of a taggable, non-Movement event or process"
            ),
            (
                "Clear reference to a conscious effort by a party to intentionally introduce challenges that hinder"
                " the occurrence \nof a detectable, non-movement-related event or process"
            ),
            (
                "Clear reference to a conscious effort by a party to intentionally introduce challenges, which may not"
                " be strictly \nphysical in nature, that hinder the occurrence of a detectable, non-movement-related"
                " event or process."
            ),
            (
                "Clear reference to a specific entity intentionally creating challenges, which may not be purely"
                " physical in \nnature, that hinder the occurrence of a distinguishable, non-movement-related event or"
                " process."
            ),
            (
                "Clear reference to a specific entity intentionally creating impediments, which may not be purely"
                " physical in \nnature, that make it more challenging to facilitate a detectable, non-movement-related"
                " occurrence or process."
            ),
        ]
    },
    "wikievents_ee_disaster": {
        "en": [
            (
                "Vehicular collision or the crashing of any type of vehicle. Explicit mentions of an outbreak of a"
                " disease in an area, \nregion, or country. A damaging fire or explosion, either natural or caused but"
                " without the intent of attacking"
            ),
            (
                "Automobile collisions or the crashing of any kind of vehicle. Clear references to an epidemic of a"
                " disease breaking \nout in a specific location, region, or nation. A harmful fire or explosion,"
                " either occurring naturally or caused \nunintentionally, without any intent of attacking"
            ),
            (
                "Automotive collision, which involves the crashing of any kind of vehicle. Clear references to an"
                " epidemic of a \ndisease breaking out in a specific location, such as a region or country. A"
                " destructive fire or explosion, either resulting \nfrom natural causes or caused by accident without"
                " any intention of attacking"
            ),
            (
                "Automobile collision, which refers to the impact or crash of any kind of vehicle. Clear indications"
                " of an epidemic \nbreakout of an illness within a specific location, region, or nation. A harmful"
                " fire or blast, either resulting from natural \ncauses or human error, but not intended as an act of"
                " aggression."
            ),
            (
                "Car accidents or the crashing of any kind of vehicle. Clear references to an epidemic breaking out in"
                " a specific \nlocation, region, or nation. Destructive fires or explosions, either naturally"
                " occurring or caused unintentionally, \nwithout the intention of launching an attack"
            ),
        ]
    },
    "wikievents_ee_generic_crime": {
        "en": [
            (
                "For use as the crime argument in a Justice event when the crime event is not taggable as any other"
                " annotation ontology \nevent type."
            ),
            (
                "This statement serves as the foundation for the crime aspect in a Justice-centered event, when the"
                " crime event \nitself cannot be categorized under any other classification within the annotation"
                " ontology."
            ),
            (
                "This statement serves as the basis for the crime argument in a Justice event when the crime event"
                " cannot be \ncategorized under any other annotation ontology event type."
            ),
            (
                "This statement serves as the premise for a Justice event when the crime event cannot be categorized"
                " under any other \nannotation ontology event type."
            ),
            (
                "This statement serves as the basis for the crime argument in a Justice event when the crime event"
                " cannot be \ncategorized under any other type of annotation within an ontology."
            ),
        ]
    },
    "wikievents_ee_justice": {
        "en": [
            (
                "Any legal proceedings, such as a court case, a trial, or a hearing. Arresting, detaining, punishing,"
                " or otherwise \ndealing with a criminal or non-criminal person or entity."
            ),
            (
                "Any legal actions, including court cases, trials, or hearings, that involve apprehending,"
                " imprisoning, \npenalizing, or otherwise handling a criminal or non-criminal individual or"
                " organization."
            ),
            (
                "Any legal actions, including court cases, trials, or hearings, that involve apprehending, confining,"
                " \npenalizing, or otherwise managing a criminal or non-criminal individual or organization."
            ),
            (
                "Any legal actions, including court cases, trials, or hearings, that involve apprehending, confining,"
                " \npenalizing, or otherwise handling a criminal or non-criminal individual or entity."
            ),
            (
                "Any legal actions, including court cases, trials, and hearings, that involve apprehending,"
                " imprisoning, \npenalizing, or otherwise handling a criminal or non-criminal individual or entity."
            ),
        ]
    },
    "wikievents_ee_life": {
        "en": [
            (
                "An person or animal takes a substance into their body. The death of a person or animal. A person or"
                " animal experiencing \nphysical harm due to sickness or illness. A person or animal is infected with"
                " a pathogen. The physical injuring of a person or \nanimal."
            ),
            (
                "A person or animal ingests a substance, resulting in the demise of a person or animal, or physical"
                " harm caused by \nillness or disease. Additionally, a person or animal may be infected with a"
                " pathogen, ultimately leading to physical \ninjury."
            ),
            (
                "A person or animal ingests a substance, which can result in the death of an individual or animal, or"
                " cause them to \nexperience physical harm due to illness or sickness. Additionally, a person or"
                " animal may become infected with a pathogen, \nleading to physical injury."
            ),
            (
                "When an individual, whether human or animal, introduces a substance into their system, it can lead to"
                " unintended \nconsequences such as death, physical harm caused by illness or injury, or even pathogen"
                " infection, which can result in injury to \nthe person or animal."
            ),
            (
                "A person or animal ingests a substance, resulting in the demise of an individual or animal. A person"
                " or animal suffers \nfrom bodily harm due to illness or disease. A person or animal becomes afflicted"
                " with a pathogen. The act of physically \ninjuring a person or animal."
            ),
        ]
    },
    "wikievents_ee_medical": {
        "en": [
            (
                "A determination of the disease or medical condition that explains a person's or animal's symptoms and"
                " signs. Any \nkind of medical treatment or intervention for a person or an animal, often following a"
                " diagnosis of a medical condition. A \nperson or an animal is inoculated against a disease."
            ),
            (
                "A conclusive assessment of the illness or health issue that accounts for an individual's or animal's"
                " symptoms and \nindications. The administration of any therapeutic intervention or treatment method"
                " to a person or an animal, frequently \nsubsequent to a diagnosis of a medical condition. An"
                " individual or an animal is protected against an illness through \ninoculation."
            ),
            (
                "An assessment that identifies the illness or health issue responsible for an individual's or animal's"
                " symptoms and \nindications, often leading to a medical treatment or intervention. The process of"
                " protecting a person or animal from a specific \ndisease through immunization."
            ),
            (
                "An assessment that identifies the illness or health issue responsible for an individual's or animal's"
                " symptoms and \nindications, often leading to a corresponding medical treatment or intervention. The"
                " process of protecting a person or animal \nagainst a specific disease through the administration of"
                " a vaccine or other preventative measure."
            ),
            (
                "A conclusion regarding the illness or health issue that accounts for an individual's or animal's"
                " indications and \nmanifestations. The administration of any form of medical therapy or intervention"
                " to a person or animal, frequently subsequent to a \ndiagnosis of a medical condition. An individual"
                " or animal is protected against an illness through immunization."
            ),
        ]
    },
    "wikievents_ee_movement_transport": {
        "en": [
            (
                "Physical movement or transportation of a person or artifact between places, includes putting and"
                " placing objects \nin locations."
            ),
            (
                'Text: "The relocation of individuals or objects from one location to another, which involves'
                " positioning and \ndepositing items in various places."
            ),
            (
                "The process of relocating individuals or items from one location to another, which encompasses the"
                " act of \npositioning and depositing objects in various places."
            ),
            (
                "The action of relocating an individual or object from one location to another, which also involves"
                " positioning and \ndepositing items in various places."
            ),
            (
                "The action of relocating an individual or object from one location to another, which also involves"
                " positioning and \ndepositing items in specific places."
            ),
        ]
    },
    "wikievents_ee_personnel": {
        "en": [
            (
                "A person changing from one position to another in the same organization, with no indication of"
                " relative levels of the \npositions. A person starting or stoping working in a position."
            ),
            (
                "An individual shifting from one role to another within the same organization, without specifying the"
                " hierarchical \nrank of the positions. A person beginning or ceasing to work in a particular"
                " position."
            ),
            (
                "An individual transitioning between roles within the same organization, without specifying the"
                " hierarchical \nrelationship between the positions. Additionally, a person beginning or ending their"
                " tenure in a particular position."
            ),
            (
                "A person shifting from one role to another within the same organization, without specifying the"
                " hierarchical rank \nof the positions. A person who begins or ends their tenure in a particular"
                " position."
            ),
            (
                "An individual shifting from one role to another within the same organization, without specifying the"
                " positions' \nhierarchical relationship. A person beginning or ending their tenure in a particular"
                " position."
            ),
        ]
    },
    "wikievents_ee_transaction": {
        "en": [
            "A transaction involving the transfer of money, goods, or services between two or more parties.",
            "A transaction entailing the movement of funds, items, or services between two or more parties.",
            (
                "A transaction entails the exchange of funds, items, or services between two or more individuals or"
                " entities."
            ),
            "A transaction entails the exchange of funds, items, or services between a minimum of two parties.",
            (
                "A transaction is a process that involves the exchange of money, goods, or services between two or"
                " more individuals or \nentities."
            ),
        ]
    },
    "wikievents_eae_damage_destroy": {
        "en": [
            "The damaging or destruction of a thing.",
            "The act of causing harm or ruin to an object.",
            "The act of causing harm or ruin to an object.",
            "The act of causing harm or ruin to an object.",
            "The act of causing harm or ruin to an object.",
        ]
    },
    "wikievents_eae_disable_defuse": {
        "en": [
            (
                "Disabling the expected functioning of a mechanical device, or software, for example defusing or"
                " removing a fuse \nfrom an explosive."
            ),
            (
                "Temporarily neutralizing the intended operation of a mechanical apparatus or software, such as"
                " deactivating or \nremoving the fuse from an explosive."
            ),
            (
                "Temporarily impairing the typical operation of a mechanical appliance or computer program, such as"
                " deactivating \nor removing the fuse from an explosive."
            ),
            (
                "Impairing the typical operation of a mechanical apparatus or software, such as disarming or"
                " extracting a fuse from \nan explosive."
            ),
            (
                "Temporarily impairing the normal operation of a mechanical apparatus or software, such as disarming"
                " or extracting \na fuse from an explosive device."
            ),
        ]
    },
    "wikievents_eae_dismantle": {
        "en": [
            "Disassembling an artifact in such a way that it could be reassembled.",
            "Dismantling a complex item so that it can be put back together again.",
            "Dismantling a complex item so that it can be put back together again.",
            "Dismantling a complex item so that it can be put back together again.",
            "Dismantling a complex item so that it can be put back together again.",
        ]
    },
    "wikievents_eae_manufacture_assemble": {
        "en": [
            (
                "Physical action of building, assembling, manufacturing, putting things together, mixing, etc. (does"
                " not include \nthe creation of intellectual property)."
            ),
            (
                'Text: "The process of constructing, assembling, producing, joining components, blending, and so on'
                " (excludes the \ndevelopment of intellectual property)."
            ),
            (
                'Text: "The process of constructing, assembling, producing, joining components, blending, and so on,'
                " excluding \nthe generation of intellectual property."
            ),
            (
                'Text: "The process of constructing, assembling, producing, joining components, blending materials,'
                " and so on, \nexcluding the generation of intellectual property."
            ),
            (
                'Text: "The process of constructing, assembling, producing, joining components, blending materials,'
                " and so on, \nexcluding the generation of intellectual property."
            ),
        ]
    },
    "wikievents_eae_identify_categorize": {
        "en": [
            (
                "The act of establishing the identity of an entity or event, or establishing the relevant category"
                " that an entity or \nevent belongs to, such as being a suspect or a target or an attack."
            ),
            (
                "The process of determining the distinct characteristics or classification of a person or occurrence,"
                " including \ndifferentiating between roles such as a suspect, a target, or the object of an attack."
            ),
            (
                "The process of determining the distinct characteristics or classification of a person or occurrence,"
                " including \nwhether they are considered a suspect, a target, or the focus of an attack."
            ),
            (
                "The process of determining the distinct characteristics or classification of a person, occurrence, or"
                " object, \nincluding whether it is considered a suspect, a target, or the subject of an attack."
            ),
            (
                "The process of determining the distinct characteristics or classification of a person or occurrence,"
                " which may \ninvolve categorizing them as a possible perpetrator, objective, or aggressor."
            ),
        ]
    },
    "wikievents_eae_inspection": {
        "en": [
            "An observation or inspection event, with any target of inspection.",
            "A monitoring or examination instance, encompassing any subject matter being scrutinized.",
            "An instance of monitoring or examination, encompassing any subject matter being scrutinized.",
            "An instance of monitoring or examination, which can encompass any subject being scrutinized.",
            "An observation or examination occasion, encompassing any objective to be inspected.",
        ]
    },
    "wikievents_eae_research": {
        "en": [
            (
                "Explicit mention of researching the answer to a question by analyzing literature or testing"
                " hypotheses through \nexperiments or reviewing opinions from experts in a field."
            ),
            (
                "Clearly stating that the answer to a question can be found by examining written works, testing ideas"
                " through \nexperiments, or consulting expert perspectives in a specific domain."
            ),
            (
                "One such approach involves conducting an in-depth exploration of a subject by examining various"
                " sources of \ninformation, such as literature, conducting experiments to test hypotheses, and seeking"
                " the insights of experts in the field."
            ),
            (
                "Specific reference to the process of investigating the response to an inquiry by examining relevant"
                " literature, \ntesting proposed explanations through experiments, or consulting the perspectives of"
                " experienced professionals in a \ngiven domain."
            ),
            (
                "Clearly stating the process of discovering a response to an inquiry by examining relevant writings,"
                " testing \nproposed assumptions through experiments, or considering perspectives from knowledgeable"
                " individuals in a given domain."
            ),
        ]
    },
    "wikievents_eae_teaching_training_learning": {
        "en": [
            (
                "The action of teaching a person or animal a particular skill or type of behavior, where the learner"
                " does not have \ncomplete mastery of the skill until the end"
            ),
            (
                "The process of instructing an individual or animal in the development of a specific skill or"
                " behavior, during which \nthe learner progresses toward full proficiency throughout the learning"
                " experience"
            ),
            (
                "The process of instructing an individual or animal in the development of a specific skill or"
                " behavior, during which \nthe learner progresses toward full proficiency throughout the learning"
                " experience"
            ),
            (
                "The process of instructing an individual or animal in the development of a specific skill or"
                " behavior, during which \nthe learner progresses toward full proficiency throughout the learning"
                " experience"
            ),
            (
                "The process of instructing an individual or animal in the development of a specific skill or"
                " behavior, in which the \nlearner does not achieve total proficiency until the conclusion"
            ),
        ]
    },
    "wikievents_eae_attack": {
        "en": [
            "An attack event, a violent physical act causing harm or damage.",
            "An assault occasion, a forceful bodily action resulting in injury or destruction.",
            "An assault occasion, a forceful physical deed leading to injury or destruction.",
            "An assault occasion, a forceful physical action resulting in injury or destruction.",
            "An act of aggression, a forceful physical action resulting in injury or destruction.",
        ]
    },
    "wikievents_eae_defeat": {
        "en": [
            "Defeat in a conflict or an election, but not a game-style competition",
            "Loss in a confrontation or an election, but not in a game-like contest",
            "Losses in battles or elections, but not in game-like contests",
            "Loss in a clash or an election, but not in a game-like contest",
            "Loss in a confrontation or an election, but not in a game-like contest",
        ]
    },
    "wikievents_eae_demonstrate": {
        "en": [
            "A demonstration, march, protest, or political gathering.",
            "A display, procession, protest, or political assembly.",
            "A display, procession, protest, or political assembly.",
            "A display, procession, protest, or political assembly.",
            "A display, procession, protest, or assemblage of a political nature.",
        ]
    },
    "wikievents_eae_contact": {
        "en": [
            (
                "An event where two or more participants communicate over any medium, where the context does not"
                " distinguish the \nnature of the contact or whether the communication is in person or remote."
            ),
            (
                "An occasion in which multiple individuals engage in communication through any means, where the"
                " circumstances do \nnot differentiate between the type of contact or whether the interaction occurs"
                " in person or at a distance."
            ),
            (
                "An interaction in which multiple individuals engage, using any mode of communication, where the"
                " circumstances do \nnot specifically identify the type of contact or whether it is conducted in"
                " person or at a distance."
            ),
            (
                "An occurrence during which multiple individuals engage in communication through any means, in which"
                " the \ncircumstances do not differentiate between the type of contact or whether the interaction is"
                " face-to-face or remote."
            ),
            (
                "An interaction in which multiple individuals engage using any means of communication, where the"
                " circumstance does \nnot differentiate between the type of contact or whether it occurs in person or"
                " at a distance."
            ),
        ]
    },
    "wikievents_eae_prevarication": {
        "en": [
            (
                "Any communication over any medium where the communicator intentionally deceives the addressee or"
                " continue an \nongoing deception, either by knowingly conveying false information, or conveying"
                " information with the intent of the \naddressee drawing the wrong conclusion, or omitting"
                " information."
            ),
            (
                "Any form of communication, across any platform, in which the sender purposefully misleads the"
                " recipient, either by \ndeliberately providing false information, by conveying information with the"
                " aim of leading the recipient to an incorrect \nconclusion, or by withholding essential information."
            ),
            (
                "Any form of communication, across any platform, in which the sender purposefully misleads the"
                " recipient, either by \ndeliberately providing false information, by transmitting data with the"
                " objective of leading the recipient to an incorrect \nconclusion, or by withholding crucial"
                " information."
            ),
            (
                "Any form of communication, across any platform, in which the sender purposefully misleads the"
                " recipient, either by \ndeliberately providing false information, by convey"
            ),
            (
                "Any form of communication that involves a sender deliberately misleading the recipient, whether by"
                " purposefully \nproviding false information, leading the recipient to an incorrect conclusion, or"
                " withholding essential information."
            ),
        ]
    },
    "wikievents_eae_request_command": {
        "en": [
            (
                "A statement making a request, appeal, command, or order, where the context does not distinguish"
                " whether the \ncommunication is one-way or two-way, or whether it is in person or remote."
            ),
            (
                "A text that conveys a message seeking an action, entreaty, direction, or requirement, without"
                " differentiating \nbetween the type of communication exchange, whether it occurs in person or at a"
                " distance."
            ),
            (
                "A communication that conveys a request, plea, directive, or instruction, which does not specify"
                " whether the \ninteraction is unidirectional or bidirectional, or whether it occurs in person or at a"
                " distance."
            ),
            (
                "A communication that conveys a request, plea, directive, or instruction, which does not specify"
                " whether the \ninteraction is unidirectional or bidirectional, or whether it occurs in person or at a"
                " distance."
            ),
            (
                "A communication that conveys a request, appeal, directive, or instruction, which does not specify"
                " whether the \ncontext is one-way or two-way, or whether it occurs in person or at a distance."
            ),
        ]
    },
    "wikievents_eae_threaten_coerce": {
        "en": [
            (
                "A statement offering threat or coercion, where the context does not distinguish whether the"
                " communication is \none-way or two-way, or whether it is in person or remote."
            ),
            (
                "A declaration that implies either intimidation or compulsion, which fails to differentiate between"
                " whether the \ncommunication is unidirectional or bidirectional, or whether it occurs in-person or at"
                " a distance."
            ),
            (
                "A declaration that implies danger or compulsion, which does not clarify whether the communication is"
                " \nunidirectional or bidirectional, or if it occurs in person or at a distance."
            ),
            (
                "A declaration that implies either coercion or intimidation, which fails to differentiate between"
                " whether the \ncommunication is unidirectional or bidirectional, or whether it occurs in-person or"
                " via remote means."
            ),
            (
                "A declaration that implicitly or explicitly implies intimidation or compulsion, which fails to"
                " differentiate \nbetween unidirectional or bidirectional communication, or whether the interaction"
                " occurs in person or at a distance."
            ),
        ]
    },
    "wikievents_eae_impede_interfere_with": {
        "en": [
            (
                "Explicit mention of an entity knowingly placing obstacles, not necessarily physical, that raise"
                " difficulties \nwith respect to the occurrence of a taggable, non-Movement event or process."
            ),
            (
                "Specific reference to a conscious effort by a party to deliberately introduce challenges that impede"
                " the \nsuccessful execution of a distinguishable, non-movement-related event or process."
            ),
            (
                "Clear reference to a conscious act by a party intentionally creating challenges that hinder the"
                " occurrence of a \ndistinguishable, non-movement incident or procedure."
            ),
            (
                "Clear reference to a conscious entity intentionally creating challenges, which may not be purely"
                " physical in \nnature, that complicate the occurrence of a detectable, non-movement-related event or"
                " process."
            ),
            (
                "Clear reference to a specific entity intentionally creating challenges, which may not be solely"
                " physical in \nnature, that hinder the occurrence of a distinguishable, non-movement-related event or"
                " process."
            ),
        ]
    },
    "wikievents_eae_crash": {
        "en": [
            "Vehicular collision or the crashing of any type of vehicle.",
            "Automotive accident or the collision of any kind of vehicle.",
            "Automotive collision, which refers to the impact of any kind of vehicle.",
            "Automobile collision, which refers to the impact of any kind of vehicle.",
            "Automotive collision, which refers to the impact of any kind of vehicle.",
        ]
    },
    "wikievents_eae_disease_outbreak": {
        "en": [
            "Explicit mentions of an outbreak of a disease in an area, region, or country.",
            "Clear references to the occurrence of a widespread disease in a specific locale, region, or nation.",
            "Clear references to the occurrence of a disease outbreak within a specific locale, region, or nation.",
            "Clear references to the occurrence of a disease epidemic within a specific region or country.",
            "Clear references to the occurrence of a widespread disease within a specific area, region, or nation.",
        ]
    },
    "wikievents_eae_fire_explosion": {
        "en": [
            "A damaging fire or explosion, either natural or caused but without the intent of attacking.",
            (
                "A destructive fire or explosion, either stemming from natural causes or resulting from an"
                " unintentional act rather \nthan a deliberate attack."
            ),
            (
                "A destructive fire or explosion, either stemming from natural causes or human error, but not"
                " resulting from a \ndeliberate attack."
            ),
            (
                "A destructive fire or explosion, either occurring naturally or caused by human action but not with"
                " the purpose of \nlaunching an attack."
            ),
            (
                "A destructive fire or explosion, either due to natural causes or caused by an external factor but not"
                " as a result of a \ndeliberate attack."
            ),
        ]
    },
    "wikievents_eae_generic_crime": {
        "en": [
            (
                "For use as the crime argument in a Justice event when the crime event is not taggable as any other"
                " annotation ontology \nevent type."
            ),
            (
                "This text is intended for use as the crime argument in a Justice event when the crime event cannot be"
                " categorized under \nany other annotation ontology event type."
            ),
            (
                "This text is intended for use in the crime argument of a Justice event, when the crime event cannot"
                " be categorized \nunder any other annotation ontology event type."
            ),
            (
                "Use this text as the basis for the crime argument in a Justice event when the crime event cannot be"
                " categorized under \nany other type of annotation within an ontology."
            ),
            (
                "This text is intended for use in a Justice event when the crime event cannot be categorized under any"
                " other annotation \nontology event type."
            ),
        ]
    },
    "wikievents_eae_acquit": {
        "en": [
            (
                "An acquit event occurs whenever a trial ends but fails to produce a conviction, including cases where"
                " the charges are \ndropped by the prosecutor."
            ),
            (
                "An acquit event transpires whenever a judicial proceeding concludes without resulting in a"
                " conviction, \nencompassing instances in which the prosecutor chooses to abandon the charges."
            ),
            (
                "An acquit event transpires whenever a judicial proceeding concludes without resulting in a"
                " conviction, \nencompassing instances in which the prosecutor chooses to abandon the charges."
            ),
            (
                "An acquit event transpires whenever a judicial proceeding concludes without resulting in a"
                " conviction, \nencompassing instances in which the prosecutor chooses to abandon the charges."
            ),
            (
                "An acquit event transpires whenever a judicial proceeding concludes without resulting in a"
                " conviction, \nencompassing instances in which the prosecutor chooses to abandon the charges."
            ),
        ]
    },
    "wikievents_eae_arrest_jail_detain": {
        "en": [
            (
                "The detention, taking hostage, kidnapping, arrest, or jailing of an individual by government or"
                " non-government \nactors."
            ),
            "The act of confining, seizing, or imprisoning a person by government or non-government entities.",
            (
                "The act of detaining, seizing, abducting, or imprisoning a person by government or non-government"
                " entities."
            ),
            "The act of imprisoning, capturing, or confining a person by government or non-government entities.",
            "The act of confining, seizing, or imprisoning a person by government or non-government entities.",
        ]
    },
    "wikievents_eae_charge_indict": {
        "en": [
            "A government actor charging, accusing, or indicting a person, organization, or GPE of a crime.",
            (
                "An official from the government levies charges, makes accusations, or issues an indictment against an"
                " individual, \ngroup, or entity for engaging in illegal activities."
            ),
            (
                "A government official levies charges, makes accusations, or issues an indictment against an"
                " individual, group, or \ncorporate entity for engaging in illegal activities."
            ),
            (
                "When a government official files charges or levies accusations against an individual, group, or"
                " entity for \nengaging in illegal activities."
            ),
            (
                "An official from the government levies charges, makes accusations, or issues an indictment against an"
                " individual, \ngroup, or government-protected entity for engaging in unlawful activities."
            ),
        ]
    },
    "wikievents_eae_convict": {
        "en": [
            "A person, organization or GPE entity is convicted whenever that entity has been found guilty of a crime.",
            (
                "An individual, group, or governmental entity is deemed guilty of an offense whenever a judgment of"
                " conviction has \nbeen rendered against that party."
            ),
            (
                "Whenever an individual, group, or entity affiliated with a government or public enterprise is judged"
                " to have \ncommitted a crime, they are considered convicted."
            ),
            (
                "An individual, group, or global political entity is deemed guilty of an offense whenever a formal"
                " judgment of \nculpability has been levied against that entity."
            ),
            (
                "An individual, organization, or governmental entity is deemed guilty and convicted when a formal"
                " judgment of \nculpability has been rendered against that entity for committing a criminal act."
            ),
        ]
    },
    "wikievents_eae_investigate_crime": {
        "en": [
            "Legal, journalistic, and other investigations of crimes.",
            "Examining criminal activities through legal, journalistic, and other methodical inquiries.",
            "Examinations of criminal activities conducted by legal, journalistic, and other professional entities.",
            "Official, news-oriented, and additional probes into felonies.",
            "Official, journalistic, and various other inquiries into criminal acts.",
        ]
    },
    "wikievents_eae_release_parole": {
        "en": [
            (
                "A release-parole event occurs whenever a state actor (GPE, organization subpart, or person"
                " representative) ends \nits custody of a person entity because the sentence has ended, because the"
                " charges are dropped, or because parole has \nbeen granted."
            ),
            (
                "A release-on-parole situation takes place whenever a state actor, such as a governmental entity,"
                " organization, or \nan individual acting on their behalf, terminates their supervision of an"
                " individual due to the completion of their \nsentence, the dismissal of charges, or the granting of"
                " parole."
            ),
            (
                "A release-on-parole situation arises whenever a state figure (such as a government, organization, or"
                " individual \nacting on their behalf) terminates their supervision of an individual because their"
                " sentence has been completed, the \ncharges against them have been dismissed, or they have been"
                " granted parole."
            ),
            (
                "A release-on-parole situation arises whenever a state agent (such as a government, organization, or"
                " individual \nacting on their behalf) terminates the custody of an individual as a result of the"
                " completion of their sentence, the \nwithdrawal of charges, or the granting of parole."
            ),
            (
                "A release-on-parole situation arises whenever a state agent (such as a government, organization, or"
                " individual \nacting on their behalf) terminates their supervision of an individual, as a result of"
                " the completion of their sentence, the \ndismissal of charges, or the granting of parole."
            ),
        ]
    },
    "wikievents_eae_sentence": {
        "en": [
            (
                "A sentence event takes place whenever the punishment, particularly incarceration, for the defendant"
                " of a trial \nevent is issued by a state actor."
            ),
            (
                "A sentence occurrence happens whenever a state representative imposes the penalty, particularly"
                " imprisonment, \non the defendant in a legal case."
            ),
            (
                "A sentence occurrence happens whenever a state official imposes a penalty, particularly imprisonment,"
                " on the \ndefendant in a legal case."
            ),
            (
                "A sentence occurrence happens whenever the penalty, particularly imprisonment, imposed by a"
                " government official \non a defendant during a trial proceeding takes place."
            ),
            (
                "A sentence occurrence happens whenever the penalty, specifically imprisonment, imposed on the"
                " defendant in a \ntrial by a state representative is carried out."
            ),
        ]
    },
    "wikievents_eae_trial_hearing": {
        "en": [
            (
                "A trial event occurs whenever a court proceeding has been initiated for the purposes of determining"
                " the guilt or \ninnocence of a person, organization or GPE accused of committing a crime. A hearing"
                " event occurs whenever a state actor (GPE, \norganization subpart, or person representative)"
                " officially gathers to discuss some criminal legal matter."
            ),
            (
                "In the context of legal proceedings, a trial event takes place whenever a court case has been"
                " initiated to ascertain \nthe culpability or innocence of an individual, organization, or governing"
                " body accused of violating the law. A hearing \nevent, on the other hand, occurs whenever a"
                " representative of a state entity, such as a government, organization, or \nindividual, convenes for"
                " the purpose of discussing a matter related to criminal law."
            ),
            (
                "In the event of a trial, a legal process is initiated to assess the guilt or innocence of an"
                " individual, organization, \nor governing body accused of violating the law. Conversely, a hearing"
                " event takes place when a representative of the \nstate, be it a government entity, organization, or"
                " individual, convenes for the purpose of discussing a matter related to \ncriminal law."
            ),
            (
                "In the context of a court case, a trial event takes place whenever a legal proceeding is initiated to"
                " ascertain the \ninnocence or guilt of an individual, organization, or other party accused of"
                " violating the law. Conversely, a hearing event \noccurs whenever a representative of a state entity,"
                " such as a government, organization, or individual, convenes for the \npurpose of discussing a matter"
                " related to criminal law."
            ),
            (
                "In a trial event, a legal process is initiated to assess the guilt or innocence of an individual,"
                " organization, or GPE \naccused of violating the law. A hearing event, on the other hand, takes place"
                " when a representative of the state, such as a GPE, \norganization, or individual, convenes for an"
                " official discussion on a criminal legal matter."
            ),
        ]
    },
    "wikievents_eae_consume": {
        "en": [
            "An person or animal takes a substance into their body.",
            "A individual or creature ingests a substance within their physique.",
            "An individual, whether human or animal, ingests a substance within their body.",
            "An individual, whether human or animal, ingests a substance within their body.",
            "A person or animal ingests a substance within their body.",
        ]
    },
    "wikievents_eae_die": {
        "en": [
            "The death of a person or animal.",
            "The loss of life for an individual or animal.",
            "The loss of life for an individual or animal.",
            "The loss of life for an individual or animal.",
            "The loss of life for an individual or animal.",
        ]
    },
    "wikievents_eae_illness": {
        "en": [
            "A person or animal experiencing physical harm due to sickness or illness.",
            "An individual or animal undergoing bodily harm as a result of disease or illness.",
            "An individual or animal undergoing bodily harm as a result of illness or sickness.",
            "An individual or creature encountering bodily distress as a result of illness or disease.",
            "A individual or creature enduring bodily harm as a result of disease or ailment.",
        ]
    },
    "wikievents_eae_infect": {
        "en": [
            "A person or animal is infected with a pathogen.",
            (
                "An individual, whether human or animal, has been exposed to and is now hosting a particular"
                " disease-causing agent."
            ),
            (
                "An individual, whether human or animal, has been exposed to and is now hosting a particular"
                " disease-causing agent."
            ),
            "An individual, either human or animal, has been exposed to a harmful microorganism or substance.",
            "An individual, either human or animal, has been exposed to a harmful agent.",
        ]
    },
    "wikievents_eae_injure": {
        "en": [
            "The physical injuring of a person or animal.",
            "The act of causing bodily harm to a person or animal.",
            "The act of causing bodily harm to a person or animal.",
            "The act of causing bodily harm to a person or animal.",
            "The act of harming a person or animal through physical means.",
        ]
    },
    "wikievents_eae_diagnosis": {
        "en": [
            (
                "A determination of the disease or medical condition that explains a person's or animal's symptoms and"
                " signs."
            ),
            (
                "The process of identifying the specific illness or health issue that accounts for an individual's or"
                " animal's \nsymptoms and indications."
            ),
            (
                "A decision as to the illness or health issue that accounts for an individual's or animal's"
                " indications and \nmanifestations."
            ),
            (
                "A conclusive assessment of the illness or health issue that accounts for an individual's or animal's"
                " symptoms and \nindications."
            ),
            (
                "A conclusive determination of the health issue, whether a disease or medical condition, that accounts"
                " for the \nobservable symptoms and indicators experienced by an individual or animal."
            ),
        ]
    },
    "wikievents_eae_intervention": {
        "en": [
            (
                "Any kind of medical treatment or intervention for a person or an animal, often following a diagnosis"
                " of a medical \ncondition."
            ),
            (
                "Any form of therapeutic care or intervention implemented for an individual, whether human or animal,"
                " frequently \nafter a medical condition has been diagnosed."
            ),
            (
                "Medical treatment or intervention, which may include various procedures or therapies, is often"
                " provided to a \nperson or an animal following a diagnosis of an illness or medical condition."
            ),
            (
                "Any type of medical procedure or remedy administered to a person or animal, frequently after a"
                " diagnosis of an \nillness or condition."
            ),
            (
                "Medical care refers to any therapeutic intervention or treatment administered to a person or animal,"
                " frequently \nafter a diagnosis of an illness or health issue."
            ),
        ]
    },
    "wikievents_eae_vaccinate": {
        "en": [
            "A person or an animal is inoculated against a disease.",
            "An individual, whether human or animal, receives an inoculation to protect against a specific illness.",
            "An individual, whether human or animal, receives a protective immunization against an illness.",
            "An individual, whether human or animal, receives an inoculation to protect against a particular illness.",
            "An individual, whether human or animal, receives an inoculation to protect against a specific disease.",
        ]
    },
    "wikievents_eae_evacuation": {
        "en": [
            (
                "The movement of a person or animal (by an agent or via their own effort) from one place to another"
                " for evacuation \npurposes."
            ),
            (
                "The action of an individual or animal being relocated (by an influencer or through their own"
                " initiative) for the \npurpose of evacuation."
            ),
            (
                "The action of an individual or animal relocating to a safe area, whether assisted by another entity"
                " or \nself-initiated, during a time of emergency."
            ),
            (
                "The action of an individual or animal being relocated (whether by an external force or their own"
                " initiative) for the \npurpose of evacuation."
            ),
            (
                "The process of relocation for an individual or animal, whether facilitated by an external agent or"
                " self-initiated, \nin order to evacuate from a given location."
            ),
        ]
    },
    "wikievents_eae_grant_allow_passage": {
        "en": [
            "Explicit mention of granting or allowing entry or exit from a location.",
            "Specific reference to granting or permitting access to or egress from a particular place.",
            "Specific reference to permitting access or egress to or from a particular place.",
            "Specific reference to permitting access or egress to or from a particular place.",
            "Specific reference to granting or permitting access to or egress from a particular place.",
        ]
    },
    "wikievents_eae_illegal_transportation": {
        "en": [
            (
                "Explicit mention of illegal physical transporting of people or things between places, such as"
                " smuggling, \ntrafficking, illegal border crossings."
            ),
            (
                "Clear reference to unlawful movement of people or objects from one location to another, including"
                " smuggling, human \ntrafficking, and illegal border passage."
            ),
            (
                'Text: "Specific reference to unlawful movement of individuals or items between locations, including'
                " smuggling, \nhuman trafficking, and unauthorized border crossings."
            ),
            (
                'Text: "Specific references to the unauthorized movement of individuals or items from one location to'
                " another, \nincluding smuggling, human trafficking, and illegal border crossing."
            ),
            (
                'Text: "Clear references to unlawful movements of people or items from one location to another,'
                " including \nsmuggling, human trafficking, and illegal border crossings."
            ),
        ]
    },
    "wikievents_eae_prevent_passage": {
        "en": [
            "Explicit mention of preventing entry or exit from a location.",
            "Clear statement regarding the prohibition of entry or exit at a specific place.",
            "Clear statement regarding the restriction of access to or egress from a specific area.",
            'Text: "Specific reference to preventing access to or egress from a particular area.',
            "Specific reference to preventing access or egress at a particular site.",
        ]
    },
    "wikievents_eae_transportation": {
        "en": [
            (
                "Physical movement or transportation of a person or artifact between places, includes putting and"
                " placing objects \nin locations."
            ),
            (
                "The process of relocating an individual or object from one location to another, which entails"
                " positioning and \ndepositing items in specific places."
            ),
            (
                "The process of relocating individuals or items from one location to another, which involves"
                " positioning and \ndepositing objects in various places."
            ),
            (
                "The process of relocating individuals or objects from one location to another, which involves"
                " positioning and \ndepositing items in specific places."
            ),
            (
                "The process of relocating individuals or objects from one location to another, which involves"
                " positioning and \ndepositing items in specific places."
            ),
        ]
    },
    "wikievents_eae_change_job_location": {
        "en": [
            "A person continuing in the same role in a different location in the same organization.",
            "An individual who retains their position while relocating to a new site within the same organization.",
            (
                "An individual who remains in the same position but transfers to a distinct area within the same"
                " organization."
            ),
            "An individual who retains their position but relocates to a new site within the same organization.",
            (
                "An individual who remains in the same position but transitions to a different site within the same"
                " organization."
            ),
        ]
    },
    "wikievents_eae_change_job_title": {
        "en": [
            (
                "A person changing from one position to another in the same organization, with no indication of"
                " relative levels of the \npositions."
            ),
            (
                "A individual transitioning between two roles within the same organization, without any indication of"
                " the \npositions' respective hierarchical rankings."
            ),
            (
                "An individual transitioning between roles within the same organization, without any indication of the"
                " \nhierarchical relationship between the positions."
            ),
            (
                "An individual transitioning between two roles within the same organization, without any indication of"
                " the \npositions' hierarchical relationship."
            ),
            (
                "An individual transferring from one role to another within the same organization, without any"
                " indication of the \npositions' hierarchical relation."
            ),
        ]
    },
    "wikievents_eae_end_position": {
        "en": [
            (
                "A person stopping working in a position, with no indication that they are changing positions within"
                " the same \norganization."
            ),
            (
                "An individual ceasing to fulfill their role in a particular position, without any sign of"
                " transitioning to another \nrole within the same organization."
            ),
            (
                "An individual who ceases to work in a particular role, without any indication that they are"
                " transitioning to a \ndifferent role within the same organization."
            ),
            (
                "An individual ceasing their duties in a specific role, without any indication that they will be"
                " transitioning to \nanother role within the same company."
            ),
            (
                "An individual ceasing their duties in a specific role, without any indication that they will be"
                " transitioning to \nanother role within the same company."
            ),
        ]
    },
    "wikievents_eae_start_position": {
        "en": [
            (
                "A person starting working in a position, with no indication that they are changing positions within"
                " the same \norganization."
            ),
            (
                "An individual commences employment in a specific role, without any sign of transitioning to another"
                " role within the \nsame company."
            ),
            (
                "An individual commences employment in a specific role, without any hint that they are transitioning"
                " to another role \nwithin the same company."
            ),
            (
                "An individual commences employment within a specific role, without any signification of transitioning"
                " to another \nrole within the same company."
            ),
            (
                "An individual commences employment in a particular role, without any suggestion that they are"
                " transitioning to a \ndifferent role within the same company."
            ),
        ]
    },
    "wikievents_eae_aid_between_governments": {
        "en": [
            (
                "A voluntary transfer of resources from one GPE to another from the perspective of governments, often"
                " with strings \nattached."
            ),
            (
                "A willing allocation of assets from one government agency to another, as seen from the viewpoint of"
                " the governments, \nfrequently accompanied by specific conditions."
            ),
            (
                "A voluntary relocation of assets from one government agency to another, as seen from the viewpoint of"
                " governments, \nfrequently accompanied by specific conditions."
            ),
            (
                "A voluntary relocation of assets from one government agency to another, as seen from the viewpoint of"
                " governments, \nfrequently accompanied by specific conditions."
            ),
            (
                "A willing allocation of assets from one government entity to another, as seen from the viewpoint of"
                " governments, \nfrequently accompanied by specific conditions."
            ),
        ]
    },
    "wikievents_eae_donation": {
        "en": [
            (
                "The voluntary provision, donation, or extension of material aid in the form of assets or commodities"
                " (nb: use \nAidBetweenGovernments instead for GPE to GPE donations)."
            ),
            (
                "The voluntary offering, contribution, or expansion of tangible support in the form of resources or"
                " goods (note: use \nAidBetweenGovernments for aid donations between governments)."
            ),
            (
                "The voluntary supply, contribution, or allocation of tangible resources in the form of assets or"
                " goods (replace \nAidBetweenGovernments with AidBetweenGovernments for government-to-government"
                " donations)."
            ),
            (
                "The voluntary contribution, gift, or allocation of tangible resources, such as assets or goods,"
                " between entities \n(e.g., governments, organizations, or individuals)."
            ),
            (
                "The voluntary offer, contribution, or allocation of resources in the form of goods or services (note:"
                " use \nAidBetweenGovernments for government-to-government donations)."
            ),
        ]
    },
    "wikievents_eae_exchange_buy_sell": {
        "en": [
            (
                "A transaction transferring or obtaining money, ownership, possession, or control of something,"
                " applicable to any \ntype, nature, or method of acquisition including barter."
            ),
            (
                "A transaction involving the transfer or acquisition of funds, property, or control, encompassing all"
                " types, \nnatures, and methods of procurement, including bartering."
            ),
            (
                "A transaction involving the exchange or acquisition of assets, including money, ownership,"
                " possession, or \ncontrol, encompassing all types, natures, and methods of procurement, such as"
                " barter."
            ),
            (
                "A transaction involving the transfer or acquisition of assets, such as money, property, or control,"
                " encompassing \nall types, natures, and methods of procurement, including barter."
            ),
            (
                "A transaction involving the transfer or acquisition of money, property, or control, encompassing all"
                " types, \nnatures, and methods of procurement, including barter."
            ),
        ]
    },
}

src/tasks/wikievents/guidelines.py
hitz-zentroa/GoLLIE
GUIDELINES = {
    "wikievents_ner_abstract": {
        "en": [
            "Abstract, non-tangible artifacts such as software (e.g., programs, tool kits, apps, e-mail),"
            " measureable\n intellectual property, contracts, etc. (nb: does not include laws, which are LAW type)\n"
        ]
    },
    "wikievents_ner_bodypart": {
        "en": ["An identifiable, living part of a human's or animal's body, such as a eye, ear, neck, leg, etc."]
    },
    "wikievents_ner_commercialproduct": {
        "en": [
            "A tangible product or article of trade for which someone pays or barters, or more generally, an"
            " artifact\n or a thing.\n"
        ]
    },
    "wikievents_ner_facility": {
        "en": [
            "A functional, primarily man-made structure. Facilities are artifacts falling under the domains of\n"
            " architecture and civil engineering, including more temporary human constructs, such as police lines"
            " and\n checkpoints.\n"
        ]
    },
    "wikievents_ner_gpe": {
        "en": [
            "Geopolitical entities such as countries, provinces, states, cities, towns, etc. GPEs are composite\n"
            " entities, consisting of a physical location, a government, and a population. All three of these"
            " elements\n must be present for an entity to be tagged as a GPE. A GPE entity may be a single"
            " geopolitical entity or\n a group.\n"
        ]
    },
    "wikievents_ner_information": {
        "en": [
            "An information object such as a field of study or topic of communication, including thoughts, opinions,\n"
            " etc.\n"
        ]
    },
    "wikievents_ner_location": {
        "en": ["Geographical entities such as geographical areas and landmasses, bodies of water."]
    },
    "wikievents_ner_medicalhealthissue": {
        "en": [
            "Any medical condition or health issue, to include everything from disease to broken bones to fever to\n"
            " general ill health, medical errors, even natural causes\n"
        ]
    },
    "wikievents_ner_money": {
        "en": [
            "A monetary payment. The extent of a Money mention includes modifying quantifiers, the amount, and the\n"
            " currency unit, all of which can be optional.\n"
        ]
    },
    "wikievents_ner_organization": {
        "en": [
            "Corporations, agencies, and other groups of people defined by an established organizational structure."
            " An\n ORG entity may be a single organization or a group. A key feature of an ORG is that it can change\n"
            " members without changing identity.\n"
        ]
    },
    "wikievents_ner_person": {
        "en": ["Person entities are limited to humans. A Person entity may be a single person or a group."]
    },
    "wikievents_ner_side_of_conflict": {
        "en": [
            "The different sides of a conflict, such as  philosophical, cultural, ideological, religious, political,\n"
            " guiding philosophical movement or group orientation.  This will encompass sides of the"
            " battle/conflict,\n sports fans when salient, and other such affiliations, in addition to religions,"
            " political parties, and\n other philosophies.\n"
        ]
    },
    "wikievents_ner_job_title": {"en": ["A person's title or job role."]},
    "wikievents_ner_numeric": {
        "en": [
            "A numerical value or non-numerical value such as an informational property such as color or make or"
            " URL.\n"
        ]
    },
    "wikievents_ner_vehicle": {
        "en": [
            "A physical device primarily designed to move an object from one location to another, by (for example)\n"
            " carrying, flying, pulling, or pushing the transported object. Vehicle entities may or may not have"
            " their\n own power source.\n"
        ]
    },
    "wikievents_ner_weapon": {
        "en": [
            "A physical device that is primarily used as an instrument for physically harming or destroying"
            " entities.\n"
        ]
    },
    "wikievents_ee_artifact_existance": {
        "en": [
            "Any event related to manufacturing, assembling, damagging, destroying, disabling, defusing, or\n"
            " dismantling an artifact or a thing.\n"
        ]
    },
    "wikievents_ee_cognitive": {
        "en": [
            "The act of establishing the identity of an entity or event, or establishing the relevant category that"
            " an\n entity or event belongs to, such as being a suspect or a target or an attack. An observation or\n"
            " inspection event, with any target of inspection. Explicit mention of researching the answer to a\n"
            " question by analyzing literature or testing hypotheses through experiments or reviewing opinions from\n"
            " experts in a field. The action of teaching a person or animal a particular skill or type of behavior,\n"
            " where the learner does not have complete mastery of the skill until the end.\n"
        ]
    },
    "wikievents_ee_conflict": {
        "en": [
            "An attack event, a violent physical act causing harm or damage. Defeat in a conflict or an election,"
            " but\n not a game-style competition. A demonstration, march, protest, or political gathering.\n"
        ]
    },
    "wikievents_ee_contact": {
        "en": [
            "An event where two or more participants communicate over any medium, where the context does not\n"
            " distinguish the nature of the contact or whether the communication is in person or remote. Any\n"
            " communication over any medium where the communicator intentionally deceives the addressee or continue"
            " an\n ongoing deception, either by knowingly conveying false information, or conveying information with"
            " the\n intent of the addressee drawing the wrong conclusion, or omitting information. A statement making"
            " a\n request, appeal, command, or order, where the context does not distinguish whether the communication"
            " is\n one-way or two-way, or whether it is in person or remote. A statement offering threat or coercion,"
            " where\n the context does not distinguish whether the communication is one-way or two-way, or whether it"
            " is in\n person or remote.\n"
        ]
    },
    "wikievents_ee_control": {
        "en": [
            "Explicit mention of an entity knowingly placing obstacles, not necessarily physical, that raise\n"
            " difficulties with respect to the occurrence of a taggable, non-Movement event or process\n"
        ]
    },
    "wikievents_ee_disaster": {
        "en": [
            "Vehicular collision or the crashing of any type of vehicle. Explicit mentions of an outbreak of a"
            " disease\n in an area, region, or country. A damaging fire or explosion, either natural or caused but"
            " without the\n intent of attacking\n"
        ]
    },
    "wikievents_ee_generic_crime": {
        "en": [
            "For use as the crime argument in a Justice event when the crime event is not taggable as any other\n"
            " annotation ontology event type.\n"
        ]
    },
    "wikievents_ee_justice": {
        "en": [
            "Any legal proceedings, such as a court case, a trial, or a hearing. Arresting, detaining, punishing, or\n"
            " otherwise dealing with a criminal or non-criminal person or entity.\n"
        ]
    },
    "wikievents_ee_life": {
        "en": [
            "An person or animal takes a substance into their body. The death of a person or animal. A person or\n"
            " animal experiencing physical harm due to sickness or illness. A person or animal is infected with a\n"
            " pathogen. The physical injuring of a person or animal.\n"
        ]
    },
    "wikievents_ee_medical": {
        "en": [
            "A determination of the disease or medical condition that explains a person's or animal's symptoms and\n"
            " signs. Any kind of medical treatment or intervention for a person or an animal, often following a\n"
            " diagnosis of a medical condition. A person or an animal is inoculated against a disease.\n"
        ]
    },
    "wikievents_ee_movement_transport": {
        "en": [
            "Physical movement or transportation of a person or artifact between places, includes putting and"
            " placing\n objects in locations.\n"
        ]
    },
    "wikievents_ee_personnel": {
        "en": [
            "A person changing from one position to another in the same organization, with no indication of relative\n"
            " levels of the positions. A person starting or stoping working in a position.\n"
        ]
    },
    "wikievents_ee_transaction": {
        "en": ["A transaction involving the transfer of money, goods, or services between two or more parties."]
    },
    "wikievents_eae_damage_destroy": {"en": ["The damaging or destruction of a thing."]},
    "wikievents_eae_disable_defuse": {
        "en": [
            "Disabling the expected functioning of a mechanical device, or software, for example defusing or"
            " removing\n a fuse from an explosive.\n"
        ]
    },
    "wikievents_eae_dismantle": {"en": ["Disassembling an artifact in such a way that it could be reassembled."]},
    "wikievents_eae_manufacture_assemble": {
        "en": [
            "Physical action of building, assembling, manufacturing, putting things together, mixing, etc. (does not\n"
            " include the creation of intellectual property).\n"
        ]
    },
    "wikievents_eae_identify_categorize": {
        "en": [
            "The act of establishing the identity of an entity or event, or establishing the relevant category that"
            " an\n entity or event belongs to, such as being a suspect or a target or an attack.\n"
        ]
    },
    "wikievents_eae_inspection": {"en": ["An observation or inspection event, with any target of inspection."]},
    "wikievents_eae_research": {
        "en": [
            "Explicit mention of researching the answer to a question by analyzing literature or testing hypotheses\n"
            " through experiments or reviewing opinions from experts in a field.\n"
        ]
    },
    "wikievents_eae_teaching_training_learning": {
        "en": [
            "The action of teaching a person or animal a particular skill or type of behavior, where the learner"
            " does\n not have complete mastery of the skill until the end\n"
        ]
    },
    "wikievents_eae_attack": {"en": ["An attack event, a violent physical act causing harm or damage."]},
    "wikievents_eae_defeat": {"en": ["Defeat in a conflict or an election, but not a game-style competition"]},
    "wikievents_eae_demonstrate": {"en": ["A demonstration, march, protest, or political gathering."]},
    "wikievents_eae_contact": {
        "en": [
            "An event where two or more participants communicate over any medium, where the context does not\n"
            " distinguish the nature of the contact or whether the communication is in person or remote.\n"
        ]
    },
    "wikievents_eae_prevarication": {
        "en": [
            "Any communication over any medium where the communicator intentionally deceives the addressee or"
            " continue\n an ongoing deception, either by knowingly conveying false information, or conveying"
            " information with the\n intent of the addressee drawing the wrong conclusion, or omitting information.\n"
        ]
    },
    "wikievents_eae_request_command": {
        "en": [
            "A statement making a request, appeal, command, or order, where the context does not distinguish whether\n"
            " the communication is one-way or two-way, or whether it is in person or remote.\n"
        ]
    },
    "wikievents_eae_threaten_coerce": {
        "en": [
            "A statement offering threat or coercion, where the context does not distinguish whether the"
            " communication\n is one-way or two-way, or whether it is in person or remote.\n"
        ]
    },
    "wikievents_eae_impede_interfere_with": {
        "en": [
            "Explicit mention of an entity knowingly placing obstacles, not necessarily physical, that raise\n"
            " difficulties with respect to the occurrence of a taggable, non-Movement event or process.\n"
        ]
    },
    "wikievents_eae_crash": {"en": ["Vehicular collision or the crashing of any type of vehicle."]},
    "wikievents_eae_disease_outbreak": {
        "en": ["Explicit mentions of an outbreak of a disease in an area, region, or country."]
    },
    "wikievents_eae_fire_explosion": {
        "en": ["A damaging fire or explosion, either natural or caused but without the intent of attacking."]
    },
    "wikievents_eae_generic_crime": {
        "en": [
            "For use as the crime argument in a Justice event when the crime event is not taggable as any other\n"
            " annotation ontology event type.\n"
        ]
    },
    "wikievents_eae_acquit": {
        "en": [
            "An acquit event occurs whenever a trial ends but fails to produce a conviction, including cases where"
            " the\n charges are dropped by the prosecutor.\n"
        ]
    },
    "wikievents_eae_arrest_jail_detain": {
        "en": [
            "The detention, taking hostage, kidnapping, arrest, or jailing of an individual by government or\n"
            " non-government actors.\n"
        ]
    },
    "wikievents_eae_charge_indict": {
        "en": ["A government actor charging, accusing, or indicting a person, organization, or GPE of a crime."]
    },
    "wikievents_eae_convict": {
        "en": [
            "A person, organization or GPE entity is convicted whenever that entity has been found guilty of a"
            " crime.\n"
        ]
    },
    "wikievents_eae_investigate_crime": {"en": ["Legal, journalistic, and other investigations of crimes."]},
    "wikievents_eae_release_parole": {
        "en": [
            "A release-parole event occurs whenever a state actor (GPE, organization subpart, or person\n"
            " representative) ends its custody of a person entity because the sentence has ended, because the"
            " charges\n are dropped, or because parole has been granted.\n"
        ]
    },
    "wikievents_eae_sentence": {
        "en": [
            "A sentence event takes place whenever the punishment, particularly incarceration, for the defendant of"
            " a\n trial event is issued by a state actor.\n"
        ]
    },
    "wikievents_eae_trial_hearing": {
        "en": [
            "A trial event occurs whenever a court proceeding has been initiated for the purposes of determining the\n"
            " guilt or innocence of a person, organization or GPE accused of committing a crime. A hearing event\n"
            " occurs whenever a state actor (GPE, organization subpart, or person representative) officially gathers\n"
            " to discuss some criminal legal matter.\n"
        ]
    },
    "wikievents_eae_consume": {"en": ["An person or animal takes a substance into their body."]},
    "wikievents_eae_die": {"en": ["The death of a person or animal."]},
    "wikievents_eae_illness": {"en": ["A person or animal experiencing physical harm due to sickness or illness."]},
    "wikievents_eae_infect": {"en": ["A person or animal is infected with a pathogen."]},
    "wikievents_eae_injure": {"en": ["The physical injuring of a person or animal."]},
    "wikievents_eae_diagnosis": {
        "en": [
            "A determination of the disease or medical condition that explains a person's or animal's symptoms and\n"
            " signs.\n"
        ]
    },
    "wikievents_eae_intervention": {
        "en": [
            "Any kind of medical treatment or intervention for a person or an animal, often following a diagnosis of"
            " a\n medical condition.\n"
        ]
    },
    "wikievents_eae_vaccinate": {"en": ["A person or an animal is inoculated against a disease."]},
    "wikievents_eae_evacuation": {
        "en": [
            "The movement of a person or animal (by an agent or via their own effort) from one place to another for\n"
            " evacuation purposes.\n"
        ]
    },
    "wikievents_eae_grant_allow_passage": {
        "en": ["Explicit mention of granting or allowing entry or exit from a location."]
    },
    "wikievents_eae_illegal_transportation": {
        "en": [
            "Explicit mention of illegal physical transporting of people or things between places, such as"
            " smuggling,\n trafficking, illegal border crossings.\n"
        ]
    },
    "wikievents_eae_prevent_passage": {"en": ["Explicit mention of preventing entry or exit from a location."]},
    "wikievents_eae_transportation": {
        "en": [
            "Physical movement or transportation of a person or artifact between places, includes putting and"
            " placing\n objects in locations.\n"
        ]
    },
    "wikievents_eae_change_job_location": {
        "en": ["A person continuing in the same role in a different location in the same organization."]
    },
    "wikievents_eae_change_job_title": {
        "en": [
            "A person changing from one position to another in the same organization, with no indication of relative\n"
            " levels of the positions.\n"
        ]
    },
    "wikievents_eae_end_position": {
        "en": [
            "A person stopping working in a position, with no indication that they are changing positions within the\n"
            " same organization.\n"
        ]
    },
    "wikievents_eae_start_position": {
        "en": [
            "A person starting working in a position, with no indication that they are changing positions within the\n"
            " same organization.\n"
        ]
    },
    "wikievents_eae_aid_between_governments": {
        "en": [
            "A voluntary transfer of resources from one GPE to another from the perspective of governments, often"
            " with\n strings attached.\n"
        ]
    },
    "wikievents_eae_donation": {
        "en": [
            "The voluntary provision, donation, or extension of material aid in the form of assets or commodities"
            " (nb:\n use AidBetweenGovernments instead for GPE to GPE donations).\n"
        ]
    },
    "wikievents_eae_exchange_buy_sell": {
        "en": [
            "A transaction transferring or obtaining money, ownership, possession, or control of something,"
            " applicable\n to any type, nature, or method of acquisition including barter.\n"
        ]
    },
}

EXAMPLES = {
    "wikievents_ner_abstract_examples": {
        "en": [
            "documents",
            "technology",
            "Proposals",
            "report",
            "research",
            "evidence",
            "data",
            "databases",
            "Twitter",
            "death penalty",
        ]
    },
    "wikievents_ner_bodypart_examples": {
        "en": [
            "tatoo",
            "limbs",
            "the bodies",
            "neck",
            "ear",
            "face",
            "leg",
            "shoulder",
            "thumb",
            "body",
        ]
    },
    "wikievents_ner_commercialproduct_examples": {
        "en": [
            "Arduino",
            "ID",
            "food",
            "bottle cap",
            "backpack",
            "package",
            "Servo",
            "medicine",
            "components",
            "windows",
        ]
    },
    "wikievents_ner_facility_examples": {
        "en": [
            "where",
            "prison",
            "station",
            "home",
            "hospital",
            "building",
            "pub",
            "house",
            "apartment",
            "school",
        ]
    },
    "wikievents_ner_gpe_examples": {
        "en": [
            "Boston",
            "U.S",
            "U.S.",
            "Venezuela",
            "country",
            "Pakistan",
            "United States",
            "city",
            "US",
            "Venezuelan",
        ]
    },
    "wikievents_ner_information_examples": {
        "en": [
            "conference",
            "images",
            "report",
            "statement",
            "name",
            "article",
            "said",
            "defense",
            "forensic analysis",
            "Photos",
        ]
    },
    "wikievents_ner_location_examples": {
        "en": [
            "area",
            "line",
            "world",
            "areas",
            "where",
            "scene",
            "site",
            "region",
            "neighborhood",
            "ground",
        ]
    },
    "wikievents_ner_medicalhealthissue_examples": {
        "en": [
            "disease",
            "fever",
            "autism spectrum disorder",
            "weak",
            "coronavirus",
            "hiccup",
            "broken ribs",
            "broken scapula",
            "severe dengue",
            "coma",
        ]
    },
    "wikievents_ner_money_examples": {
        "en": [
            "cash",
            "$925",
            "$2,775",
            "$5.2 million",
            "$400",
            "$97",
            "$20.99",
            "$41.28",
            "$44",
            "¥500,000",
        ]
    },
    "wikievents_ner_organization_examples": {
        "en": [
            "police",
            "government",
            "FBI",
            "IRA",
            "military",
            "group",
            "court",
            "Taliban",
            "it",
            "its",
        ]
    },
    "wikievents_ner_person_examples": {
        "en": [
            "he",
            "his",
            "who",
            "people",
            "they",
            "He",
            "I",
            "their",
            "Tsarnaev",
            "him",
        ]
    },
    "wikievents_ner_side_of_conflict_examples": {"en": ["church", "officers", "supporters", "radical Islam"]},
    "wikievents_ner_job_title_examples": {
        "en": [
            "President",
            "president",
            "Attorney General",
            "opposition activist",
            "Prime Minister",
            "consultant",
            "law enforcement",
            "psychiatrist",
            "spokesman",
            "organizer",
        ]
    },
    "wikievents_ner_numeric_examples": {
        "en": [
            "two",
            "three",
            "2",
            "five",
            "Tuesday",
            "1",
            "15",
            "2011",
            "48",
            "3",
        ]
    },
    "wikievents_ner_vehicle_examples": {
        "en": [
            "car",
            "drone",
            "drones",
            "vehicle",
            "boat",
            "aircraft",
            "it",
            "train",
            "truck",
            "vehicles",
        ]
    },
    "wikievents_ner_weapon_examples": {
        "en": [
            "bomb",
            "bombs",
            "drone",
            "explosives",
            "device",
            "drones",
            "devices",
            "it",
            "explosive",
            "that",
        ]
    },
    "wikievents_ee_artifact_existance_examples": {
        "en": [
            "damage",
            "damaged",
            "destroyed",
            "make",
            "wrecked",
            "made",
            "DESTRUCTION",
            "build",
            "burn",
            "destruction",
        ]
    },
    "wikievents_ee_cognitive_examples": {
        "en": [
            "identified",
            "found",
            "discovered",
            "investigation",
            "searched",
            "investigate",
            "inspections",
            "researching",
            "saw",
            "identifying",
        ]
    },
    "wikievents_ee_conflict_examples": {
        "en": [
            "attack",
            "bombing",
            "attacks",
            "explosion",
            "bombings",
            "exploded",
            "detonated",
            "blast",
            "went off",
            "shot",
        ]
    },
    "wikievents_ee_contact_examples": {
        "en": [
            "said",
            "told",
            "reported",
            "wrote",
            "say",
            "called",
            "met",
            "says",
            "released",
            "reports",
        ]
    },
    "wikievents_ee_control_examples": {
        "en": [
            "embargo",
            "intercepted",
            "banned",
            "disrupted",
            "combat",
            "stop",
            "blocks",
            "curfew",
            "stopping",
        ]
    },
    "wikievents_ee_disaster_examples": {
        "en": [
            "crashed",
            "crash",
            "Chikungunya",
            "crashing",
            "outbreak",
            "disease",
            "virus",
        ]
    },
    "wikievents_ee_generic_crime_examples": {
        "en": [
            "crime",
            "assassination",
            "kidnappings",
            "AWOL",
            "criminal",
            "robbery",
            "crimes",
            "murder",
            "possession",
            "pornography",
        ]
    },
    "wikievents_ee_justice_examples": {
        "en": [
            "arrested",
            "investigation",
            "charged",
            "convicted",
            "sentenced",
            "arrest",
            "charges",
            "accused",
            "trial",
            "custody",
        ]
    },
    "wikievents_ee_life_examples": {
        "en": [
            "killed",
            "injured",
            "wounded",
            "killing",
            "died",
            "dead",
            "injuring",
            "deaths",
            "injuries",
            "kill",
        ]
    },
    "wikievents_ee_medical_examples": {
        "en": [
            "treated",
            "treatment",
            "surgery",
            "hospitalized",
            "helped",
            "amputations",
            "medication",
            "rushed",
        ]
    },
    "wikievents_ee_transport_examples": {
        "en": [
            "moved",
            "evacuated",
            "travel",
            "fenced off",
            "driven",
            "fled",
            "traveled",
            "carrying",
            "returned",
            "brought",
        ]
    },
    "wikievents_ee_personnel_examples": {
        "en": [
            "discharged",
            "discharge",
            "relinquished",
            "becoming",
            "inauguration",
            "fire",
            "is",
        ]
    },
    "wikievents_ee_transaction_examples": {
        "en": [
            "bought",
            "buy",
            "purchase",
            "purchased",
            "sold",
            "selling",
            "buys",
            "purchases",
            "sells",
            "purchasing",
        ]
    },
}

src/tasks/wikievents/guidelines_gold.py
hitz-zentroa/GoLLIE
from __future__ import annotations

import json
import os
from argparse import ArgumentParser
from collections import defaultdict
from copy import deepcopy
from typing import Iterable, Tuple

from tqdm.auto import tqdm


class WikiEventsPreprocessor:
    def __init__(self) -> WikiEventsPreprocessor:
        self._nlp = None

    def _sent_tokenize(self: WikiEventsPreprocessor, text: str, start_pos: int) -> Iterable[Tuple[int, str]]:
        if not self._nlp:
            import spacy

            self._nlp = spacy.load("en_core_web_sm")

        for sent in self._nlp(text).sents:
            yield [sent[0].idx + start_pos, sent.text]

    @staticmethod
    def _find_subsentence(offset, sentences):
        sentences_ = sentences.copy() + [(sentences[-1][0] + len(sentences[-1][1]) + 1, "")]
        return next((i - 1 for i, (idx, _) in enumerate(sentences_) if offset < idx), -1)

    def __call__(self, input_path: str) -> Iterable[dict]:
        """A function that loads and converts the WikiEvents dataset into sentence level."""
        with open(input_path, "rt") as data_f:
            for data_line in tqdm(data_f):
                instance = json.loads(data_line)

                entities = {entity["id"]: entity for entity in instance["entity_mentions"]}
                tokens = [token for sentence in instance["sentences"] for token in sentence[0]]

                sub_sentence_information = defaultdict(dict)
                all_sub_sentences = [
                    list(self._sent_tokenize(text, tokens[0][1])) for tokens, text in instance["sentences"]
                ]

                for i, sentences in enumerate(all_sub_sentences):
                    for j, sent in enumerate(sentences):
                        sub_sentence_information[f"{i}-{j}"]["text"] = sent[-1]

                for value in entities.values():
                    sent_idx = value["sent_idx"]
                    sub_sent_idx = self._find_subsentence(tokens[value["start"]][1], all_sub_sentences[sent_idx])

                    sentence = sub_sentence_information[f"{sent_idx}-{sub_sent_idx}"]["text"]

                    new_value = deepcopy(value)
                    new_value["start"] = tokens[value["start"]][1] - all_sub_sentences[sent_idx][sub_sent_idx][0]
                    # Fix start if needed
                    _shift = sentence[new_value["start"] :].find(new_value["text"])
                    if _shift == -1:
                        continue
                    elif _shift > 0:
                        new_value["start"] += _shift

                    # new_value["end"] = tokens[value["end"] - 1][-1] - all_sub_sentences[sent_idx][sub_sent_idx][0]
                    new_value["end"] = new_value["start"] + len(new_value["text"])

                    assert (
                        sentence[new_value["start"] : new_value["end"]] == new_value["text"]
                    ), f"{sentence[new_value['start']:new_value['end']]}|{new_value['text']}"

                    if "entity_mentions" not in sub_sentence_information[f"{sent_idx}-{sub_sent_idx}"]:
                        sub_sentence_information[f"{sent_idx}-{sub_sent_idx}"]["entity_mentions"] = []

                    sub_sentence_information[f"{sent_idx}-{sub_sent_idx}"]["entity_mentions"].append(new_value)

                for event in instance["event_mentions"]:
                    sent_idx = event["trigger"]["sent_idx"]
                    sub_sent_idx = self._find_subsentence(
                        tokens[event["trigger"]["start"]][1], all_sub_sentences[sent_idx]
                    )

                    sentence = sub_sentence_information[f"{sent_idx}-{sub_sent_idx}"]["text"]

                    new_value = deepcopy(event)
                    new_value["trigger"]["start"] = (
                        tokens[event["trigger"]["start"]][1] - all_sub_sentences[sent_idx][sub_sent_idx][0]
                    )
                    # Fix start if needed
                    _shift = sentence[new_value["trigger"]["start"] :].find(new_value["trigger"]["text"])
                    if _shift == -1:
                        continue
                    elif _shift > 0:
                        new_value["trigger"]["start"] += _shift
                    # new_value["trigger"]["end"] = (
                    #     tokens[event["trigger"]["end"] - 1][-1] - all_sub_sentences[sent_idx][sub_sent_idx][0]
                    # )
                    new_value["trigger"]["end"] = new_value["trigger"]["start"] + len(new_value["trigger"]["text"])

                    sent_entities = (
                        {
                            ent["id"]
                            for ent in sub_sentence_information[f"{sent_idx}-{sub_sent_idx}"]["entity_mentions"]
                        }
                        if "entity_mentions" in sub_sentence_information[f"{sent_idx}-{sub_sent_idx}"]
                        else {}
                    )

                    for argument in new_value["arguments"]:
                        if argument["entity_id"] not in sent_entities:
                            argument["role"] = f"[OOR]_{argument['role']}"

                    assert (
                        sentence[new_value["trigger"]["start"] : new_value["trigger"]["end"]]
                        == new_value["trigger"]["text"]
                    ), f"{sentence[new_value['trigger']['start']:new_value['trigger']['end']]}|{new_value['trigger']['text']}"

                    if "event_mentions" not in sub_sentence_information[f"{sent_idx}-{sub_sent_idx}"]:
                        sub_sentence_information[f"{sent_idx}-{sub_sent_idx}"]["event_mentions"] = []

                    sub_sentence_information[f"{sent_idx}-{sub_sent_idx}"]["event_mentions"].append(new_value)

                for i, sentences in enumerate(all_sub_sentences):
                    for j, _ in enumerate(sentences):
                        info = sub_sentence_information[f"{i}-{j}"]
                        info["doc_id"] = f"{instance['doc_id']}_{i}_{j}"

                        if "entity_mentions" not in info and "event_mentions" not in info and len(info["text"]) <= 25:
                            continue

                        if "entity_mentions" not in info:
                            info["entity_mentions"] = []
                        if "event_mentions" not in info:
                            info["event_mentions"] = []

                        yield info


def main(args):
    processor = WikiEventsPreprocessor()

    for filename in os.listdir(args.dir_path):
        full_name = os.path.join(args.dir_path, filename)
        if os.path.isfile(full_name) and "sentence" not in filename:
            output_file = os.path.join(args.dir_path, filename.replace(".jsonl", ".sentence.jsonl"))
            with open(output_file, "w", encoding="utf-8") as f:
                for inst in processor(full_name):
                    print(json.dumps(inst, ensure_ascii=False), file=f)


if __name__ == "__main__":
    parser = ArgumentParser("WikiEvents preprocessing")

    parser.add_argument("-d", "--dir_path", type=str, default="data/wikievents/")

    args = parser.parse_args()
    main(args)

src/tasks/wikievents/preprocess_wikievents.py
hitz-zentroa/GoLLIE
from typing import Dict, List, Type

from src.tasks.utils_typing import Entity, Event, dataclass


"""Entity definitions

The entity defitions are extracted from the KAIROS project guidelines.
"""


@dataclass
class Abstract(Entity):
    """{wikievents_ner_abstract}"""

    span: str  # {wikievents_ner_abstract_examples}


@dataclass
class BodyPart(Entity):
    """{wikievents_ner_bodypart}"""

    span: str  # {wikievents_ner_bodypart_examples}


@dataclass
class CommercialProduct(Entity):
    """{wikievents_ner_commercialproduct}"""

    span: str  # {wikievents_ner_commercialproduct_examples}


@dataclass
class Facility(Entity):
    """{wikievents_ner_facility}"""

    span: str  # {wikievents_ner_facility_examples}


@dataclass
class GPE(Entity):
    """{wikievents_ner_gpe}"""

    span: str  # {wikievents_ner_gpe_examples}


@dataclass
class Information(Entity):
    """{wikievents_ner_information}"""

    span: str  # {wikievents_ner_information_examples}


@dataclass
class Location(Entity):
    """{wikievents_ner_location}"""

    span: str  # {wikievents_ner_location_examples}


@dataclass
class MedicalHealthIssue(Entity):
    """{wikievents_ner_medicalhealthissue}"""

    span: str  # {wikievents_ner_medicalhealthissue_examples}


@dataclass
class Money(Entity):
    """{wikievents_ner_money}"""

    span: str  # {wikievents_ner_money_examples}


@dataclass
class Organization(Entity):
    """{wikievents_ner_organization}"""

    span: str  # {wikievents_ner_organization_examples}


@dataclass
class Person(Entity):
    """{wikievents_ner_person}"""

    span: str  # {wikievents_ner_person_examples}


@dataclass
class SideOfConflict(Entity):
    """{wikievents_ner_side_of_conflict}"""

    span: str  # {wikievents_ner_side_of_conflict_examples}


@dataclass
class JobTitle(Entity):
    """{wikievents_ner_job_title}"""

    span: str  # {wikievents_ner_job_title_examples}


@dataclass
class Numeric(Entity):
    """{wikievents_ner_numeric}"""

    span: str  # {wikievents_ner_numeric_examples}


@dataclass
class Vehicle(Entity):
    """{wikievents_ner_vehicle}"""

    span: str  # {wikievents_ner_vehicle_examples}


@dataclass
class Weapon(Entity):
    """{wikievents_ner_weapon}"""

    span: str  # {wikievents_ner_weapon_examples}


ENTITY_DEFINITIONS: List[Type] = [
    Abstract,
    BodyPart,
    CommercialProduct,
    Facility,
    GPE,
    Information,
    Location,
    MedicalHealthIssue,
    Money,
    Organization,
    Person,
    SideOfConflict,
    JobTitle,
    Numeric,
    Vehicle,
    Weapon,
]


@dataclass
class ArtifactExistanceEvent(Event):
    """{wikievents_ee_artifact_existance}"""

    mention: str
    """The text span that triggers the event.
    {wikievents_ee_artifact_existance_examples}
    """


@dataclass
class CognitiveEvent(Event):
    """{wikievents_ee_cognitive}"""

    mention: str
    """The text span that triggers the event.
    {wikievents_ee_cognitive_examples}
    """


@dataclass
class ConflictEvent(Event):
    """{wikievents_ee_conflict}"""

    mention: str
    """The text span that triggers the event.
    {wikievents_ee_conflict_examples}
    """


@dataclass
class ContactEvent(Event):
    """{wikievents_ee_contact}"""

    mention: str
    """The text span that triggers the event.
    {wikievents_ee_contact_examples}
    """


@dataclass
class ControlEvent(Event):
    """{wikievents_ee_control}"""

    mention: str
    """The text span that triggers the event.
    {wikievents_ee_control_examples}
    """


@dataclass
class DisasterEvent(Event):
    """{wikievents_ee_disaster}"""

    mention: str
    """The text span that triggers the event.
    {wikievents_ee_disaster_examples}
    """


@dataclass
class GenericCrimeEvent(Event):
    """{wikievents_ee_generic_crime}"""

    mention: str
    """The text span that triggers the event.
    {wikievents_ee_generic_crime_examples}
    """


@dataclass
class JusticeEvent(Event):
    """{wikievents_ee_justice}"""

    mention: str
    """The text span that triggers the event.
    {wikievents_ee_justice_examples}
    """


@dataclass
class LifeEvent(Event):
    """{wikievents_ee_life}"""

    mention: str
    """The text span that triggers the event.
    {wikievents_ee_life_examples}
    """


@dataclass
class MedicalEvent(Event):
    """{wikievents_ee_medical}"""

    mention: str
    """The text span that triggers the event.
    {wikievents_ee_medical_examples}
    """


@dataclass
class MovementTransportEvent(Event):
    """{wikievents_ee_movement_transport}"""

    mention: str
    """The text span that triggers the event.
    {wikievents_ee_transport_examples}
    """


@dataclass
class PersonnelEvent(Event):
    """{wikievents_ee_personnel}"""

    mention: str
    """The text span that triggers the event.
    {wikievents_ee_personnel_examples}
    """


@dataclass
class TransactionEvent(Event):
    """{wikievents_ee_transaction}"""

    mention: str
    """The text span that triggers the event.
    {wikievents_ee_transaction_examples}
    """


COARSE_EVENT_DEFINITIONS: List[Type] = [
    ArtifactExistanceEvent,
    CognitiveEvent,
    ConflictEvent,
    ContactEvent,
    ControlEvent,
    DisasterEvent,
    GenericCrimeEvent,
    JusticeEvent,
    LifeEvent,
    MedicalEvent,
    MovementTransportEvent,
    PersonnelEvent,
    TransactionEvent,
]


@dataclass
class DamageDestroy(ArtifactExistanceEvent):
    """{wikievents_eae_damage_destroy}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    agent: List[str]  # The person damaging the artifact
    artifact: List[str]  # The artifact being damaged
    instrument: List[str]  # The object used to damage the artifact
    place: List[str]  # Where the event occurred


@dataclass
class DisableDefuse(ArtifactExistanceEvent):
    """{wikievents_eae_disable_defuse}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    agent: List[str]  # The person disabling, defusing or deactivating the artifact
    artifact: List[str]  # The artifact being disabled
    instrument: List[str]  # The object used to disable, defuse or deactivating the artifact
    place: List[str]  # Where the event occurred


@dataclass
class Dismantle(ArtifactExistanceEvent):
    """{wikievents_eae_dismantle}"""

    mention: str  # # The text span that most clearly expresses (triggers) the event
    agent: List[str]  # The person dismantling the artifact
    artifact: List[str]  # The artifact being dismantled into components
    components: List[str]  # The components that originally composed the artifact
    instrument: List[str]  # The object used to dismantle the artifact
    place: List[str]  # Where the event occurred


@dataclass
class ManufactureAssemble(ArtifactExistanceEvent):
    """{wikievents_eae_manufacture_assemble}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    agent: List[str]  # The person manufacturing or assembling the artifact
    artifact: List[str]  # The artifact being manufactured or assembled
    components: List[str]  # The components that compose the artifact
    place: List[str]  # Where the event occurred


@dataclass
class IdentifyCategorize(CognitiveEvent):
    """{wikievents_eae_identify_categorize}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    agent: List[str]  # The person identifying the object
    object: List[str]  # The identifyied object
    role: List[str]  # The role or categorization of the object
    place: List[str]  # Where the event occurred


@dataclass
class Inspection(CognitiveEvent):
    """{wikievents_eae_inspection}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    agent: List[str]  # The person observing the entity
    entity: List[str]  # The entity being observed
    instrument: List[str]  # The instrument used to observe the entity
    place: List[str]  # Where the event occurred


@dataclass
class Research(CognitiveEvent):
    """{wikievents_eae_research}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    agent: List[str]  # The person conducting the research
    subject: List[str]  # The subject being researched
    place: List[str]  # Where the event occurred


@dataclass
class TeachingTrainingLearning(CognitiveEvent):
    """{wikievents_eae_teaching_training_learning}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    agent: List[str]  # The person who teaches or trains
    patient: List[str]  # The person who is being taught or trained


@dataclass
class Attack(ConflictEvent):
    """{wikievents_eae_attack}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    agent: List[str]  # The entity attacking (the attacker)
    target: List[str]  # The entity being attacked (the victim)
    instrument: List[str]  # The object used to attack
    explosive_device: List[str]  # The explosive device used to attack (if any)
    place: List[str]  # Where the event occurred


@dataclass
class Defeat(ConflictEvent):
    """{wikievents_eae_defeat}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    victor: List[str]  # The entity defeating (the defeater)
    defeated: List[str]  # The entity being defeated (the defeated)
    place: List[str]  # Where the event occurred


@dataclass
class Demonstrate(ConflictEvent):
    """{wikievents_eae_demonstrate}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    demonstrator: List[str]  # The entity demonstrating
    regulator: List[str]  #  The entity regulating the demonstration
    topic: List[str]  # The topic of the demonstration
    target: List[str]  # The entity being demonstrated against


@dataclass
class Contact(ContactEvent):
    """{wikievents_eae_contact}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    participants: List[str]  # The entities in contact
    instrument: List[str]  # The object used to contact
    topic: List[str]  # The topic of the contact
    place: List[str]  # Where the event occurred


@dataclass
class RequestCommand(ContactEvent):
    """{wikievents_eae_request_command}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    communicator: List[str]  # The entity communicating
    receptor: List[str]  # The entity receiving the communication
    topic: List[str]  # The topic of the communication
    place: List[str]  # Where the event occurred


@dataclass
class ThreatenCoerce(ContactEvent):
    """{wikievents_eae_threaten_coerce}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    communicator: List[str]  # The entity threatening
    recipient: List[str]  # The entity being threatened
    place: List[str]  # Where the event occurred


@dataclass
class ImpedeInterfereWith(ControlEvent):
    """{wikievents_eae_impede_interfere_with}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    impeder: List[str]  # The entity impeding
    place: List[str]  # Where the event occurred


@dataclass
class Crash(DisasterEvent):
    """{wikievents_eae_crash}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    vehicle: List[str]  # The vehicle being driven
    object: List[str]  # The object crashed into
    place: List[str]  # Where the event occurred


@dataclass
class DiseaseOutbreak(DisasterEvent):
    """{wikievents_eae_disease_outbreak}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    disease: List[str]  # The disease being outbreak
    victim: List[str]  # The entity being infected
    place: List[str]  # Where the event occurred


@dataclass
class GenericCrime(GenericCrimeEvent):
    """{wikievents_eae_generic_crime}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    perpetrator: List[str]  # The entity perpetrating the crime
    victim: List[str]  # The entity being perpetrator
    place: List[str]  # Where the event occurred


@dataclass
class Acquit(JusticeEvent):
    """{wikievents_eae_acquit}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    defendant: List[str]  # The entity being acquitted


@dataclass
class ArrestJailDetain(JusticeEvent):
    """{wikievents_eae_arrest_jail_detain}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    jailer: List[str]  # The entity jailing
    detainee: List[str]  # The entity being detained
    place: List[str]  # Where the event occurred


@dataclass
class ChargeIndict(JusticeEvent):
    """{wikievents_eae_charge_indict}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    defendant: List[str]  # The entity being charged or indicted
    prosecutor: List[str]  # The entity prosecuting
    judge_court: List[str]  # The entity charging or indicting
    crime: List[str]  # The crime being charged
    place: List[str]  # Where the event occurred


@dataclass
class Convict(JusticeEvent):
    """{wikievents_eae_convict}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    judge_court: List[str]  # The entity convicting
    defendant: List[str]  # The entity being convicted
    place: List[str]  # Where the event occurred


@dataclass
class InvestigateCrime(JusticeEvent):
    """{wikievents_eae_investigate_crime}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    investigator: List[str]  # The entity investigating
    defendant: List[str]  # The entity being investigated
    crime: List[str]  # The crime being investigated
    place: List[str]  # Where the event occurred


@dataclass
class ReleaseParole(JusticeEvent):
    """{wikievents_eae_release_parole}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    judge_court: List[str]  # The entity releasing
    defendant: List[str]  # The entity being released
    place: List[str]  # Where the event occurred


@dataclass
class Sentence(JusticeEvent):
    """{wikievents_eae_sentence}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    judge_court: List[str]  # The entity sentence
    defendant: List[str]  # The entity being sentence
    place: List[str]  # Where the event occurred


@dataclass
class TrialHearing(JusticeEvent):
    """{wikievents_eae_trial_hearing}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    prosecutor: List[str]  # The entity prosecuting
    defendant: List[str]  # The entity being prosecuted
    judge_court: List[str]  # The entity trialing
    place: List[str]  # Where the event occurred


@dataclass
class Die(LifeEvent):
    """{wikievents_eae_die}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    victim: List[str]  # The entity being killed
    killer: List[str]  # The entity killing the victim
    place: List[str]  # Where the event occurred


@dataclass
class Infect(LifeEvent):
    """{wikievents_eae_infect}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    victim: List[str]  # The entity being infected


@dataclass
class Injure(LifeEvent):
    """{wikievents_eae_injure}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    victim: List[str]  # The entity being injured
    injurer: List[str]  #  The entity injuring the victim
    instrument: List[str]  # The instrument used to injure the victim
    body_part: List[str]  # The body part injured
    place: List[str]  # Where the event occurred


@dataclass
class Intervention(MedicalEvent):
    """{wikievents_eae_intervention}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    treater: List[str]  # The entity performing the intervention
    patient: List[str]  # The entity being intervened upon
    place: List[str]  # Where the event occurred


@dataclass
class Evacuation(MovementTransportEvent):
    """{wikievents_eae_evacuation}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    transporter: List[str]  # The entity transporting
    passenger_artifact: List[str]  # The entity being transported
    origin: List[str]  # The origin of the event
    destination: List[str]  # The destination of the event


@dataclass
class IllegalTransportation(MovementTransportEvent):
    """{wikievents_eae_illegal_transportation}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    transporter: List[str]  # The entity transporting
    passenger_artifact: List[str]  # The entity being transported
    vehicle: List[str]  # The vehicle used
    origin: List[str]  # The origin of the event
    destination: List[str]  # The destination of the event


@dataclass
class PreventPassage(MovementTransportEvent):
    """{wikievents_eae_prevent_passage}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    transporter: List[str]  # The entity transporting
    passenger_artifact: List[str]  # The entity being transported
    vehicle: List[str]  # The vehicle used
    preventer: List[str]  # The transportation preventer
    origin: List[str]  # The origin of the event
    destination: List[str]  # The destination of the event


@dataclass
class Transportation(MovementTransportEvent):
    """{wikievents_eae_transportation}"""

    mention: str  # The text span that most clearly expresses (triggers) the event
    transporter: List[str]  # The entity transporting
    passenger_artifact: List[str]  # The entity being transported
    vehicle: List[str]  # The vehicle used
    origin: List[str]  # The origin of the event
    destination: List[str]  # The destination of the event


@dataclass
class EndPosition(PersonnelEvent):
    """{wikievents_eae_end_position}"""

    mention: str  # The text span that most clearl expresses (triggers) the event
    employee: List[str]  # The person ending the job
    organization: List[str]  # Place of employment


@dataclass
class StartPosition(PersonnelEvent):
    """{wikievents_eae_start_position}"""

    mention: str  # The text span that most clearl expresses (triggers) the event
    employee: List[str]  # The person starting the job
    organization: List[str]  # Place of employment
    position: List[str]  # The position started
    place: List[str]  # Where the event occurred


@dataclass
class Donation(TransactionEvent):
    """{wikievents_eae_donation}"""

    mention: str  # The text span that most clearl expresses (triggers) the event
    giver: List[str]  # The entity donating
    recipient: List[str]  # The entity receiving the money
    artifact_money: List[str]  # The artifact or money given


@dataclass
class ExchangeBuySell(TransactionEvent):
    """{wikievents_eae_exchange_buy_sell}"""

    mention: str  # The text span that most clearl expresses (triggers) the event
    giver: List[str]  # The entity selling
    recipient: List[str]  # The entity buying
    adquired_entity: List[str]  # The entity being bought or sold
    payment_barter: List[str]  # The artifact or money given


EVENT_DEFINITIONS: List[Type] = [
    DamageDestroy,
    DisableDefuse,
    Dismantle,
    ManufactureAssemble,
    IdentifyCategorize,
    Inspection,
    Research,
    TeachingTrainingLearning,
    Attack,
    Defeat,
    Demonstrate,
    Contact,
    RequestCommand,
    ThreatenCoerce,
    ImpedeInterfereWith,
    Crash,
    DiseaseOutbreak,
    GenericCrime,
    Acquit,
    ArrestJailDetain,
    ChargeIndict,
    Convict,
    InvestigateCrime,
    ReleaseParole,
    Sentence,
    TrialHearing,
    Die,
    Infect,
    Injure,
    Intervention,
    Evacuation,
    IllegalTransportation,
    PreventPassage,
    Transportation,
    EndPosition,
    StartPosition,
    Donation,
    ExchangeBuySell,
]

FINE_TO_COARSE_EVENTS: Dict[Type, Type] = {_def: _def.__base__ for _def in EVENT_DEFINITIONS}
COARSE_TO_FINE_EVENTS: Dict[Type, List[Type]] = {}
for fine, coarse in FINE_TO_COARSE_EVENTS.items():
    if coarse not in COARSE_TO_FINE_EVENTS:
        COARSE_TO_FINE_EVENTS[coarse] = []
    COARSE_TO_FINE_EVENTS[coarse].append(fine)

src/tasks/wikievents/prompts.py
hitz-zentroa/GoLLIE
from typing import Dict, List, Type

from src.tasks.utils_scorer import EventScorer, SpanScorer
from src.tasks.utils_typing import Entity
from src.tasks.wikievents.prompts import (
    COARSE_EVENT_DEFINITIONS,
    ENTITY_DEFINITIONS,
    EVENT_DEFINITIONS,
)


class WikiEventsEntityScorer(SpanScorer):
    """WikiEvents Entity identification and classification scorer."""

    valid_types: List[Type] = ENTITY_DEFINITIONS

    def __call__(self, reference: List[Entity], predictions: List[Entity]) -> Dict[str, Dict[str, float]]:
        output = super().__call__(reference, predictions)
        return {"entities": output["spans"]}


class WikiEventsEventScorer(EventScorer):
    """WikiEvents Event and argument classification scorer."""

    valid_types: List[Type] = COARSE_EVENT_DEFINITIONS


class WikiEventsEventArgumentScorer(EventScorer):
    """WikiEvents Event and argument classification scorer."""

    valid_types: List[Type] = EVENT_DEFINITIONS

src/tasks/wikievents/scorer.py
hitz-zentroa/GoLLIE
from . import data_loader, prompts, scorer


__all__ = ["data_loader", "prompts", "scorer"]

src/tasks/wnut/__init__.py
hitz-zentroa/GoLLIE
from typing import Dict, List, Tuple, Type, Union

from src.tasks.label_encoding import rewrite_labels
from src.tasks.wnut.guidelines import GUIDELINES
from src.tasks.wnut.guidelines_gold import EXAMPLES
from src.tasks.wnut.prompts import ENTITY_DEFINITIONS, Corporation, CreativeWork, Group, Location, Person, Product

from ..utils_data import DatasetLoader, Sampler
from ..utils_typing import Entity


def get_bwnut_hf(
    split: str,
    ENTITY_TO_CLASS_MAPPING: Dict[str, Type[Entity]],
) -> Tuple[List[List[str]], List[List[Entity]]]:
    """
    Get the Wnut dataset from the huggingface datasets library
    Args:
        split (str): The path_or_split to load. Can be one of `train`, `validation` or `test`.
    Returns:
        (List[str],List[Entity]): The text and the entities
    """
    from datasets import load_dataset

    dataset = load_dataset("wnut_17")
    # Dirty fix to prevent errors reading the labels
    label_names = dataset["train"].features["ner_tags"].feature.names
    i = label_names.index("B-creative-work")
    label_names[i] = "B-creative_work"
    i = label_names.index("I-creative-work")
    label_names[i] = "I-creative_work"

    id2label = dict(enumerate(label_names))
    dataset_sentences: List[List[str]] = []
    dataset_entities: List[List[Entity]] = []

    for example in dataset[split]:
        words = example["tokens"]
        # Ensure IOB2 encoding
        labels = rewrite_labels(labels=[id2label[label] for label in example["ner_tags"]], encoding="iob2")

        # Get labeled word spans
        spans = []
        for i, label in enumerate(labels):
            if label == "O":
                continue
            elif label.startswith("B-"):
                spans.append([label[2:], i, i + 1])
            elif label.startswith("I-"):
                spans[-1][2] += 1
            else:
                raise ValueError(f"Found an unexpected label: {label}")

        # Get entities
        entities = []
        for label, start, end in spans:
            entities.append(ENTITY_TO_CLASS_MAPPING[label](span=" ".join(words[start:end])))

        dataset_sentences.append(words)
        dataset_entities.append(entities)

    return dataset_sentences, dataset_entities


class WnutDatasetLoader(DatasetLoader):
    """
    A `DatasetLoader` for the Wnut dataset.

    Args:
        split (`str`):
            The split to load. Can be one of `train`, `validation` or `test`.

    Raises:
        `ValueError`:
            raised when a not defined value found.
    """

    ENTITY_TO_CLASS_MAPPING = None

    def __init__(self, path_or_split: str, **kwargs) -> None:
        self.ENTITY_TO_CLASS_MAPPING = {
            "corporation": Corporation,
            "creative_work": CreativeWork,
            "group": Group,
            "location": Location,
            "person": Person,
            "product": Product,
        }

        self.elements = {}

        dataset_words, dataset_entities = get_bwnut_hf(
            split=path_or_split,
            ENTITY_TO_CLASS_MAPPING=self.ENTITY_TO_CLASS_MAPPING,
        )

        for id, (words, entities) in enumerate(zip(dataset_words, dataset_entities)):
            self.elements[id] = {
                "id": id,
                "doc_id": id,
                "text": " ".join(words),
                "entities": entities,
                "gold": entities,
            }


class WnutSampler(Sampler):
    """
    A data `Sampler` for the Wnut dataset.

    Args:
        dataset_loader (`WnutDatasetLoader`):
            The dataset loader that contains the data information.
        task (`str`, optional):
            The task to sample. It must be one of the following: NER, VER, RE, EE.
            Defaults to `None`.
        split (`str`, optional):
            The path_or_split to sample. It must be one of the following: "train", "dev" or
            "test". Depending on the path_or_split the sampling strategy differs. Defaults to
            `"train"`.
        parallel_instances (`Union[int, Tuple[int, int]]`, optional):
            The number of sentences sampled in parallel. Options:

                * **`int`**: The amount of elements that will be sampled in parallel.
                * **`tuple`**: The range of elements that will be sampled in parallel.

            Defaults to 1.
        max_guidelines (`int`, optional):
            The number of guidelines to append to the example at the same time. If `-1`
            is given then all the guidelines are appended. Defaults to `-1`.
        guideline_dropout (`float`, optional):
            The probability to dropout a guideline definition for the given example. This
            is only applied on training. Defaults to `0.0`.
        seed (`float`, optional):
            The seed to sample the examples. Defaults to `0`.
        prompt_template (`str`, optional):
            The path to the prompt template. Defaults to `"templates/prompt.txt"`.
        ensure_positives_on_train (bool, optional):
            Whether to ensure that the guidelines of annotated examples are not removed.
            Defaults to `False`.
        dataset_name (str, optional):
            The name of the dataset. Defaults to `None`.
        scorer (`str`, optional):
           The scorer class import string. Defaults to `None`.
        sample_only_gold_guidelines (`bool`, optional):
            Whether to sample only guidelines of present annotations. Defaults to `False`.
    """

    def __init__(
        self,
        dataset_loader: WnutDatasetLoader,
        task: str = None,
        split: str = "train",
        parallel_instances: Union[int, Tuple[int, int]] = 1,
        max_guidelines: int = -1,
        guideline_dropout: float = 0.0,
        seed: float = 0,
        prompt_template: str = "templates/prompt.txt",
        ensure_positives_on_train: bool = False,
        dataset_name: str = None,
        scorer: str = None,
        sample_only_gold_guidelines: bool = False,
        **kwargs,
    ) -> None:
        assert task in [
            "NER",
        ], f"Wnut only supports NER task. {task} is not supported."

        task_definitions, task_target = {
            "NER": (ENTITY_DEFINITIONS, "entities"),
        }[task]

        super().__init__(
            dataset_loader=dataset_loader,
            task=task,
            split=split,
            parallel_instances=parallel_instances,
            max_guidelines=max_guidelines,
            guideline_dropout=guideline_dropout,
            seed=seed,
            prompt_template=prompt_template,
            ensure_positives_on_train=ensure_positives_on_train,
            sample_only_gold_guidelines=sample_only_gold_guidelines,
            dataset_name=dataset_name,
            scorer=scorer,
            task_definitions=task_definitions,
            task_target=task_target,
            definitions=GUIDELINES,
            examples=EXAMPLES,
            **kwargs,
        )

src/tasks/wnut/data_loader.py
hitz-zentroa/GoLLIE
GUIDELINES = {
    "wnut_person": {
        "en": [
            (
                "Names of people (e.g. Virginia Wade). Don\u2019t mark people that don\u2019t have their own name."
                " Include punctuation in the \nmiddle of names. Fictional people can be included, as long as"
                " they\u2019re referred to by name (e.g. Harry Potter)."
            ),
            (
                "For example, list the names of individuals (such as Virginia Wade). Do not include people who do not"
                " have their own \ndistinct name. Ensure that punctuation is placed within the names. Feel free to"
                " incorporate fictional characters, as long as \nthey are identified by their name (like Harry"
                " Potter)."
            ),
            (
                "List the names of individuals (for example, Virginia Wade). Do not include individuals who do not"
                " have their own \nunique name. Ensure that punctuation is present within the middle of names. Feel"
                " free to incorporate fictional \ncharacters, as long as they are identified by their name (for"
                " example, Harry Potter)."
            ),
            (
                "For example, list the names of individuals such as Virginia Wade. Exclude people who aren't"
                " specifically \nidentified by their own name. Incorporate punctuation within the names of"
                " individuals. Feel free to include fictional \ncharacters like Harry Potter, as long as they are"
                " referred to by their proper name."
            ),
            (
                "For example, list the names of individuals such as Virginia Wade. Exclude individuals who aren't"
                " specifically \nnamed. Incorporate punctuation within the names of people. You may also include"
                " fictional characters like Harry Potter, \nas long as they are identified by their names."
            ),
        ]
    },
    "wnut_location": {
        "en": [
            (
                "Names that are locations (e.g. France). Don\u2019t mark locations that don\u2019t have their own"
                " name. Include punctuation in \nthe mid\u0002dle of names. Fictional locations can be included, as"
                " long as they\u2019re referred to by name (e.g. Hog\u0002warts)."
            ),
            (
                "Names that are locations, such as France, should be included. Exclude locations that lack their own"
                " name. Ensure \nthat punctuation is part of the middle of names. Fictional locations like Hogwarts,"
                " if referred to by name, can also be \nincluded."
            ),
            (
                "Geographical designations (for example, France). Do not indicate places that lack a distinct name."
                " Incorporate \npunctuation within the center of names. Fictional locations can be incorporated, as"
                " long as they are mentioned by name (e.g. \nHogwarts)."
            ),
            (
                "Names that denote specific locations (e.g., France). Do not mark locations that lack their own name."
                " Ensure that \npunctuation is included in the middle of names. Fictional locations can be"
                " incorporated, as long as they are referred to by name \n(e.g., Hogwarts)."
            ),
            (
                "Names associated with specific locations (for example, France). Do not mark locations that do not"
                " have their own \ndesignated name. Ensure that punctuation is included in the middle of such names."
                " Feel free to incorporate fictional locations \nas well, as long as they are referenced by name"
                " (e.g., Hogwarts)."
            ),
        ]
    },
    "wnut_corporation": {
        "en": [
            (
                "Names of corporations (e.g. Google). Don\u2019t mark locations that don\u2019t have their own name."
                " Include punctuation in \nthe middle of names."
            ),
            (
                "List the names of corporations (for example, Google). Do not highlight locations that lack their own"
                " specific name. \nEnsure that punctuation is included in the middle of the names."
            ),
            (
                "List the names of corporations (for example, Google). Do not indicate locations that do not have"
                " their own specific \nname. Ensure that punctuation is included in the middle of the company names."
            ),
            (
                "List the names of corporations (for example, Google). Do not highlight locations that do not have"
                " their own distinct \nname. Ensure that punctuation is properly included within the middle of names."
            ),
            (
                'Rephrased Text: "List the names of corporations (for example, Google). Do not mark locations that'
                " lack their own \nspecific name. Ensure that punctuation is included in the middle of the corporation"
                " names."
            ),
        ]
    },
    "wnut_product": {
        "en": [
            (
                "Name of products (e.g. iPhone). Don\u2019t mark products that don\u2019t have their own name. Include"
                " punctuation in the \nmiddle of names. Fictional products can be included, as long as they\u2019re"
                " referred to by name (e.g. Everlasting \nGob\u0002stopper). It\u2019s got to be something you can"
                " touch, and it\u2019s got to be the official name."
            ),
            (
                "Specify the names of items (for example, iPhone). Don't enumerate products that lack distinct names."
                " Integrate \npunctuation within the midst of names. Include imaginary products, provided they are"
                " denoted by name (for instance, \nEverlasting Gob\u0002stopper). It must be a tangible object and"
                " encompass the official title."
            ),
            (
                "Entries should consist of specific product names (for example, iPhone). Exclude items that lack their"
                " own \ndesignated name. Incorporate appropriate punctuation within product names. Feel free to"
                " include fictitious items, as long as \nthey have an identifiable name (e.g., Everlasting"
                " Gobstopper). The focus should be on tangible items with official \nnames."
            ),
            (
                "Provide the names of items (for example, iPhone). Don't include products that don't have distinct"
                " names. \nIncorporate punctuation within the names. Feel free to include fictional products, as long"
                " as they have designated names (for \nexample, Everlasting Gob\u0002stopper). The objective is to"
                " identify tangible items with official names."
            ),
            (
                "Specify the names of items (for example, iPhone). Do not mention products that lack distinct names."
                " Incorporate \npunctuation within the names of products. Fictional products can be included, as long"
                " as they are identified by their respective \nnames (for example, Everlasting Gob\u0002stopper). The"
                " items must be tangible, and the names provided should be the official \nones."
            ),
        ]
    },
    "wnut_creativework": {
        "en": [
            (
                "Names of creative works (e.g. Bohemian Rhapsody). Include punctuation in the middle of names. The"
                " work should be \ncreated by a human, and referred to by its specific name."
            ),
            (
                "Title examples for imaginative productions (such as Bohemian Rhapsody). Incorporate punctuation"
                " within the \nnames. The piece should be crafted by a person, and identified by its distinctive name."
            ),
            (
                "Cite examples of creative works, such as 'Bohemian Rhapsody.' Maintain punctuation within the titles."
                " These works \nmust be produced by individuals and identified by their specific names."
            ),
            (
                "List the titles of original creations (for example, Bohemian Rhapsody). Incorporate punctuation"
                " within the names \nof these works. These works should be produced by human authors and identified by"
                " their distinctive names."
            ),
            (
                "Example titles of artistic productions (such as Bohemian Rhapsody). Incorporate punctuation within"
                " the titles. \nThese works should be crafted by individuals and identified by their unique names."
            ),
        ]
    },
    "wnut_group": {
        "en": [
            (
                "Names of groups (e.g. Nirvana, San Diego Padres). Don\u2019t mark groups that don\u2019t have a"
                " specific, unique name, or \ncompanies (which should be marked corporation)."
            ),
            (
                "Example categories include bands like Nirvana and sports teams such as the San Diego Padres. Exclude"
                " groups without \na distinct, unique name or corporations, which should be designated as"
                " 'corporation'."
            ),
            (
                "Example group names include Nirvana and the San Diego Padres. Exclude any groups without a distinct,"
                " unique name, as \nwell as companies, which should be labeled as corporations."
            ),
            (
                "For example, list the names of bands (like Nirvana) and sports teams (such as the San Diego Padres)."
                " However, exclude \nany groups without a distinct, unique name, as well as businesses, which should"
                " be labeled as corporations."
            ),
            (
                "For example, list the names of groups (such as Nirvana or San Diego Padres). Exclude any groups"
                " without a distinct, \nunique name or companies (which should be labeled as corporations)."
            ),
        ]
    },
}

src/tasks/wnut/guidelines.py
hitz-zentroa/GoLLIE
GUIDELINES = {
    "wnut_person": {
        "en": [
            "Names of people (e.g. Virginia Wade). Don’t mark people that don’t have their own name. Include"
            " punctuation in the middle of names. Fictional people can be included, as long as they’re referred to by"
            " name (e.g. Harry Potter)."
        ]
    },
    "wnut_location": {
        "en": [
            "Names that are locations (e.g. France). Don’t mark locations that don’t have their own name. Include"
            " punctuation in the middle of names. Fictional locations can be included, as long as they’re referred to"
            " by name (e.g. Hogwarts)."
        ]
    },
    "wnut_corporation": {
        "en": [
            "Names of corporations (e.g. Google). Don’t mark locations that don’t have their own name. Include"
            " punctuation in the middle of names."
        ]
    },
    "wnut_product": {
        "en": [
            "Name of products (e.g. iPhone). Don’t mark products that don’t have their own name. Include punctuation"
            " in the middle of names. Fictional products can be included, as long as they’re referred to by name (e.g."
            " Everlasting Gobstopper). It’s got to be something you can touch, and it’s got to be the official name."
        ]
    },
    "wnut_creativework": {
        "en": [
            "Names of creative works (e.g. Bohemian Rhapsody). Include punctuation in the middle of names. The work"
            " should be created by a human, and referred to by its specific name."
        ]
    },
    "wnut_group": {
        "en": [
            "Names of groups (e.g. Nirvana, San Diego Padres). Don’t mark groups that don’t have a specific, unique"
            " name, or companies (which should be marked corporation)."
        ]
    },
}

EXAMPLES = {
    "wnut_corporation_examples": {
        "en": [
            "Twitter",
            "twitter",
            "YouTube",
            "facebook",
            "Facebook",
            "Ustream",
            "Playboy",
            "Pepsi",
            "Yahoo",
            "iTunes",
        ]
    },
    "wnut_creativework_examples": {
        "en": [
            "The Town",
            "TheDeAndreWay",
            "The DeAndre Way",
            "Winter 's Bone",
            "Friday Night Lights",
            "Gigli",
            "True Blood",
            "Private Practice",
            "Danger Days : The True Lives Of The Fabulous Killjoys",
            "Indiesent Exposure",
        ]
    },
    "wnut_group_examples": {
        "en": [
            "Eagles",
            "FDA",
            "Raw Illest",
            "Storm",
            "Cowboys",
            "Ohio State",
            "Eyeopener",
            "JV soccer",
            "sharks",
            "30stm",
        ]
    },
    "wnut_location_examples": {
        "en": ["Chicago", "GA", "NYC", "UK", "London", "DC", "NY", "MD", "Hawaii", "VISIONS LOUNGE"]
    },
    "wnut_person_examples": {
        "en": [
            "Justin Bieber",
            "God",
            "Lindsay Lohan",
            "Pope",
            "Justin",
            "Scooter Braun",
            "JFK",
            "Soulja Boy",
            "Allah",
            "justin bieber",
        ]
    },
    "wnut_product_examples": {
        "en": [
            "iPad",
            "iPhone",
            "HTC Desire Z",
            "Coke",
            "Xbox 360",
            "Pumpkin Moonshine",
            "Club Penguin",
            "Flip MinoHD",
            "vodka",
            "ipod",
        ]
    },
}

src/tasks/wnut/guidelines_gold.py
hitz-zentroa/GoLLIE
from typing import List

from ..utils_typing import Entity, dataclass


"""Entity definitions

The entity definitions are derived from the official WNUT17 guidelines:
https://aclanthology.org/W17-4418.pdf


"""


@dataclass
class Person(Entity):
    """{wnut_person}"""

    span: str  # {wnut_person_examples}


@dataclass
class Location(Entity):
    """{wnut_location}"""

    span: str  # {wnut_location_examples}


@dataclass
class Corporation(Entity):
    """{wnut_corporation}"""

    span: str  # {wnut_corporation_examples}


@dataclass
class Product(Entity):
    """{wnut_product}"""

    span: str  # {wnut_product_examples}


@dataclass
class CreativeWork(Entity):
    """{wnut_creativework}"""

    span: str  # {wnut_creativework_examples}


@dataclass
class Group(Entity):
    """{wnut_group}"""

    span: str  # {wnut_group_examples}


ENTITY_DEFINITIONS: List[Entity] = [Person, Location, Corporation, Product, CreativeWork, Group]

# __all__ = list(map(str, [*ENTITY_DEFINITIONS]))

src/tasks/wnut/prompts.py
hitz-zentroa/GoLLIE
from typing import Dict, List, Type

from src.tasks.utils_scorer import SpanScorer
from src.tasks.utils_typing import Entity
from src.tasks.wnut.prompts import ENTITY_DEFINITIONS


class WnutEntityScorer(SpanScorer):
    """CoNLL03 Entity identification and classification scorer."""

    valid_types: List[Type] = ENTITY_DEFINITIONS

    def __call__(self, reference: List[Entity], predictions: List[Entity]) -> Dict[str, Dict[str, float]]:
        output = super().__call__(reference, predictions)
        return {"entities": output["spans"]}

src/tasks/wnut/scorer.py
hitz-zentroa/GoLLIE
import json
import os
import unittest


class TestDataLoaders(unittest.TestCase):
    @unittest.skipIf(not os.path.exists("data/ace05/english.sentence.json"), "No ACE data available")
    def test_ACE(self):
        from src.tasks.ace.data_loader import ACEDatasetLoader, ACESampler

        with open("configs/data_configs/ace_config.json") as f:
            config = json.load(f)
        if isinstance(config["seed"], list):
            config["seed"] = 0
            config["label_noise_prob"] = 0.0

        dataloader = ACEDatasetLoader("data/ace05/english.sentence.json", group_by="sentence")

        _ = list(ACESampler(dataloader, task="RE", **config, **config["task_configuration"]["RE"]))

        # TODO: Implement a better TEST

    @unittest.skipIf(not os.path.exists("data/casie/data.jsonl"), "No CASIE data available")
    def test_CASIE(self):
        from src.tasks.casie.data_loader import CASIEDatasetLoader, CASIESampler

        with open("configs/data_configs/casie_config.json") as f:
            config = json.load(f)

        dataloader = CASIEDatasetLoader("data/casie/data.dev.jsonl")

        _ = list(
            CASIESampler(
                dataloader, task="EAE", remove_guidelines=True, **config, **config["task_configuration"]["EE"]
            )
        )

    @unittest.skipIf(not os.path.exists("data/wikievents/train.sentence.jsonl"), "No WikiEvents data available")
    def test_WikiEvents(self):
        from src.tasks.wikievents.data_loader import WikiEventsDatasetLoader, WikiEventsSampler

        with open("configs/data_configs/wikievents_config.json") as f:
            config = json.load(f)
        if isinstance(config["seed"], list):
            config["seed"] = 0
            config["label_noise_prob"] = 0.0

        dataloader = WikiEventsDatasetLoader("data/wikievents/train.sentence.jsonl", group_by="sentence")

        _ = list(WikiEventsSampler(dataloader, task="EAE", **config, **config["task_configuration"]["EAE"]))

        # TODO: Implement a better TEST

    @unittest.skipIf(not os.path.exists("data/rams/dev.jsonlines"), "No RAMS data available")
    def test_RAMS(self):
        from src.tasks.rams.data_loader import RAMSDatasetLoader, RAMSSampler

        with open("configs/data_configs/rams_config.json") as f:
            config = json.load(f)
        if isinstance(config["seed"], list):
            config["seed"] = 0
            config["label_noise_prob"] = 0.0

        dataloader = RAMSDatasetLoader("data/rams/dev.jsonlines")

        _ = list(RAMSSampler(dataloader, task="EAE", **config, **config["task_configuration"]["EAE"]))

        # TODO: Implement a better TEST

    @unittest.skipIf(not os.path.exists("data/tacred/dev.json"), "No TACRED data available")
    def test_TACRED(self):
        from src.tasks.tacred.data_loader import TACREDDatasetLoader, TACREDSampler

        with open("configs/data_configs/tacred_config.json") as f:
            config = json.load(f)
        if isinstance(config["seed"], list):
            config["seed"] = 0
            config["label_noise_prob"] = 0.0

        dataloader = TACREDDatasetLoader("data/tacred/train.json")[:10]

        _ = list(TACREDSampler(dataloader, task="SF", **config, **config["task_configuration"]["SF"]))

    @unittest.skipIf(not os.path.exists("data/conll/en.conll.train.tsv"), "No CoNLL data available")
    def test_CoNLL03(self):
        from src.tasks.conll03.data_loader import CoNLL03Sampler, CoNLLDatasetLoader
        from src.tasks.conll03.prompts import Miscellaneous, Organization, Person

        with open("configs/data_configs/conll03_config.json") as f:
            config = json.load(f)
        if isinstance(config["seed"], list):
            config["seed"] = 0
            config["label_noise_prob"] = 0.0

        config["task_configuration"] = {
            "NER": {
                "parallel_instances": 1,
                "max_guidelines": -1,
                "guideline_dropout": 0,
                "scorer": "src.tasks.conll03.scorer.CoNLL03EntityScorer",
            }
        }

        config["include_misc"] = True
        dataloader = CoNLLDatasetLoader("validation", **config)

        _ = list(CoNLL03Sampler(dataloader, task="NER", **config, **config["task_configuration"]["NER"]))

        self.assertDictEqual(
            {
                "id": 0,
                "doc_id": 0,
                "text": "CRICKET - LEICESTERSHIRE TAKE OVER AT TOP AFTER INNINGS VICTORY .",
                "entities": [Organization(span="LEICESTERSHIRE")],
                "gold": [Organization(span="LEICESTERSHIRE")],
            },
            dataloader[0],
        )

        self.assertDictEqual(
            {
                "id": 2,
                "doc_id": 2,
                "text": (
                    "West Indian all-rounder Phil Simmons took four for 38 on Friday as Leicestershire beat Somerset"
                    " by an innings and 39 runs in two days to take over at the head of the county championship ."
                ),
                "entities": [
                    Miscellaneous(span="West Indian"),
                    Person(span="Phil Simmons"),
                    Organization(span="Leicestershire"),
                    Organization(span="Somerset"),
                ],
                "gold": [
                    Miscellaneous(span="West Indian"),
                    Person(span="Phil Simmons"),
                    Organization(span="Leicestershire"),
                    Organization(span="Somerset"),
                ],
            },
            dataloader[2],
        )

        config["include_misc"] = False
        dataloader = CoNLLDatasetLoader("validation", **config)

        _ = list(CoNLL03Sampler(dataloader, task="NER", **config, **config["task_configuration"]["NER"]))

        self.assertDictEqual(
            {
                "id": 0,
                "doc_id": 0,
                "text": "CRICKET - LEICESTERSHIRE TAKE OVER AT TOP AFTER INNINGS VICTORY .",
                "entities": [Organization(span="LEICESTERSHIRE")],
                "gold": [Organization(span="LEICESTERSHIRE")],
            },
            dataloader[0],
        )

        self.assertDictEqual(
            {
                "id": 2,
                "doc_id": 2,
                "text": (
                    "West Indian all-rounder Phil Simmons took four for 38 on Friday as Leicestershire beat Somerset"
                    " by an innings and 39 runs in two days to take over at the head of the county championship ."
                ),
                "entities": [
                    Person(span="Phil Simmons"),
                    Organization(span="Leicestershire"),
                    Organization(span="Somerset"),
                ],
                "gold": [
                    Person(span="Phil Simmons"),
                    Organization(span="Leicestershire"),
                    Organization(span="Somerset"),
                ],
            },
            dataloader[2],
        )

src/tests/test_dataloaders.py
hitz-zentroa/GoLLIE
import json
import os
import random
import tempfile
import unittest

from src.dataset.dataset import CollieDataset
from transformers import PreTrainedTokenizerBase


def get_dataset(
    tokenizer: PreTrainedTokenizerBase,
    is_encoder_decoder: bool,
    inference: bool,
    prompt_loss_weight: float,
    num_epochs: int = -1,
    max_length: int = 2048,
) -> (CollieDataset, str, str):
    text = """@dataclass
class EnergyAndInfrastructureEvent:
    \"\"\"This class is used to instantiate events that involve Chinese energy and infrastructure projects.\"\"\"
    meeting_attendees: Union[List[str], None] # Persons or organizations that attended the meeting.
    meeting_location: Union[List[str], None] # Location where the meeting happened.
    meeting_topic: Union[List[str], None] # Topic discussed on the meeting
    project_location: Union[List[str], None] # Location of the project
    project_name: Union[List[str], None] # Name of the project

# This is the sentence to analyze
sentence = "The Chinese and Rongovian delegations met at the sidelines of the Berlin Development Futures conference to discuss Rongovia's proposed Pangean Reunification Facility.

# The following list contains the events instances that happens in the sentence defined above
result = [
    EnergyAndInfrastructureEvent(
        meeting_attendees=["Chinese", "Rongovian"],
        meeting_location=["Berlin"],
        meeting_topic=["Pangean Reunification Facility"],
        project_location=["Rongovia"],
        project_name=["Pangean Reunification Facility"]
    ),
]"""

    prompt = """@dataclass
class EnergyAndInfrastructureEvent:
    \"\"\"This class is used to instantiate events that involve Chinese energy and infrastructure projects.\"\"\"
    meeting_attendees: Union[List[str], None] # Persons or organizations that attended the meeting.
    meeting_location: Union[List[str], None] # Location where the meeting happened.
    meeting_topic: Union[List[str], None] # Topic discussed on the meeting
    project_location: Union[List[str], None] # Location of the project
    project_name: Union[List[str], None] # Name of the project

# This is the sentence to analyze
sentence = "The Chinese and Rongovian delegations met at the sidelines of the Berlin Development Futures conference to discuss Rongovia's proposed Pangean Reunification Facility.

# The following list contains the events instances that happens in the sentence defined above
result ="""
    result = """[
    EnergyAndInfrastructureEvent(
        meeting_attendees=["Chinese", "Rongovian"],
        meeting_location=["Berlin"],
        meeting_topic=["Pangean Reunification Facility"],
        project_location=["Rongovia"],
        project_name=["Pangean Reunification Facility"]
    ),
]"""
    if num_epochs == -1:
        with tempfile.TemporaryDirectory() as tmpdirname:
            with open(os.path.join(tmpdirname, "tmp.ee.train.jsonl"), "w", encoding="utf8") as f:
                print(json.dumps({"text": text}, ensure_ascii=False), file=f)

            dataset = CollieDataset(
                tokenizer=tokenizer,
                dataset_path=os.path.join(tmpdirname, "tmp.ee.train.jsonl"),
                is_encoder_decoder=is_encoder_decoder,
                max_length=max_length,
                inference=inference,
                prompt_loss_weight=prompt_loss_weight,
                num_workers=1,
            )

    else:
        # List of random integers with len = num_epochs
        random_seeds = random.sample(range(0, 100000), num_epochs)
        with tempfile.TemporaryDirectory() as tmpdirname:
            for epoch in random_seeds:
                with open(os.path.join(tmpdirname, f"tmp.ee.train.{epoch}.jsonl"), "w", encoding="utf8") as f:
                    print(json.dumps({"text": text}, ensure_ascii=False), file=f)

            dataset = CollieDataset(
                tokenizer=tokenizer,
                dataset_path=os.path.join(tmpdirname, "tmp.ee.train.jsonl"),
                is_encoder_decoder=is_encoder_decoder,
                max_length=max_length,
                inference=inference,
                prompt_loss_weight=prompt_loss_weight,
                num_workers=1,
            )

    return dataset, prompt, result


class TestCollieDataset(unittest.TestCase):
    def test_add_eos(self):
        from transformers import AutoTokenizer

        tokenizer = AutoTokenizer.from_pretrained(
            (
                "/gaueko1/hizkuntza-ereduak/LLaMA/lm/huggingface/7B/"
                if os.path.exists("/gaueko1/hizkuntza-ereduak/LLaMA/lm/huggingface/7B/")
                else "EleutherAI/gpt-neo-125m"
            ),
            add_eos_token=True,
            use_fast=True,
        )

        simple_sentence = "This is a sentence to test if the tokenizer adds eos token."
        simple_sentence_ids = tokenizer(simple_sentence, add_special_tokens=True)
        if simple_sentence_ids["input_ids"][-1] != tokenizer.eos_token_id:
            simple_sentence_ids["input_ids"].append(tokenizer.eos_token_id)
            simple_sentence_ids["attention_mask"].append(1)
            # print(simple_sentence_ids)

        self.assertEqual(
            tokenizer.decode(simple_sentence_ids.input_ids, skip_special_tokens=True),
            simple_sentence,
        )

        self.assertEqual(simple_sentence_ids.input_ids[-1], tokenizer.eos_token_id)

        if tokenizer.pad_token_id is None:
            tokenizer.pad_token_id = tokenizer.unk_token_id

        dataset, _, _ = get_dataset(
            tokenizer=tokenizer,
            is_encoder_decoder=False,
            inference=False,
            prompt_loss_weight=0.05,
        )

        # Check if every instance `input_ids` has `eos_token_id`
        self.assertTrue(
            all(inst.input_ids[-1] == tokenizer.eos_token_id for inst in dataset),
            "There are `input_ids` without `eos_token_ids` at the end.",
        )

        # Check if every instance labels has `eos_token_id`
        self.assertTrue(
            all(inst.labels[-1] == tokenizer.eos_token_id for inst in dataset),
            "There are `labels` without `eos_token_ids` at the end.",
        )

        # Check that at inference we don't have eos token
        dataset, _, _ = get_dataset(
            tokenizer=tokenizer,
            is_encoder_decoder=False,
            inference=True,
            prompt_loss_weight=0.05,
        )

        # Check if every instance `input_ids` does not have `eos_token_id`
        self.assertTrue(
            all(inst.input_ids[-1] != tokenizer.eos_token_id for inst in dataset),
            "There are `input_ids` without `eos_token_ids` at the end.",
        )

    def test_inference_token_ids(self):
        from transformers import AutoTokenizer

        tokenizer = AutoTokenizer.from_pretrained(
            (
                "/gaueko1/hizkuntza-ereduak/LLaMA/lm/huggingface/7B/"
                if os.path.exists("/gaueko1/hizkuntza-ereduak/LLaMA/lm/huggingface/7B/")
                else "EleutherAI/gpt-neo-125m"
            ),
            add_eos_token=True,
        )

        tokenizer.padding_side = "left"

        if tokenizer.pad_token_id is None:
            tokenizer.pad_token_id = tokenizer.unk_token_id

        text1 = """@dataclass
        class EnergyAndInfrastructureEvent:
            \"\"\"This class is used to instantiate events that involve Chinese energy and infrastructure projects.\"\"\"
            meeting_attendees: Union[List[str], None] # Persons or organizations that attended the meeting.
            meeting_location: Union[List[str], None] # Location where the meeting happened.
            meeting_topic: Union[List[str], None] # Topic discussed on the meeting
            project_location: Union[List[str], None] # Location of the project
            project_name: Union[List[str], None] # Name of the project

        # This is the sentence to analyze
        sentence = "The Chinese and Rongovian delegations met at the sidelines of the Berlin Development Futures conference to discuss Rongovia's proposed Pangean Reunification Facility.

        # The following list contains the events instances that happens in the sentence defined above
        result = [
            EnergyAndInfrastructureEvent(
                meeting_attendees=["Chinese", "Rongovian"],
                meeting_location=["Berlin"],
                meeting_topic=["Pangean Reunification Facility"],
                project_location=["Rongovia"],
                project_name=["Pangean Reunification Facility"]
            ),
        ]"""

        text2 = """@dataclass
        class EnergyAndInfrastructureEvent:
            \"\"\"This class is used to instantiate events that involve Chinese energy and infrastructure projects.\"\"\"
            meeting_attendees: Union[List[str], None] # Persons or organizations that attended the meeting.
            meeting_location: Union[List[str], None] # Location where the meeting happened.
            meeting_topic: Union[List[str], None] # Topic discussed on the meeting
            project_location: Union[List[str], None] # Location of the project
            project_name: Union[List[str], None] # Name of the project

        # This is the sentence to analyze
        sentence = "The Spanish and French delegations met at the sidelines of the Berlin Development Futures conference to discuss Rongovia's proposed Pangean Reunification Facility.

        # The following list contains the events instances that happens in the sentence defined above
        result = []"""

        with tempfile.TemporaryDirectory() as tmpdirname:
            with open(os.path.join(tmpdirname, "tmp.ee.train.jsonl"), "w", encoding="utf8") as f:
                print(json.dumps({"text": text2}, ensure_ascii=False), file=f)
                print(json.dumps({"text": text1}, ensure_ascii=False), file=f)

            dataset = CollieDataset(
                tokenizer=tokenizer,
                dataset_path=os.path.join(tmpdirname, "tmp.ee.train.jsonl"),
                is_encoder_decoder=False,
                max_length=2048,
                inference=True,
                prompt_loss_weight=0.05,
            )

            self.assertEqual(len(dataset), 2)

            # Check if every instance of `input_ids` end with the same token, so at inference the first token
            # of the prompt is the same for all instances
            self.assertTrue(
                all(inst.input_ids[-1] == dataset[0].input_ids[-1] for inst in dataset),
                "The last token of the `input_ids` is not the same for all instances at inference.",
            )

    def test_encoder(self):
        from transformers import AutoTokenizer

        tokenizer = AutoTokenizer.from_pretrained(
            (
                "/gaueko1/hizkuntza-ereduak/LLaMA/lm/huggingface/7B/"
                if os.path.exists("/gaueko1/hizkuntza-ereduak/LLaMA/lm/huggingface/7B/")
                else "EleutherAI/gpt-neo-125m"
            ),
            add_eos_token=True,
        )

        if tokenizer.pad_token_id is None:
            tokenizer.pad_token_id = tokenizer.unk_token_id

        # Test Train
        dataset, prompt, result = get_dataset(
            tokenizer=tokenizer,
            is_encoder_decoder=False,
            inference=False,
            prompt_loss_weight=0.05,
        )

        model_input = dataset[0]["input_ids"]
        labels = dataset[0]["labels"]
        self.assertEqual(model_input, labels)
        self.assertEqual(
            tokenizer.decode(model_input, skip_special_tokens=True, clean_up_tokenization_spaces=False),
            prompt + " " + result,
        )

        # Test Inference
        dataset, prompt, result = get_dataset(
            tokenizer=tokenizer,
            is_encoder_decoder=False,
            inference=True,
            prompt_loss_weight=0.05,
        )

        model_input = dataset[0]["input_ids"]
        self.assertFalse("labels" in dataset[0])
        self.assertEqual(
            tokenizer.decode(model_input, skip_special_tokens=True, clean_up_tokenization_spaces=False),
            prompt,
        )

    """
    We do not support encoder-decoder models yet. T5/mT5/FlanT5 lack the representation
    for '\n' or multiple whitespaces so they cannot be used with CoLLIE prompt encoding.


    def test_encoder_decoder(self):
        from transformers import AutoTokenizer

        tokenizer = AutoTokenizer.from_pretrained("google/mt5-small")
        if tokenizer.pad_token_id is None:
            tokenizer.pad_token_id = tokenizer.unk_token_id
        if tokenizer.decode(tokenizer.encode("\n", add_special_tokens=False)) != "\n":
            #T5 does not have a newline token, so we add one
            tokenizer.add_tokens("\n")
        # Test Train
        dataset, prompt, result = get_dataset(
            tokenizer=tokenizer,
            is_encoder_decoder=True,
            pad_to_max_length=False,
            inference=False,
        )

        model_input = dataset[0]["input_ids"]
        labels = dataset[0]["labels"]
        self.assertEqual(
            tokenizer.decode(
                model_input, skip_special_tokens=True, clean_up_tokenization_spaces=False
            ),
            prompt,
        )
        self.assertEqual(
            tokenizer.decode(
                labels, skip_special_tokens=True, clean_up_tokenization_spaces=False
            ),
            result,
        )

        # Test Inference
        dataset, prompt, result = get_dataset(
            tokenizer=tokenizer,
            is_encoder_decoder=True,
            pad_to_max_length=False,
            inference=True,
        )

        model_input = dataset[0]["input_ids"]
        labels = dataset[0]["labels"]
        self.assertEqual(
            tokenizer.decode(
                model_input, skip_special_tokens=True, clean_up_tokenization_spaces=False
            ),
            prompt,
        )
        self.assertEqual(
            tokenizer.decode(
                labels, skip_special_tokens=True, clean_up_tokenization_spaces=False
            ),
            result,
        )

        #  Test train with pad_to_max_length
        dataset, prompt, result = get_dataset(
            tokenizer=tokenizer,
            is_encoder_decoder=True,
            pad_to_max_length=True,
            inference=False,
        )

        model_input = dataset[0]["input_ids"]
        labels = dataset[0]["labels"]
        self.assertEqual(
            tokenizer.decode(
                model_input, skip_special_tokens=True, clean_up_tokenization_spaces=False
            ),
            prompt,
        )
        self.assertEqual(
            tokenizer.decode(
                labels, skip_special_tokens=True, clean_up_tokenization_spaces=False
            ),
            result,
        )
        self.assertTrue(len(model_input) == 2048)
    """

    def test_dataloader(self):
        from torch.utils.data import DataLoader

        from src.dataset.dataset import DataCollatorForCoLLIE
        from transformers import AutoTokenizer

        tokenizer = AutoTokenizer.from_pretrained(
            (
                "/gaueko1/hizkuntza-ereduak/LLaMA/lm/huggingface/7B/"
                if os.path.exists("/gaueko1/hizkuntza-ereduak/LLaMA/lm/huggingface/7B/")
                else "EleutherAI/gpt-neo-125m"
            ),
            add_eos_token=True,
        )

        tokenizer.padding_side = "left"

        if tokenizer.pad_token_id is None:
            tokenizer.pad_token_id = tokenizer.unk_token_id

        # Test Train
        dataset, prompt, result = get_dataset(
            tokenizer=tokenizer,
            is_encoder_decoder=False,
            inference=False,
            prompt_loss_weight=0.05,
        )

        datacollator = DataCollatorForCoLLIE(
            tokenizer,
            pad_to_multiple_of=2048,
            return_tensors="pt",
            padding=True,
            label_pad_token_id=-100,
        )

        dataloder = DataLoader(dataset, batch_size=1, collate_fn=datacollator, shuffle=False)
        batch = list(dataloder)[0]

        model_input = batch["input_ids"][0].tolist()
        labels = batch["labels"][0].tolist()
        self.assertEqual(
            tokenizer.decode(model_input, skip_special_tokens=True, clean_up_tokenization_spaces=False),
            prompt + " " + result,
        )
        self.assertEqual(
            tokenizer.decode(
                [x for x in labels if x != -100],
                skip_special_tokens=True,
                clean_up_tokenization_spaces=False,
            ),
            prompt + " " + result,
        )
        self.assertEqual(model_input[0], tokenizer.pad_token_id)
        self.assertEqual(labels[0], -100)

        datacollator = DataCollatorForCoLLIE(
            tokenizer,
            pad_to_multiple_of=2048,
            return_tensors="pt",
            padding=True,
            label_pad_token_id=tokenizer.pad_token_id,
        )

        dataloder = DataLoader(dataset, batch_size=1, collate_fn=datacollator, shuffle=False)
        batch = list(dataloder)[0]

        model_input = batch["input_ids"][0].tolist()
        labels = batch["labels"][0].tolist()
        self.assertEqual(model_input, labels)
        self.assertEqual(
            tokenizer.decode(model_input, skip_special_tokens=True, clean_up_tokenization_spaces=False),
            prompt + " " + result,
        )

        self.assertEqual(model_input[0], tokenizer.pad_token_id)
        self.assertEqual(labels[0], tokenizer.pad_token_id)

    def test_weight_loss_mask(self):
        import numpy as np
        from torch.utils.data import DataLoader

        from src.dataset.dataset import DataCollatorForCoLLIE
        from transformers import AutoTokenizer

        tokenizer = AutoTokenizer.from_pretrained(
            (
                "/gaueko1/hizkuntza-ereduak/LLaMA/lm/huggingface/7B/"
                if os.path.exists("/gaueko1/hizkuntza-ereduak/LLaMA/lm/huggingface/7B/")
                else "EleutherAI/gpt-neo-125m"
            ),
            add_eos_token=True,
        )

        tokenizer.padding_side = "left"

        if tokenizer.pad_token_id is None:
            tokenizer.pad_token_id = tokenizer.unk_token_id

        # Padding = Max Length , Ignore pad token for loss = True
        for prompt_loss_weight in [0.0, 0.05, 0.2, 0.5]:
            # Test Train
            dataset, prompt, result = get_dataset(
                tokenizer=tokenizer,
                is_encoder_decoder=False,
                inference=False,
                prompt_loss_weight=prompt_loss_weight,
            )

            datacollator = DataCollatorForCoLLIE(
                tokenizer,
                pad_to_multiple_of=2048,
                return_tensors="pt",
                padding="max_length",
                label_pad_token_id=-100,
            )

            dataloder = DataLoader(dataset, batch_size=1, collate_fn=datacollator, shuffle=False)
            batch = list(dataloder)[0]

            model_input = batch["input_ids"][0].tolist()
            labels = batch["labels"][0].tolist()
            loss_weights_mask = batch["loss_weight_mask"][0].tolist()

            self.assertEqual(
                tokenizer.decode(model_input, skip_special_tokens=True, clean_up_tokenization_spaces=False),
                prompt + " " + result,
            )
            self.assertEqual(
                tokenizer.decode(
                    [x for x in labels if x != -100],
                    skip_special_tokens=True,
                    clean_up_tokenization_spaces=False,
                ),
                prompt + " " + result,
            )

            prompt_tokens = tokenizer(
                text=prompt,
                max_length=2048,
                truncation=True,
                padding=False,
                return_tensors=None,
                add_special_tokens=True,
            )["input_ids"]

            # Remove the last token if it is an eos token
            if prompt_tokens[-1] == tokenizer.eos_token_id:
                prompt_tokens = prompt_tokens[:-1]

            result_tokens = tokenizer(
                text=result,
                max_length=2048,
                truncation=True,
                padding=False,
                return_tensors=None,
                add_special_tokens=True,
            )["input_ids"]

            if result_tokens[-1] != tokenizer.eos_token_id:
                result_tokens = result_tokens + [tokenizer.eos_token_id]

            num_pad_tokens = len(labels) - len(prompt_tokens) - len(result_tokens)

            # Test that all pad tokens are 0 in loss_weights_mask
            np.testing.assert_almost_equal(
                loss_weights_mask[:num_pad_tokens],
                [0.0] * num_pad_tokens,
            )

            # Test that all result tokens are 1.0 in loss_weights_mask
            np.testing.assert_almost_equal(
                loss_weights_mask[num_pad_tokens + len(prompt_tokens) :],
                [1.0] * len(result_tokens),
            )

            prompt_tokens_loss = sum(loss_weights_mask[num_pad_tokens : num_pad_tokens + len(prompt_tokens)])
            result_tokens_loss = sum(loss_weights_mask[num_pad_tokens + len(prompt_tokens) :])
            total_loss = prompt_tokens_loss + result_tokens_loss

            # print(f"Prompt loss weight: {prompt_loss_weight}")
            # print(f"Prompt loss: {prompt_tokens_loss}")
            # print(f"Result loss: {result_tokens_loss}")
            # print(f"Total loss: {total_loss}")
            # print()

            # Test that the loss of the prompt tokens is prompt_loss_weight of the total loss
            self.assertAlmostEqual(
                prompt_tokens_loss / total_loss,
                prompt_loss_weight,
                places=5,
            )

            # Test that the loss of the result tokens is (1 - prompt_loss_weight) of the total loss
            self.assertAlmostEqual(
                result_tokens_loss / total_loss,
                1 - prompt_loss_weight,
                places=5,
            )

        # Padding = True , Ignore pad token for loss = True
        for prompt_loss_weight in [0.0, 0.05, 0.2, 0.5]:
            # Test Train
            dataset, prompt, result = get_dataset(
                tokenizer=tokenizer,
                is_encoder_decoder=False,
                inference=False,
                prompt_loss_weight=prompt_loss_weight,
            )

            datacollator = DataCollatorForCoLLIE(
                tokenizer,
                pad_to_multiple_of=2048,
                return_tensors="pt",
                padding=True,
                label_pad_token_id=-100,
            )

            dataloder = DataLoader(dataset, batch_size=1, collate_fn=datacollator, shuffle=False)
            batch = list(dataloder)[0]

            model_input = batch["input_ids"][0].tolist()
            labels = batch["labels"][0].tolist()
            loss_weights_mask = batch["loss_weight_mask"][0].tolist()

            self.assertEqual(
                tokenizer.decode(model_input, skip_special_tokens=True, clean_up_tokenization_spaces=False),
                prompt + " " + result,
            )
            self.assertEqual(
                tokenizer.decode(
                    [x for x in labels if x != -100],
                    skip_special_tokens=True,
                    clean_up_tokenization_spaces=False,
                ),
                prompt + " " + result,
            )

            prompt_tokens = tokenizer(
                text=prompt,
                max_length=2048,
                truncation=True,
                padding=False,
                return_tensors=None,
                add_special_tokens=True,
            )["input_ids"]

            # Remove the last token if it is an eos token
            if prompt_tokens[-1] == tokenizer.eos_token_id:
                prompt_tokens = prompt_tokens[:-1]

            result_tokens = tokenizer(
                text=result,
                max_length=2048,
                truncation=True,
                padding=False,
                return_tensors=None,
                add_special_tokens=True,
            )["input_ids"]

            if result_tokens[-1] != tokenizer.eos_token_id:
                result_tokens = result_tokens + [tokenizer.eos_token_id]

            num_pad_tokens = len(labels) - len(prompt_tokens) - len(result_tokens)

            # Test that all pad tokens are 0 in loss_weights_mask
            np.testing.assert_almost_equal(
                loss_weights_mask[:num_pad_tokens],
                [0.0] * num_pad_tokens,
            )

            # Test that all result tokens are 1.0 in loss_weights_mask
            np.testing.assert_almost_equal(
                loss_weights_mask[num_pad_tokens + len(prompt_tokens) :],
                [1.0] * len(result_tokens),
            )

            prompt_tokens_loss = sum(loss_weights_mask[num_pad_tokens : num_pad_tokens + len(prompt_tokens)])
            result_tokens_loss = sum(loss_weights_mask[num_pad_tokens + len(prompt_tokens) :])
            total_loss = prompt_tokens_loss + result_tokens_loss

            # Test that the loss of the prompt tokens is prompt_loss_weight of the total loss
            self.assertAlmostEqual(
                prompt_tokens_loss / total_loss,
                prompt_loss_weight,
                places=5,
            )

            # Test that the loss of the result tokens is (1 - prompt_loss_weight) of the total loss
            self.assertAlmostEqual(
                result_tokens_loss / total_loss,
                1 - prompt_loss_weight,
                places=5,
            )

        # Padding = "Max len" , Ignore pad token for loss = False
        for prompt_loss_weight in [0.0, 0.05, 0.2, 0.5]:
            # Test Train
            dataset, prompt, result = get_dataset(
                tokenizer=tokenizer,
                is_encoder_decoder=False,
                inference=False,
                prompt_loss_weight=prompt_loss_weight,
            )

            datacollator = DataCollatorForCoLLIE(
                tokenizer,
                pad_to_multiple_of=2048,
                return_tensors="pt",
                padding="max_length",
                label_pad_token_id=tokenizer.pad_token_id,
            )

            dataloder = DataLoader(dataset, batch_size=1, collate_fn=datacollator, shuffle=False)
            batch = list(dataloder)[0]

            model_input = batch["input_ids"][0].tolist()
            labels = batch["labels"][0].tolist()
            loss_weights_mask = batch["loss_weight_mask"][0].tolist()

            self.assertEqual(
                tokenizer.decode(model_input, skip_special_tokens=True, clean_up_tokenization_spaces=False),
                prompt + " " + result,
            )
            self.assertEqual(
                tokenizer.decode(
                    [x for x in labels if x != -100],
                    skip_special_tokens=True,
                    clean_up_tokenization_spaces=False,
                ),
                prompt + " " + result,
            )

            prompt_tokens = tokenizer(
                text=prompt,
                max_length=2048,
                truncation=True,
                padding=False,
                return_tensors=None,
                add_special_tokens=True,
            )["input_ids"]

            # Remove the last token if it is an eos token
            if prompt_tokens[-1] == tokenizer.eos_token_id:
                prompt_tokens = prompt_tokens[:-1]

            result_tokens = tokenizer(
                text=result,
                max_length=2048,
                truncation=True,
                padding=False,
                return_tensors=None,
                add_special_tokens=True,
            )["input_ids"]

            if result_tokens[-1] != tokenizer.eos_token_id:
                result_tokens = result_tokens + [tokenizer.eos_token_id]

            num_pad_tokens = len(labels) - len(prompt_tokens) - len(result_tokens)

            # Test that all pad tokens are 1.0 in loss_weights_mask
            np.testing.assert_almost_equal(
                loss_weights_mask[:num_pad_tokens],
                [1.0] * num_pad_tokens,
            )

            # Test that all result tokens are 1.0 in loss_weights_mask
            np.testing.assert_almost_equal(
                loss_weights_mask[num_pad_tokens + len(prompt_tokens) :],
                [1.0] * len(result_tokens),
            )

            prompt_tokens_loss = sum(loss_weights_mask[num_pad_tokens : num_pad_tokens + len(prompt_tokens)])
            result_tokens_loss = sum(loss_weights_mask[num_pad_tokens + len(prompt_tokens) :])
            total_loss = prompt_tokens_loss + result_tokens_loss

            # Test that the loss of the prompt tokens is prompt_loss_weight of the total loss
            self.assertAlmostEqual(
                prompt_tokens_loss / total_loss,
                prompt_loss_weight,
                places=5,
            )

            # Test that the loss of the result tokens is (1 - prompt_loss_weight) of the total loss
            self.assertAlmostEqual(
                result_tokens_loss / total_loss,
                1 - prompt_loss_weight,
                places=5,
            )

    def test_dataset_rotation(self):
        from src.trainer import ConcatDataset
        from transformers import AutoTokenizer

        tokenizer = AutoTokenizer.from_pretrained(
            (
                "/gaueko1/hizkuntza-ereduak/LLaMA/lm/huggingface/7B/"
                if os.path.exists("/gaueko1/hizkuntza-ereduak/LLaMA/lm/huggingface/7B/")
                else "EleutherAI/gpt-neo-125m"
            ),
            add_eos_token=True,
        )

        tokenizer.padding_side = "left"

        if tokenizer.pad_token_id is None:
            tokenizer.pad_token_id = tokenizer.unk_token_id

        text1 = """@dataclass
        class EnergyAndInfrastructureEvent:
            \"\"\"This class is used to instantiate events that involve Chinese energy and infrastructure projects.\"\"\"
            meeting_attendees: Union[List[str], None] # Persons or organizations that attended the meeting.
            meeting_location: Union[List[str], None] # Location where the meeting happened.
            meeting_topic: Union[List[str], None] # Topic discussed on the meeting
            project_location: Union[List[str], None] # Location of the project
            project_name: Union[List[str], None] # Name of the project

        # This is the sentence to analyze
        sentence = "The Chinese and Rongovian delegations met at the sidelines of the Berlin Development Futures conference to discuss Rongovia's proposed Pangean Reunification Facility.

        # The following list contains the events instances that happens in the sentence defined above
        result = [
            EnergyAndInfrastructureEvent(
                meeting_attendees=["Chinese", "Rongovian"],
                meeting_location=["Berlin"],
                meeting_topic=["Pangean Reunification Facility"],
                project_location=["Rongovia"],
                project_name=["Pangean Reunification Facility"]
            ),
        ]"""

        prompt1 = """@dataclass
        class EnergyAndInfrastructureEvent:
            \"\"\"This class is used to instantiate events that involve Chinese energy and infrastructure projects.\"\"\"
            meeting_attendees: Union[List[str], None] # Persons or organizations that attended the meeting.
            meeting_location: Union[List[str], None] # Location where the meeting happened.
            meeting_topic: Union[List[str], None] # Topic discussed on the meeting
            project_location: Union[List[str], None] # Location of the project
            project_name: Union[List[str], None] # Name of the project

        # This is the sentence to analyze
        sentence = "The Chinese and Rongovian delegations met at the sidelines of the Berlin Development Futures conference to discuss Rongovia's proposed Pangean Reunification Facility.

        # The following list contains the events instances that happens in the sentence defined above
        result ="""
        result1 = """[
            EnergyAndInfrastructureEvent(
                meeting_attendees=["Chinese", "Rongovian"],
                meeting_location=["Berlin"],
                meeting_topic=["Pangean Reunification Facility"],
                project_location=["Rongovia"],
                project_name=["Pangean Reunification Facility"]
            ),
        ]"""

        text2 = """@dataclass
        class EnergyAndInfrastructureEvent:
            \"\"\"This class is used to instantiate events that involve Chinese energy and infrastructure projects.\"\"\"
            meeting_attendees: Union[List[str], None] # Persons or organizations that attended the meeting.
            meeting_location: Union[List[str], None] # Location where the meeting happened.
            meeting_topic: Union[List[str], None] # Topic discussed on the meeting
            project_location: Union[List[str], None] # Location of the project
            project_name: Union[List[str], None] # Name of the project

        # This is the sentence to analyze
        sentence = "The Spanish and French delegations met at the sidelines of the Berlin Development Futures conference to discuss Rongovia's proposed Pangean Reunification Facility.

        # The following list contains the events instances that happens in the sentence defined above
        result = [
            EnergyAndInfrastructureEvent(
                meeting_attendees=["Spanish", "French"],
                meeting_location=["Berlin"],
                meeting_topic=["Pangean Reunification Facility"],
                project_location=["Rongovia"],
                project_name=["Pangean Reunification Facility"]
            ),
        ]"""

        prompt2 = """@dataclass
        class EnergyAndInfrastructureEvent:
            \"\"\"This class is used to instantiate events that involve Chinese energy and infrastructure projects.\"\"\"
            meeting_attendees: Union[List[str], None] # Persons or organizations that attended the meeting.
            meeting_location: Union[List[str], None] # Location where the meeting happened.
            meeting_topic: Union[List[str], None] # Topic discussed on the meeting
            project_location: Union[List[str], None] # Location of the project
            project_name: Union[List[str], None] # Name of the project

        # This is the sentence to analyze
        sentence = "The Spanish and French delegations met at the sidelines of the Berlin Development Futures conference to discuss Rongovia's proposed Pangean Reunification Facility.

        # The following list contains the events instances that happens in the sentence defined above
        result ="""
        result2 = """[
            EnergyAndInfrastructureEvent(
                meeting_attendees=["Spanish", "French"],
                meeting_location=["Berlin"],
                meeting_topic=["Pangean Reunification Facility"],
                project_location=["Rongovia"],
                project_name=["Pangean Reunification Facility"]
            ),
        ]"""

        with tempfile.TemporaryDirectory() as tmpdirname:
            with open(os.path.join(tmpdirname, "tmp.ee.train.8.jsonl"), "w", encoding="utf8") as f:
                print(json.dumps({"text": text1}, ensure_ascii=False), file=f)
            with open(os.path.join(tmpdirname, "tmp.ee.train.42.jsonl"), "w", encoding="utf8") as f:
                print(json.dumps({"text": text2}, ensure_ascii=False), file=f)
                print(json.dumps({"text": text1}, ensure_ascii=False), file=f)

            dataset1 = CollieDataset(
                tokenizer=tokenizer,
                dataset_path=os.path.join(tmpdirname, "tmp.ee.train.jsonl"),
                is_encoder_decoder=False,
                max_length=2048,
                inference=False,
                prompt_loss_weight=0.05,
            )

        dataset3, prompt3, result3 = get_dataset(
            tokenizer=tokenizer,
            is_encoder_decoder=False,
            inference=False,
            prompt_loss_weight=0.05,
        )

        train_dataset = ConcatDataset([dataset1, dataset3])

        self.assertEqual(len(train_dataset), 3)

        model_input = train_dataset[0]["input_ids"]
        labels = train_dataset[0]["labels"]
        self.assertEqual(model_input, labels)
        self.assertEqual(
            tokenizer.decode(model_input, skip_special_tokens=True, clean_up_tokenization_spaces=False),
            prompt1 + " " + result1,
        )
        self.assertNotEqual(
            tokenizer.decode(model_input, skip_special_tokens=True, clean_up_tokenization_spaces=False),
            prompt2 + " " + result2,
        )
        self.assertNotEqual(
            tokenizer.decode(model_input, skip_special_tokens=True, clean_up_tokenization_spaces=False),
            prompt3 + " " + result3,
        )

        model_input = train_dataset[1]["input_ids"]
        labels = train_dataset[1]["labels"]
        self.assertEqual(model_input, labels)
        self.assertEqual(
            tokenizer.decode(model_input, skip_special_tokens=True, clean_up_tokenization_spaces=False),
            prompt1 + " " + result1,
        )
        self.assertNotEqual(
            tokenizer.decode(model_input, skip_special_tokens=True, clean_up_tokenization_spaces=False),
            prompt2 + " " + result2,
        )
        self.assertNotEqual(
            tokenizer.decode(model_input, skip_special_tokens=True, clean_up_tokenization_spaces=False),
            prompt3 + " " + result3,
        )

        model_input = train_dataset[2]["input_ids"]
        labels = train_dataset[2]["labels"]
        self.assertEqual(model_input, labels)
        self.assertEqual(
            tokenizer.decode(model_input, skip_special_tokens=True, clean_up_tokenization_spaces=False),
            prompt3 + " " + result3,
        )

        train_dataset.rotate_split()
        train_dataset.rotate_split()
        train_dataset.rotate_split()
        train_dataset.rotate_split()
        train_dataset.rotate_split()

        self.assertEqual(len(train_dataset), 3)
        model_input = train_dataset[0]["input_ids"]
        labels = train_dataset[0]["labels"]
        self.assertEqual(model_input, labels)
        self.assertEqual(
            tokenizer.decode(model_input, skip_special_tokens=True, clean_up_tokenization_spaces=False),
            prompt2 + " " + result2,
        )
        self.assertNotEqual(
            tokenizer.decode(model_input, skip_special_tokens=True, clean_up_tokenization_spaces=False),
            prompt1 + " " + result1,
        )

        self.assertNotEqual(
            tokenizer.decode(model_input, skip_special_tokens=True, clean_up_tokenization_spaces=False),
            prompt3 + " " + result3,
        )

        model_input = train_dataset[1]["input_ids"]
        labels = train_dataset[1]["labels"]
        self.assertEqual(model_input, labels)
        self.assertEqual(
            tokenizer.decode(model_input, skip_special_tokens=True, clean_up_tokenization_spaces=False),
            prompt1 + " " + result1,
        )
        self.assertNotEqual(
            tokenizer.decode(model_input, skip_special_tokens=True, clean_up_tokenization_spaces=False),
            prompt2 + " " + result2,
        )
        self.assertNotEqual(
            tokenizer.decode(model_input, skip_special_tokens=True, clean_up_tokenization_spaces=False),
            prompt3 + " " + result3,
        )

        model_input = train_dataset[2]["input_ids"]
        labels = train_dataset[2]["labels"]
        self.assertEqual(model_input, labels)
        self.assertEqual(
            tokenizer.decode(model_input, skip_special_tokens=True, clean_up_tokenization_spaces=False),
            prompt3 + " " + result3,
        )

src/tests/test_dataset.py
hitz-zentroa/GoLLIE
import unittest

import black


def to_str(x):
    return black.format_str(x.__repr__(), mode=black.Mode())


class TestEvaluate(unittest.TestCase):
    def test_annotation_list(self):
        from src.tasks.ace.prompts import GPE, Location, Organization, Person
        from src.tasks.utils_typing import AnnotationList

        text = "Peter went to Hitz-zentroa where Carlos works, after visiting Donosti, Tokyo."

        annotations = [
            Person("Peter"),
            Organization("Hitz-zentroa"),
            Person("carlos"),
            Location("Tokyo"),
            Location("Donosti"),
            GPE("UE"),
        ]

        annotations = AnnotationList.from_output(
            to_str(annotations),
            task_module="src.tasks.ace.prompts",
            text=text,
            filter_hallucinations=True,
        )

        self.assertListEqual(
            annotations,
            AnnotationList(
                [
                    Person("Peter"),
                    Organization("Hitz-zentroa"),
                    Person("carlos"),
                    Location("Tokyo"),
                    Location("Donosti"),
                ]
            ),
        )

    def test_entity_hallucination(self):
        from src.tasks.ace.prompts import (
            GPE,
            Location,
            Organization,
            Person,
        )
        from src.tasks.utils_typing import AnnotationList

        unlabelled_sentence = "Peter was born in Donosti. He married Carlos on May 18th."
        predictions = [
            Person("Peter"),
            Organization("Hitz-zentroa"),
            Person("carlos"),
            Location("Tokyo"),
            Location("Donosti"),
            GPE("UE"),
        ]

        predictions = AnnotationList.from_output(str(predictions), task_module="src.tasks.ace.prompts")
        filtered_predictions = predictions.filter_hallucinations(text=unlabelled_sentence)

        self.assertListEqual(
            filtered_predictions,
            AnnotationList(
                [
                    Person("Peter"),
                    Person("carlos"),
                    Location("Donosti"),
                ]
            ),
        )

    def test_type_hallucination(self):
        from src.tasks.ace.prompts import (
            Location,
            Person,
        )
        from src.tasks.utils_typing import AnnotationList

        unlabelled_sentence = "Peter was born in Donosti. He married Carlos on May 18th."
        predictions = (
            "["
            + "Person('Peter'), "
            + "Organization('Hitz-zentroa'), "
            + "Person('carlos'), "
            + "Location('Tokyo'), "
            + "Location('Donosti'), "
            + "ASDF('UE'), ] "
        )

        predictions = AnnotationList.from_output(predictions, task_module="src.tasks.ace.prompts")
        filtered_predictions = predictions.filter_hallucinations(text=unlabelled_sentence)

        self.assertListEqual(
            filtered_predictions,
            AnnotationList(
                [
                    Person("Peter"),
                    Person("carlos"),
                    Location("Donosti"),
                ]
            ),
        )

    def test_relation_hallucination(self):
        from src.tasks.ace.prompts import (
            Family,
            Geographical,
            Located,
        )
        from src.tasks.utils_typing import AnnotationList

        unlabelled_sentence = "Peter was born in Donosti. He married Carlos on May 18th."
        location = Located("Peter", "Donosti")
        family = Family("Peter", "carlos")
        geographical = Geographical("Donosti", "Spain")
        predictions = [
            location,
            family,
            geographical,
        ]

        predictions = AnnotationList.from_output(str(predictions), task_module="src.tasks.ace.prompts")
        filtered_predictions = predictions.filter_hallucinations(text=unlabelled_sentence)

        self.assertListEqual(
            filtered_predictions,
            AnnotationList(
                [
                    location,
                    family,
                ]
            ),
        )

    def test_event_hallucination(self):
        from src.tasks.ace.prompts import BeBorn, Marry
        from src.tasks.utils_typing import AnnotationList

        text = "Peter was born in Donosti. He married Carlos on May 18th."

        predictions = str(
            [
                BeBorn("born", person=[], time=[], place=[]),
                Marry(
                    "married",
                    person=[
                        "Peter",
                    ],
                    time=["May 18th"],
                    place=["Tokyo"],
                ),
            ]
        )
        predictions = AnnotationList.from_output(
            predictions, task_module="src.tasks.ace.prompts", text=text, filter_hallucinations=True
        )

        self.assertListEqual(
            predictions,
            AnnotationList(
                [
                    BeBorn("born", person=[], time=[], place=[]),
                    Marry(
                        "married",
                        person=[
                            "Peter",
                        ],
                        time=["May 18th"],
                        place=[],
                    ),
                ]
            ),
        )

src/tests/test_evaluate.py
hitz-zentroa/GoLLIE
import unittest

from src.tasks.label_encoding import rewrite_labels


class TestToIOB(unittest.TestCase):
    def test_from_iob(self):
        text_in = (
            "word1 O\n"
            "word2 I-LOC\n"
            "word3 I-LOC\n"
            "word4 I-LOC\n"
            "wordnl  O\n"
            "word5 I-LOC\n"
            "word6 O\n"
            "word7 I-PER\n"
            "word8 B-PER\n"
            "word9 I-LOC\n"
            "word10 O\n"
            "word11 I-LOC\n"
            "word12 I-ORG\n"
            "word13 I-LOC\n"
            "word14 O\n"
            "word15 O\n"
            "word16 O\n"
            "word17 O\n"
            "word18 I-PER\n"
            "word19 O\n"
        )

        in_labels = [label for _, label in [line.split() for line in text_in.splitlines()]]
        gold_labels = [label for _, label in [line.split() for line in text_in.splitlines()]]

        rewritten_labels = rewrite_labels(labels=in_labels, encoding="iob")

        self.assertEqual(gold_labels, rewritten_labels)

    def test_from_corrupted_iob(self):
        text_in = (
            "word1 O\n"
            "word2 I-LOC\n"
            "word3 I-LOC\n"
            "word4 I-LOC\n"
            "wordnl  O\n"
            "word5 I-LOC\n"
            "word6 O\n"
            "word7 B-PER\n"
            "word8 B-PER\n"
            "word9 I-LOC\n"
            "word10 O\n"
            "word11 I-LOC\n"
            "word12 B-ORG\n"
            "word13 I-LOC\n"
            "word14 O\n"
            "word15 O\n"
            "word16 O\n"
            "word17 O\n"
            "word18 B-PER\n"
            "word19 O\n"
        )

        text_out = (
            "word1 O\n"
            "word2 I-LOC\n"
            "word3 I-LOC\n"
            "word4 I-LOC\n"
            "wordnl  O\n"
            "word5 I-LOC\n"
            "word6 O\n"
            "word7 I-PER\n"
            "word8 B-PER\n"
            "word9 I-LOC\n"
            "word10 O\n"
            "word11 I-LOC\n"
            "word12 I-ORG\n"
            "word13 I-LOC\n"
            "word14 O\n"
            "word15 O\n"
            "word16 O\n"
            "word17 O\n"
            "word18 I-PER\n"
            "word19 O\n"
        )

        in_labels = [label for _, label in [line.split() for line in text_in.splitlines()]]
        gold_labels = [label for _, label in [line.split() for line in text_out.splitlines()]]

        rewritten_labels = rewrite_labels(labels=in_labels, encoding="iob")

        self.assertEqual(gold_labels, rewritten_labels)

    def test_from_iob2(self):
        text_in = (
            "word1 O\n"
            "word2 B-LOC\n"
            "word3 I-LOC\n"
            "word4 I-LOC\n"
            "wordnl  O\n"
            "word5 B-LOC\n"
            "word6 O\n"
            "word7 B-PER\n"
            "word8 B-PER\n"
            "word9 B-LOC\n"
            "word10 O\n"
            "word11 B-LOC\n"
            "word12 B-ORG\n"
            "word13 B-LOC\n"
            "word14 O\n"
            "word15 O\n"
            "word16 O\n"
            "word17 O\n"
            "word18 B-PER\n"
            "word19 I-PER\n"
        )

        text_out = (
            "word1 O\n"
            "word2 I-LOC\n"
            "word3 I-LOC\n"
            "word4 I-LOC\n"
            "wordnl  O\n"
            "word5 I-LOC\n"
            "word6 O\n"
            "word7 I-PER\n"
            "word8 B-PER\n"
            "word9 I-LOC\n"
            "word10 O\n"
            "word11 I-LOC\n"
            "word12 I-ORG\n"
            "word13 I-LOC\n"
            "word14 O\n"
            "word15 O\n"
            "word16 O\n"
            "word17 O\n"
            "word18 I-PER\n"
            "word19 I-PER\n"
        )

        in_labels = [label for _, label in [line.split() for line in text_in.splitlines()]]
        gold_labels = [label for _, label in [line.split() for line in text_out.splitlines()]]

        rewritten_labels = rewrite_labels(labels=in_labels, encoding="iob")

        self.assertEqual(gold_labels, rewritten_labels)

    def test_from_bilou(self):
        text_in = (
            "word1 O\n"
            "word2 B-LOC\n"
            "word3 I-LOC\n"
            "word4 L-LOC\n"
            "word5 U-LOC\n"
            "word6 O\n"
            "word7 U-PER\n"
            "word8 U-PER\n"
            "word9 U-LOC\n"
            "word10 O\n"
            "word11 U-LOC\n"
            "word12 U-ORG\n"
            "word13 U-LOC\n"
            "word14 O\n"
            "word15 O\n"
            "word16 O\n"
            "word17 O\n"
            "word18 B-PER\n"
            "word19 L-PER\n"
        )

        text_out = (
            "word1 O\n"
            "word2 I-LOC\n"
            "word3 I-LOC\n"
            "word4 I-LOC\n"
            "word5 B-LOC\n"
            "word6 O\n"
            "word7 I-PER\n"
            "word8 B-PER\n"
            "word9 I-LOC\n"
            "word10 O\n"
            "word11 I-LOC\n"
            "word12 I-ORG\n"
            "word13 I-LOC\n"
            "word14 O\n"
            "word15 O\n"
            "word16 O\n"
            "word17 O\n"
            "word18 I-PER\n"
            "word19 I-PER\n"
        )

        in_labels = [label for _, label in [line.split() for line in text_in.splitlines()]]
        [label for _, label in [line.split() for line in text_out.splitlines()]]

        rewrite_labels(labels=in_labels, encoding="iob")


class TestToIOB2(unittest.TestCase):
    def test_from_iob(self):
        text_in = (
            "word1 O\n"
            "word2 I-LOC\n"
            "word3 I-LOC\n"
            "word4 I-LOC\n"
            "wordnl  O\n"
            "word5 I-LOC\n"
            "word6 O\n"
            "word7 I-PER\n"
            "word8 B-PER\n"
            "word9 I-LOC\n"
            "word10 O\n"
            "word11 I-LOC\n"
            "word12 I-ORG\n"
            "word13 I-LOC\n"
            "word14 O\n"
            "word15 O\n"
            "word16 O\n"
            "word17 O\n"
            "word18 I-PER\n"
            "word19 I-PER\n"
        )

        text_out = (
            "word1 O\n"
            "word2 B-LOC\n"
            "word3 I-LOC\n"
            "word4 I-LOC\n"
            "wordnl  O\n"
            "word5 B-LOC\n"
            "word6 O\n"
            "word7 B-PER\n"
            "word8 B-PER\n"
            "word9 B-LOC\n"
            "word10 O\n"
            "word11 B-LOC\n"
            "word12 B-ORG\n"
            "word13 B-LOC\n"
            "word14 O\n"
            "word15 O\n"
            "word16 O\n"
            "word17 O\n"
            "word18 B-PER\n"
            "word19 I-PER\n"
        )

        in_labels = [label for _, label in [line.split() for line in text_in.splitlines()]]
        gold_labels = [label for _, label in [line.split() for line in text_out.splitlines()]]

        rewritten_labels = rewrite_labels(labels=in_labels, encoding="iob2")

        self.assertEqual(gold_labels, rewritten_labels)

    def test_from_iob2(self):
        text_in = (
            "word1 O\n"
            "word2 B-LOC\n"
            "word3 I-LOC\n"
            "word4 I-LOC\n"
            "word5 B-LOC\n"
            "word6 O\n"
            "word7 B-PER\n"
            "word8 B-PER\n"
            "word9 B-LOC\n"
            "word10 O\n"
            "word11 B-LOC\n"
            "word12 B-ORG\n"
            "word13 B-LOC\n"
            "word14 O\n"
            "word15 O\n"
            "word16 O\n"
            "word17 O\n"
            "word18 B-PER\n"
            "word19 I-PER\n"
        )

        in_labels = [label for _, label in [line.split() for line in text_in.splitlines()]]
        gold_labels = [label for _, label in [line.split() for line in text_in.splitlines()]]

        rewritten_labels = rewrite_labels(labels=in_labels, encoding="iob2")

        self.assertEqual(gold_labels, rewritten_labels)

    def test_from_bilou(self):
        text_in = (
            "word1 O\n"
            "word2 B-LOC\n"
            "word3 I-LOC\n"
            "word4 L-LOC\n"
            "word5 U-LOC\n"
            "word6 O\n"
            "word7 U-PER\n"
            "word8 U-PER\n"
            "word9 U-LOC\n"
            "word10 O\n"
            "word11 U-LOC\n"
            "word12 U-ORG\n"
            "word13 U-LOC\n"
            "word14 O\n"
            "word15 O\n"
            "word16 O\n"
            "word17 O\n"
            "word18 B-PER\n"
            "word19 L-PER\n"
        )

        text_out = (
            "word1 O\n"
            "word2 B-LOC\n"
            "word3 I-LOC\n"
            "word4 I-LOC\n"
            "word5 B-LOC\n"
            "word6 O\n"
            "word7 B-PER\n"
            "word8 B-PER\n"
            "word9 B-LOC\n"
            "word10 O\n"
            "word11 B-LOC\n"
            "word12 B-ORG\n"
            "word13 B-LOC\n"
            "word14 O\n"
            "word15 O\n"
            "word16 O\n"
            "word17 O\n"
            "word18 B-PER\n"
            "word19 I-PER\n"
        )

        in_labels = [label for _, label in [line.split() for line in text_in.splitlines()]]
        gold_labels = [label for _, label in [line.split() for line in text_out.splitlines()]]

        rewritten_labels = rewrite_labels(labels=in_labels, encoding="iob2")

        self.assertEqual(gold_labels, rewritten_labels)

    def test_from_corrupted_iob2(self):
        text_in = (
            "word1 O\n"
            "word2 B-LOC\n"
            "word3 I-LOC\n"
            "word4 L-LOC\n"
            "word5 B-LOC\n"
            "word6 O\n"
            "word7 B-PER\n"
            "word8 U-PER\n"
            "word9 B-LOC\n"
            "word10 O\n"
            "word11 B-LOC\n"
            "word12 I-ORG\n"
            "word13 I-LOC\n"
            "word14 O\n"
            "word15 O\n"
            "word16 O\n"
            "word17 O\n"
            "word18 B-PER\n"
            "word19 L-PER\n"
        )

        text_out = (
            "word1 O\n"
            "word2 B-LOC\n"
            "word3 I-LOC\n"
            "word4 I-LOC\n"
            "word5 B-LOC\n"
            "word6 O\n"
            "word7 B-PER\n"
            "word8 B-PER\n"
            "word9 B-LOC\n"
            "word10 O\n"
            "word11 B-LOC\n"
            "word12 B-ORG\n"
            "word13 B-LOC\n"
            "word14 O\n"
            "word15 O\n"
            "word16 O\n"
            "word17 O\n"
            "word18 B-PER\n"
            "word19 I-PER\n"
        )

        in_labels = [label for _, label in [line.split() for line in text_in.splitlines()]]
        gold_labels = [label for _, label in [line.split() for line in text_out.splitlines()]]

        rewritten_labels = rewrite_labels(labels=in_labels, encoding="iob2")

        self.assertEqual(gold_labels, rewritten_labels)


class TestToBILOU(unittest.TestCase):
    def test_from_iob(self):
        text_in = (
            "word1 O\n"
            "word2 I-LOC\n"
            "word3 I-LOC\n"
            "word4 I-LOC\n"
            "wordnl  O\n"
            "word5 I-LOC\n"
            "word6 O\n"
            "word7 I-PER\n"
            "word8 B-PER\n"
            "word9 I-LOC\n"
            "word10 O\n"
            "word11 I-LOC\n"
            "word12 I-ORG\n"
            "word13 I-LOC\n"
            "word14 O\n"
            "word15 O\n"
            "word16 O\n"
            "word17 O\n"
            "word18 I-PER\n"
            "word19 I-PER\n"
        )

        text_out = (
            "word1 O\n"
            "word2 B-LOC\n"
            "word3 I-LOC\n"
            "word4 L-LOC\n"
            "wordnl  O\n"
            "word5 U-LOC\n"
            "word6 O\n"
            "word7 U-PER\n"
            "word8 U-PER\n"
            "word9 U-LOC\n"
            "word10 O\n"
            "word11 U-LOC\n"
            "word12 U-ORG\n"
            "word13 U-LOC\n"
            "word14 O\n"
            "word15 O\n"
            "word16 O\n"
            "word17 O\n"
            "word18 B-PER\n"
            "word19 L-PER\n"
        )

        in_labels = [label for _, label in [line.split() for line in text_in.splitlines()]]
        gold_labels = [label for _, label in [line.split() for line in text_out.splitlines()]]

        rewritten_labels = rewrite_labels(labels=in_labels, encoding="bilou")

        self.assertEqual(gold_labels, rewritten_labels)

    def test_from_iob2(self):
        text_in = (
            "word1 O\n"
            "word2 B-LOC\n"
            "word3 I-LOC\n"
            "word4 I-LOC\n"
            "word5 B-LOC\n"
            "word6 O\n"
            "word7 B-PER\n"
            "word8 B-PER\n"
            "word9 B-LOC\n"
            "word10 O\n"
            "word11 B-LOC\n"
            "word12 B-ORG\n"
            "word13 B-LOC\n"
            "word14 O\n"
            "word15 O\n"
            "word16 O\n"
            "word17 O\n"
            "word18 B-PER\n"
            "word19 I-PER\n"
            "word20 B-PER\n"
            "word21 O\n"
        )

        text_out = (
            "word1 O\n"
            "word2 B-LOC\n"
            "word3 I-LOC\n"
            "word4 L-LOC\n"
            "word5 U-LOC\n"
            "word6 O\n"
            "word7 U-PER\n"
            "word8 U-PER\n"
            "word9 U-LOC\n"
            "word10 O\n"
            "word11 U-LOC\n"
            "word12 U-ORG\n"
            "word13 U-LOC\n"
            "word14 O\n"
            "word15 O\n"
            "word16 O\n"
            "word17 O\n"
            "word18 B-PER\n"
            "word19 L-PER\n"
            "word20 U-PER\n"
            "word21 O\n"
        )

        in_labels = [label for _, label in [line.split() for line in text_in.splitlines()]]
        gold_labels = [label for _, label in [line.split() for line in text_out.splitlines()]]

        rewritten_labels = rewrite_labels(labels=in_labels, encoding="bilou")

        self.assertEqual(gold_labels, rewritten_labels)

    def test_from_bilou(self):
        text_in = (
            "word1 O\n"
            "word2 B-LOC\n"
            "word3 I-LOC\n"
            "word4 L-LOC\n"
            "word5 U-LOC\n"
            "word6 O\n"
            "word7 U-PER\n"
            "word8 U-PER\n"
            "word9 U-LOC\n"
            "word10 O\n"
            "word11 U-LOC\n"
            "word12 U-ORG\n"
            "word13 U-LOC\n"
            "word14 O\n"
            "word15 O\n"
            "word16 O\n"
            "word17 O\n"
            "word18 B-PER\n"
            "word19 L-PER\n"
        )

        in_labels = [label for _, label in [line.split() for line in text_in.splitlines()]]
        gold_labels = [label for _, label in [line.split() for line in text_in.splitlines()]]

        rewritten_labels = rewrite_labels(labels=in_labels, encoding="bilou")

        self.assertEqual(gold_labels, rewritten_labels)

    def test_from_corrupted_bilou(self):
        text_in = (
            "word1 O\n"
            "word2 B-LOC\n"
            "word3 L-LOC\n"
            "word4 L-LOC\n"
            "word5 U-LOC\n"
            "word6 O\n"
            "word7 U-PER\n"
            "word8 U-PER\n"
            "word9 U-LOC\n"
            "word10 O\n"
            "word11 U-LOC\n"
            "word12 B-ORG\n"
            "word13 U-LOC\n"
            "word14 O\n"
            "word15 O\n"
            "word16 O\n"
            "word17 B-PER\n"
            "word18 L-PER\n"
            "word19 L-PER\n"
            "word20 B-PER\n"
            "word21 L-LOC\n"
            "word22 L-PER\n"
            "word23 B-PER\n"
            "word24 I-PER\n"
            "word25 I-PER\n"
        )

        text_out = (
            "word1 O\n"
            "word2 B-LOC\n"
            "word3 I-LOC\n"
            "word4 L-LOC\n"
            "word5 U-LOC\n"
            "word6 O\n"
            "word7 U-PER\n"
            "word8 U-PER\n"
            "word9 U-LOC\n"
            "word10 O\n"
            "word11 U-LOC\n"
            "word12 U-ORG\n"
            "word13 U-LOC\n"
            "word14 O\n"
            "word15 O\n"
            "word16 O\n"
            "word17 B-PER\n"
            "word18 I-PER\n"
            "word19 L-PER\n"
            "word20 U-PER\n"
            "word21 U-LOC\n"
            "word22 U-PER\n"
            "word23 B-PER\n"
            "word24 I-PER\n"
            "word25 L-PER\n"
        )

        in_labels = [label for _, label in [line.split() for line in text_in.splitlines()]]
        gold_labels = [label for _, label in [line.split() for line in text_out.splitlines()]]

        rewritten_labels = rewrite_labels(labels=in_labels, encoding="bilou")

        self.assertEqual(gold_labels, rewritten_labels)

src/tests/test_label_encoding.py
hitz-zentroa/GoLLIE
import unittest


class TestEntityPrompts(unittest.TestCase):
    def test_ACE(self):
        pass

src/tests/test_prompts.py
hitz-zentroa/GoLLIE
import unittest


class TestACEScorers(unittest.TestCase):
    def test_entity_scorer(self):
        from src.tasks.ace.prompts import (
            GPE,
            Location,
            Organization,
            Person,
        )
        from src.tasks.ace.scorer import ACEEntityScorer

        reference = [
            Person("Peter"),
            Organization("OpenAI"),
            Person("Carlos"),
            Person("Peter"),
            Location("Tokyo"),
        ]
        predictions = [
            Person("Peter"),
            Organization("Google"),
            Person("Carlos"),
            GPE("Tokyo"),
        ]

        scorer = ACEEntityScorer()

        # TP = 2 (Person("Peter"), Person("Carlos"))
        # FP = 2 (Organization("Google"), GPE("Tokyo"))
        # FN = 3 (Organization("OpenAI"), Person("Peter"), Location("Tokyo"))
        # precision -> 2 / 4 = 0.5
        # recall -> 2 / 5 = 0.4
        # F1 -> 2 * 0.5 * 0.4 / (0.5 + 0.4) = 0.5

        scorer_result = scorer(reference=reference, predictions=predictions)["entities"]
        scorer_result.pop("class_scores")

        self.assertDictEqual(
            scorer_result,
            {"precision": 0.5, "recall": 0.4, "f1-score": 0.4444444444444445},
        )

    def test_event_scorer(self):
        from src.tasks.ace.prompts import (
            BeBorn,
            Marry,
        )
        from src.tasks.ace.scorer import ACEEventScorer

        reference = [
            BeBorn("born", person=["Peter"], time=[], place=["Donosti"]),
            Marry("married", person=["Peter", "Carlos"], time=["May 18th"], place=[]),
        ]
        predictions = [
            BeBorn("birth", person=[], time=[], place=[]),
            Marry(
                "married",
                person=[
                    "Peter",
                ],
                time=["May 18th"],
                place=["Tokyo"],
            ),
        ]

        scorer = ACEEventScorer()

        # Events
        # TP = 1 (Marry("married"))
        # FP = 1 (BeBorn("birth"))
        # FN = 1 (BeBorn("born"))
        # precision -> 1 / 2 = 0.5
        # recall -> 1 / 2 = 0.5
        # F1 -> 2* 0.5 * 0.5 / (0.5 + 0.5) = 0.5

        scorer_results = scorer(reference=reference, predictions=predictions)["events"]
        if "class_scores" in scorer_results:
            scorer_results.pop("class_scores")

        self.assertDictEqual(
            scorer_results,
            {"precision": 0.5, "recall": 0.5, "f1-score": 0.5},
        )
        # Arguments
        # TP = 2 (Marry("married", person=["Peter"], time=["May 18th"]))
        # FP = 1 (Marry("married", time=["Tokyo"]))
        # FN = 3 (BeBorn("born", person=["Peter"], place=["Donosti"]),
        #         Marry("married", person=["Carlos"]))
        # precision -> 2 / 3 = 0.6666666666666666
        # recall -> 2 / 5 = 0.4
        # F1 -> 2 * (0.6666666666666666 * 0.4) / (0.6666666666666666 + 0.4) = 0.5

        scorer_results = scorer(reference=reference, predictions=predictions)["arguments"]
        if "class_scores" in scorer_results:
            scorer_results.pop("class_scores")

        self.assertDictEqual(
            scorer_results,
            {"precision": 0.6666666666666666, "recall": 0.4, "f1-score": 0.5},
        )


class TestRAMSScorers(unittest.TestCase):
    def test_event_scorer(self):
        from src.tasks.rams.prompts import TransferOwnership, Yield
        from src.tasks.rams.scorer import RAMSEventScorer

        scorer = RAMSEventScorer()

        reference = [
            TransferOwnership(
                mention="bought",
                subtype=None,
                giver=["commodities trader"],
                recipient=[],
                beneficiary=[],
                artifact=["appointment"],
                preventer=[],
                place=[],
            ),
            Yield(
                mention="surrendered",
                subtype=None,
                agent=["Japan"],
                recipient=[],
                place=[],
                origin=[],
                destination=[],
            ),
        ]

        # Perfect alingment
        scorer_results = scorer(reference=reference, predictions=reference)["arguments"]
        if "class_scores" in scorer_results:
            scorer_results.pop("class_scores")

        self.assertDictEqual(
            scorer_results,
            {"precision": 1.0, "recall": 1.0, "f1-score": 1.0},
        )

        # No alignment
        scorer_results = scorer(reference=reference, predictions=[])["arguments"]
        if "class_scores" in scorer_results:
            scorer_results.pop("class_scores")

        self.assertDictEqual(
            scorer_results,
            {"precision": 0.0, "recall": 0.0, "f1-score": 0.0},
        )

        predictions = [
            TransferOwnership(
                mention="bought",
                subtype=None,
                giver=[],
                recipient=["commodities trader"],
                beneficiary=[],
                artifact=["appointment"],
                preventer=[],
                place=[],
            ),
            Yield(
                mention="surrendered",
                subtype=None,
                agent=["Japan"],
                recipient=[],
                place=[],
                origin=[],
                destination=[],
            ),
        ]

        # Partially aligned
        # TP = 2 (artifact("apointment"), agent("Japan"))
        # FP = 1 (recipient("commodities trader"))
        # FN = 1 (giver("commodities trader"))

        # precision -> 2 / 3 = 0.6666666666666666
        # recall -> 2 / 3 = 0.6666666666666666
        # F1 -> 2 * (0.6666666666666666 * 0.6666666666666666) / (0.6666666666666666 + 0.6666666666666666) = 0.6666666666666666

        # No alignment

        scorer_results = scorer(reference=reference, predictions=predictions)["arguments"]
        if "class_scores" in scorer_results:
            scorer_results.pop("class_scores")

        self.assertDictEqual(
            scorer_results,
            {
                "precision": 0.6666666666666666,
                "recall": 0.6666666666666666,
                "f1-score": 0.6666666666666666,
            },
        )


class TestTACREDScorers(unittest.TestCase):
    def test_template_scorer(self):
        from src.tasks.tacred.prompts import PersonTemplate
        from src.tasks.tacred.scorer import TACREDTemplateScorer
        from src.tasks.utils_typing import Name, String, Value

        scorer = TACREDTemplateScorer()

        reference = [
            PersonTemplate(
                query="Peter",
                city_of_birth=Name("Lasarte"),
                age=[Value("26")],
                title=[String("Researcher"), String("Student")],
            )
        ]

        # Assert the typing constrains are satisfied
        for template in reference:
            template.assert_typing_constraints()

        # Perfect alingment
        scorer_results = scorer(reference=reference, predictions=reference)["slots"]
        if "class_scores" in scorer_results:
            scorer_results.pop("class_scores")

        self.assertDictEqual(
            scorer_results,
            {"precision": 1.0, "recall": 1.0, "f1-score": 1.0},
        )

        # No alignment
        scorer_results = scorer(reference=reference, predictions=[])["slots"]
        if "class_scores" in scorer_results:
            scorer_results.pop("class_scores")

        self.assertDictEqual(
            scorer_results,
            {"precision": 0.0, "recall": 0.0, "f1-score": 0.0},
        )

        predictions = [
            PersonTemplate(query="Peter", city_of_birth=Name("Lasarte"), age=Value("26"), origin=Name("Spain"))
        ]

        # Partially aligned
        # TP = 2 (city_of_birth=Name("Lasarte"), age=Value("26"))
        # FP = 1 (origin=Name("Spain"))
        # FN = 2 (title=[String("Researcher"), String("Student"))

        # precision -> 2 / 3 = 0.6666666666666666
        # recall -> 2 / 4 = 0.5
        # F1 -> 2 * (0.6666666666666666 * 0.5) / (0.6666666666666666 + 0.5) = 0.5714285714285715

        # No alignment
        scorer_results = scorer(reference=reference, predictions=predictions)["slots"]
        if "class_scores" in scorer_results:
            scorer_results.pop("class_scores")

        self.assertDictEqual(
            scorer_results,
            {
                "precision": 0.6666666666666666,
                "recall": 0.5,
                "f1-score": 0.5714285714285715,
            },
        )


class TestCASIEScorers(unittest.TestCase):
    def test_event_scorer(self):
        from src.tasks.casie.prompts_ed import VulnerabilityDiscover
        from src.tasks.casie.scorer import CASIEEventScorer

        scorer = CASIEEventScorer(allow_partial_match=True)

        reference = [VulnerabilityDiscover(mention="vulnerability found"), VulnerabilityDiscover(mention="reported")]

        # Perfect alingment
        scorer_results = scorer(reference=reference, predictions=reference)["events"]
        if "class_scores" in scorer_results:
            scorer_results.pop("class_scores")

        self.assertDictEqual(
            scorer_results,
            {"precision": 1.0, "recall": 1.0, "f1-score": 1.0},
        )

        # No alignment

        scorer_results = scorer(reference=reference, predictions=[])["events"]
        if "class_scores" in scorer_results:
            scorer_results.pop("class_scores")
        self.assertDictEqual(
            scorer_results,
            {"precision": 0.0, "recall": 0.0, "f1-score": 0.0},
        )

        predictions = [VulnerabilityDiscover(mention="vulnerability"), VulnerabilityDiscover(mention="discovered")]

        scorer_results = scorer(reference=reference, predictions=predictions)["events"]
        if "class_scores" in scorer_results:
            scorer_results.pop("class_scores")

        self.assertDictEqual(
            scorer_results,
            {
                "precision": 1.0,
                "recall": 1.0,
                "f1-score": 1.0,
            },
        )

    def test_event_argument_scorer(self):
        from src.tasks.casie.prompts_eae import VulnerabilityDiscover
        from src.tasks.casie.scorer import CASIEEventArgumentScorer

        scorer = CASIEEventArgumentScorer(allow_partial_match=True)

        reference = [
            VulnerabilityDiscover(
                mention="says",
                cve=["CVE-2018-12799", "CVE-2018-12808"],
                used_for=["lead to arbitrary code execution"],
                discoverer=["The tech giant"],
                supported_platform=[],
                vulnerability=[
                    "an out of bounds write issue",
                    "the security flaws",
                    "an untrusted pointer dereference problem",
                ],
                vulnerable_system=[],
                system_owner=[],
                system_version=[],
                time=[],
            ),
        ]

        # Perfect alingment
        scorer_results = scorer(reference=reference, predictions=reference)["arguments"]
        if "class_scores" in scorer_results:
            scorer_results.pop("class_scores")

        self.assertDictEqual(
            scorer_results,
            {"precision": 1.0, "recall": 1.0, "f1-score": 1.0},
        )

        # No alignment
        scorer_results = scorer(reference=reference, predictions=[])["arguments"]
        if "class_scores" in scorer_results:
            scorer_results.pop("class_scores")

        self.assertDictEqual(
            scorer_results,
            {"precision": 0.0, "recall": 0.0, "f1-score": 0.0},
        )

        predictions = [
            VulnerabilityDiscover(
                mention="says",
                cve="CVE-2018-12799",
                used_for=["arbitrary code execution"],
                discoverer=["tech giant"],
                supported_platform=[],
                vulnerability=[
                    "an out of bounds write issue",
                    "the security flaws",
                    "an untrusted pointer",
                ],
                vulnerable_system=[],
                system_owner=[],
                system_version=[],
                time=[],
            ),
        ]
        for pred in predictions:
            pred.assert_typing_constraints()

        import rich

        rich.print(predictions)
        rich.print(reference)

        # Partially aligned
        # TP = 6
        #   - cve(""CVE-2018-12799"")
        #   - used_for("lead to arbitrary code execution")
        #   - discoverer("The tech giant")
        #   - vulnerability("an out of bounds write issue")
        #   - vulnerability("the security flaws")
        #   - vulnerability("an untrusted pointer dereference problem")
        # FP = 0
        # FN = 0

        # precision -> 6 / 6 = 1.0
        # recall -> 6 / 7 = 0.8571428571428571
        # F1 -> 2 * (0.8571428571428571 * 1.0) / (0.8571428571428571 + 1.0) = 0.9285714285714286
        scorer_results = scorer(reference=reference, predictions=predictions)["arguments"]
        if "class_scores" in scorer_results:
            scorer_results.pop("class_scores")
        self.assertDictEqual(
            scorer_results,
            {
                "precision": 1.0,
                "recall": 0.8571428571428571,
                "f1-score": 0.923076923076923,
            },
        )

src/tests/test_scorers.py
hitz-zentroa/GoLLIE
import os
import unittest

import psutil

from src.dataset.dataset import DataCollatorForCoLLIE
from src.model.load_model import load_model
from transformers.testing_utils import require_bitsandbytes, require_torch_gpu


class TestCollieTrainer(unittest.TestCase):
    @unittest.skipIf(
        not os.path.exists("/gaueko1/hizkuntza-ereduak/LLaMA/lm/huggingface/7B"),
        "No LLaMA model available",
    )
    @require_torch_gpu
    @require_bitsandbytes
    def test_save(self):
        from tempfile import TemporaryDirectory

        from src.config import ModelArguments
        from src.trainer import CollieTrainer
        from transformers import Seq2SeqTrainingArguments

        with TemporaryDirectory() as tmpdirname:
            model_args = ModelArguments(
                model_name_or_path="/gaueko1/hizkuntza-ereduak/LLaMA/lm/huggingface/7B/",
                lora_weights_name_or_path=None,
                torch_dtype="auto",
                use_lora=True,
                quantization=4,
            )

            training_args = Seq2SeqTrainingArguments(
                output_dir=tmpdirname,
            )

            model, _ = load_model(
                inference=False,
                model_weights_name_or_path=model_args.model_name_or_path,
                quantization=model_args.quantization,
                use_lora=model_args.use_lora,
                lora_target_modules=model_args.lora_target_modules,
                torch_dtype=model_args.torch_dtype,
            )

            trainer = CollieTrainer(model=model, args=training_args)
            trainer.save_model()

            self.assertTrue(
                os.path.exists(os.path.join(tmpdirname, "adapter_model.bin")),
                "If PEFT is used, then the adapter must be saved.",
            )
            self.assertFalse(
                any(elem.startswith("pytorch_model") for elem in os.listdir(tmpdirname)),
                "If PEFT is used, then the model should not be saved.",
            )

        with TemporaryDirectory() as tmpdirname:
            model_args = ModelArguments(
                model_name_or_path="/gaueko1/hizkuntza-ereduak/LLaMA/lm/huggingface/7B/",
                lora_weights_name_or_path=None,
                torch_dtype="auto",
                use_lora=False,
                quantization=None,
            )

            training_args = Seq2SeqTrainingArguments(
                output_dir=tmpdirname,
            )

            model, _ = load_model(
                inference=False,
                model_weights_name_or_path=model_args.model_name_or_path,
                quantization=model_args.quantization,
                use_lora=model_args.use_lora,
                lora_target_modules=model_args.lora_target_modules,
                torch_dtype=model_args.torch_dtype,
            )

            trainer = CollieTrainer(model=model, args=training_args)
            trainer.save_model()

            self.assertFalse(
                os.path.exists(os.path.join(tmpdirname, "adapter_model.bin")),
                "If PEFT is not used, then no adapter must be saved.",
            )
            self.assertTrue(
                any(elem.startswith("pytorch_model") for elem in os.listdir(tmpdirname)),
                "If PEFT is not used, then the model have to be saved.",
            )

    @unittest.skipIf(psutil.virtual_memory().total / (1 << 30) < 15, "Not enough RAM available to run the test")
    def test_compute_loss(self):
        from tempfile import TemporaryDirectory

        from torch.utils.data import DataLoader

        from src.model.load_model import load_model
        from src.tests.test_dataset import get_dataset
        from src.trainer import CollieTrainer
        from transformers import Seq2SeqTrainingArguments, Trainer

        with TemporaryDirectory() as tmpdirname:
            training_args = Seq2SeqTrainingArguments(
                output_dir=tmpdirname,
            )

            model, tokenizer = load_model(
                inference=False,
                model_weights_name_or_path="EleutherAI/gpt-neo-125m",
                quantization=None,
                use_lora=False,
                lora_target_modules=None,
                torch_dtype="float32",
            )

            collie_trainer = CollieTrainer(model=model, args=training_args)
            collie_trainer.model = collie_trainer.model.to("cpu")
            trainer = Trainer(model=model, args=training_args)
            trainer.model = trainer.model.to("cpu")

            dataset, prompt, result = get_dataset(
                tokenizer=tokenizer,
                is_encoder_decoder=False,
                inference=False,
                prompt_loss_weight=0.0,
                max_length=tokenizer.model_max_length,
            )

            datacollator = DataCollatorForCoLLIE(
                tokenizer,
                pad_to_multiple_of=tokenizer.model_max_length,
                return_tensors="pt",
                padding=True,
                label_pad_token_id=-100,
            )

            dataloader = DataLoader(dataset, batch_size=1, collate_fn=datacollator, shuffle=False)
            inputs = list(dataloader)[0]
            inputs = {key: value.to("cpu") for key, value in inputs.items()}
            not_result_mask = inputs["loss_weight_mask"] < 1.0

            collie_loss = collie_trainer.compute_loss(model=model, inputs=inputs, return_outputs=False)

            inputs = list(dataloader)[0]
            inputs = {key: value.to("cpu") for key, value in inputs.items() if key not in ["loss_weight_mask"]}
            inputs["labels"][not_result_mask] = -100

            model_loss = model(**inputs).loss
            original_loss = trainer.compute_loss(model=model, inputs=inputs, return_outputs=False)

            self.assertAlmostEqual(model_loss.item(), original_loss.item(), delta=1e-5)
            self.assertAlmostEqual(collie_loss.item(), original_loss.item(), delta=1e-5)

    @unittest.skipIf(psutil.virtual_memory().total / (1 << 30) < 15, "Not enough RAM available to run the test")
    def test_compute_loss_lora(self):
        from tempfile import TemporaryDirectory

        from torch.utils.data import DataLoader

        from src.model.load_model import load_model
        from src.tests.test_dataset import get_dataset
        from src.trainer import CollieTrainer
        from transformers import Seq2SeqTrainingArguments, Trainer

        with TemporaryDirectory() as tmpdirname:
            training_args = Seq2SeqTrainingArguments(
                output_dir=tmpdirname,
            )

            model, tokenizer = load_model(
                inference=False,
                model_weights_name_or_path="EleutherAI/gpt-neo-125m",
                quantization=None,
                use_lora=True,
                torch_dtype="float32",
            )

            collie_trainer = CollieTrainer(model=model, args=training_args)
            collie_trainer.model = collie_trainer.model.to("cpu")
            trainer = Trainer(model=model, args=training_args)
            trainer.model = trainer.model.to("cpu")

            dataset, prompt, result = get_dataset(
                tokenizer=tokenizer,
                is_encoder_decoder=False,
                inference=False,
                prompt_loss_weight=0.0,
                max_length=tokenizer.model_max_length,
            )

            datacollator = DataCollatorForCoLLIE(
                tokenizer,
                pad_to_multiple_of=tokenizer.model_max_length,
                return_tensors="pt",
                padding=True,
                label_pad_token_id=-100,
            )

            dataloader = DataLoader(dataset, batch_size=1, collate_fn=datacollator, shuffle=False)
            inputs = list(dataloader)[0]
            inputs = {key: value.to("cpu") for key, value in inputs.items()}
            not_result_mask = inputs["loss_weight_mask"] < 1.0

            collie_loss = collie_trainer.compute_loss(model=model, inputs=inputs, return_outputs=False)

            inputs = list(dataloader)[0]
            inputs = {key: value.to("cpu") for key, value in inputs.items() if key not in ["loss_weight_mask"]}
            inputs["labels"][not_result_mask] = -100

            model_loss = model(**inputs).loss
            original_loss = trainer.compute_loss(model=model, inputs=inputs, return_outputs=False)

            self.assertAlmostEqual(model_loss.item(), original_loss.item(), delta=1e-5)
            self.assertAlmostEqual(collie_loss.item(), original_loss.item(), delta=1e-5)

src/tests/test_trainer.py
hitz-zentroa/GoLLIE
import bisect
import logging
import os
from typing import Callable, Dict, List, Optional, Tuple, Union

import torch
import torch.nn as nn
from rich.progress import Progress, SpinnerColumn, TimeElapsedColumn
from torch.utils.data import Dataset
from torch.utils.data.dataset import Iterable, IterableDataset, T_co

from src.config import ModelArguments
from src.dataset.dataset import CollieDataset
from transformers import (
    DataCollator,
    PreTrainedModel,
    PreTrainedTokenizerBase,
    Seq2SeqTrainer,
    Seq2SeqTrainingArguments,
    TrainerControl,
    TrainerState,
    TrainingArguments,
)
from transformers.modeling_utils import unwrap_model
from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
from transformers.trainer import TRAINING_ARGS_NAME, logger
from transformers.trainer_callback import TrainerCallback
from transformers.trainer_utils import EvalPrediction, has_length
from transformers.utils import SAFE_WEIGHTS_NAME, WEIGHTS_NAME, is_safetensors_available


if is_safetensors_available():
    import safetensors.torch


class RichProgressCallback(TrainerCallback):
    """
    A [`TrainerCallback`] that displays the progress of training or evaluation.
    """

    def __init__(self):
        self.training_bar = None
        self.training_task = None
        self.prediction_bar = None
        self.prediction_task = None

    def on_train_begin(self, args, state, control, **kwargs):
        if state.is_local_process_zero:
            self.training_bar = Progress(
                SpinnerColumn(),
                *Progress.get_default_columns(),
                TimeElapsedColumn(),
                auto_refresh=False,
            )

            self.training_bar.start()
            self.training_task = self.training_bar.add_task("[cyan]Training: ", total=state.max_steps)
        self.current_step = 0

    def on_step_end(self, args, state, control, **kwargs):
        if state.is_local_process_zero:
            self.training_bar.update(
                self.training_task,
                advance=state.global_step - self.current_step,
                refresh=True,
            )
            self.current_step = state.global_step

    def on_prediction_step(self, args, state, control, eval_dataloader=None, **kwargs):
        if state.is_local_process_zero and has_length(eval_dataloader):
            if self.prediction_bar is None:
                self.prediction_bar = Progress(
                    SpinnerColumn(),
                    *Progress.get_default_columns(),
                    TimeElapsedColumn(),
                    auto_refresh=False,
                )
                self.prediction_bar.start()
                self.prediction_task = self.prediction_bar.add_task("[cyan]Predicting: ", total=len(eval_dataloader))
            self.prediction_bar.update(self.prediction_task, advance=1, refresh=True)

    def on_evaluate(self, args, state, control, **kwargs):
        if state.is_local_process_zero:
            if self.prediction_bar is not None:
                self.prediction_bar.stop()
            self.prediction_bar = None
            self.prediction_task = None

    def on_predict(self, args, state, control, **kwargs):
        if state.is_local_process_zero:
            if self.prediction_bar is not None:
                self.prediction_bar.stop()
            self.prediction_bar = None
            self.prediction_task = None

    def on_log(self, args, state, control, logs=None, **kwargs):
        pass

    def on_train_end(self, args, state, control, **kwargs):
        if state.is_local_process_zero:
            self.training_bar.stop()
            self.training_bar = None
            self.training_task = None


class CollieTrainer(Seq2SeqTrainer):
    """
    The CollieTrainer is an adaptation of the 🤗 Transformers Trainer.

    Args:
        model ([`PreTrainedModel`] or `torch.nn.Module`, *optional*):
            The model to train, evaluate or use for predictions. If not provided, a `model_init` must be passed.
            <Tip>
            [`Trainer`] is optimized to work with the [`PreTrainedModel`] provided by the library. You can still use
            your own models defined as `torch.nn.Module` as long as they work the same way as the 🤗 Transformers
            models.
            </Tip>
        args ([`TrainingArguments`], *optional*):
            The arguments to tweak for training. Will default to a basic instance of [`TrainingArguments`] with the
            `output_dir` set to a directory named *tmp_trainer* in the current directory if not provided.
        data_collator (`DataCollator`, *optional*):
            The function to use to form a batch from a list of elements of `train_dataset` or `eval_dataset`. Will
            default to [`default_data_collator`] if no `tokenizer` is provided, an instance of
            [`DataCollatorWithPadding`] otherwise.
        train_dataset (`torch.utils.data.Dataset` or `torch.utils.data.IterableDataset`, *optional*):
            The dataset to use for training. If it is a [`~datasets.Dataset`], columns not accepted by the
            `model.forward()` method are automatically removed.
            Note that if it's a `torch.utils.data.IterableDataset` with some randomization and you are training in a
            distributed fashion, your iterable dataset should either use a internal attribute `generator` that is a
            `torch.Generator` for the randomization that must be identical on all processes (and the Trainer will
            manually set the seed of this `generator` at each epoch) or have a `set_epoch()` method that internally
            sets the seed of the RNGs used.
        eval_dataset (Union[`torch.utils.data.Dataset`, Dict[str, `torch.utils.data.Dataset`]), *optional*):
             The dataset to use for evaluation. If it is a [`~datasets.Dataset`], columns not accepted by the
             `model.forward()` method are automatically removed. If it is a dictionary, it will evaluate on each
             dataset prepending the dictionary key to the metric name.
        tokenizer ([`PreTrainedTokenizerBase`], *optional*):
            The tokenizer used to preprocess the data. If provided, will be used to automatically pad the inputs to the
            maximum length when batching inputs, and it will be saved along the model to make it easier to rerun an
            interrupted training or reuse the fine-tuned model.
        model_init (`Callable[[], PreTrainedModel]`, *optional*):
            A function that instantiates the model to be used. If provided, each call to [`~Trainer.train`] will start
            from a new instance of the model as given by this function.
            The function may have zero argument, or a single one containing the optuna/Ray Tune/SigOpt trial object, to
            be able to choose different architectures according to hyper parameters (such as layer count, sizes of
            inner layers, dropout probabilities etc).
        compute_metrics (`Callable[[EvalPrediction], Dict]`, *optional*):
            The function that will be used to compute metrics at evaluation. Must take a [`EvalPrediction`] and return
            a dictionary string to metric values.
        callbacks (List of [`TrainerCallback`], *optional*):
            A list of callbacks to customize the training loop. Will add those to the list of default callbacks
            detailed in [here](callback).
            If you want to remove one of the default callbacks used, use the [`Trainer.remove_callback`] method.
        optimizers (`Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`, *optional*): A tuple
            containing the optimizer and the scheduler to use. Will default to an instance of [`AdamW`] on your model
            and a scheduler given by [`get_linear_schedule_with_warmup`] controlled by `args`.
        preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`, *optional*):
            A function that preprocess the logits right before caching them at each evaluation step. Must take two
            tensors, the logits and the labels, and return the logits once processed as desired. The modifications made
            by this function will be reflected in the predictions received by `compute_metrics`.
            Note that the labels (second parameter) will be `None` if the dataset does not have them.
    Important attributes:
        - **model** -- Always points to the core model. If using a transformers model, it will be a [`PreTrainedModel`]
          subclass.
        - **model_wrapped** -- Always points to the most external model in case one or more other modules wrap the
          original model. This is the model that should be used for the forward pass. For example, under `DeepSpeed`,
          the inner model is wrapped in `DeepSpeed` and then again in `torch.nn.DistributedDataParallel`. If the inner
          model hasn't been wrapped, then `self.model_wrapped` is the same as `self.model`.
        - **is_model_parallel** -- Whether or not a model has been switched to a model parallel mode (different from
          data parallelism, this means some of the model layers are split on different GPUs).
        - **place_model_on_device** -- Whether or not to automatically place the model on the device - it will be set
          to `False` if model parallel or deepspeed is used, or if the default
          `TrainingArguments.place_model_on_device` is overridden to return `False` .
        - **is_in_train** -- Whether or not a model is currently running `train` (e.g. when `evaluate` is called while
          in `train`)
    """

    def __init__(
        self,
        model: Union[PreTrainedModel, nn.Module] = None,
        args: TrainingArguments = None,
        data_collator: Optional[DataCollator] = None,
        train_dataset: Optional[Dataset] = None,
        eval_dataset: Optional[Union[Dataset, Dict[str, Dataset]]] = None,
        tokenizer: Optional[PreTrainedTokenizerBase] = None,
        model_init: Optional[Callable[[], PreTrainedModel]] = None,
        compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None,
        callbacks: Optional[List[TrainerCallback]] = None,
        optimizers: Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (
            None,
            None,
        ),
        preprocess_logits_for_metrics: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None,
    ):
        if callbacks is None:
            callbacks = [RotateDatasetCallback()]
        else:
            callbacks.append(RotateDatasetCallback())

        self.first_train_batch = True

        # HuggingFace mad TrainingArguments inmutable and therefore the next function crashes
        # We made TrainingArguments mutable again
        TrainingArguments.__setattr__ = object.__setattr__

        # Ensure that the values are floats
        args.set_optimizer(
            name=args.optim,
            learning_rate=float(args.learning_rate),
            weight_decay=float(args.weight_decay),
            beta1=float(args.adam_beta1),
            beta2=float(args.adam_beta2),
            epsilon=float(args.adam_epsilon),
            args=args.optim_args,
        )

        super().__init__(
            model=model,
            args=args,
            data_collator=data_collator,
            train_dataset=train_dataset,
            eval_dataset=eval_dataset,
            tokenizer=None,  # We don't want to save the tokenizer with the model or use it for padding
            model_init=model_init,
            compute_metrics=compute_metrics,
            callbacks=callbacks,
            optimizers=optimizers,
            preprocess_logits_for_metrics=preprocess_logits_for_metrics,
        )
        # Change the tqdm progress callback with `RichProgressCallback`
        # _prev_progress_callback = self.pop_callback(ProgressCallback)
        # if _prev_progress_callback:
        #    self.add_callback(RichProgressCallback)

        if tokenizer is not None:
            # We want the tokenizer to decode the first training batch for debugging purposes
            self.tokenizer = tokenizer
        else:
            self.tokenizer = None

    def compute_loss(self, model, inputs, return_outputs=False):
        """
        How the loss is computed by Trainer. By default, all models return the loss in the first element.
        Subclass and override for custom behavior.
        """

        if "labels" in inputs:
            labels = inputs.pop("labels")
        else:
            raise ValueError("You should supply a labels key to compute the loss")

        if "loss_weight_mask" in inputs:
            loss_weight_mask = inputs.pop("loss_weight_mask")
        else:
            raise ValueError("You should supply a loss_weight_mask key to compute the loss")

        # Print first batch of training data for debugging
        if self.first_train_batch:
            self.first_train_batch = False
            print_input_ids = inputs["input_ids"][:8].clone().detach().cpu()
            print_attention_mask = inputs["attention_mask"][:8].clone().detach().cpu()
            print_labels = labels[:8].clone().detach().cpu()
            print_loss_weight_mask = loss_weight_mask[:8].clone().detach().cpu()

            print("*** First batch of training data ***")
            print("-- input_ids --")
            if self.tokenizer is not None:
                print_input_ids[print_input_ids == -100] = self.tokenizer.pad_token_id
                print(self.tokenizer.batch_decode(print_input_ids))
            else:
                print(print_input_ids.tolist())
            print("-- attention_mask --")
            print(print_attention_mask.tolist())
            print("-- labels --")
            if self.tokenizer is not None:
                print_labels[print_labels == -100] = self.tokenizer.pad_token_id
                print(self.tokenizer.batch_decode(print_labels))
            else:
                print(print_labels[:8].tolist())
            print("-- loss_weight_mask --")
            print(print_loss_weight_mask.tolist())
            print()

        outputs = model(**inputs, use_cache=False)

        logits = outputs["logits"] if isinstance(outputs, dict) else outputs[0]

        model_name = unwrap_model(model)._get_name()
        if model_name in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES.values() or model_name == "PeftModelForCausalLM":
            logits = logits[..., :-1, :].contiguous()
            labels = labels[..., 1:].contiguous()
            loss_weight_mask = loss_weight_mask[..., 1:].contiguous()

        logits = logits.view(-1, logits.size(-1))
        labels = labels.view(-1)
        loss_weight_mask = loss_weight_mask.view(-1)
        loss_fct = nn.CrossEntropyLoss(reduction="none", ignore_index=-100)

        loss = loss_fct(logits, labels)
        loss = torch.sum(loss * loss_weight_mask) / torch.sum(loss_weight_mask)

        return (loss, outputs) if return_outputs else loss

    # Modify the Seq2SeqTrainer from transformers to only save the LoRA weights if we are using a LoRA model
    # Original trainer saves the full state dict. It doesn't make sense for us to create a full copy
    # of LLaMA weights each time we save a checkpoint since we do not modify them.
    def _save(self, output_dir: Optional[str] = None, state_dict=None):
        # If we are executing this function, we are the process zero, so we don't check for that.
        output_dir = output_dir if output_dir is not None else self.args.output_dir
        os.makedirs(output_dir, exist_ok=True)
        logger.info(f"Saving model checkpoint to {output_dir}")
        # Save a trained model and configuration using `save_pretrained()`.
        # They can then be reloaded using `from_pretrained()

        # Find out if the model is a LoRA Peft Model
        # try:
        #     from peft import PeftModel, LoraModel

        #     if isinstance(unwrap_model(self.model), PeftModel):
        #         if isinstance(unwrap_model(self.model).base_model, LoraModel):
        #             unwrap_model(self.model).save_pretrained(
        #                 output_dir,
        #             )
        #             return
        # except ImportError:
        #     pass

        try:
            from peft import PeftModel
        except ImportError:
            PeftModel = None

        if not isinstance(self.model, PreTrainedModel) and not (PeftModel and isinstance(self.model, PeftModel)):
            if state_dict is None:
                state_dict = self.model.state_dict()

            if isinstance(unwrap_model(self.model), PreTrainedModel):
                unwrap_model(self.model).save_pretrained(
                    output_dir,
                    state_dict=state_dict,
                    safe_serialization=self.args.save_safetensors,
                )
            else:
                logger.info("Trainer.model is not a `PreTrainedModel`, only saving its state dict.")
                if self.args.save_safetensors:
                    safetensors.torch.save_file(state_dict, os.path.join(output_dir, SAFE_WEIGHTS_NAME))
                else:
                    torch.save(state_dict, os.path.join(output_dir, WEIGHTS_NAME))
        else:
            self.model.save_pretrained(
                output_dir,
                state_dict=state_dict,
                safe_serialization=self.args.save_safetensors,
            )

        if self.tokenizer is not None:
            self.tokenizer.save_pretrained(output_dir)

        # Good practice: save your training arguments together with the trained model
        torch.save(self.args, os.path.join(output_dir, TRAINING_ARGS_NAME))


class ConcatDataset(Dataset[T_co]):
    r"""Dataset as a concatenation of multiple datasets.

    This class is useful to assemble different existing datasets.

    Args:
        datasets (sequence): List of datasets to be concatenated
    """
    datasets: List[CollieDataset]
    cumulative_sizes: List[int]

    @staticmethod
    def cumsum(sequence):
        r, s = [], 0
        for e in sequence:
            l = len(e)
            r.append(l + s)
            s += l
        return r

    def __init__(self, datasets: Iterable[CollieDataset]) -> None:
        super().__init__()
        self.datasets = list(datasets)
        assert len(self.datasets) > 0, "datasets should not be an empty iterable"  # type: ignore[arg-type]
        for d in self.datasets:
            assert not isinstance(d, IterableDataset), "ConcatDataset does not support IterableDataset"
        self.cumulative_sizes = self.cumsum(self.datasets)

    def __len__(self):
        return self.cumulative_sizes[-1]

    def __getitem__(self, idx):
        if idx < 0:
            if -idx > len(self):
                raise ValueError("absolute value of index should not exceed dataset length")
            idx = len(self) + idx
        dataset_idx = bisect.bisect_right(self.cumulative_sizes, idx)
        if dataset_idx == 0:
            sample_idx = idx
        else:
            sample_idx = idx - self.cumulative_sizes[dataset_idx - 1]
        return self.datasets[dataset_idx][sample_idx]

    @property
    def cummulative_sizes(self):
        logging.warning("cummulative_sizes attribute is renamed to cumulative_sizes", DeprecationWarning, stacklevel=2)
        return self.cumulative_sizes

    def rotate_split(self):
        for x in self.datasets:
            x.rotate_split()
        self.cumulative_sizes = self.cumsum(self.datasets)


class RotateDatasetCallback(TrainerCallback):
    def on_epoch_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
        if "train_dataloader" in kwargs:
            kwargs["train_dataloader"].dataset.rotate_split()
        else:
            logging.warning("No train_dataloader in kwargs. Skipping rotate_split()")


def get_correct_torch_dtype(
    quantization: int,
    model_args: ModelArguments,
    training_args: Seq2SeqTrainingArguments,
) -> "str":
    """
    Returns the correct torch dtype based on the model and training arguments (if quantization is enabled).

    Args:
        quantization (`int`, optional):
            '4' or '8' for 4 bits or 8 bits quantization or None for 16/32bits training. Defaults to `None`.
        model_args (:class:`~transformers.ModelArguments`):
            The model arguments.
        training_args (:class:`~transformers.Seq2SeqTrainingArguments`):
            The training arguments.

    Returns:
        :obj:`str`: The correct torch dtype.
    """

    if isinstance(quantization, str):
        quantization = int(quantization)

    if quantization in [4, 8]:
        if training_args.fp16:
            if model_args.torch_dtype in ["auto", None]:
                logging.warning(
                    "Quantification and fp16 are enabled, but torch_dtype is not set. Setting torch_dtype to float16."
                )

            elif model_args.torch_dtype != "float16":
                logging.warning(
                    f"Quantification and fp16 are enabled, but torch_dtype is set to {model_args.torch_dtype}. "
                    "This can cause issues. We will override torch_dtype to float16."
                )
            return "float16"

        elif training_args.bf16:
            if model_args.torch_dtype in ["auto", None]:
                logging.warning(
                    "Quantification and bf16 are enabled, but torch_dtype is not set. Setting torch_dtype to bfloat16."
                )
            elif model_args.torch_dtype != "bfloat16":
                logging.warning(
                    f"Quantification and bf16 are enabled, but torch_dtype is set to {model_args.torch_dtype}. "
                    "This can cause issues. We will override torch_dtype to bfloat16."
                )
            return "bfloat16"

    return model_args.torch_dtype

src/trainer.py
hitz-zentroa/GoLLIE