notebooks/Create Custom Task.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": [
    "# Custom Tasks with GoLLIE\n",
    "\n",
    "This notebook is an example of how to run Custom Tasks with GoLLIE. This notebook covers:\n",
    "\n",
    "- How to define the guidelines for a task\n",
    "- How to load GoLLIE\n",
    "- How to generate model inputs\n",
    "- How to parse the output\n",
    "- How to implement a scorer and evaluate the output\n",
    "\n",
    "You can modify this notebook to run any task 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": 23,
   "id": "9ed51491",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append(\"../\") # Add the GoLLIE base directory to sys path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "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 as jinja2Template\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": 25,
   "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",
      "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"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "22558d19f46842f89184ff8fc6576c9f",
       "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": "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",
    "The following guidelines have been defined for this example. They were not part of the pre-training dataset. Therefore, we will run GOLLIE in zero-shot settings using unseen labels.\n",
    "\n",
    "We will use the `Generic` class, which is a versatile class that allows for the implementation of any task you want. However, since the model has never seen the Generic label during training, we will rename it to Template, which is recognized by the model (as it was used in the Tacred dataset).\n",
    "\n",
    "We will define two classes: `Launcher` and `Mission`. Each class will have a definition and a set of slots that the model needs to fill. Each slot also requires a type definition and a short description, which can include examples. For instance, for the `Launcher` class, we define three slots:\n",
    "\n",
    "- The `mention`, which will be the name of the Launcher vehicle and should be a string.\n",
    "- The `space_company` that operated the vehicle, which will also be a string.\n",
    "- The `crew`, which is defined as a list of astronauts. Therefore, GoLLIE will fill this slot with a list of strings.\n",
    "\n",
    "💡 Be creative and try to define your own guidelines to test GoLLIE!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "dd5d262f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "\n",
    "from src.tasks.utils_typing import dataclass\n",
    "from src.tasks.utils_typing import Generic as Template\n",
    "\n",
    "\"\"\"\n",
    "Entity definitions\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class Launcher(Template):\n",
    "    \"\"\"Refers to a vehicle designed primarily to transport payloads from the Earth's \n",
    "    surface to space. Launchers can carry various payloads, including satellites, \n",
    "    crewed spacecraft, and cargo, into various orbits or even beyond Earth's orbit. \n",
    "    They are usually multi-stage vehicles that use rocket engines for propulsion.\"\"\"\n",
    "\n",
    "    mention: str  \n",
    "    \"\"\"\n",
    "    The name of the launcher vehicle. \n",
    "    Such as: \"Sturn V\", \"Atlas V\", \"Soyuz\", \"Ariane 5\"\n",
    "    \"\"\"\n",
    "    space_company: str # The company that operates the launcher. Such as: \"Blue origin\", \"ESA\", \"Boeing\", \"ISRO\", \"Northrop Grumman\", \"Arianespace\"\n",
    "    crew: List[str] # Names of the crew members boarding the Launcher. Such as: \"Neil Armstrong\", \"Michael Collins\", \"Buzz Aldrin\"\n",
    "    \n",
    "\n",
    "@dataclass\n",
    "class Mission(Template):\n",
    "    \"\"\"Any planned or accomplished journey beyond Earth's atmosphere with specific objectives, \n",
    "    either crewed or uncrewed. It includes missions to satellites, the International \n",
    "    Space Station (ISS), other celestial bodies, and deep space.\"\"\"\n",
    "    \n",
    "    mention: str\n",
    "    \"\"\"\n",
    "    The name of the mission. \n",
    "    Such as: \"Apollo 11\", \"Artemis\", \"Mercury\"\n",
    "    \"\"\"\n",
    "    date: str # The start date of the mission\n",
    "    departure: str # The place from which the vehicle will be launched. Such as: \"Florida\", \"Houston\", \"French Guiana\"\n",
    "    destination: str # The place or planet to which the launcher will be sent. Such as \"Moon\", \"low-orbit\", \"Saturn\"\n",
    "\n",
    "\n",
    "ENTITY_DEFINITIONS: List[Template] = [\n",
    "    Launcher,\n",
    "    Mission,\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": 27,
   "id": "d4736a37",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"guidelines.py\",\"w\",encoding=\"utf8\") as python_guidelines:\n",
    "    print(cell_txt,file=python_guidelines)\n",
    "\n",
    "from guidelines import *"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0ac3db6f",
   "metadata": {},
   "source": [
    "We use inspect.getsource to get the guidelines as a string"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "89454475",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['@dataclass\\nclass Launcher(Template):\\n    \"\"\"Refers to a vehicle designed primarily to transport payloads from the Earth\\'s \\n    surface to space. Launchers can carry various payloads, including satellites, \\n    crewed spacecraft, and cargo, into various orbits or even beyond Earth\\'s orbit. \\n    They are usually multi-stage vehicles that use rocket engines for propulsion.\"\"\"\\n\\n    mention: str  \\n    \"\"\"\\n    The name of the launcher vehicle. \\n    Such as: \"Sturn V\", \"Atlas V\", \"Soyuz\", \"Ariane 5\"\\n    \"\"\"\\n    space_company: str # The company that operates the launcher. Such as: \"Blue origin\", \"ESA\", \"Boeing\", \"ISRO\", \"Northrop Grumman\", \"Arianespace\"\\n    crew: List[str] # Names of the crew members boarding the Launcher. Such as: \"Neil Armstrong\", \"Michael Collins\", \"Buzz Aldrin\"\\n',\n",
       " '@dataclass\\nclass Mission(Template):\\n    \"\"\"Any planned or accomplished journey beyond Earth\\'s atmosphere with specific objectives, \\n    either crewed or uncrewed. It includes missions to satellites, the International \\n    Space Station (ISS), other celestial bodies, and deep space.\"\"\"\\n    \\n    mention: str\\n    \"\"\"\\n    The name of the mission. \\n    Such as: \"Apollo 11\", \"Artemis\", \"Mercury\"\\n    \"\"\"\\n    date: str # The start date of the mission\\n    departure: str # The place from which the vehicle will be launched. Such as: \"Florida\", \"Houston\", \"French Guiana\"\\n    destination: str # The place or planet to which the launcher will be sent. Such as \"Moon\", \"low-orbit\", \"Saturn\"\\n']"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "guidelines = [inspect.getsource(definition) for definition in ENTITY_DEFINITIONS]\n",
    "guidelines"
   ]
  },
  {
   "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": 29,
   "id": "dfb92535",
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"The Ares 3 mission to Mars is scheduled for 2032. The Starship rocket build by SpaceX will take off from Boca Chica, carrying the astronauts Max Rutherford, Elena Soto, and Jake Martinez.\"\n",
    "gold = [\n",
    "    Launcher(mention=\"Starship\",space_company=\"SpaceX\",crew=[\"Max Rutherford\",\"Elena Soto\",\"Jake Martinez\"]),\n",
    "    Mission(mention=\"Ares 3\",date=\"2032\",departure=\"Boca Chica\",destination=\"Mars\")\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "90501322",
   "metadata": {},
   "source": [
    "## Filling a template\n",
    "\n",
    "We need to define a template. For this task, we will include only the class definitions and the text to be annotated. However, you can design different templates to incorporate more information (for example, event triggers, as demonstrated in the Event Extraction notebook).\n",
    "\n",
    "We will 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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "9b365413",
   "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",
    "\"\"\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "0f54034f",
   "metadata": {},
   "outputs": [],
   "source": [
    "template = jinja2Template(template_txt)\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": 32,
   "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": 33,
   "id": "fa5f3106",
   "metadata": {
    "scrolled": false
   },
   "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\">Launcher</span><span style=\"font-weight: bold\">(</span>Template<span style=\"font-weight: bold\">)</span>:\n",
       "    <span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"Refers to a vehicle designed primarily to transport payloads from the Earth's\n",
       "    surface to space. Launchers can carry various payloads, including satellites,\n",
       "    crewed spacecraft, and cargo, into various orbits or even beyond Earth's orbit.\n",
       "    They are usually multi-stage vehicles that use rocket engines for propulsion.<span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"\n",
       "\n",
       "    mention: str\n",
       "    <span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"\n",
       "    The name of the launcher vehicle. \n",
       "    Such as: <span style=\"color: #008000; text-decoration-color: #008000\">\"Sturn V\"</span>, <span style=\"color: #008000; text-decoration-color: #008000\">\"Atlas V\"</span>, <span style=\"color: #008000; text-decoration-color: #008000\">\"Soyuz\"</span>, <span style=\"color: #008000; text-decoration-color: #008000\">\"Ariane 5\"</span>\n",
       "    <span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"\n",
       "    space_company: str  # The company that operates the launcher. Such as: <span style=\"color: #008000; text-decoration-color: #008000\">\"Blue origin\"</span>, <span style=\"color: #008000; text-decoration-color: #008000\">\"ESA\"</span>, <span style=\"color: #008000; text-decoration-color: #008000\">\"Boeing\"</span>, <span style=\"color: #008000; text-decoration-color: #008000\">\"ISRO\"</span>, \n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">\"Northrop Grumman\"</span>, <span style=\"color: #008000; text-decoration-color: #008000\">\"Arianespace\"</span>\n",
       "    crew: List<span style=\"font-weight: bold\">[</span>\n",
       "        str\n",
       "    <span style=\"font-weight: bold\">]</span>  # Names of the crew members boarding the Launcher. Such as: <span style=\"color: #008000; text-decoration-color: #008000\">\"Neil Armstrong\"</span>, <span style=\"color: #008000; text-decoration-color: #008000\">\"Michael Collins\"</span>, <span style=\"color: #008000; text-decoration-color: #008000\">\"Buzz </span>\n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">Aldrin\"</span>\n",
       "\n",
       "\n",
       "@dataclass\n",
       "class <span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">Mission</span><span style=\"font-weight: bold\">(</span>Template<span style=\"font-weight: bold\">)</span>:\n",
       "    <span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"Any planned or accomplished journey beyond Earth's atmosphere with specific objectives,\n",
       "    either crewed or uncrewed. It includes missions to satellites, the International\n",
       "    Space Station <span style=\"font-weight: bold\">(</span>ISS<span style=\"font-weight: bold\">)</span>, other celestial bodies, and deep space.<span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"\n",
       "\n",
       "    mention: str\n",
       "    <span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"\n",
       "    The name of the mission. \n",
       "    Such as: <span style=\"color: #008000; text-decoration-color: #008000\">\"Apollo 11\"</span>, <span style=\"color: #008000; text-decoration-color: #008000\">\"Artemis\"</span>, <span style=\"color: #008000; text-decoration-color: #008000\">\"Mercury\"</span>\n",
       "    <span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"\n",
       "    date: str  # The start date of the mission\n",
       "    departure: str  # The place from which the vehicle will be launched. Such as: <span style=\"color: #008000; text-decoration-color: #008000\">\"Florida\"</span>, <span style=\"color: #008000; text-decoration-color: #008000\">\"Houston\"</span>, <span style=\"color: #008000; text-decoration-color: #008000\">\"French </span>\n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">Guiana\"</span>\n",
       "    destination: str  # The place or planet to which the launcher will be sent. Such as <span style=\"color: #008000; text-decoration-color: #008000\">\"Moon\"</span>, <span style=\"color: #008000; text-decoration-color: #008000\">\"low-orbit\"</span>, \n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">\"Saturn\"</span>\n",
       "\n",
       "\n",
       "# This is the text to analyze\n",
       "text = <span style=\"color: #008000; text-decoration-color: #008000\">\"The Ares 3 mission to Mars is scheduled for 2032. The Starship rocket build by SpaceX will take off from </span>\n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">Boca Chica, carrying the astronauts Max Rutherford, Elena Soto, and Jake Martinez.\"</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\">Launcher</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\">\"Starship\"</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">space_company</span>=<span style=\"color: #008000; text-decoration-color: #008000\">\"SpaceX\"</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">crew</span>=<span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">\"Max Rutherford\"</span>, <span style=\"color: #008000; text-decoration-color: #008000\">\"Elena Soto\"</span>, <span style=\"color: #008000; text-decoration-color: #008000\">\"Jake Martinez\"</span><span style=\"font-weight: bold\">]</span>,\n",
       "    <span style=\"font-weight: bold\">)</span>,\n",
       "    <span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">Mission</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\">\"Ares 3\"</span>, <span style=\"color: #808000; text-decoration-color: #808000\">date</span>=<span style=\"color: #008000; text-decoration-color: #008000\">\"2032\"</span>, <span style=\"color: #808000; text-decoration-color: #808000\">departure</span>=<span style=\"color: #008000; text-decoration-color: #008000\">\"Boca Chica\"</span>, <span style=\"color: #808000; text-decoration-color: #808000\">destination</span>=<span style=\"color: #008000; text-decoration-color: #008000\">\"Mars\"</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;35mLauncher\u001b[0m\u001b[1m(\u001b[0mTemplate\u001b[1m)\u001b[0m:\n",
       "    \u001b[32m\"\"\u001b[0m\"Refers to a vehicle designed primarily to transport payloads from the Earth's\n",
       "    surface to space. Launchers can carry various payloads, including satellites,\n",
       "    crewed spacecraft, and cargo, into various orbits or even beyond Earth's orbit.\n",
       "    They are usually multi-stage vehicles that use rocket engines for propulsion.\u001b[32m\"\"\u001b[0m\"\n",
       "\n",
       "    mention: str\n",
       "    \u001b[32m\"\"\u001b[0m\"\n",
       "    The name of the launcher vehicle. \n",
       "    Such as: \u001b[32m\"Sturn V\"\u001b[0m, \u001b[32m\"Atlas V\"\u001b[0m, \u001b[32m\"Soyuz\"\u001b[0m, \u001b[32m\"Ariane 5\"\u001b[0m\n",
       "    \u001b[32m\"\"\u001b[0m\"\n",
       "    space_company: str  # The company that operates the launcher. Such as: \u001b[32m\"Blue origin\"\u001b[0m, \u001b[32m\"ESA\"\u001b[0m, \u001b[32m\"Boeing\"\u001b[0m, \u001b[32m\"ISRO\"\u001b[0m, \n",
       "\u001b[32m\"Northrop Grumman\"\u001b[0m, \u001b[32m\"Arianespace\"\u001b[0m\n",
       "    crew: List\u001b[1m[\u001b[0m\n",
       "        str\n",
       "    \u001b[1m]\u001b[0m  # Names of the crew members boarding the Launcher. Such as: \u001b[32m\"Neil Armstrong\"\u001b[0m, \u001b[32m\"Michael Collins\"\u001b[0m, \u001b[32m\"Buzz \u001b[0m\n",
       "\u001b[32mAldrin\"\u001b[0m\n",
       "\n",
       "\n",
       "@dataclass\n",
       "class \u001b[1;35mMission\u001b[0m\u001b[1m(\u001b[0mTemplate\u001b[1m)\u001b[0m:\n",
       "    \u001b[32m\"\"\u001b[0m\"Any planned or accomplished journey beyond Earth's atmosphere with specific objectives,\n",
       "    either crewed or uncrewed. It includes missions to satellites, the International\n",
       "    Space Station \u001b[1m(\u001b[0mISS\u001b[1m)\u001b[0m, other celestial bodies, and deep space.\u001b[32m\"\"\u001b[0m\"\n",
       "\n",
       "    mention: str\n",
       "    \u001b[32m\"\"\u001b[0m\"\n",
       "    The name of the mission. \n",
       "    Such as: \u001b[32m\"Apollo 11\"\u001b[0m, \u001b[32m\"Artemis\"\u001b[0m, \u001b[32m\"Mercury\"\u001b[0m\n",
       "    \u001b[32m\"\"\u001b[0m\"\n",
       "    date: str  # The start date of the mission\n",
       "    departure: str  # The place from which the vehicle will be launched. Such as: \u001b[32m\"Florida\"\u001b[0m, \u001b[32m\"Houston\"\u001b[0m, \u001b[32m\"French \u001b[0m\n",
       "\u001b[32mGuiana\"\u001b[0m\n",
       "    destination: str  # The place or planet to which the launcher will be sent. Such as \u001b[32m\"Moon\"\u001b[0m, \u001b[32m\"low-orbit\"\u001b[0m, \n",
       "\u001b[32m\"Saturn\"\u001b[0m\n",
       "\n",
       "\n",
       "# This is the text to analyze\n",
       "text = \u001b[32m\"The Ares 3 mission to Mars is scheduled for 2032. The Starship rocket build by SpaceX will take off from \u001b[0m\n",
       "\u001b[32mBoca Chica, carrying the astronauts Max Rutherford, Elena Soto, and Jake Martinez.\"\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;35mLauncher\u001b[0m\u001b[1m(\u001b[0m\n",
       "        \u001b[33mmention\u001b[0m=\u001b[32m\"Starship\"\u001b[0m,\n",
       "        \u001b[33mspace_company\u001b[0m=\u001b[32m\"SpaceX\"\u001b[0m,\n",
       "        \u001b[33mcrew\u001b[0m=\u001b[1m[\u001b[0m\u001b[32m\"Max Rutherford\"\u001b[0m, \u001b[32m\"Elena Soto\"\u001b[0m, \u001b[32m\"Jake Martinez\"\u001b[0m\u001b[1m]\u001b[0m,\n",
       "    \u001b[1m)\u001b[0m,\n",
       "    \u001b[1;35mMission\u001b[0m\u001b[1m(\u001b[0m\u001b[33mmention\u001b[0m=\u001b[32m\"Ares\u001b[0m\u001b[32m 3\"\u001b[0m, \u001b[33mdate\u001b[0m=\u001b[32m\"2032\"\u001b[0m, \u001b[33mdeparture\u001b[0m=\u001b[32m\"Boca\u001b[0m\u001b[32m Chica\"\u001b[0m, \u001b[33mdestination\u001b[0m=\u001b[32m\"Mars\"\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": 34,
   "id": "19eabf38",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt, _ = formated_text.split(\"result =\")\n",
    "prompt = prompt + \"result =\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0620be2",
   "metadata": {},
   "source": [
    "Tokenize the input sentence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "f13c79d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_input = tokenizer(prompt, add_special_tokens=True, return_tensors=\"pt\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dff91970",
   "metadata": {},
   "source": [
    "Remove the `eos` token from the input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "ee9d69ce",
   "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": 37,
   "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 3.24 s, sys: 29.3 ms, total: 3.27 s\n",
      "Wall time: 3.33 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": "983f4a2a",
   "metadata": {},
   "source": [
    "### Print the results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "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\">Mission</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\">\"Ares 3\"</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">date</span>=<span style=\"color: #008000; text-decoration-color: #008000\">\"2032\"</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">departure</span>=<span style=\"color: #008000; text-decoration-color: #008000\">\"Boca Chica\"</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">destination</span>=<span style=\"color: #008000; text-decoration-color: #008000\">\"Mars\"</span>,\n",
       "    <span style=\"font-weight: bold\">)</span>,\n",
       "    <span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">Launcher</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\">\"Starship\"</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">space_company</span>=<span style=\"color: #008000; text-decoration-color: #008000\">\"SpaceX\"</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">crew</span>=<span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">\"Max Rutherford\"</span>, <span style=\"color: #008000; text-decoration-color: #008000\">\"Elena Soto\"</span>, <span style=\"color: #008000; text-decoration-color: #008000\">\"Jake Martinez\"</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;35mMission\u001b[0m\u001b[1m(\u001b[0m\n",
       "        \u001b[33mmention\u001b[0m=\u001b[32m\"Ares\u001b[0m\u001b[32m 3\"\u001b[0m,\n",
       "        \u001b[33mdate\u001b[0m=\u001b[32m\"2032\"\u001b[0m,\n",
       "        \u001b[33mdeparture\u001b[0m=\u001b[32m\"Boca\u001b[0m\u001b[32m Chica\"\u001b[0m,\n",
       "        \u001b[33mdestination\u001b[0m=\u001b[32m\"Mars\"\u001b[0m,\n",
       "    \u001b[1m)\u001b[0m,\n",
       "    \u001b[1;35mLauncher\u001b[0m\u001b[1m(\u001b[0m\n",
       "        \u001b[33mmention\u001b[0m=\u001b[32m\"Starship\"\u001b[0m,\n",
       "        \u001b[33mspace_company\u001b[0m=\u001b[32m\"SpaceX\"\u001b[0m,\n",
       "        \u001b[33mcrew\u001b[0m=\u001b[1m[\u001b[0m\u001b[32m\"Max Rutherford\"\u001b[0m, \u001b[32m\"Elena Soto\"\u001b[0m, \u001b[32m\"Jake Martinez\"\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": "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": 39,
   "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>\n",
       "    <span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">Mission</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\">'Ares 3'</span>, <span style=\"color: #808000; text-decoration-color: #808000\">date</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'2032'</span>, <span style=\"color: #808000; text-decoration-color: #808000\">departure</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'Boca Chica'</span>, <span style=\"color: #808000; text-decoration-color: #808000\">destination</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'Mars'</span><span style=\"font-weight: bold\">)</span>,\n",
       "    <span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">Launcher</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\">'Starship'</span>, <span style=\"color: #808000; text-decoration-color: #808000\">space_company</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'SpaceX'</span>, <span style=\"color: #808000; text-decoration-color: #808000\">crew</span>=<span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">'Max Rutherford'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'Elena Soto'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'Jake Martinez'</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;35mMission\u001b[0m\u001b[1m(\u001b[0m\u001b[33mmention\u001b[0m=\u001b[32m'Ares 3'\u001b[0m, \u001b[33mdate\u001b[0m=\u001b[32m'2032'\u001b[0m, \u001b[33mdeparture\u001b[0m=\u001b[32m'Boca Chica'\u001b[0m, \u001b[33mdestination\u001b[0m=\u001b[32m'Mars'\u001b[0m\u001b[1m)\u001b[0m,\n",
       "    \u001b[1;35mLauncher\u001b[0m\u001b[1m(\u001b[0m\u001b[33mmention\u001b[0m=\u001b[32m'Starship'\u001b[0m, \u001b[33mspace_company\u001b[0m=\u001b[32m'SpaceX'\u001b[0m, \u001b[33mcrew\u001b[0m=\u001b[1m[\u001b[0m\u001b[32m'Max Rutherford'\u001b[0m, \u001b[32m'Elena Soto'\u001b[0m, \u001b[32m'Jake Martinez'\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\"\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": 40,
   "id": "bd039309",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "guidelines.Mission"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(result[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "2703325b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Ares 3'"
      ]
     },
     "execution_count": 41,
     "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 Named Entity Recognition, we will use the `SpanScorer` 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": 42,
   "id": "d44e1f46",
   "metadata": {},
   "outputs": [],
   "source": [
    "from src.tasks.utils_scorer import TemplateScorer\n",
    "\n",
    "class MyScorer(TemplateScorer):\n",
    "    \"\"\"Compute the F1 score for Generic Task\"\"\"\n",
    "\n",
    "    valid_types: List[Type] = ENTITY_DEFINITIONS"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9907289",
   "metadata": {},
   "source": [
    "### Instanciate the scorer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "93d6c8dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "scorer = MyScorer()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ccb11ce1",
   "metadata": {},
   "source": [
    "### Compute F1 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "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\">'templates'</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\">'Launcher'</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=\"color: #008000; text-decoration-color: #008000\">'Mission'</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\">'slots'</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'templates'\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'Launcher'\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[32m'Mission'\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'slots'\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": "3d86d47e",
   "metadata": {},
   "source": [
    "GoLLIE has successfully labeled a sentence using a set of labels that were not part of the pretraining dataset 🎉🎉🎉\n",
    "\n",
    "GoLLIE will perform well on labels with well-defined and clearly bounded guidelines. \n",
    "\n",
    "Please share your cool experiments with us; we'd love to see what everyone is doing 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
}