notebooks/Event Extraction.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8b2afe3b",
   "metadata": {},
   "source": [
    "<img src=\"../assets/CoLLIE_blue.png\" alt=\"GoLLIE\" width=\"200\"/>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "003b73a2",
   "metadata": {},
   "source": [
    "# Event Extraction with GoLLIE\n",
    "\n",
    "This notebook provides an example of how to conduct Event Extraction using GoLLIE.\n",
    "\n",
    "In the paper, to compare GoLLIE with the previous state-of-the-art (SOTA), we divided the Event Extraction task into two sub-tasks: **Event Extraction** (EE) and **Event Argument Extraction** (EAE). The former focuses on detecting event instances in a text, while the latter, given an event instance, predicts its attributes, such as the persons involved and their roles. Additionally, GoLLIE is capable of performing **End-to-End Event Extraction**. In this notebook, we will demonstrate all three scenarios. This notebook covers:\n",
    "\n",
    "- How to define guidelines for a task\n",
    "- How to load GoLLIE\n",
    "- How to conduct Event Extraction (EE) using GoLLIE\n",
    "- How to perform Event Argument Extraction (EAE) with GoLLIE\n",
    "- How to execute end-to-end Event Extraction with GoLLIE\n",
    "- How to implement a scorer and evaluate the results\n",
    "\n",
    "You can modify this notebook to run any Event Extraction task you want"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b015c64",
   "metadata": {},
   "source": [
    "### Import requeriments\n",
    "\n",
    "See the requeriments.txt file in the main directory to install the required dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 168,
   "id": "9ed51491",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append(\"../\") # Add the GoLLIE base directory to sys path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 169,
   "id": "28ff498a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import rich\n",
    "import logging\n",
    "from src.model.load_model import load_model\n",
    "import black\n",
    "import inspect\n",
    "from jinja2 import Template\n",
    "import tempfile\n",
    "from src.tasks.utils_typing import AnnotationList\n",
    "logging.basicConfig(level=logging.INFO)\n",
    "from typing import Dict, List, Type"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "004626bb",
   "metadata": {},
   "source": [
    "## Load GoLLIE\n",
    "\n",
    "We will load GOLLIE-7B from the huggingface-hub.\n",
    "You can use the function AutoModelForCausalLM.from_pretrained if you prefer it. However, we provide a handy load_model function with many functionalities already implemented that will assist you in reproducing our results.\n",
    "\n",
    "Please note that setting use_flash_attention=True is mandatory. Our flash attention implementation has small numerical differences compared to the attention implementation in Huggingface. Using use_flash_attention=False will result in the model producing inferior results. Flash attention requires an available CUDA GPU. Running GOLLIE pre-trained models on a CPU is not supported. We plan to address this in future releases.\n",
    "\n",
    "- Set force_auto_device_map=True to automatically load the model on available GPUs.\n",
    "- Set quantization=4 if the model doesn't fit in your GPU memory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "cb841c40",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:Loading model model from HiTZ/GoLLIE-7B\n",
      "WARNING:root:Using auto device map, we will split the model across GPUs and CPU to fit the model in memory.\n",
      "INFO:root:We will load the model using the following device map: auto and max_memory: None\n",
      "Loading the tokenizer from the `special_tokens_map.json` and the `added_tokens.json` will be removed in `transformers 5`,  it is kept for forward compatibility, but it is recommended to update your `tokenizer_config.json` by uploading it again. You will see the new `added_tokens_decoder` attribute that will store the relevant information.\n",
      "INFO:root:Loading model with dtype: torch.bfloat16\n",
      "WARNING:root:Model HiTZ/GoLLIE-7B is an decoder-only model. We will load it as a CausalLM model.\n",
      "WARNING:root:Using Flash Attention for LLaMA model.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      ">>>> Flash Attention installed\n",
      ">>>> Flash RoPE installed\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0aa48f07d3f04d88978cf218bc79c97b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:Model dtype: torch.bfloat16\n",
      "INFO:root:Total model memory footprint: 13477.101762 MB\n",
      "INFO:root:Quantization is enabled, we will not merge LoRA layers into the model. Inference will be slower.\n"
     ]
    }
   ],
   "source": [
    "model, tokenizer = load_model(\n",
    "    inference=True,\n",
    "    model_weights_name_or_path=\"HiTZ/GoLLIE-7B\",\n",
    "    quantization=None,\n",
    "    use_lora=False,\n",
    "    force_auto_device_map=True,\n",
    "    use_flash_attention=True,\n",
    "    torch_dtype=\"bfloat16\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff31e54f",
   "metadata": {},
   "source": [
    "# Event Extraction"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a662bf3",
   "metadata": {},
   "source": [
    "## Define the guidelines\n",
    "\n",
    "First, we will define the labels and guidelines for the task. We will represent them as Python classes.\n",
    "\n",
    "We will use the ACE05 dataset guidelines for this example. You can find more information about the dataset and the guidelines here: https://www.ldc.upenn.edu/sites/www.ldc.upenn.edu/files/english-entities-guidelines-v6.6.pdf\n",
    "\n",
    "💡 Be creative and try to define your own guidelines to test GoLLIE!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 170,
   "id": "b3381513",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "\n",
    "from src.tasks.utils_typing import Event, dataclass\n",
    "\n",
    "# The following lines describe the task definition\n",
    "@dataclass\n",
    "class JusticeEvent(Event):\n",
    "    \"\"\"A JusticeEvent refers to any judicial action such as: arresting, jailing, releasing, granting parole, trial\n",
    "    starting, hearing, charging, indicting, suing, convicting, sentencing, fine, executing, extraditing, adquiting,\n",
    "    appealing or pardoning a Person entity.\"\"\"\n",
    "\n",
    "    mention: str\n",
    "    \"\"\"The text span that most clearly expresses the event.\n",
    "        Such as: \"parole\", \"sued\", \"sentenced\", \"appeal\", \"charged\" \n",
    "    \"\"\"\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class PersonnelEvent(Event):\n",
    "    \"\"\"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, gets nominated or is\n",
    "    elected for a position in a Organization.\"\"\"\n",
    "\n",
    "    mention: str\n",
    "    \"\"\"The text span that most clearly expresses the event.\n",
    "        Such as: \"won\", \"appoint\", \"retired\", \"fired\", \"appointed\" \n",
    "    \"\"\"\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class BusinessEvent(Event):\n",
    "    \"\"\"A BusinessEvent refers to actions related to Organizations such as: creating, merging, declaring bankruptcy or\n",
    "    ending organizations (including government agencies).\"\"\"\n",
    "\n",
    "    mention: str\n",
    "    \"\"\"The text span that most clearly expresses the event.\n",
    "        Such as: \"started\", \"open\", \"create\", \"closing\", \"merged\" \n",
    "    \"\"\"\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class LifeEvent(Event):\n",
    "    \"\"\"A LifeEvent occurs whenever a Person Entity borns, dies, gets married, divorced or gets injured.\"\"\"\n",
    "\n",
    "    mention: str\n",
    "    \"\"\"The text span that most clearly expresses the event.\n",
    "        Such as: \"wounded\", \"divorce\", \"birth\", \"born\", \"marriage\" \n",
    "    \"\"\"\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class MovementEvent(Event):\n",
    "    \"\"\"A TransportEvent occurs whenever an Artifact (Weapon or Vehicle) or a Person is moved from one Place (GPE, Facility,\n",
    "    Location) to another. This event requires the explicit mention of the Artifact or Person.\n",
    "    \"\"\"\n",
    "\n",
    "    mention: str\n",
    "    \"\"\"The text span that most clearly expresses the event.\n",
    "        Such as: \"travel\", \"arrived\", \"going\", \"moving\", \"take\" \n",
    "    \"\"\"\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class TransactionEvent(Event):\n",
    "    \"\"\"A TransactionEvent refers to buying, selling, loaning, borrowing, giving, or receving of Artifacts or\n",
    "    Organizations; or giving, receiving, borrowing, or lending Money.\"\"\"\n",
    "\n",
    "    mention: str\n",
    "    \"\"\"The text span that most clearly expresses the event.\n",
    "        Such as: \"donate\", \"received\", \"paying\", \"seize\", \"contributions\" \n",
    "    \"\"\"\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class ContactEvent(Event):\n",
    "    \"\"\"A ContactEvent occurs whenever two or more entities (persons or organization's representatives) come together at\n",
    "    a single location and interact with one another face-to-face or directly enages in discussion via written or\n",
    "    telephone communication.\"\"\"\n",
    "\n",
    "    mention: str\n",
    "    \"\"\"The text span that most clearly expresses the event.\n",
    "        Such as: \"meetings\", \"conference\", \"talked\", \"met\", \"letters\" \n",
    "    \"\"\"\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class ConflictEvent(Event):\n",
    "    \"\"\"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, ...).\n",
    "    \"\"\"\n",
    "\n",
    "    mention: str\n",
    "    \"\"\"The text span that most clearly expresses the event.\n",
    "        Such as: \"terrorism\", \"combat\", \"hit\", \"fight\", \"bombing\" \n",
    "    \"\"\"\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "EVENTS_DEFINITIONS: List[Event] = [\n",
    "    LifeEvent,\n",
    "    MovementEvent,\n",
    "    TransactionEvent,\n",
    "    BusinessEvent,\n",
    "    ConflictEvent,\n",
    "    ContactEvent,\n",
    "    PersonnelEvent,\n",
    "    JusticeEvent,\n",
    "]\n",
    "    \n",
    "if __name__ == \"__main__\":\n",
    "    cell_txt = In[-1]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40f8ef55",
   "metadata": {},
   "source": [
    "### Print the guidelines to guidelines.py\n",
    "\n",
    "Due to IPython limitations, we must write the content of the previous cell to a file and then import the content from that file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 171,
   "id": "d4736a37",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"guidelines_ee.py\",\"w\",encoding=\"utf8\") as python_guidelines:\n",
    "    print(cell_txt,file=python_guidelines)\n",
    "\n",
    "import guidelines_ee as guidelines_ee"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0ac3db6f",
   "metadata": {},
   "source": [
    "We use inspect.getsource to get the guidelines as a string"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 172,
   "id": "89454475",
   "metadata": {},
   "outputs": [],
   "source": [
    "guidelines = [inspect.getsource(definition) for definition in guidelines_ee.EVENTS_DEFINITIONS]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2bd26b7",
   "metadata": {},
   "source": [
    "## Define input sentence\n",
    "\n",
    "Here we define the input sentence and the gold labels.\n",
    "\n",
    "You can define and empy list as gold labels if you don't have gold annotations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 173,
   "id": "4d1f1c8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"The Times reported that Vivendi Universal was negotiating to sell its flagship theme parks to the New York investment firm, Blackstone Group, for 8 billion as the first step toward dismantling its entertainment empire.\"\n",
    "gold = [\n",
    "    guidelines_ee.TransactionEvent(mention='sell'),\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "90501322",
   "metadata": {},
   "source": [
    "## Filling a template\n",
    "\n",
    "For EE we will use the following prompt template.\n",
    "We use Jinja templates, which are easy to implement and exceptionally fast. For more information, visit: https://jinja.palletsprojects.com/en/3.1.x/api/#high-level-api.\n",
    "\n",
    "```Python\n",
    "# The following lines describe the task definition\n",
    "{%- for definition in guidelines %}\n",
    "{{ definition }}\n",
    "{%- endfor %}\n",
    "\n",
    "# This is the text to analyze\n",
    "text = {{ text.__repr__() }}\n",
    "\n",
    "# The annotation instances that take place in the text above are listed here\n",
    "result = [\n",
    "{%- for ann in annotations %}\n",
    "    {{ ann }},\n",
    "{%- endfor %}\n",
    "]\n",
    "\n",
    "```\n",
    "\n",
    "This template is stored in `templates/prompt.txt`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 174,
   "id": "0f54034f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read template\n",
    "with open(\"../templates/prompt.txt\", \"rt\") as f:\n",
    "    template = Template(f.read())\n",
    "# Fill the template\n",
    "formated_text = template.render(guidelines=guidelines, text=text, annotations=gold, gold=gold)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "886dbc0c",
   "metadata": {},
   "source": [
    "### Black Code Formatter\n",
    "\n",
    "We use the Black Code Formatter to automatically unify all the prompts to the same format. \n",
    "\n",
    "https://github.com/psf/black"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 175,
   "id": "c8994924",
   "metadata": {},
   "outputs": [],
   "source": [
    "black_mode = black.Mode()\n",
    "formated_text = black.format_str(formated_text, mode=black_mode)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80aa5c0e",
   "metadata": {},
   "source": [
    "### Print the filled and formatted template"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 176,
   "id": "fa5f3106",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"># The following lines describe the task definition\n",
       "@dataclass\n",
       "class <span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">LifeEvent</span><span style=\"font-weight: bold\">(</span>Event<span style=\"font-weight: bold\">)</span>:\n",
       "    <span style=\"color: #008000; text-decoration-color: #008000\">\"\"\"A LifeEvent occurs whenever a Person Entity borns, dies, gets married, divorced or gets injured.\"\"\"</span>\n",
       "\n",
       "    mention: str\n",
       "    <span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"The text span that most clearly expresses the event.\n",
       "        Such as: <span style=\"color: #008000; text-decoration-color: #008000\">\"wounded\"</span>, <span style=\"color: #008000; text-decoration-color: #008000\">\"divorce\"</span>, <span style=\"color: #008000; text-decoration-color: #008000\">\"birth\"</span>, <span style=\"color: #008000; text-decoration-color: #008000\">\"born\"</span>, <span style=\"color: #008000; text-decoration-color: #008000\">\"marriage\"</span> \n",
       "    <span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"\n",
       "\n",
       "\n",
       "@dataclass\n",
       "class <span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">MovementEvent</span><span style=\"font-weight: bold\">(</span>Event<span style=\"font-weight: bold\">)</span>:\n",
       "    <span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"A TransportEvent occurs whenever an Artifact <span style=\"font-weight: bold\">(</span>Weapon or Vehicle<span style=\"font-weight: bold\">)</span> or a Person is moved from one Place <span style=\"font-weight: bold\">(</span>GPE, \n",
       "Facility,\n",
       "    Location<span style=\"font-weight: bold\">)</span> to another. This event requires the explicit mention of the Artifact or Person.\n",
       "    <span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"\n",
       "\n",
       "    mention: str\n",
       "    <span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"The text span that most clearly expresses the event.\n",
       "        Such as: <span style=\"color: #008000; text-decoration-color: #008000\">\"travel\"</span>, <span style=\"color: #008000; text-decoration-color: #008000\">\"arrived\"</span>, <span style=\"color: #008000; text-decoration-color: #008000\">\"going\"</span>, <span style=\"color: #008000; text-decoration-color: #008000\">\"moving\"</span>, <span style=\"color: #008000; text-decoration-color: #008000\">\"take\"</span> \n",
       "    <span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"\n",
       "\n",
       "\n",
       "@dataclass\n",
       "class <span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">TransactionEvent</span><span style=\"font-weight: bold\">(</span>Event<span style=\"font-weight: bold\">)</span>:\n",
       "    <span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"A TransactionEvent refers to buying, selling, loaning, borrowing, giving, or receving of Artifacts or\n",
       "    Organizations; or giving, receiving, borrowing, or lending Money.<span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"\n",
       "\n",
       "    mention: str\n",
       "    <span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"The text span that most clearly expresses the event.\n",
       "        Such as: <span style=\"color: #008000; text-decoration-color: #008000\">\"donate\"</span>, <span style=\"color: #008000; text-decoration-color: #008000\">\"received\"</span>, <span style=\"color: #008000; text-decoration-color: #008000\">\"paying\"</span>, <span style=\"color: #008000; text-decoration-color: #008000\">\"seize\"</span>, <span style=\"color: #008000; text-decoration-color: #008000\">\"contributions\"</span> \n",
       "    <span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"\n",
       "\n",
       "\n",
       "@dataclass\n",
       "class <span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">BusinessEvent</span><span style=\"font-weight: bold\">(</span>Event<span style=\"font-weight: bold\">)</span>:\n",
       "    <span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"A BusinessEvent refers to actions related to Organizations such as: creating, merging, declaring bankruptcy \n",
       "or\n",
       "    ending organizations <span style=\"font-weight: bold\">(</span>including government agencies<span style=\"font-weight: bold\">)</span>.<span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"\n",
       "\n",
       "    mention: str\n",
       "    <span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"The text span that most clearly expresses the event.\n",
       "        Such as: <span style=\"color: #008000; text-decoration-color: #008000\">\"started\"</span>, <span style=\"color: #008000; text-decoration-color: #008000\">\"open\"</span>, <span style=\"color: #008000; text-decoration-color: #008000\">\"create\"</span>, <span style=\"color: #008000; text-decoration-color: #008000\">\"closing\"</span>, <span style=\"color: #008000; text-decoration-color: #008000\">\"merged\"</span> \n",
       "    <span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"\n",
       "\n",
       "\n",
       "@dataclass\n",
       "class <span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">ConflictEvent</span><span style=\"font-weight: bold\">(</span>Event<span style=\"font-weight: bold\">)</span>:\n",
       "    <span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"A ConflictEvent refers to either violent physical acts causing harm or damage, but are not covered by Life \n",
       "events\n",
       "    <span style=\"font-weight: bold\">(</span>conflicts, clashes, fighting, gunfire, <span style=\"color: #808000; text-decoration-color: #808000\">...</span><span style=\"font-weight: bold\">)</span> or demonstrations <span style=\"font-weight: bold\">(</span>protests, sit-ins, strikes, riots, <span style=\"color: #808000; text-decoration-color: #808000\">...</span><span style=\"font-weight: bold\">)</span>.\n",
       "    <span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"\n",
       "\n",
       "    mention: str\n",
       "    <span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"The text span that most clearly expresses the event.\n",
       "        Such as: <span style=\"color: #008000; text-decoration-color: #008000\">\"terrorism\"</span>, <span style=\"color: #008000; text-decoration-color: #008000\">\"combat\"</span>, <span style=\"color: #008000; text-decoration-color: #008000\">\"hit\"</span>, <span style=\"color: #008000; text-decoration-color: #008000\">\"fight\"</span>, <span style=\"color: #008000; text-decoration-color: #008000\">\"bombing\"</span> \n",
       "    <span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"\n",
       "\n",
       "\n",
       "@dataclass\n",
       "class <span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">ContactEvent</span><span style=\"font-weight: bold\">(</span>Event<span style=\"font-weight: bold\">)</span>:\n",
       "    <span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"A ContactEvent occurs whenever two or more entities <span style=\"font-weight: bold\">(</span>persons or organization's representatives<span style=\"font-weight: bold\">)</span> come \n",
       "together at\n",
       "    a single location and interact with one another face-to-face or directly enages in discussion via written or\n",
       "    telephone communication.<span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"\n",
       "\n",
       "    mention: str\n",
       "    <span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"The text span that most clearly expresses the event.\n",
       "        Such as: <span style=\"color: #008000; text-decoration-color: #008000\">\"meetings\"</span>, <span style=\"color: #008000; text-decoration-color: #008000\">\"conference\"</span>, <span style=\"color: #008000; text-decoration-color: #008000\">\"talked\"</span>, <span style=\"color: #008000; text-decoration-color: #008000\">\"met\"</span>, <span style=\"color: #008000; text-decoration-color: #008000\">\"letters\"</span> \n",
       "    <span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"\n",
       "\n",
       "\n",
       "@dataclass\n",
       "class <span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">PersonnelEvent</span><span style=\"font-weight: bold\">(</span>Event<span style=\"font-weight: bold\">)</span>:\n",
       "    <span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"A PersonnelEvent occurs when a Person entity changes its job position <span style=\"font-weight: bold\">(</span>JobTitle entity<span style=\"font-weight: bold\">)</span> with respect an\n",
       "    Organization entity. It includes when a person starts working, ends working, changes offices within, gets \n",
       "nominated or is\n",
       "    elected for a position in a Organization.<span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"\n",
       "\n",
       "    mention: str\n",
       "    <span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"The text span that most clearly expresses the event.\n",
       "        Such as: <span style=\"color: #008000; text-decoration-color: #008000\">\"won\"</span>, <span style=\"color: #008000; text-decoration-color: #008000\">\"appoint\"</span>, <span style=\"color: #008000; text-decoration-color: #008000\">\"retired\"</span>, <span style=\"color: #008000; text-decoration-color: #008000\">\"fired\"</span>, <span style=\"color: #008000; text-decoration-color: #008000\">\"appointed\"</span> \n",
       "    <span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"\n",
       "\n",
       "\n",
       "@dataclass\n",
       "class <span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">JusticeEvent</span><span style=\"font-weight: bold\">(</span>Event<span style=\"font-weight: bold\">)</span>:\n",
       "    <span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"A JusticeEvent refers to any judicial action such as: arresting, jailing, releasing, granting parole, trial\n",
       "    starting, hearing, charging, indicting, suing, convicting, sentencing, fine, executing, extraditing, adquiting,\n",
       "    appealing or pardoning a Person entity.<span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"\n",
       "\n",
       "    mention: str\n",
       "    <span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"The text span that most clearly expresses the event.\n",
       "        Such as: <span style=\"color: #008000; text-decoration-color: #008000\">\"parole\"</span>, <span style=\"color: #008000; text-decoration-color: #008000\">\"sued\"</span>, <span style=\"color: #008000; text-decoration-color: #008000\">\"sentenced\"</span>, <span style=\"color: #008000; text-decoration-color: #008000\">\"appeal\"</span>, <span style=\"color: #008000; text-decoration-color: #008000\">\"charged\"</span> \n",
       "    <span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"\n",
       "\n",
       "\n",
       "# This is the text to analyze\n",
       "text = <span style=\"color: #008000; text-decoration-color: #008000\">\"The Times reported that Vivendi Universal was negotiating to sell its flagship theme parks to the New York </span>\n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">investment firm, Blackstone Group, for 8 billion as the first step toward dismantling its entertainment empire.\"</span>\n",
       "\n",
       "# The annotation instances that take place in the text above are listed here\n",
       "result = <span style=\"font-weight: bold\">[</span>\n",
       "    <span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">TransactionEvent</span><span style=\"font-weight: bold\">(</span><span style=\"color: #808000; text-decoration-color: #808000\">mention</span>=<span style=\"color: #008000; text-decoration-color: #008000\">\"sell\"</span><span style=\"font-weight: bold\">)</span>,\n",
       "<span style=\"font-weight: bold\">]</span>\n",
       "\n",
       "</pre>\n"
      ],
      "text/plain": [
       "# The following lines describe the task definition\n",
       "@dataclass\n",
       "class \u001B[1;35mLifeEvent\u001B[0m\u001B[1m(\u001B[0mEvent\u001B[1m)\u001B[0m:\n",
       "    \u001B[32m\"\"\"A LifeEvent occurs whenever a Person Entity borns, dies, gets married, divorced or gets injured.\"\"\"\u001B[0m\n",
       "\n",
       "    mention: str\n",
       "    \u001B[32m\"\"\u001B[0m\"The text span that most clearly expresses the event.\n",
       "        Such as: \u001B[32m\"wounded\"\u001B[0m, \u001B[32m\"divorce\"\u001B[0m, \u001B[32m\"birth\"\u001B[0m, \u001B[32m\"born\"\u001B[0m, \u001B[32m\"marriage\"\u001B[0m \n",
       "    \u001B[32m\"\"\u001B[0m\"\n",
       "\n",
       "\n",
       "@dataclass\n",
       "class \u001B[1;35mMovementEvent\u001B[0m\u001B[1m(\u001B[0mEvent\u001B[1m)\u001B[0m:\n",
       "    \u001B[32m\"\"\u001B[0m\"A TransportEvent occurs whenever an Artifact \u001B[1m(\u001B[0mWeapon or Vehicle\u001B[1m)\u001B[0m or a Person is moved from one Place \u001B[1m(\u001B[0mGPE, \n",
       "Facility,\n",
       "    Location\u001B[1m)\u001B[0m to another. This event requires the explicit mention of the Artifact or Person.\n",
       "    \u001B[32m\"\"\u001B[0m\"\n",
       "\n",
       "    mention: str\n",
       "    \u001B[32m\"\"\u001B[0m\"The text span that most clearly expresses the event.\n",
       "        Such as: \u001B[32m\"travel\"\u001B[0m, \u001B[32m\"arrived\"\u001B[0m, \u001B[32m\"going\"\u001B[0m, \u001B[32m\"moving\"\u001B[0m, \u001B[32m\"take\"\u001B[0m \n",
       "    \u001B[32m\"\"\u001B[0m\"\n",
       "\n",
       "\n",
       "@dataclass\n",
       "class \u001B[1;35mTransactionEvent\u001B[0m\u001B[1m(\u001B[0mEvent\u001B[1m)\u001B[0m:\n",
       "    \u001B[32m\"\"\u001B[0m\"A TransactionEvent refers to buying, selling, loaning, borrowing, giving, or receving of Artifacts or\n",
       "    Organizations; or giving, receiving, borrowing, or lending Money.\u001B[32m\"\"\u001B[0m\"\n",
       "\n",
       "    mention: str\n",
       "    \u001B[32m\"\"\u001B[0m\"The text span that most clearly expresses the event.\n",
       "        Such as: \u001B[32m\"donate\"\u001B[0m, \u001B[32m\"received\"\u001B[0m, \u001B[32m\"paying\"\u001B[0m, \u001B[32m\"seize\"\u001B[0m, \u001B[32m\"contributions\"\u001B[0m \n",
       "    \u001B[32m\"\"\u001B[0m\"\n",
       "\n",
       "\n",
       "@dataclass\n",
       "class \u001B[1;35mBusinessEvent\u001B[0m\u001B[1m(\u001B[0mEvent\u001B[1m)\u001B[0m:\n",
       "    \u001B[32m\"\"\u001B[0m\"A BusinessEvent refers to actions related to Organizations such as: creating, merging, declaring bankruptcy \n",
       "or\n",
       "    ending organizations \u001B[1m(\u001B[0mincluding government agencies\u001B[1m)\u001B[0m.\u001B[32m\"\"\u001B[0m\"\n",
       "\n",
       "    mention: str\n",
       "    \u001B[32m\"\"\u001B[0m\"The text span that most clearly expresses the event.\n",
       "        Such as: \u001B[32m\"started\"\u001B[0m, \u001B[32m\"open\"\u001B[0m, \u001B[32m\"create\"\u001B[0m, \u001B[32m\"closing\"\u001B[0m, \u001B[32m\"merged\"\u001B[0m \n",
       "    \u001B[32m\"\"\u001B[0m\"\n",
       "\n",
       "\n",
       "@dataclass\n",
       "class \u001B[1;35mConflictEvent\u001B[0m\u001B[1m(\u001B[0mEvent\u001B[1m)\u001B[0m:\n",
       "    \u001B[32m\"\"\u001B[0m\"A ConflictEvent refers to either violent physical acts causing harm or damage, but are not covered by Life \n",
       "events\n",
       "    \u001B[1m(\u001B[0mconflicts, clashes, fighting, gunfire, \u001B[33m...\u001B[0m\u001B[1m)\u001B[0m or demonstrations \u001B[1m(\u001B[0mprotests, sit-ins, strikes, riots, \u001B[33m...\u001B[0m\u001B[1m)\u001B[0m.\n",
       "    \u001B[32m\"\"\u001B[0m\"\n",
       "\n",
       "    mention: str\n",
       "    \u001B[32m\"\"\u001B[0m\"The text span that most clearly expresses the event.\n",
       "        Such as: \u001B[32m\"terrorism\"\u001B[0m, \u001B[32m\"combat\"\u001B[0m, \u001B[32m\"hit\"\u001B[0m, \u001B[32m\"fight\"\u001B[0m, \u001B[32m\"bombing\"\u001B[0m \n",
       "    \u001B[32m\"\"\u001B[0m\"\n",
       "\n",
       "\n",
       "@dataclass\n",
       "class \u001B[1;35mContactEvent\u001B[0m\u001B[1m(\u001B[0mEvent\u001B[1m)\u001B[0m:\n",
       "    \u001B[32m\"\"\u001B[0m\"A ContactEvent occurs whenever two or more entities \u001B[1m(\u001B[0mpersons or organization's representatives\u001B[1m)\u001B[0m come \n",
       "together at\n",
       "    a single location and interact with one another face-to-face or directly enages in discussion via written or\n",
       "    telephone communication.\u001B[32m\"\"\u001B[0m\"\n",
       "\n",
       "    mention: str\n",
       "    \u001B[32m\"\"\u001B[0m\"The text span that most clearly expresses the event.\n",
       "        Such as: \u001B[32m\"meetings\"\u001B[0m, \u001B[32m\"conference\"\u001B[0m, \u001B[32m\"talked\"\u001B[0m, \u001B[32m\"met\"\u001B[0m, \u001B[32m\"letters\"\u001B[0m \n",
       "    \u001B[32m\"\"\u001B[0m\"\n",
       "\n",
       "\n",
       "@dataclass\n",
       "class \u001B[1;35mPersonnelEvent\u001B[0m\u001B[1m(\u001B[0mEvent\u001B[1m)\u001B[0m:\n",
       "    \u001B[32m\"\"\u001B[0m\"A PersonnelEvent occurs when a Person entity changes its job position \u001B[1m(\u001B[0mJobTitle entity\u001B[1m)\u001B[0m with respect an\n",
       "    Organization entity. It includes when a person starts working, ends working, changes offices within, gets \n",
       "nominated or is\n",
       "    elected for a position in a Organization.\u001B[32m\"\"\u001B[0m\"\n",
       "\n",
       "    mention: str\n",
       "    \u001B[32m\"\"\u001B[0m\"The text span that most clearly expresses the event.\n",
       "        Such as: \u001B[32m\"won\"\u001B[0m, \u001B[32m\"appoint\"\u001B[0m, \u001B[32m\"retired\"\u001B[0m, \u001B[32m\"fired\"\u001B[0m, \u001B[32m\"appointed\"\u001B[0m \n",
       "    \u001B[32m\"\"\u001B[0m\"\n",
       "\n",
       "\n",
       "@dataclass\n",
       "class \u001B[1;35mJusticeEvent\u001B[0m\u001B[1m(\u001B[0mEvent\u001B[1m)\u001B[0m:\n",
       "    \u001B[32m\"\"\u001B[0m\"A JusticeEvent refers to any judicial action such as: arresting, jailing, releasing, granting parole, trial\n",
       "    starting, hearing, charging, indicting, suing, convicting, sentencing, fine, executing, extraditing, adquiting,\n",
       "    appealing or pardoning a Person entity.\u001B[32m\"\"\u001B[0m\"\n",
       "\n",
       "    mention: str\n",
       "    \u001B[32m\"\"\u001B[0m\"The text span that most clearly expresses the event.\n",
       "        Such as: \u001B[32m\"parole\"\u001B[0m, \u001B[32m\"sued\"\u001B[0m, \u001B[32m\"sentenced\"\u001B[0m, \u001B[32m\"appeal\"\u001B[0m, \u001B[32m\"charged\"\u001B[0m \n",
       "    \u001B[32m\"\"\u001B[0m\"\n",
       "\n",
       "\n",
       "# This is the text to analyze\n",
       "text = \u001B[32m\"The Times reported that Vivendi Universal was negotiating to sell its flagship theme parks to the New York \u001B[0m\n",
       "\u001B[32minvestment firm, Blackstone Group, for 8 billion as the first step toward dismantling its entertainment empire.\"\u001B[0m\n",
       "\n",
       "# The annotation instances that take place in the text above are listed here\n",
       "result = \u001B[1m[\u001B[0m\n",
       "    \u001B[1;35mTransactionEvent\u001B[0m\u001B[1m(\u001B[0m\u001B[33mmention\u001B[0m=\u001B[32m\"sell\"\u001B[0m\u001B[1m)\u001B[0m,\n",
       "\u001B[1m]\u001B[0m\n",
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "rich.print(formated_text)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a6165c0d",
   "metadata": {},
   "source": [
    "## Prepare model inputs\n",
    "\n",
    "We remove everything after `result =` to run inference with the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 177,
   "id": "19eabf38",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt, _ = formated_text.split(\"result =\")\n",
    "prompt = prompt + \"result =\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "395394ae",
   "metadata": {},
   "source": [
    "Tokenize the input sentence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 178,
   "id": "f13c79d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_input = tokenizer(prompt, add_special_tokens=True, return_tensors=\"pt\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bef49edd",
   "metadata": {},
   "source": [
    "Remove the `eos` token from the input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 179,
   "id": "6dfc622f",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_input[\"input_ids\"] = model_input[\"input_ids\"][:, :-1]\n",
    "model_input[\"attention_mask\"] = model_input[\"attention_mask\"][:, :-1]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6718528",
   "metadata": {},
   "source": [
    "## Run GoLLIE\n",
    "\n",
    "We generate the predictions using GoLLIE.\n",
    "\n",
    "We use `num_beams=1` and `do_sample=False` in our exmperiments. But feel free to experiment with differen decoding strategies 😊"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 180,
   "id": "b6f95263",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:2 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 600 ms, sys: 5.04 ms, total: 606 ms\n",
      "Wall time: 608 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "model_ouput = model.generate(\n",
    "    **model_input.to(model.device),\n",
    "    max_new_tokens=128,\n",
    "    do_sample=False,\n",
    "    min_new_tokens=0,\n",
    "    num_beams=1,\n",
    "    num_return_sequences=1,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "983f4a2a",
   "metadata": {},
   "source": [
    "### Print the results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 181,
   "id": "31808b61",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Answer 0\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">[</span>\n",
       "    <span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">TransactionEvent</span><span style=\"font-weight: bold\">(</span><span style=\"color: #808000; text-decoration-color: #808000\">mention</span>=<span style=\"color: #008000; text-decoration-color: #008000\">\"sell\"</span><span style=\"font-weight: bold\">)</span>,\n",
       "<span style=\"font-weight: bold\">]</span>\n",
       "\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001B[1m[\u001B[0m\n",
       "    \u001B[1;35mTransactionEvent\u001B[0m\u001B[1m(\u001B[0m\u001B[33mmention\u001B[0m=\u001B[32m\"sell\"\u001B[0m\u001B[1m)\u001B[0m,\n",
       "\u001B[1m]\u001B[0m\n",
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "for y, x in enumerate(model_ouput):\n",
    "    print(f\"Answer {y}\")\n",
    "    rich.print(tokenizer.decode(x,skip_special_tokens=True).split(\"result = \")[-1])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95e2013f",
   "metadata": {},
   "source": [
    "## Parse the output\n",
    "\n",
    "The output is a Python list of instances, we can execute it  🤯\n",
    "\n",
    "We define the AnnotationList class to parse the output with a single line of code. The `AnnotationList.from_output` function filters any label that we did not define (hallucinations) to prevent getting an `undefined class` error. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 182,
   "id": "5d66fb9e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">[</span><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">TransactionEvent</span><span style=\"font-weight: bold\">(</span><span style=\"color: #808000; text-decoration-color: #808000\">mention</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'sell'</span><span style=\"font-weight: bold\">)]</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001B[1m[\u001B[0m\u001B[1;35mTransactionEvent\u001B[0m\u001B[1m(\u001B[0m\u001B[33mmention\u001B[0m=\u001B[32m'sell'\u001B[0m\u001B[1m)\u001B[0m\u001B[1m]\u001B[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "result = AnnotationList.from_output(\n",
    "    tokenizer.decode(model_ouput[0],skip_special_tokens=True).split(\"result = \")[-1],\n",
    "    task_module=\"guidelines_ee\"\n",
    "    )\n",
    "rich.print(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27586bd8",
   "metadata": {},
   "source": [
    "Labels are an instance of the defined classes:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 183,
   "id": "bd039309",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "guidelines_ee.TransactionEvent"
      ]
     },
     "execution_count": 183,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(result[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 184,
   "id": "2703325b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'sell'"
      ]
     },
     "execution_count": 184,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result[0].mention"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34c52537",
   "metadata": {},
   "source": [
    "# Evaluate the result\n",
    "\n",
    "Finally, we will evaluate the outputs from the model."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "715535b7",
   "metadata": {},
   "source": [
    "First, we define an Scorer, for EE, we will use the `EventScorer` class.\n",
    "\n",
    "We need to define the `valid_types` for the scorer, which will be the labels that we have defined. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 185,
   "id": "d44e1f46",
   "metadata": {},
   "outputs": [],
   "source": [
    "from src.tasks.utils_scorer import EventScorer\n",
    "\n",
    "\n",
    "class MyScorer(EventScorer):\n",
    "    \"\"\"Event scorer.\"\"\"\n",
    "\n",
    "    valid_types: List[Type] = guidelines_ee.EVENTS_DEFINITIONS"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9907289",
   "metadata": {},
   "source": [
    "### Instanciate the scorer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 186,
   "id": "1a0825db",
   "metadata": {},
   "outputs": [],
   "source": [
    "scorer = MyScorer()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ccb11ce1",
   "metadata": {},
   "source": [
    "### Compute F1 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 187,
   "id": "da72e844",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">{</span>\n",
       "    <span style=\"color: #008000; text-decoration-color: #008000\">'events'</span>: <span style=\"font-weight: bold\">{</span>\n",
       "        <span style=\"color: #008000; text-decoration-color: #008000\">'precision'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.0</span>,\n",
       "        <span style=\"color: #008000; text-decoration-color: #008000\">'recall'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.0</span>,\n",
       "        <span style=\"color: #008000; text-decoration-color: #008000\">'f1-score'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.0</span>,\n",
       "        <span style=\"color: #008000; text-decoration-color: #008000\">'class_scores'</span>: <span style=\"font-weight: bold\">{</span>\n",
       "            <span style=\"color: #008000; text-decoration-color: #008000\">'TransactionEvent'</span>: <span style=\"font-weight: bold\">{</span>\n",
       "                <span style=\"color: #008000; text-decoration-color: #008000\">'tp'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,\n",
       "                <span style=\"color: #008000; text-decoration-color: #008000\">'total_pos'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,\n",
       "                <span style=\"color: #008000; text-decoration-color: #008000\">'total_pre'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,\n",
       "                <span style=\"color: #008000; text-decoration-color: #008000\">'precision'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.0</span>,\n",
       "                <span style=\"color: #008000; text-decoration-color: #008000\">'recall'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.0</span>,\n",
       "                <span style=\"color: #008000; text-decoration-color: #008000\">'f1-score'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.0</span>\n",
       "            <span style=\"font-weight: bold\">}</span>\n",
       "        <span style=\"font-weight: bold\">}</span>\n",
       "    <span style=\"font-weight: bold\">}</span>,\n",
       "    <span style=\"color: #008000; text-decoration-color: #008000\">'arguments'</span>: <span style=\"font-weight: bold\">{</span><span style=\"color: #008000; text-decoration-color: #008000\">'precision'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'recall'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'f1-score'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0</span><span style=\"font-weight: bold\">}</span>\n",
       "<span style=\"font-weight: bold\">}</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001B[1m{\u001B[0m\n",
       "    \u001B[32m'events'\u001B[0m: \u001B[1m{\u001B[0m\n",
       "        \u001B[32m'precision'\u001B[0m: \u001B[1;36m1.0\u001B[0m,\n",
       "        \u001B[32m'recall'\u001B[0m: \u001B[1;36m1.0\u001B[0m,\n",
       "        \u001B[32m'f1-score'\u001B[0m: \u001B[1;36m1.0\u001B[0m,\n",
       "        \u001B[32m'class_scores'\u001B[0m: \u001B[1m{\u001B[0m\n",
       "            \u001B[32m'TransactionEvent'\u001B[0m: \u001B[1m{\u001B[0m\n",
       "                \u001B[32m'tp'\u001B[0m: \u001B[1;36m1\u001B[0m,\n",
       "                \u001B[32m'total_pos'\u001B[0m: \u001B[1;36m1\u001B[0m,\n",
       "                \u001B[32m'total_pre'\u001B[0m: \u001B[1;36m1\u001B[0m,\n",
       "                \u001B[32m'precision'\u001B[0m: \u001B[1;36m1.0\u001B[0m,\n",
       "                \u001B[32m'recall'\u001B[0m: \u001B[1;36m1.0\u001B[0m,\n",
       "                \u001B[32m'f1-score'\u001B[0m: \u001B[1;36m1.0\u001B[0m\n",
       "            \u001B[1m}\u001B[0m\n",
       "        \u001B[1m}\u001B[0m\n",
       "    \u001B[1m}\u001B[0m,\n",
       "    \u001B[32m'arguments'\u001B[0m: \u001B[1m{\u001B[0m\u001B[32m'precision'\u001B[0m: \u001B[1;36m0.0\u001B[0m, \u001B[32m'recall'\u001B[0m: \u001B[1;36m0.0\u001B[0m, \u001B[32m'f1-score'\u001B[0m: \u001B[1;36m0.0\u001B[0m\u001B[1m}\u001B[0m\n",
       "\u001B[1m}\u001B[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "scorer_results = scorer(reference=[gold],predictions=[result])\n",
    "rich.print(scorer_results)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1519861",
   "metadata": {},
   "source": [
    "GoLLIE has successfully predicted the `TransactionEvent` in the sentence. The argument F1-score is 0.0 because we are only predicting event instances. We will predict the arguments in the next step."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b7ee0f2",
   "metadata": {},
   "source": [
    "# Event Argument Extraction"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc602bf3",
   "metadata": {},
   "source": [
    "In this step, we will predict the arguments for the `TransactionEvent`. `TransactionEvent` is a coarse category that encompasses two fine-grained categories: `TransferMoney` and `TransferOwnership`. Each of these fine-grained categories has distinct arguments. First, we define both classes and their respective arguments."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 188,
   "id": "8c773cbc",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "\n",
    "from src.tasks.utils_typing import Event, dataclass\n",
    "from src.tasks.ace.prompts import TransactionEvent\n",
    "\n",
    "# The following lines describe the task definition\n",
    "@dataclass\n",
    "class TransferMoney(TransactionEvent):\n",
    "    \"\"\"TransferMoney Events refer to the giving, receiving, borrowing, or lending money when it is not in the context of\n",
    "    purchasing something. The canonical examples are: (1) people giving money to organizations (and getting nothing tangible in\n",
    "    return); and (2) organizations lending money to people or other orgs.\"\"\"\n",
    "\n",
    "    mention: str  # The text span that most clearly expresses (triggers) the event\n",
    "    giver: List[str]  # The donating agent\n",
    "    recipient: List[str]  # The recipient agent\n",
    "    beneficiary: List[str]  # The agent that benefits from the transfer\n",
    "    money: List[str]  # The amount given, donated or loaned\n",
    "    time: List[str]  # When the amount is transferred\n",
    "    place: List[str]  # Where the transation takes place\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class TransferOwnership(TransactionEvent):\n",
    "    \"\"\"TransferOwnership Events refer to the buying, selling, loaning, borrowing, giving, or receiving of artifacts or\n",
    "    organizations.\"\"\"\n",
    "\n",
    "    mention: str  # The text span that most clearly expresses (triggers) the event\n",
    "    buyer: List[str]  # The buying agent\n",
    "    seller: List[str]  # The selling agent\n",
    "    beneficiary: List[str]  # The agent that benefits from the transaction\n",
    "    artifact: List[str]  # The item or Organization that was bought or sold\n",
    "    price: List[str]  # The sale price of the artifact\n",
    "    time: List[str]  # When the sale takes place\n",
    "    place: List[str]  # Where the sale takes place\n",
    "\n",
    "\n",
    "\n",
    "EVENTS_EAE_DEFINITIONS: List[Event] = [\n",
    "    TransferMoney,\n",
    "    TransferOwnership\n",
    "]\n",
    "    \n",
    "if __name__ == \"__main__\":\n",
    "    cell_txt = In[-1]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86facef3",
   "metadata": {},
   "source": [
    "As before, we write the content of the previous cell to a file and then import its content. We export the guidelines to a separate file and load them under a different name; otherwise, Python will not load them correctly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 189,
   "id": "676baa37",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"guidelines_eae.py\",\"w\",encoding=\"utf8\") as python_guidelines:\n",
    "    print(cell_txt,file=python_guidelines)\n",
    "    \n",
    "import guidelines_eae as guidelines_eae"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 190,
   "id": "2421253b",
   "metadata": {},
   "outputs": [],
   "source": [
    "guidelines = [inspect.getsource(definition) for definition in guidelines_eae.EVENTS_EAE_DEFINITIONS]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84c8899e",
   "metadata": {},
   "source": [
    "Here we define the input sentence and the gold labels.\n",
    "\n",
    "You can define and empy list as gold labels if you don't have gold annotations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 202,
   "id": "0be963db",
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"The Times reported that Vivendi Universal was negotiating to sell its flagship theme parks to the New York investment firm, Blackstone Group, for 8 billion as the first step toward dismantling its entertainment empire.\"\n",
    "gold = [\n",
    "    guidelines_eae.TransferOwnership(\n",
    "        mention='sell', \n",
    "        buyer=['Blackstone Group'], \n",
    "        seller=['Vivendi Universal'], \n",
    "        beneficiary=[], \n",
    "        artifact=['theme parks'], \n",
    "        price=[\"8 billion\"], \n",
    "        time=[], \n",
    "        place=[]\n",
    "    )\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d130dfe2",
   "metadata": {},
   "source": [
    "## Filling a template\n",
    "\n",
    "For EAE, we will use a distinct template. This template incorporates both the event category and the trigger. As a result, GoLLIE only needs to predict the arguments corresponding to the specified events.\n",
    "\n",
    "We use Jinja templates, which are easy to implement and exceptionally fast. For more information, visit: https://jinja.palletsprojects.com/en/3.1.x/api/#high-level-api.\n",
    "\n",
    "```Python\n",
    "# The following lines describe the task definition\n",
    "{%- for definition in guidelines %}\n",
    "{{ definition }}\n",
    "{%- endfor %}\n",
    "\n",
    "# This is the text to analyze\n",
    "text = {{ text.__repr__() }}\n",
    "\n",
    "# The list called result contains the instances for the following events according to the guidelines above:\n",
    "{%- for ann in gold %}\n",
    "#    - \"{{ann.mention}}\" triggers a {{ann.__class__.__name__}} event.\n",
    "{%- endfor %}\n",
    "# \n",
    "result = [\n",
    "{%- for ann in annotations %}\n",
    "    {{ ann }},\n",
    "{%- endfor %}\n",
    "]\n",
    "\n",
    "```\n",
    "\n",
    "This template is stored in `templates/prompt_ace_eae.txt`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 203,
   "id": "1f339eb7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read template\n",
    "with open(\"../templates/prompt_ace_eae.txt\", \"rt\") as f:\n",
    "    template = Template(f.read())\n",
    "# Fill the template\n",
    "formated_text = template.render(guidelines=guidelines, text=text, annotations=gold, gold=gold)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a70f053",
   "metadata": {},
   "source": [
    "### Black Code Formatter\n",
    "\n",
    "We use the Black Code Formatter to automatically unify all the prompts to the same format. \n",
    "\n",
    "https://github.com/psf/black"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 204,
   "id": "0cc64f05",
   "metadata": {},
   "outputs": [],
   "source": [
    "black_mode = black.Mode()\n",
    "formated_text = black.format_str(formated_text, mode=black_mode)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 205,
   "id": "67907a32",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"># The following lines describe the task definition\n",
       "@dataclass\n",
       "class <span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">TransferMoney</span><span style=\"font-weight: bold\">(</span>TransactionEvent<span style=\"font-weight: bold\">)</span>:\n",
       "    <span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"TransferMoney Events refer to the giving, receiving, borrowing, or lending money when it is not in the \n",
       "context of\n",
       "    purchasing something. The canonical examples are: <span style=\"font-weight: bold\">(</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span><span style=\"font-weight: bold\">)</span> people giving money to organizations <span style=\"font-weight: bold\">(</span>and getting nothing\n",
       "tangible in\n",
       "    return<span style=\"font-weight: bold\">)</span>; and <span style=\"font-weight: bold\">(</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span><span style=\"font-weight: bold\">)</span> organizations lending money to people or other orgs.<span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"\n",
       "\n",
       "    mention: str  # The text span that most clearly expresses <span style=\"font-weight: bold\">(</span>triggers<span style=\"font-weight: bold\">)</span> the event\n",
       "    giver: List  # The donating agent\n",
       "    recipient: List  # The recipient agent\n",
       "    beneficiary: List  # The agent that benefits from the transfer\n",
       "    money: List  # The amount given, donated or loaned\n",
       "    time: List  # When the amount is transferred\n",
       "    place: List  # Where the transation takes place\n",
       "\n",
       "\n",
       "@dataclass\n",
       "class <span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">TransferOwnership</span><span style=\"font-weight: bold\">(</span>TransactionEvent<span style=\"font-weight: bold\">)</span>:\n",
       "    <span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"TransferOwnership Events refer to the buying, selling, loaning, borrowing, giving, or receiving of artifacts\n",
       "or\n",
       "    organizations.<span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"\n",
       "\n",
       "    mention: str  # The text span that most clearly expresses <span style=\"font-weight: bold\">(</span>triggers<span style=\"font-weight: bold\">)</span> the event\n",
       "    buyer: List  # The buying agent\n",
       "    seller: List  # The selling agent\n",
       "    beneficiary: List  # The agent that benefits from the transaction\n",
       "    artifact: List  # The item or Organization that was bought or sold\n",
       "    price: List  # The sale price of the artifact\n",
       "    time: List  # When the sale takes place\n",
       "    place: List  # Where the sale takes place\n",
       "\n",
       "\n",
       "# This is the text to analyze\n",
       "text = <span style=\"color: #008000; text-decoration-color: #008000\">\"The Times reported that Vivendi Universal was negotiating to sell its flagship theme parks to the New York </span>\n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">investment firm, Blackstone Group, for 8 billion as the first step toward dismantling its entertainment empire.\"</span>\n",
       "\n",
       "# The list called result contains the instances for the following events according to the guidelines above:\n",
       "#    - <span style=\"color: #008000; text-decoration-color: #008000\">\"sell\"</span> triggers a TransferOwnership event.\n",
       "#\n",
       "result = <span style=\"font-weight: bold\">[</span>\n",
       "    <span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">TransferOwnership</span><span style=\"font-weight: bold\">(</span>\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">mention</span>=<span style=\"color: #008000; text-decoration-color: #008000\">\"sell\"</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">buyer</span>=<span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">\"Blackstone Group\"</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">seller</span>=<span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">\"Vivendi Universal\"</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">beneficiary</span>=<span style=\"font-weight: bold\">[]</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">artifact</span>=<span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">\"theme parks\"</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">price</span>=<span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">\"8 billion\"</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">time</span>=<span style=\"font-weight: bold\">[]</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">place</span>=<span style=\"font-weight: bold\">[]</span>,\n",
       "    <span style=\"font-weight: bold\">)</span>,\n",
       "<span style=\"font-weight: bold\">]</span>\n",
       "\n",
       "</pre>\n"
      ],
      "text/plain": [
       "# The following lines describe the task definition\n",
       "@dataclass\n",
       "class \u001B[1;35mTransferMoney\u001B[0m\u001B[1m(\u001B[0mTransactionEvent\u001B[1m)\u001B[0m:\n",
       "    \u001B[32m\"\"\u001B[0m\"TransferMoney Events refer to the giving, receiving, borrowing, or lending money when it is not in the \n",
       "context of\n",
       "    purchasing something. The canonical examples are: \u001B[1m(\u001B[0m\u001B[1;36m1\u001B[0m\u001B[1m)\u001B[0m people giving money to organizations \u001B[1m(\u001B[0mand getting nothing\n",
       "tangible in\n",
       "    return\u001B[1m)\u001B[0m; and \u001B[1m(\u001B[0m\u001B[1;36m2\u001B[0m\u001B[1m)\u001B[0m organizations lending money to people or other orgs.\u001B[32m\"\"\u001B[0m\"\n",
       "\n",
       "    mention: str  # The text span that most clearly expresses \u001B[1m(\u001B[0mtriggers\u001B[1m)\u001B[0m the event\n",
       "    giver: List  # The donating agent\n",
       "    recipient: List  # The recipient agent\n",
       "    beneficiary: List  # The agent that benefits from the transfer\n",
       "    money: List  # The amount given, donated or loaned\n",
       "    time: List  # When the amount is transferred\n",
       "    place: List  # Where the transation takes place\n",
       "\n",
       "\n",
       "@dataclass\n",
       "class \u001B[1;35mTransferOwnership\u001B[0m\u001B[1m(\u001B[0mTransactionEvent\u001B[1m)\u001B[0m:\n",
       "    \u001B[32m\"\"\u001B[0m\"TransferOwnership Events refer to the buying, selling, loaning, borrowing, giving, or receiving of artifacts\n",
       "or\n",
       "    organizations.\u001B[32m\"\"\u001B[0m\"\n",
       "\n",
       "    mention: str  # The text span that most clearly expresses \u001B[1m(\u001B[0mtriggers\u001B[1m)\u001B[0m the event\n",
       "    buyer: List  # The buying agent\n",
       "    seller: List  # The selling agent\n",
       "    beneficiary: List  # The agent that benefits from the transaction\n",
       "    artifact: List  # The item or Organization that was bought or sold\n",
       "    price: List  # The sale price of the artifact\n",
       "    time: List  # When the sale takes place\n",
       "    place: List  # Where the sale takes place\n",
       "\n",
       "\n",
       "# This is the text to analyze\n",
       "text = \u001B[32m\"The Times reported that Vivendi Universal was negotiating to sell its flagship theme parks to the New York \u001B[0m\n",
       "\u001B[32minvestment firm, Blackstone Group, for 8 billion as the first step toward dismantling its entertainment empire.\"\u001B[0m\n",
       "\n",
       "# The list called result contains the instances for the following events according to the guidelines above:\n",
       "#    - \u001B[32m\"sell\"\u001B[0m triggers a TransferOwnership event.\n",
       "#\n",
       "result = \u001B[1m[\u001B[0m\n",
       "    \u001B[1;35mTransferOwnership\u001B[0m\u001B[1m(\u001B[0m\n",
       "        \u001B[33mmention\u001B[0m=\u001B[32m\"sell\"\u001B[0m,\n",
       "        \u001B[33mbuyer\u001B[0m=\u001B[1m[\u001B[0m\u001B[32m\"Blackstone Group\"\u001B[0m\u001B[1m]\u001B[0m,\n",
       "        \u001B[33mseller\u001B[0m=\u001B[1m[\u001B[0m\u001B[32m\"Vivendi Universal\"\u001B[0m\u001B[1m]\u001B[0m,\n",
       "        \u001B[33mbeneficiary\u001B[0m=\u001B[1m[\u001B[0m\u001B[1m]\u001B[0m,\n",
       "        \u001B[33martifact\u001B[0m=\u001B[1m[\u001B[0m\u001B[32m\"theme parks\"\u001B[0m\u001B[1m]\u001B[0m,\n",
       "        \u001B[33mprice\u001B[0m=\u001B[1m[\u001B[0m\u001B[32m\"8 billion\"\u001B[0m\u001B[1m]\u001B[0m,\n",
       "        \u001B[33mtime\u001B[0m=\u001B[1m[\u001B[0m\u001B[1m]\u001B[0m,\n",
       "        \u001B[33mplace\u001B[0m=\u001B[1m[\u001B[0m\u001B[1m]\u001B[0m,\n",
       "    \u001B[1m)\u001B[0m,\n",
       "\u001B[1m]\u001B[0m\n",
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "rich.print(formated_text)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8da7029e",
   "metadata": {},
   "source": [
    "As observed in the previous cell, there's an additional comment preceding the `result` variable that indicates the event trigger and the category of the event."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2826b4b0",
   "metadata": {},
   "source": [
    "## Prepare model inputs\n",
    "\n",
    "We remove everything after `result =` to run inference with the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 206,
   "id": "3ab8bc2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt, _ = formated_text.split(\"result =\")\n",
    "prompt = prompt + \"result =\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "806c242c",
   "metadata": {},
   "source": [
    "Tokenize the input sentence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 207,
   "id": "0f74763b",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_input = tokenizer(prompt, add_special_tokens=True, return_tensors=\"pt\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5497f326",
   "metadata": {},
   "source": [
    "Remove the `eos` token from the input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 208,
   "id": "28dc5528",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_input[\"input_ids\"] = model_input[\"input_ids\"][:, :-1]\n",
    "model_input[\"attention_mask\"] = model_input[\"attention_mask\"][:, :-1]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d8174a9",
   "metadata": {},
   "source": [
    "## Run GoLLIE\n",
    "\n",
    "We generate the predictions using GoLLIE.\n",
    "\n",
    "We use `num_beams=1` and `do_sample=False` in our exmperiments. But feel free to experiment with differen decoding strategies 😊"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 209,
   "id": "d5b88b6d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:2 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 2.49 s, sys: 6.28 ms, total: 2.5 s\n",
      "Wall time: 2.51 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "model_ouput = model.generate(\n",
    "    **model_input.to(model.device),\n",
    "    max_new_tokens=128,\n",
    "    do_sample=False,\n",
    "    min_new_tokens=0,\n",
    "    num_beams=1,\n",
    "    num_return_sequences=1,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ff452e1",
   "metadata": {},
   "source": [
    "### Print the results\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 210,
   "id": "07d142d2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Answer 0\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">[</span>\n",
       "    <span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">TransferOwnership</span><span style=\"font-weight: bold\">(</span>\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">mention</span>=<span style=\"color: #008000; text-decoration-color: #008000\">\"sell\"</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">buyer</span>=<span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">\"Blackstone Group\"</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">seller</span>=<span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">\"Vivendi Universal\"</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">beneficiary</span>=<span style=\"font-weight: bold\">[]</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">artifact</span>=<span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">\"theme parks\"</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">price</span>=<span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">\"8 billion\"</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">time</span>=<span style=\"font-weight: bold\">[]</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">place</span>=<span style=\"font-weight: bold\">[]</span>,\n",
       "    <span style=\"font-weight: bold\">)</span>,\n",
       "<span style=\"font-weight: bold\">]</span>\n",
       "\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001B[1m[\u001B[0m\n",
       "    \u001B[1;35mTransferOwnership\u001B[0m\u001B[1m(\u001B[0m\n",
       "        \u001B[33mmention\u001B[0m=\u001B[32m\"sell\"\u001B[0m,\n",
       "        \u001B[33mbuyer\u001B[0m=\u001B[1m[\u001B[0m\u001B[32m\"Blackstone Group\"\u001B[0m\u001B[1m]\u001B[0m,\n",
       "        \u001B[33mseller\u001B[0m=\u001B[1m[\u001B[0m\u001B[32m\"Vivendi Universal\"\u001B[0m\u001B[1m]\u001B[0m,\n",
       "        \u001B[33mbeneficiary\u001B[0m=\u001B[1m[\u001B[0m\u001B[1m]\u001B[0m,\n",
       "        \u001B[33martifact\u001B[0m=\u001B[1m[\u001B[0m\u001B[32m\"theme parks\"\u001B[0m\u001B[1m]\u001B[0m,\n",
       "        \u001B[33mprice\u001B[0m=\u001B[1m[\u001B[0m\u001B[32m\"8 billion\"\u001B[0m\u001B[1m]\u001B[0m,\n",
       "        \u001B[33mtime\u001B[0m=\u001B[1m[\u001B[0m\u001B[1m]\u001B[0m,\n",
       "        \u001B[33mplace\u001B[0m=\u001B[1m[\u001B[0m\u001B[1m]\u001B[0m,\n",
       "    \u001B[1m)\u001B[0m,\n",
       "\u001B[1m]\u001B[0m\n",
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "for y, x in enumerate(model_ouput):\n",
    "    print(f\"Answer {y}\")\n",
    "    rich.print(tokenizer.decode(x,skip_special_tokens=True).split(\"result = \")[-1])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f32756b1",
   "metadata": {},
   "source": [
    "## Parse the output\n",
    "\n",
    "The output is a Python list of instances, we can execute it  🤯\n",
    "\n",
    "We define the AnnotationList class to parse the output with a single line of code. The `AnnotationList.from_output` function filters any label that we did not define (hallucinations) to prevent getting an `undefined class` error. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 211,
   "id": "95fdee69",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">[</span>\n",
       "    <span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">TransferOwnership</span><span style=\"font-weight: bold\">(</span><span style=\"color: #808000; text-decoration-color: #808000\">mention</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'sell'</span>, <span style=\"color: #808000; text-decoration-color: #808000\">buyer</span>=<span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">'Blackstone Group'</span><span style=\"font-weight: bold\">]</span>, <span style=\"color: #808000; text-decoration-color: #808000\">seller</span>=<span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">'Vivendi Universal'</span><span style=\"font-weight: bold\">]</span>, <span style=\"color: #808000; text-decoration-color: #808000\">beneficiary</span>=<span style=\"font-weight: bold\">[]</span>, \n",
       "<span style=\"color: #808000; text-decoration-color: #808000\">artifact</span>=<span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">'theme parks'</span><span style=\"font-weight: bold\">]</span>, <span style=\"color: #808000; text-decoration-color: #808000\">price</span>=<span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">'8 billion'</span><span style=\"font-weight: bold\">]</span>, <span style=\"color: #808000; text-decoration-color: #808000\">time</span>=<span style=\"font-weight: bold\">[]</span>, <span style=\"color: #808000; text-decoration-color: #808000\">place</span>=<span style=\"font-weight: bold\">[])</span>\n",
       "<span style=\"font-weight: bold\">]</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001B[1m[\u001B[0m\n",
       "    \u001B[1;35mTransferOwnership\u001B[0m\u001B[1m(\u001B[0m\u001B[33mmention\u001B[0m=\u001B[32m'sell'\u001B[0m, \u001B[33mbuyer\u001B[0m=\u001B[1m[\u001B[0m\u001B[32m'Blackstone Group'\u001B[0m\u001B[1m]\u001B[0m, \u001B[33mseller\u001B[0m=\u001B[1m[\u001B[0m\u001B[32m'Vivendi Universal'\u001B[0m\u001B[1m]\u001B[0m, \u001B[33mbeneficiary\u001B[0m=\u001B[1m[\u001B[0m\u001B[1m]\u001B[0m, \n",
       "\u001B[33martifact\u001B[0m=\u001B[1m[\u001B[0m\u001B[32m'theme parks'\u001B[0m\u001B[1m]\u001B[0m, \u001B[33mprice\u001B[0m=\u001B[1m[\u001B[0m\u001B[32m'8 billion'\u001B[0m\u001B[1m]\u001B[0m, \u001B[33mtime\u001B[0m=\u001B[1m[\u001B[0m\u001B[1m]\u001B[0m, \u001B[33mplace\u001B[0m=\u001B[1m[\u001B[0m\u001B[1m]\u001B[0m\u001B[1m)\u001B[0m\n",
       "\u001B[1m]\u001B[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "result = AnnotationList.from_output(\n",
    "    tokenizer.decode(model_ouput[0],skip_special_tokens=True).split(\"result = \")[-1],\n",
    "    task_module=\"guidelines_eae\"\n",
    "    )\n",
    "rich.print(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "92ede72c",
   "metadata": {},
   "source": [
    "Labels are an instance of the defined classes:\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 212,
   "id": "5e14e92d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "guidelines_eae.TransferOwnership"
      ]
     },
     "execution_count": 212,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(result[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 213,
   "id": "1f234a7f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'sell'"
      ]
     },
     "execution_count": 213,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result[0].mention"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c3b6328",
   "metadata": {},
   "source": [
    "# Evaluate the result\n",
    "\n",
    "Finally, we will evaluate the outputs from the model.\n",
    "\n",
    "First, we define an Scorer, for EE, we will use the `EventScorer` class.\n",
    "\n",
    "We need to define the `valid_types` for the scorer, which will be the labels that we have defined. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 214,
   "id": "8c278ebc",
   "metadata": {},
   "outputs": [],
   "source": [
    "from src.tasks.utils_scorer import EventScorer\n",
    "\n",
    "\n",
    "class MyScorer(EventScorer):\n",
    "    \"\"\"Event scorer.\"\"\"\n",
    "\n",
    "    valid_types: List[Type] = guidelines_eae.EVENTS_EAE_DEFINITIONS"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a6d58301",
   "metadata": {},
   "source": [
    "### Compute F1 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 215,
   "id": "6c393d61",
   "metadata": {},
   "outputs": [],
   "source": [
    "scorer = MyScorer()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 216,
   "id": "f0d0d932",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">{</span>\n",
       "    <span style=\"color: #008000; text-decoration-color: #008000\">'events'</span>: <span style=\"font-weight: bold\">{</span>\n",
       "        <span style=\"color: #008000; text-decoration-color: #008000\">'precision'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.0</span>,\n",
       "        <span style=\"color: #008000; text-decoration-color: #008000\">'recall'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.0</span>,\n",
       "        <span style=\"color: #008000; text-decoration-color: #008000\">'f1-score'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.0</span>,\n",
       "        <span style=\"color: #008000; text-decoration-color: #008000\">'class_scores'</span>: <span style=\"font-weight: bold\">{</span>\n",
       "            <span style=\"color: #008000; text-decoration-color: #008000\">'TransferOwnership'</span>: <span style=\"font-weight: bold\">{</span>\n",
       "                <span style=\"color: #008000; text-decoration-color: #008000\">'tp'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,\n",
       "                <span style=\"color: #008000; text-decoration-color: #008000\">'total_pos'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,\n",
       "                <span style=\"color: #008000; text-decoration-color: #008000\">'total_pre'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,\n",
       "                <span style=\"color: #008000; text-decoration-color: #008000\">'precision'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.0</span>,\n",
       "                <span style=\"color: #008000; text-decoration-color: #008000\">'recall'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.0</span>,\n",
       "                <span style=\"color: #008000; text-decoration-color: #008000\">'f1-score'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.0</span>\n",
       "            <span style=\"font-weight: bold\">}</span>\n",
       "        <span style=\"font-weight: bold\">}</span>\n",
       "    <span style=\"font-weight: bold\">}</span>,\n",
       "    <span style=\"color: #008000; text-decoration-color: #008000\">'arguments'</span>: <span style=\"font-weight: bold\">{</span><span style=\"color: #008000; text-decoration-color: #008000\">'precision'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.0</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'recall'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.0</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'f1-score'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.0</span><span style=\"font-weight: bold\">}</span>\n",
       "<span style=\"font-weight: bold\">}</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001B[1m{\u001B[0m\n",
       "    \u001B[32m'events'\u001B[0m: \u001B[1m{\u001B[0m\n",
       "        \u001B[32m'precision'\u001B[0m: \u001B[1;36m1.0\u001B[0m,\n",
       "        \u001B[32m'recall'\u001B[0m: \u001B[1;36m1.0\u001B[0m,\n",
       "        \u001B[32m'f1-score'\u001B[0m: \u001B[1;36m1.0\u001B[0m,\n",
       "        \u001B[32m'class_scores'\u001B[0m: \u001B[1m{\u001B[0m\n",
       "            \u001B[32m'TransferOwnership'\u001B[0m: \u001B[1m{\u001B[0m\n",
       "                \u001B[32m'tp'\u001B[0m: \u001B[1;36m1\u001B[0m,\n",
       "                \u001B[32m'total_pos'\u001B[0m: \u001B[1;36m1\u001B[0m,\n",
       "                \u001B[32m'total_pre'\u001B[0m: \u001B[1;36m1\u001B[0m,\n",
       "                \u001B[32m'precision'\u001B[0m: \u001B[1;36m1.0\u001B[0m,\n",
       "                \u001B[32m'recall'\u001B[0m: \u001B[1;36m1.0\u001B[0m,\n",
       "                \u001B[32m'f1-score'\u001B[0m: \u001B[1;36m1.0\u001B[0m\n",
       "            \u001B[1m}\u001B[0m\n",
       "        \u001B[1m}\u001B[0m\n",
       "    \u001B[1m}\u001B[0m,\n",
       "    \u001B[32m'arguments'\u001B[0m: \u001B[1m{\u001B[0m\u001B[32m'precision'\u001B[0m: \u001B[1;36m1.0\u001B[0m, \u001B[32m'recall'\u001B[0m: \u001B[1;36m1.0\u001B[0m, \u001B[32m'f1-score'\u001B[0m: \u001B[1;36m1.0\u001B[0m\u001B[1m}\u001B[0m\n",
       "\u001B[1m}\u001B[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "scorer_results = scorer(reference=[gold],predictions=[result])\n",
    "rich.print(scorer_results)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7528296e",
   "metadata": {},
   "source": [
    "GoLLIE has successfully predicted the event arguments for the `TransferOwnership` event in the sentence. In the EAE task, we can anticipate the F1 score for the Events to always be `1.0`, as we are providing the model with the gold triggers and classes. While it's possible to implement the EE and EAE tasks as a pipeline, let's venture into something even more exciting! Next, we will explore end-to-end event extraction with GoLLIE.   🚀🚀🚀\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6fed103",
   "metadata": {},
   "source": [
    "# END-to-END Event Extraction"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "061830c0",
   "metadata": {},
   "source": [
    "In the EAE task, we provided the model with the gold triggers and classes. Now, let's modify the template to remove that information.\n",
    "\n",
    "We will use the same guidelines as in the EAE task, meaning that we will define the `TransferMoney` and `TransferOwnership` events. It's worth noting that the ACE05 dataset includes 32 fine-grained event classes. But incorporating all of them would result in an huge prompt with a lot of tokens. While GoLLIE can natively handle inputs of up to `4096` tokens (which can be expanded using RoPE scaling to `64000` tokens or even more), we will keep this example simple by focusing on only 2 event classes. Feel free to experiment by adding more event definitions!! \n",
    "\n",
    "It's important to mention that GoLLIE has not been pre-trained for end-to-end Event Extraction tasks. Nonetheless, its robust generalization capabilities make such a task feasible."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 217,
   "id": "a1d0ba0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "template_txt =(\n",
    "\"\"\"# The following lines describe the task definition\n",
    "{%- for definition in guidelines %}\n",
    "{{ definition }}\n",
    "{%- endfor %}\n",
    "\n",
    "# This is the text to analyze\n",
    "text = {{ text.__repr__() }}\n",
    "\n",
    "# The annotation instances that take place in the text above are listed here\n",
    "result = [\n",
    "{%- for ann in annotations %}\n",
    "    {{ ann }},\n",
    "{%- endfor %}\n",
    "]\n",
    "\"\"\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 218,
   "id": "ca411094",
   "metadata": {},
   "outputs": [],
   "source": [
    "template = Template(template_txt)\n",
    "# Fill the template\n",
    "formated_text = template.render(guidelines=guidelines, text=text, annotations=gold, gold=gold)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 219,
   "id": "268a4e4f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"># The following lines describe the task definition\n",
       "@dataclass\n",
       "class <span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">TransferMoney</span><span style=\"font-weight: bold\">(</span>TransactionEvent<span style=\"font-weight: bold\">)</span>:\n",
       "    <span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"TransferMoney Events refer to the giving, receiving, borrowing, or lending money when it is not in the \n",
       "context of\n",
       "    purchasing something. The canonical examples are: <span style=\"font-weight: bold\">(</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span><span style=\"font-weight: bold\">)</span> people giving money to organizations <span style=\"font-weight: bold\">(</span>and getting nothing\n",
       "tangible in\n",
       "    return<span style=\"font-weight: bold\">)</span>; and <span style=\"font-weight: bold\">(</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span><span style=\"font-weight: bold\">)</span> organizations lending money to people or other orgs.<span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"\n",
       "\n",
       "    mention: str  # The text span that most clearly expresses <span style=\"font-weight: bold\">(</span>triggers<span style=\"font-weight: bold\">)</span> the event\n",
       "    giver: List  # The donating agent\n",
       "    recipient: List  # The recipient agent\n",
       "    beneficiary: List  # The agent that benefits from the transfer\n",
       "    money: List  # The amount given, donated or loaned\n",
       "    time: List  # When the amount is transferred\n",
       "    place: List  # Where the transation takes place\n",
       "\n",
       "\n",
       "@dataclass\n",
       "class <span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">TransferOwnership</span><span style=\"font-weight: bold\">(</span>TransactionEvent<span style=\"font-weight: bold\">)</span>:\n",
       "    <span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"TransferOwnership Events refer to the buying, selling, loaning, borrowing, giving, or receiving of artifacts\n",
       "or\n",
       "    organizations.<span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"\n",
       "\n",
       "    mention: str  # The text span that most clearly expresses <span style=\"font-weight: bold\">(</span>triggers<span style=\"font-weight: bold\">)</span> the event\n",
       "    buyer: List  # The buying agent\n",
       "    seller: List  # The selling agent\n",
       "    beneficiary: List  # The agent that benefits from the transaction\n",
       "    artifact: List  # The item or Organization that was bought or sold\n",
       "    price: List  # The sale price of the artifact\n",
       "    time: List  # When the sale takes place\n",
       "    place: List  # Where the sale takes place\n",
       "\n",
       "\n",
       "# This is the text to analyze\n",
       "text = <span style=\"color: #008000; text-decoration-color: #008000\">\"The Times reported that Vivendi Universal was negotiating to sell its flagship theme parks to the New York </span>\n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">investment firm, Blackstone Group, for 8 billion as the first step toward dismantling its entertainment empire.\"</span>\n",
       "\n",
       "# The annotation instances that take place in the text above are listed here\n",
       "result = <span style=\"font-weight: bold\">[</span>\n",
       "    <span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">TransferOwnership</span><span style=\"font-weight: bold\">(</span>\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">mention</span>=<span style=\"color: #008000; text-decoration-color: #008000\">\"sell\"</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">buyer</span>=<span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">\"Blackstone Group\"</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">seller</span>=<span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">\"Vivendi Universal\"</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">beneficiary</span>=<span style=\"font-weight: bold\">[]</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">artifact</span>=<span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">\"theme parks\"</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">price</span>=<span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">\"8 billion\"</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">time</span>=<span style=\"font-weight: bold\">[]</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">place</span>=<span style=\"font-weight: bold\">[]</span>,\n",
       "    <span style=\"font-weight: bold\">)</span>,\n",
       "<span style=\"font-weight: bold\">]</span>\n",
       "\n",
       "</pre>\n"
      ],
      "text/plain": [
       "# The following lines describe the task definition\n",
       "@dataclass\n",
       "class \u001B[1;35mTransferMoney\u001B[0m\u001B[1m(\u001B[0mTransactionEvent\u001B[1m)\u001B[0m:\n",
       "    \u001B[32m\"\"\u001B[0m\"TransferMoney Events refer to the giving, receiving, borrowing, or lending money when it is not in the \n",
       "context of\n",
       "    purchasing something. The canonical examples are: \u001B[1m(\u001B[0m\u001B[1;36m1\u001B[0m\u001B[1m)\u001B[0m people giving money to organizations \u001B[1m(\u001B[0mand getting nothing\n",
       "tangible in\n",
       "    return\u001B[1m)\u001B[0m; and \u001B[1m(\u001B[0m\u001B[1;36m2\u001B[0m\u001B[1m)\u001B[0m organizations lending money to people or other orgs.\u001B[32m\"\"\u001B[0m\"\n",
       "\n",
       "    mention: str  # The text span that most clearly expresses \u001B[1m(\u001B[0mtriggers\u001B[1m)\u001B[0m the event\n",
       "    giver: List  # The donating agent\n",
       "    recipient: List  # The recipient agent\n",
       "    beneficiary: List  # The agent that benefits from the transfer\n",
       "    money: List  # The amount given, donated or loaned\n",
       "    time: List  # When the amount is transferred\n",
       "    place: List  # Where the transation takes place\n",
       "\n",
       "\n",
       "@dataclass\n",
       "class \u001B[1;35mTransferOwnership\u001B[0m\u001B[1m(\u001B[0mTransactionEvent\u001B[1m)\u001B[0m:\n",
       "    \u001B[32m\"\"\u001B[0m\"TransferOwnership Events refer to the buying, selling, loaning, borrowing, giving, or receiving of artifacts\n",
       "or\n",
       "    organizations.\u001B[32m\"\"\u001B[0m\"\n",
       "\n",
       "    mention: str  # The text span that most clearly expresses \u001B[1m(\u001B[0mtriggers\u001B[1m)\u001B[0m the event\n",
       "    buyer: List  # The buying agent\n",
       "    seller: List  # The selling agent\n",
       "    beneficiary: List  # The agent that benefits from the transaction\n",
       "    artifact: List  # The item or Organization that was bought or sold\n",
       "    price: List  # The sale price of the artifact\n",
       "    time: List  # When the sale takes place\n",
       "    place: List  # Where the sale takes place\n",
       "\n",
       "\n",
       "# This is the text to analyze\n",
       "text = \u001B[32m\"The Times reported that Vivendi Universal was negotiating to sell its flagship theme parks to the New York \u001B[0m\n",
       "\u001B[32minvestment firm, Blackstone Group, for 8 billion as the first step toward dismantling its entertainment empire.\"\u001B[0m\n",
       "\n",
       "# The annotation instances that take place in the text above are listed here\n",
       "result = \u001B[1m[\u001B[0m\n",
       "    \u001B[1;35mTransferOwnership\u001B[0m\u001B[1m(\u001B[0m\n",
       "        \u001B[33mmention\u001B[0m=\u001B[32m\"sell\"\u001B[0m,\n",
       "        \u001B[33mbuyer\u001B[0m=\u001B[1m[\u001B[0m\u001B[32m\"Blackstone Group\"\u001B[0m\u001B[1m]\u001B[0m,\n",
       "        \u001B[33mseller\u001B[0m=\u001B[1m[\u001B[0m\u001B[32m\"Vivendi Universal\"\u001B[0m\u001B[1m]\u001B[0m,\n",
       "        \u001B[33mbeneficiary\u001B[0m=\u001B[1m[\u001B[0m\u001B[1m]\u001B[0m,\n",
       "        \u001B[33martifact\u001B[0m=\u001B[1m[\u001B[0m\u001B[32m\"theme parks\"\u001B[0m\u001B[1m]\u001B[0m,\n",
       "        \u001B[33mprice\u001B[0m=\u001B[1m[\u001B[0m\u001B[32m\"8 billion\"\u001B[0m\u001B[1m]\u001B[0m,\n",
       "        \u001B[33mtime\u001B[0m=\u001B[1m[\u001B[0m\u001B[1m]\u001B[0m,\n",
       "        \u001B[33mplace\u001B[0m=\u001B[1m[\u001B[0m\u001B[1m]\u001B[0m,\n",
       "    \u001B[1m)\u001B[0m,\n",
       "\u001B[1m]\u001B[0m\n",
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "black_mode = black.Mode()\n",
    "formated_text = black.format_str(formated_text, mode=black_mode)\n",
    "rich.print(formated_text)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7c7ec97",
   "metadata": {},
   "source": [
    "As observed in the previous cell, we no longer provide the model with the gold trigger or event class. Therefore, GoLLIE will attempt to predict the correct trigger, event class, and arguments."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 220,
   "id": "e5ba6d02",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt, _ = formated_text.split(\"result =\")\n",
    "prompt = prompt + \"result =\"\n",
    "model_input = tokenizer(prompt, add_special_tokens=True, return_tensors=\"pt\")\n",
    "model_input[\"input_ids\"] = model_input[\"input_ids\"][:, :-1]\n",
    "model_input[\"attention_mask\"] = model_input[\"attention_mask\"][:, :-1]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d8c8558",
   "metadata": {},
   "source": [
    "Let's run the model 👀"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 221,
   "id": "517024d2",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:2 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 2.47 s, sys: 7.25 ms, total: 2.48 s\n",
      "Wall time: 2.49 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "model_ouput = model.generate(\n",
    "    **model_input.to(model.device),\n",
    "    max_new_tokens=128,\n",
    "    do_sample=False,\n",
    "    min_new_tokens=0,\n",
    "    num_beams=1,\n",
    "    num_return_sequences=1,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4442528e",
   "metadata": {},
   "source": [
    "And parse the output 🤞"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 224,
   "id": "f0fabd71",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">[</span>\n",
       "    <span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">TransferOwnership</span><span style=\"font-weight: bold\">(</span><span style=\"color: #808000; text-decoration-color: #808000\">mention</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'sell'</span>, <span style=\"color: #808000; text-decoration-color: #808000\">buyer</span>=<span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">'Blackstone Group'</span><span style=\"font-weight: bold\">]</span>, <span style=\"color: #808000; text-decoration-color: #808000\">seller</span>=<span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">'Vivendi Universal'</span><span style=\"font-weight: bold\">]</span>, <span style=\"color: #808000; text-decoration-color: #808000\">beneficiary</span>=<span style=\"font-weight: bold\">[]</span>, \n",
       "<span style=\"color: #808000; text-decoration-color: #808000\">artifact</span>=<span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">'theme parks'</span><span style=\"font-weight: bold\">]</span>, <span style=\"color: #808000; text-decoration-color: #808000\">price</span>=<span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">'8 billion'</span><span style=\"font-weight: bold\">]</span>, <span style=\"color: #808000; text-decoration-color: #808000\">time</span>=<span style=\"font-weight: bold\">[]</span>, <span style=\"color: #808000; text-decoration-color: #808000\">place</span>=<span style=\"font-weight: bold\">[])</span>\n",
       "<span style=\"font-weight: bold\">]</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001B[1m[\u001B[0m\n",
       "    \u001B[1;35mTransferOwnership\u001B[0m\u001B[1m(\u001B[0m\u001B[33mmention\u001B[0m=\u001B[32m'sell'\u001B[0m, \u001B[33mbuyer\u001B[0m=\u001B[1m[\u001B[0m\u001B[32m'Blackstone Group'\u001B[0m\u001B[1m]\u001B[0m, \u001B[33mseller\u001B[0m=\u001B[1m[\u001B[0m\u001B[32m'Vivendi Universal'\u001B[0m\u001B[1m]\u001B[0m, \u001B[33mbeneficiary\u001B[0m=\u001B[1m[\u001B[0m\u001B[1m]\u001B[0m, \n",
       "\u001B[33martifact\u001B[0m=\u001B[1m[\u001B[0m\u001B[32m'theme parks'\u001B[0m\u001B[1m]\u001B[0m, \u001B[33mprice\u001B[0m=\u001B[1m[\u001B[0m\u001B[32m'8 billion'\u001B[0m\u001B[1m]\u001B[0m, \u001B[33mtime\u001B[0m=\u001B[1m[\u001B[0m\u001B[1m]\u001B[0m, \u001B[33mplace\u001B[0m=\u001B[1m[\u001B[0m\u001B[1m]\u001B[0m\u001B[1m)\u001B[0m\n",
       "\u001B[1m]\u001B[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "result = AnnotationList.from_output(\n",
    "    tokenizer.decode(model_ouput[0],skip_special_tokens=True).split(\"result = \")[-1],\n",
    "    task_module=\"guidelines_eae\"\n",
    "    )\n",
    "rich.print(result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 225,
   "id": "e25cc91c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">{</span>\n",
       "    <span style=\"color: #008000; text-decoration-color: #008000\">'events'</span>: <span style=\"font-weight: bold\">{</span>\n",
       "        <span style=\"color: #008000; text-decoration-color: #008000\">'precision'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.0</span>,\n",
       "        <span style=\"color: #008000; text-decoration-color: #008000\">'recall'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.0</span>,\n",
       "        <span style=\"color: #008000; text-decoration-color: #008000\">'f1-score'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.0</span>,\n",
       "        <span style=\"color: #008000; text-decoration-color: #008000\">'class_scores'</span>: <span style=\"font-weight: bold\">{</span>\n",
       "            <span style=\"color: #008000; text-decoration-color: #008000\">'TransferOwnership'</span>: <span style=\"font-weight: bold\">{</span>\n",
       "                <span style=\"color: #008000; text-decoration-color: #008000\">'tp'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,\n",
       "                <span style=\"color: #008000; text-decoration-color: #008000\">'total_pos'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,\n",
       "                <span style=\"color: #008000; text-decoration-color: #008000\">'total_pre'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,\n",
       "                <span style=\"color: #008000; text-decoration-color: #008000\">'precision'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.0</span>,\n",
       "                <span style=\"color: #008000; text-decoration-color: #008000\">'recall'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.0</span>,\n",
       "                <span style=\"color: #008000; text-decoration-color: #008000\">'f1-score'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.0</span>\n",
       "            <span style=\"font-weight: bold\">}</span>\n",
       "        <span style=\"font-weight: bold\">}</span>\n",
       "    <span style=\"font-weight: bold\">}</span>,\n",
       "    <span style=\"color: #008000; text-decoration-color: #008000\">'arguments'</span>: <span style=\"font-weight: bold\">{</span><span style=\"color: #008000; text-decoration-color: #008000\">'precision'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.0</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'recall'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.0</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'f1-score'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.0</span><span style=\"font-weight: bold\">}</span>\n",
       "<span style=\"font-weight: bold\">}</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001B[1m{\u001B[0m\n",
       "    \u001B[32m'events'\u001B[0m: \u001B[1m{\u001B[0m\n",
       "        \u001B[32m'precision'\u001B[0m: \u001B[1;36m1.0\u001B[0m,\n",
       "        \u001B[32m'recall'\u001B[0m: \u001B[1;36m1.0\u001B[0m,\n",
       "        \u001B[32m'f1-score'\u001B[0m: \u001B[1;36m1.0\u001B[0m,\n",
       "        \u001B[32m'class_scores'\u001B[0m: \u001B[1m{\u001B[0m\n",
       "            \u001B[32m'TransferOwnership'\u001B[0m: \u001B[1m{\u001B[0m\n",
       "                \u001B[32m'tp'\u001B[0m: \u001B[1;36m1\u001B[0m,\n",
       "                \u001B[32m'total_pos'\u001B[0m: \u001B[1;36m1\u001B[0m,\n",
       "                \u001B[32m'total_pre'\u001B[0m: \u001B[1;36m1\u001B[0m,\n",
       "                \u001B[32m'precision'\u001B[0m: \u001B[1;36m1.0\u001B[0m,\n",
       "                \u001B[32m'recall'\u001B[0m: \u001B[1;36m1.0\u001B[0m,\n",
       "                \u001B[32m'f1-score'\u001B[0m: \u001B[1;36m1.0\u001B[0m\n",
       "            \u001B[1m}\u001B[0m\n",
       "        \u001B[1m}\u001B[0m\n",
       "    \u001B[1m}\u001B[0m,\n",
       "    \u001B[32m'arguments'\u001B[0m: \u001B[1m{\u001B[0m\u001B[32m'precision'\u001B[0m: \u001B[1;36m1.0\u001B[0m, \u001B[32m'recall'\u001B[0m: \u001B[1;36m1.0\u001B[0m, \u001B[32m'f1-score'\u001B[0m: \u001B[1;36m1.0\u001B[0m\u001B[1m}\u001B[0m\n",
       "\u001B[1m}\u001B[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from src.tasks.utils_scorer import EventScorer\n",
    "\n",
    "\n",
    "class MyScorer(EventScorer):\n",
    "    \"\"\"Event scorer.\"\"\"\n",
    "\n",
    "    valid_types: List[Type] = guidelines_eae.EVENTS_EAE_DEFINITIONS\n",
    "\n",
    "scorer = MyScorer()\n",
    "scorer_results = scorer(reference=[gold],predictions=[result])\n",
    "rich.print(scorer_results)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc27a0a2",
   "metadata": {},
   "source": [
    "GoLLIE has successfully performed end-to-end event extraction, even though it was not pre-trained for this task. The predicted trigger, event class, and all the arguments are correct. This demonstrates the strong generalization capabilities of the model 🎉🎉🎉\n",
    "\n",
    "\n",
    "We have not extensively tested GoLLIE's capabilities to perform end-to-end event extraction in zero-shot settings (unseen event types). However, we plan to do so in the near future. Feel free to venture into the unknown, and please share your exciting experiments with us; we'd love to see what everyone is achieving with GoLLIE!\n",
    "- [@iker_garciaf](https://twitter.com/iker_garciaf)\n",
    "- [@osainz59](https://twitter.com/osainz59)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "collie",
   "language": "python",
   "name": "collie"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}