notebooks/Relation 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": [
    "# Relation Extraction with GoLLIE\n",
    "\n",
    "This notebook is an example of how to run Relation Extraction 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",
    "In this notebook, we will demonstrate how to perform basic relation extraction with two string arguments per Relation class. However, GoLLIE can handle more complex relations with multiple arguments. Please refer to the `Create Custom Task` notebook if you wish to undertake more advanced relation extraction tasks 🔥🔥🔥."
   ]
  },
  {
   "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": 30,
   "id": "9ed51491",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append(\"../\") # Add the GoLLIE base directory to sys path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "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": "1d2171e8854a433c9975d5f5c36ca8cf",
       "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",
    " For this demonstration, we will define two Relations extracted from the ACE05 guidelines: 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": 32,
   "id": "b3381513",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "\n",
    "from src.tasks.utils_typing import Relation, dataclass\n",
    "\n",
    "\"\"\"\n",
    "Relation definitions\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class PhysicalRelation(Relation):\n",
    "    \"\"\"The Physical Relation captures the physical location relation of entities such as: \n",
    "    a Person entity located in a Facility, Location or GPE; or two entities that are near, \n",
    "    but neither entity is a part of the other or located in/at the other.\"\"\"\n",
    "\n",
    "    arg1: str\n",
    "    arg2: str\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class PersonalSocialRelation(Relation):\n",
    "    \"\"\"The Personal-Social Relation describe the relationship between people. Both arguments must be entities\n",
    "       of type Person. Please note: The arguments of these Relations are not ordered. The Relations are\n",
    "       symmetric.\"\"\"\n",
    "\n",
    "    arg1: str\n",
    "    arg2: str\n",
    "\n",
    "\n",
    "\n",
    "ENTITY_DEFINITIONS: List[Relation] = [\n",
    "    PhysicalRelation,\n",
    "    PersonalSocialRelation,\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": 33,
   "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": 34,
   "id": "89454475",
   "metadata": {},
   "outputs": [],
   "source": [
    "guidelines = [inspect.getsource(definition) for definition in ENTITY_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": 10,
   "id": "4d1f1c8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"Ana and Mary are sisters. Mary was at the supermarket while Ana was at home.\"\n",
    "gold = [\n",
    "    PersonalSocialRelation(arg1=\"Ana\",arg2=\"Mary\"),\n",
    "    PhysicalRelation(arg1=\"Mary\",arg2=\"supermarket\"),\n",
    "    PhysicalRelation(arg1=\"Ana\",arg2=\"home\"),\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "90501322",
   "metadata": {},
   "source": [
    "## Filling a template\n",
    "\n",
    "For Relation Extraction 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": 35,
   "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": 36,
   "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": 37,
   "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\">PhysicalRelation</span><span style=\"font-weight: bold\">(</span>Relation<span style=\"font-weight: bold\">)</span>:\n",
       "    <span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"The Physical Relation captures the physical location relation of entities such as:\n",
       "    a Person entity located in a Facility, Location or GPE; or two entities that are near,\n",
       "    but neither entity is a part of the other or located in/at the other.<span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"\n",
       "\n",
       "    arg1: str\n",
       "    arg2: str\n",
       "\n",
       "\n",
       "@dataclass\n",
       "class <span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">PersonalSocialRelation</span><span style=\"font-weight: bold\">(</span>Relation<span style=\"font-weight: bold\">)</span>:\n",
       "    <span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"The Personal-Social Relation describe the relationship between people. Both arguments must be entities\n",
       "    of type Person. Please note: The arguments of these Relations are not ordered. The Relations are\n",
       "    symmetric.<span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"\n",
       "\n",
       "    arg1: str\n",
       "    arg2: str\n",
       "\n",
       "\n",
       "# This is the text to analyze\n",
       "text = <span style=\"color: #008000; text-decoration-color: #008000\">\"Ana and Mary are sisters. Mary was at the supermarket while Ana was at home.\"</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\">PersonalSocialRelation</span><span style=\"font-weight: bold\">(</span><span style=\"color: #808000; text-decoration-color: #808000\">arg1</span>=<span style=\"color: #008000; text-decoration-color: #008000\">\"Ana\"</span>, <span style=\"color: #808000; text-decoration-color: #808000\">arg2</span>=<span style=\"color: #008000; text-decoration-color: #008000\">\"Mary\"</span><span style=\"font-weight: bold\">)</span>,\n",
       "    <span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">PhysicalRelation</span><span style=\"font-weight: bold\">(</span><span style=\"color: #808000; text-decoration-color: #808000\">arg1</span>=<span style=\"color: #008000; text-decoration-color: #008000\">\"Mary\"</span>, <span style=\"color: #808000; text-decoration-color: #808000\">arg2</span>=<span style=\"color: #008000; text-decoration-color: #008000\">\"supermarket\"</span><span style=\"font-weight: bold\">)</span>,\n",
       "    <span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">PhysicalRelation</span><span style=\"font-weight: bold\">(</span><span style=\"color: #808000; text-decoration-color: #808000\">arg1</span>=<span style=\"color: #008000; text-decoration-color: #008000\">\"Ana\"</span>, <span style=\"color: #808000; text-decoration-color: #808000\">arg2</span>=<span style=\"color: #008000; text-decoration-color: #008000\">\"home\"</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;35mPhysicalRelation\u001B[0m\u001B[1m(\u001B[0mRelation\u001B[1m)\u001B[0m:\n",
       "    \u001B[32m\"\"\u001B[0m\"The Physical Relation captures the physical location relation of entities such as:\n",
       "    a Person entity located in a Facility, Location or GPE; or two entities that are near,\n",
       "    but neither entity is a part of the other or located in/at the other.\u001B[32m\"\"\u001B[0m\"\n",
       "\n",
       "    arg1: str\n",
       "    arg2: str\n",
       "\n",
       "\n",
       "@dataclass\n",
       "class \u001B[1;35mPersonalSocialRelation\u001B[0m\u001B[1m(\u001B[0mRelation\u001B[1m)\u001B[0m:\n",
       "    \u001B[32m\"\"\u001B[0m\"The Personal-Social Relation describe the relationship between people. Both arguments must be entities\n",
       "    of type Person. Please note: The arguments of these Relations are not ordered. The Relations are\n",
       "    symmetric.\u001B[32m\"\"\u001B[0m\"\n",
       "\n",
       "    arg1: str\n",
       "    arg2: str\n",
       "\n",
       "\n",
       "# This is the text to analyze\n",
       "text = \u001B[32m\"Ana and Mary are sisters. Mary was at the supermarket while Ana was at home.\"\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;35mPersonalSocialRelation\u001B[0m\u001B[1m(\u001B[0m\u001B[33marg1\u001B[0m=\u001B[32m\"Ana\"\u001B[0m, \u001B[33marg2\u001B[0m=\u001B[32m\"Mary\"\u001B[0m\u001B[1m)\u001B[0m,\n",
       "    \u001B[1;35mPhysicalRelation\u001B[0m\u001B[1m(\u001B[0m\u001B[33marg1\u001B[0m=\u001B[32m\"Mary\"\u001B[0m, \u001B[33marg2\u001B[0m=\u001B[32m\"supermarket\"\u001B[0m\u001B[1m)\u001B[0m,\n",
       "    \u001B[1;35mPhysicalRelation\u001B[0m\u001B[1m(\u001B[0m\u001B[33marg1\u001B[0m=\u001B[32m\"Ana\"\u001B[0m, \u001B[33marg2\u001B[0m=\u001B[32m\"home\"\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": 38,
   "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": 39,
   "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": 40,
   "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": 41,
   "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 2.07 s, sys: 5.39 ms, total: 2.08 s\n",
      "Wall time: 2.09 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": 42,
   "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\">PersonalSocialRelation</span><span style=\"font-weight: bold\">(</span><span style=\"color: #808000; text-decoration-color: #808000\">arg1</span>=<span style=\"color: #008000; text-decoration-color: #008000\">\"Ana\"</span>, <span style=\"color: #808000; text-decoration-color: #808000\">arg2</span>=<span style=\"color: #008000; text-decoration-color: #008000\">\"Mary\"</span><span style=\"font-weight: bold\">)</span>,\n",
       "    <span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">PhysicalRelation</span><span style=\"font-weight: bold\">(</span><span style=\"color: #808000; text-decoration-color: #808000\">arg1</span>=<span style=\"color: #008000; text-decoration-color: #008000\">\"Mary\"</span>, <span style=\"color: #808000; text-decoration-color: #808000\">arg2</span>=<span style=\"color: #008000; text-decoration-color: #008000\">\"supermarket\"</span><span style=\"font-weight: bold\">)</span>,\n",
       "    <span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">PhysicalRelation</span><span style=\"font-weight: bold\">(</span><span style=\"color: #808000; text-decoration-color: #808000\">arg1</span>=<span style=\"color: #008000; text-decoration-color: #008000\">\"Ana\"</span>, <span style=\"color: #808000; text-decoration-color: #808000\">arg2</span>=<span style=\"color: #008000; text-decoration-color: #008000\">\"home\"</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;35mPersonalSocialRelation\u001B[0m\u001B[1m(\u001B[0m\u001B[33marg1\u001B[0m=\u001B[32m\"Ana\"\u001B[0m, \u001B[33marg2\u001B[0m=\u001B[32m\"Mary\"\u001B[0m\u001B[1m)\u001B[0m,\n",
       "    \u001B[1;35mPhysicalRelation\u001B[0m\u001B[1m(\u001B[0m\u001B[33marg1\u001B[0m=\u001B[32m\"Mary\"\u001B[0m, \u001B[33marg2\u001B[0m=\u001B[32m\"supermarket\"\u001B[0m\u001B[1m)\u001B[0m,\n",
       "    \u001B[1;35mPhysicalRelation\u001B[0m\u001B[1m(\u001B[0m\u001B[33marg1\u001B[0m=\u001B[32m\"Ana\"\u001B[0m, \u001B[33marg2\u001B[0m=\u001B[32m\"home\"\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": 43,
   "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\">PersonalSocialRelation</span><span style=\"font-weight: bold\">(</span><span style=\"color: #808000; text-decoration-color: #808000\">arg1</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'Ana'</span>, <span style=\"color: #808000; text-decoration-color: #808000\">arg2</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'Mary'</span><span style=\"font-weight: bold\">)</span>,\n",
       "    <span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">PhysicalRelation</span><span style=\"font-weight: bold\">(</span><span style=\"color: #808000; text-decoration-color: #808000\">arg1</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'Mary'</span>, <span style=\"color: #808000; text-decoration-color: #808000\">arg2</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'supermarket'</span><span style=\"font-weight: bold\">)</span>,\n",
       "    <span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">PhysicalRelation</span><span style=\"font-weight: bold\">(</span><span style=\"color: #808000; text-decoration-color: #808000\">arg1</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'Ana'</span>, <span style=\"color: #808000; text-decoration-color: #808000\">arg2</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'home'</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;35mPersonalSocialRelation\u001B[0m\u001B[1m(\u001B[0m\u001B[33marg1\u001B[0m=\u001B[32m'Ana'\u001B[0m, \u001B[33marg2\u001B[0m=\u001B[32m'Mary'\u001B[0m\u001B[1m)\u001B[0m,\n",
       "    \u001B[1;35mPhysicalRelation\u001B[0m\u001B[1m(\u001B[0m\u001B[33marg1\u001B[0m=\u001B[32m'Mary'\u001B[0m, \u001B[33marg2\u001B[0m=\u001B[32m'supermarket'\u001B[0m\u001B[1m)\u001B[0m,\n",
       "    \u001B[1;35mPhysicalRelation\u001B[0m\u001B[1m(\u001B[0m\u001B[33marg1\u001B[0m=\u001B[32m'Ana'\u001B[0m, \u001B[33marg2\u001B[0m=\u001B[32m'home'\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": 44,
   "id": "bd039309",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "guidelines.PersonalSocialRelation"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(result[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "2703325b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Ana'"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result[0].arg1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "7349d97f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Mary'"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result[0].arg2"
   ]
  },
  {
   "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 Relation Extraction, we will use the `RelationScorer` 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": 47,
   "id": "d44e1f46",
   "metadata": {},
   "outputs": [],
   "source": [
    "from src.tasks.utils_scorer import RelationScorer\n",
    "\n",
    "class MyScorer(RelationScorer):\n",
    "    \"\"\"Compute the F1 score for Relation Extraction\"\"\"\n",
    "\n",
    "    valid_types: List[Type] = ENTITY_DEFINITIONS\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9907289",
   "metadata": {},
   "source": [
    "### Instanciate the scorer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "93d6c8dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "scorer = MyScorer()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ccb11ce1",
   "metadata": {},
   "source": [
    "### Compute F1 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "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\">'relations'</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\">'PersonalSocialRelation'</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\">'PhysicalRelation'</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\">2</span>,\n",
       "                <span style=\"color: #008000; text-decoration-color: #008000\">'total_pos'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>,\n",
       "                <span style=\"color: #008000; text-decoration-color: #008000\">'total_pre'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</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=\"font-weight: bold\">}</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001B[1m{\u001B[0m\n",
       "    \u001B[32m'relations'\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'PersonalSocialRelation'\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'PhysicalRelation'\u001B[0m: \u001B[1m{\u001B[0m\n",
       "                \u001B[32m'tp'\u001B[0m: \u001B[1;36m2\u001B[0m,\n",
       "                \u001B[32m'total_pos'\u001B[0m: \u001B[1;36m2\u001B[0m,\n",
       "                \u001B[32m'total_pre'\u001B[0m: \u001B[1;36m2\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[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": "fc27a0a2",
   "metadata": {},
   "source": [
    "GoLLIE has successfully labeled the sententence with the defined relations 🎉🎉🎉\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
}