coref_mt5/colabs/coreference_decoder.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "Copyright 2023 Google LLC.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");"
      ],
      "metadata": {
        "id": "ZMBdysAcyTdm"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "```\n",
        "# Copyright 2022 The Google Research Authors.\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     http://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License.\n",
        "\n",
        "```"
      ],
      "metadata": {
        "id": "MneR1qP_uFsR"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Decoder of the paper\n",
        "# [Coreference Resolution through a seq2seq Transition-Based System](https://arxiv.org/abs/2211.12142)\n",
        "\n",
        "\n",
        "\n",
        "```\n",
        "@misc{https://doi.org/10.48550/arxiv.2211.12142,\n",
        "  url = {https://arxiv.org/abs/2211.12142},\n",
        "  author = {Bohnet, Bernd and Alberti, Chris and Collins, Michael},\n",
        "  title = {Coreference Resolution through a seq2seq Transition-Based System},\n",
        "  publisher = {TACL},\n",
        "  year = {2023},\n",
        "}\n",
        "\n",
        "```\n",
        "\n",
        "\n",
        "\n",
        "# Requirements\n",
        "\n",
        "\n",
        "1.   Please start a mt5 sever with the checkpoint [Coref-mT5-XXL model](https://console.cloud.google.com/storage/browser/gresearch/correference_seq2seq/checkpoint_1140000)\n",
        "2.   For instance, via rpc-calls send and get results. Not included in this colab.\n",
        "\n",
        "**Important !!**\n",
        "\n",
        "**The python predictor_fn needs to be implemented. This code is missing.**\n",
        "**This method calls a prediction serverl.**"
      ],
      "metadata": {
        "id": "gx8Wa9xY8VUb"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Imports and libs"
      ],
      "metadata": {
        "id": "J2-8TD74M5Ck"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import functools\n",
        "import multiprocessing\n",
        "import time\n",
        "\n",
        "import tensorflow as tf\n",
        "\n",
        "import nltk\n",
        "nltk.download('stopwords')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "zQu8X1LnBuYQ",
        "outputId": "cfb910f8-33fe-4ca1-8fc2-2228ffb7eec2"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "[nltk_data] Downloading package stopwords to /root/nltk_data...\n",
            "[nltk_data]   Unzipping corpora/stopwords.zip.\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "True"
            ]
          },
          "metadata": {},
          "execution_count": 1
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!git clone https://github.com/huggingface/transformers.git\n",
        "!pip install ./transformers\n",
        "!pip install sentencepiece"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "uX4rUKIYISeE",
        "outputId": "1c5c1d97-7f95-47c5-afd6-1725bc2621e8"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Cloning into 'transformers'...\n",
            "remote: Enumerating objects: 132356, done.\u001b[K\n",
            "remote: Counting objects: 100% (698/698), done.\u001b[K\n",
            "remote: Compressing objects: 100% (238/238), done.\u001b[K\n",
            "remote: Total 132356 (delta 434), reused 623 (delta 391), pack-reused 131658\u001b[K\n",
            "Receiving objects: 100% (132356/132356), 127.53 MiB | 17.40 MiB/s, done.\n",
            "Resolving deltas: 100% (99808/99808), done.\n",
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Processing ./transformers\n",
            "  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n",
            "  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n",
            "  Preparing metadata (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
            "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.8/dist-packages (from transformers==4.27.0.dev0) (2022.6.2)\n",
            "Requirement already satisfied: tqdm>=4.27 in /usr/local/lib/python3.8/dist-packages (from transformers==4.27.0.dev0) (4.64.1)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.8/dist-packages (from transformers==4.27.0.dev0) (3.9.0)\n",
            "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.8/dist-packages (from transformers==4.27.0.dev0) (23.0)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.8/dist-packages (from transformers==4.27.0.dev0) (2.25.1)\n",
            "Collecting tokenizers!=0.11.3,<0.14,>=0.11.1\n",
            "  Downloading tokenizers-0.13.2-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (7.6 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m7.6/7.6 MB\u001b[0m \u001b[31m49.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.8/dist-packages (from transformers==4.27.0.dev0) (1.22.4)\n",
            "Collecting huggingface-hub<1.0,>=0.11.0\n",
            "  Downloading huggingface_hub-0.12.1-py3-none-any.whl (190 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m190.3/190.3 KB\u001b[0m \u001b[31m19.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.8/dist-packages (from transformers==4.27.0.dev0) (6.0)\n",
            "Requirement already satisfied: typing-extensions>=3.7.4.3 in /usr/local/lib/python3.8/dist-packages (from huggingface-hub<1.0,>=0.11.0->transformers==4.27.0.dev0) (4.5.0)\n",
            "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.8/dist-packages (from requests->transformers==4.27.0.dev0) (1.26.14)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.8/dist-packages (from requests->transformers==4.27.0.dev0) (2022.12.7)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.8/dist-packages (from requests->transformers==4.27.0.dev0) (2.10)\n",
            "Requirement already satisfied: chardet<5,>=3.0.2 in /usr/local/lib/python3.8/dist-packages (from requests->transformers==4.27.0.dev0) (4.0.0)\n",
            "Building wheels for collected packages: transformers\n",
            "  Building wheel for transformers (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for transformers: filename=transformers-4.27.0.dev0-py3-none-any.whl size=6670994 sha256=30d7756c8df193784b72c85f5cc63193122642b04fb04a5d11b0880cdcc3a043\n",
            "  Stored in directory: /tmp/pip-ephem-wheel-cache-zb9ebigj/wheels/15/57/14/0d2873a0295966ca166ea9d9225761a50cce27e4d6b0341fcc\n",
            "Successfully built transformers\n",
            "Installing collected packages: tokenizers, huggingface-hub, transformers\n",
            "Successfully installed huggingface-hub-0.12.1 tokenizers-0.13.2 transformers-4.27.0.dev0\n",
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Collecting sentencepiece\n",
            "  Downloading sentencepiece-0.1.97-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.3 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.3/1.3 MB\u001b[0m \u001b[31m34.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hInstalling collected packages: sentencepiece\n",
            "Successfully installed sentencepiece-0.1.97\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from transformers import T5Tokenizer\n",
        "tokenizer_mt5 =  T5Tokenizer.from_pretrained('google/mt5-xxl')\n",
        "tokenizer_mt5.encode('Hello world')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "vyo1dJooIWIQ",
        "outputId": "8cb309d7-3d63-46bd-89c9-aeeb8c160e9a"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[30273, 4836, 1]"
            ]
          },
          "metadata": {},
          "execution_count": 11
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# NEEDS IMPLEMENTATION depending on the infrastructure\n",
        "def predictor_fn(batches):\n",
        "  \"\"\"The method takes a list of batches and processes each batch.\n",
        "\n",
        "    This could be also implemented as a loop that sends each batch to a server.\n",
        "    The demo uses list of batches but always only a single batch\n",
        "    with one sentence is included for this the demo.\n",
        "\n",
        "    For the efficient processing of a larger document set the ability to process\n",
        "    documents in parallel is kept.\n",
        "\n",
        "  Args:\n",
        "    batches: list[list[str]] list of batches wiht each bach a list of documents.\n",
        "\n",
        "  Returns:\n",
        "    mt5 predictions (typically consisting of output_text and scores).\n",
        "  \"\"\"\n",
        "  raise NotImplementedError('method \"predictor_fn\" needs to be implement depending on the used infrastructure')"
      ],
      "metadata": {
        "id": "37hcziYxB8sU"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cdOYc-2-8PXr"
      },
      "outputs": [],
      "source": [
        "\n",
        "def predict_coreferences(batches, threads_to_use=1):\n",
        "  \"\"\"Predict coreferences of focus part (e.g. one sentence)\"\"\"\n",
        "\n",
        "  with multiprocessing.pool.ThreadPool(threads_to_use) as pool:\n",
        "    results = pool.map(predictor_fn, batches)\n",
        "\n",
        "  return results\n",
        "\n",
        "def extract_result_string(predictions):\n",
        "  \"\"\"Extract the results from prediction.\"\"\"\n",
        "  results = []\n",
        "  for resp in  predictions:\n",
        "    output_text = tf.make_ndarray(resp.outputs['output_text'])\n",
        "    scores = tf.make_ndarray(resp.outputs['scores'])\n",
        "\n",
        "    for text, score in zip(output_text, scores):\n",
        "      text = text[0].decode('utf-8')\n",
        "      results.append(text)\n",
        "  return results"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Extract coreferences"
      ],
      "metadata": {
        "id": "FGYBAbNCD4k1"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# @title helper\n",
        "def normalize_speaker(speaker_in):\n",
        "  \"\"\"Add '_' before and after speaker name if it does not contain it already\"\"\"\n",
        "  if speaker_in == '-' or speaker_in == '__':\n",
        "    return '_'\n",
        "\n",
        "  speaker = speaker_in.replace(' ', '_')\n",
        "  speaker = speaker.strip()\n",
        "\n",
        "  if not speaker.startswith('_'):\n",
        "    speaker = '_'+speaker\n",
        "  if not speaker.endswith('_'):\n",
        "    speaker = speaker+'_'\n",
        "  return speaker\n",
        "\n",
        "\n",
        "def match_mention_state(m, inputs, maps, position=None, debug=False, start_index=0):\n",
        "\n",
        "  if '##' in m:\n",
        "    index_num = m.index('##')\n",
        "  else:\n",
        "    if not m[0].startswith('['):\n",
        "      print('get_chains::error ## not in split', m)\n",
        "    index_num = len(m)\n",
        "\n",
        "  if ']]' in inputs:\n",
        "    end_index = inputs.index(']]')\n",
        "  elif '**' in inputs:\n",
        "    end_index = inputs.index('**')\n",
        "  else:\n",
        "    end_index = len(inputs)\n",
        "\n",
        "  # m_clean = [x for x in m if x != '##']\n",
        "  m_clean = []\n",
        "  for x in m:\n",
        "    if x != '##':\n",
        "      m_clean.append(x)\n",
        "    if x == '**':\n",
        "      break\n",
        "\n",
        "  # get context\n",
        "  context = []\n",
        "  found_num = False\n",
        "  for s in m:\n",
        "    if found_num:\n",
        "      context.append(s)\n",
        "    if '##' == s:\n",
        "      found_num = True\n",
        "\n",
        "  maps_index = 0\n",
        "  indices = []\n",
        "  for i in range(start_index, end_index):\n",
        "    maps_index = i\n",
        "    if inputs[i] == m_clean[0]:\n",
        "      if inputs[i:i+len(m_clean)] == m_clean:\n",
        "        indices.append((maps[maps_index], maps[maps_index + index_num  - 1]))\n",
        "\n",
        "        if maps[maps_index + index_num  - 1] == -1:\n",
        "          print('index negative', maps[maps_index:], ' index_num',  index_num)\n",
        "          print('index negative', inputs[i:], ' index_num',  index_num)\n",
        "          print(f'i {i} maps_index {maps_index}')\n",
        "\n",
        "\n",
        "  if len(indices) == 0:\n",
        "    print('none found match_mention', m)\n",
        "    print('inputs', inputs)\n",
        "    return []\n",
        "  elif len(indices) > 1 and debug:\n",
        "    print('match_mention: too many ', m,  indices, 'm_clean - use both')\n",
        "\n",
        "  if (-1,-1) in indices:\n",
        "    print('error for ',m, indices)\n",
        "    return []\n",
        "\n",
        "  return indices\n",
        "\n",
        "def match_link_state(link, inputs, maps, cluster_name_to_cluster,\n",
        "                     debug=True, node_wise=True):\n",
        "  link_mentions = [m.split(' ') for m in link]\n",
        "  links = []\n",
        "  if len(link_mentions) == 1 and node_wise:\n",
        "    m0 = link_mentions[0]\n",
        "    try:\n",
        "      index_m0 = match_mention_state(m0, inputs, maps, position=None)\n",
        "      links = [index_m0]\n",
        "    except Exception as e:\n",
        "      print(str(e))\n",
        "    return links\n",
        "\n",
        "\n",
        "  m0 = link_mentions[0]\n",
        "  m1 = link_mentions[1]\n",
        "\n",
        "  if debug:\n",
        "    print('match_link', m0, m1)\n",
        "\n",
        "  # invert indices\n",
        "  if m1[0].startswith('[') and len(m1[0]) > 0:\n",
        "    cluster = cluster_name_to_cluster.get(m1[0], None)\n",
        "    if cluster is not None:\n",
        "      index_m1 = [cluster[-1]]\n",
        "    else:\n",
        "      print('cluster does not exists')\n",
        "      return []\n",
        "  else:\n",
        "    index_m1 = match_mention_state(m1, inputs, maps, position=None)\n",
        "\n",
        "\n",
        "  if debug:\n",
        "    print(index_m1 ,'match' ,m1)\n",
        "\n",
        "  if len(index_m1) > 1:\n",
        "    print('index_m1', index_m1)\n",
        "\n",
        "  try:\n",
        "    index_m0 = match_mention_state(m0, inputs, maps, position=None)\n",
        "  except Exception as e:\n",
        "    print('error', str(e))\n",
        "    index_m0 = []\n",
        "\n",
        "  if debug:\n",
        "    print(index_m0 ,'match' , m0)\n",
        "\n",
        "  if len(index_m0) > 1:\n",
        "    print('index_m0', index_m0)\n",
        "\n",
        "  if len(index_m1) > 0 and len(index_m0) > 0:\n",
        "      i1 = index_m1[-1]\n",
        "      i2 = index_m0[-1]\n",
        "      links.append([i1, i2])\n",
        "\n",
        "  # use only last link\n",
        "  if len(links) > 1:\n",
        "    print('too many links, ', links, 'for link', link)\n",
        "    print('context', inputs)\n",
        "\n",
        "    return links[-1:]\n",
        "\n",
        "  return links\n",
        "\n",
        "\n",
        "def get_mentions_for_link_state(link, node_wise):\n",
        "  link_split = link.split('->')\n",
        "\n",
        "  if node_wise and len(link_split) == 1:\n",
        "    m0 = link_split[0].strip()\n",
        "    # print('link has only one mention?', link, m0)\n",
        "    return [m0]\n",
        "\n",
        "  elif len(link_split) < 2:\n",
        "    print('link has only one mention - skipping mention', link)\n",
        "    return []\n",
        "\n",
        "  if len(link_split) > 2:\n",
        "    print('link has too many mentions - using first two.', link)\n",
        "  m0 = link_split[0].strip()\n",
        "  m1 = link_split[1].strip()\n",
        "  return [m0, m1]"
      ],
      "metadata": {
        "cellView": "form",
        "id": "AaUYfXUBCvTF"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# use mt5 and large context\n",
        "\n",
        "class State(object):\n",
        "  \"\"\"Document state.\"\"\"\n",
        "\n",
        "  def __init__(self, input_document, node_wise=True, max_len_doc=3000):\n",
        "    \"\"\" Create State object to process documents.\n",
        "\n",
        "    Args:\n",
        "      input_document: dictonary with the input document.\n",
        "      node_wise: Predict mentions too.\n",
        "      max_len_doc: max sentence pieace tokens, eg. 2000 or 3000 (bit better).\n",
        "\n",
        "    \"\"\"\n",
        "    self.sentence_num = -1\n",
        "    self.clusters_num = 0\n",
        "\n",
        "    self.token_map_context, self.annotation_context = [], []\n",
        "    self.annotation_coreference_start, self.annotation_coreference_end = [], []\n",
        "    self.token_map, self.annotation = [], []\n",
        "\n",
        "    # a mention index to cluster mapping, e.g. (23, 24) -> [(23, 24), (41, 42)]\n",
        "    self.mention_index_to_cluster = {}\n",
        "\n",
        "    # the first link names the cluster, e.g. (23, 24) -> '1'\n",
        "    self.mention_index_to_cluster_name = {}\n",
        "    self.cluster_name_to_cluster = {}\n",
        "\n",
        "    self.input_document = input_document\n",
        "    # print('sentence_num', self.sentence_num)\n",
        "    self.genre = input_document['genres'][0][0]\n",
        "    self.speakers = {t: spk for (t, spk) in self.input_document['speakers']}\n",
        "\n",
        "    self.done = False\n",
        "    self.predictions_str = {}  # keep the predictions\n",
        "    self.node_wise = node_wise\n",
        "\n",
        "    self.max_len_doc = max_len_doc\n",
        "\n",
        "    # move to initial position.\n",
        "    self.extend()\n",
        "\n",
        "\n",
        "  def extend_done(self):\n",
        "    return self.done\n",
        "\n",
        "  def extend(self, prediction_str=None, use_gold_cluster=False, move=True):\n",
        "\n",
        "    # move annotation to context\n",
        "    self.token_map_context +=  self.token_map\n",
        "    self.annotation_context += self.annotation\n",
        "\n",
        "    for k in range(len(self.annotation)):\n",
        "      self.annotation_coreference_start.append([])\n",
        "      self.annotation_coreference_end.append([])\n",
        "\n",
        "    assert len(self.annotation_context)  == len(self.annotation_coreference_start)\n",
        "\n",
        "    self.annotation, self.token_map = [], []\n",
        "\n",
        "    link_found = False\n",
        "    if prediction_str is not None and not 'None [' in prediction_str:\n",
        "      links = [l for l in prediction_str.split(';;') if l != '' ]\n",
        "\n",
        "      annotation_update = []\n",
        "      for link in links:\n",
        "        link_found = True\n",
        "        link_mentions = get_mentions_for_link_state(link, self.node_wise)\n",
        "\n",
        "        if len(link_mentions) < 2 and not (self.node_wise and len(link_mentions)):\n",
        "          print('less mentions as needed skip', link_mentions)\n",
        "          continue\n",
        "        indices = match_link_state(link_mentions, self.annotation_full,\n",
        "                                   self.annotation_full_map,\n",
        "                                   self.cluster_name_to_cluster,\n",
        "                                   debug=False)\n",
        "\n",
        "        if not indices:\n",
        "          print('not found !!')\n",
        "          print('indices not found', link, indices)\n",
        "          print('self.annotation_full', self.annotation_full )\n",
        "          print('annotation + context', self.get_input_annotation())\n",
        "          continue\n",
        "\n",
        "        if True:\n",
        "          index = indices[0]\n",
        "          cluster = []\n",
        "          for mention_index in index:\n",
        "            if str(mention_index) in self.mention_index_to_cluster:\n",
        "              cluster = self.mention_index_to_cluster[str(mention_index)]\n",
        "              break\n",
        "\n",
        "\n",
        "          if not cluster:\n",
        "            self.clusters_num += 1\n",
        "            cluster_name = str(self.clusters_num)\n",
        "\n",
        "            if use_gold_cluster:  # just to evaluate on gold\n",
        "\n",
        "              for ni, cx in enumerate(self.input_document['clusters']):\n",
        "                for mx in cx:\n",
        "                  if mx in index:\n",
        "                    cluster_name = str(ni+1)\n",
        "                    break\n",
        "\n",
        "          else:\n",
        "            cluster_name = self.mention_index_to_cluster_name[str(cluster[0])]\n",
        "\n",
        "          for mention_index in index:\n",
        "            if mention_index not in cluster:\n",
        "              cluster.append(mention_index)\n",
        "              self.mention_index_to_cluster[str(mention_index)] = cluster\n",
        "              self.cluster_name_to_cluster['['+cluster_name] = cluster\n",
        "              self.mention_index_to_cluster_name[str(mention_index)] = cluster_name\n",
        "              annotation_update.append([mention_index, cluster_name])\n",
        "\n",
        "      # update the annotation\n",
        "      if True:\n",
        "        for update in annotation_update:\n",
        "          update_index = update[0]\n",
        "          update_name = update[1]\n",
        "\n",
        "          for t, coref_starts, coref_end, tid in zip(self.annotation_context,\n",
        "                                    self.annotation_coreference_start,\n",
        "                                    self.annotation_coreference_end,\n",
        "                                    self.token_map_context):\n",
        "\n",
        "\n",
        "\n",
        "            if update_index[0] == tid:\n",
        "              coref_starts.append(update)\n",
        "              coref_starts.sort( key=lambda x: x[0][1], reverse=True)\n",
        "\n",
        "\n",
        "            if update_index[1] == tid:\n",
        "              coref_end.append(']')\n",
        "\n",
        "    if move or 'None [' in prediction_str or not link_found:\n",
        "      self.sentence_num += 1\n",
        "\n",
        "      if self.sentence_num not in self.input_document['sentences']:\n",
        "        self.done = True\n",
        "        return True\n",
        "\n",
        "      tokens = self.input_document['sentences'][self.sentence_num]\n",
        "      token_map = self.input_document['token_maps'][self.sentence_num]\n",
        "      first = True\n",
        "\n",
        "\n",
        "      for tid, t in zip(token_map, tokens):\n",
        "        if first:\n",
        "          self.token_map.append(-1)\n",
        "          speaker = normalize_speaker(self.speakers[tid])\n",
        "          self.annotation.append(speaker)\n",
        "          first = False\n",
        "        self.token_map.append(tid)\n",
        "        self.annotation.append(t)\n",
        "\n",
        "    if self.sentence_num not in self.predictions_str:\n",
        "      self.predictions_str[self.sentence_num] = ''\n",
        "\n",
        "    if prediction_str is not None:\n",
        "      self.predictions_str[self.sentence_num] += prediction_str\n",
        "\n",
        "    return False\n",
        "\n",
        "  def input_annotation(self):\n",
        "\n",
        "    self.annotation_full = ['coref:', self.genre]\n",
        "    self.annotation_full_map = [-1, -1]\n",
        "    for t, coref_starts, coref_end, tid in zip(self.annotation_context,\n",
        "                                  self.annotation_coreference_start,\n",
        "                                  self.annotation_coreference_end,\n",
        "                                  self.token_map_context):\n",
        "\n",
        "      for coref_start in coref_starts:\n",
        "        coref_name = coref_start[-1]\n",
        "\n",
        "        self.annotation_full.append('[' + coref_name)\n",
        "        self.annotation_full_map.append(-1)\n",
        "\n",
        "      self.annotation_full.append(t)\n",
        "      self.annotation_full_map.append(tid)\n",
        "\n",
        "      for end in coref_end:\n",
        "        coref_name = end[-1]\n",
        "        self.annotation_full.append(coref_name)\n",
        "        self.annotation_full_map.append(-1)\n",
        "\n",
        "    self.annotation_full += ['|'] + self.annotation\n",
        "    self.annotation_full_map += [-1] + self.token_map\n",
        "    self.annotation_full += ['**']\n",
        "    self.annotation_full_map += [-1]\n",
        "\n",
        "\n",
        "  def encode(self, annotation_str):\n",
        "    return tokenizer_mt5.encode(annotation_str)\n",
        "\n",
        "  def get_input_annotation(self, context_right=True):\n",
        "\n",
        "    self.input_annotation()\n",
        "    annotation_str = ' '.join(self.annotation_full)\n",
        "\n",
        "    enc = self.encode(annotation_str)\n",
        "    shorten = len(enc) > self.max_len_doc\n",
        "\n",
        "    while len(enc) > self.max_len_doc:   # inefficient ...\n",
        "      self.annotation_context = self.annotation_context[1:]\n",
        "      self.token_map_context = self.token_map_context[1:]\n",
        "      self.annotation_coreference_start = self.annotation_coreference_start[1:]\n",
        "      self.annotation_coreference_end = self.annotation_coreference_end[1:]\n",
        "\n",
        "      self.input_annotation()\n",
        "      annotation_str = ' '.join(self.annotation_full)\n",
        "      enc = self.encode(annotation_str)\n",
        "\n",
        "    last_token_id = self.annotation_full_map[-2]  # the last one is **\n",
        "    self.annotation_context_right = []\n",
        "\n",
        "    if not shorten and context_right:\n",
        "      sentence_num = self.sentence_num\n",
        "      total_len = len(enc)\n",
        "\n",
        "      while True:\n",
        "        sentence_num += 1\n",
        "        if sentence_num not in self.input_document['sentences']:\n",
        "          break\n",
        "\n",
        "        first = True\n",
        "        annotation_context_next = []\n",
        "\n",
        "        for t, tid in zip(self.input_document['sentences'][sentence_num], self.input_document['token_maps'][sentence_num]):\n",
        "          if first:\n",
        "            speaker = normalize_speaker(self.speakers[tid])\n",
        "            annotation_context_next.append(speaker)\n",
        "            first = False\n",
        "          annotation_context_next.append(t)\n",
        "\n",
        "        annotation_context_right = self.annotation_context_right + annotation_context_next\n",
        "        enc = self.encode(' '.join(annotation_context_right))\n",
        "\n",
        "        if (len(enc) + total_len) > self.max_len_doc:\n",
        "          break\n",
        "        self.annotation_context_right = annotation_context_right\n",
        "      if self.annotation_context_right:\n",
        "        annotation_str = annotation_str + ' ' + ' '.join(self.annotation_context_right)\n",
        "\n",
        "    enc = self.encode(annotation_str)\n",
        "    if len(enc) > self.max_len_doc:\n",
        "      print('warning: document too long', len(enc))\n",
        "\n",
        "    return annotation_str\n"
      ],
      "metadata": {
        "id": "_zSQBdAXELkR"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title function def to reate a input state\n",
        "tokenizer_nltk = nltk.WordPunctTokenizer()\n",
        "\n",
        "\n",
        "def create_document(document: str, title: str = 'not_named'):\n",
        "  \"\"\"Creates a datastructure with a title and uses nltk for tokenization.\n",
        "\n",
        "  Args:\n",
        "    document: sentences separated with newline ('\\n').\n",
        "    title: the name of the document.\n",
        "\n",
        "  Returns:\n",
        "    dict with sentences, maps to token-ids, speakers, and genres.\n",
        "  \"\"\"\n",
        "  input_document = {\n",
        "      'doc_key': title,\n",
        "      'sentences': {},\n",
        "      'token_maps': {},\n",
        "      'speakers': [],\n",
        "      'genres': []\n",
        "  }\n",
        "\n",
        "  tid = 0\n",
        "  for k, sentence in enumerate(document.split('\\n')):\n",
        "    input_document['sentences'][k] = tokenizer_nltk.tokenize(text=sentence)\n",
        "    input_document['token_maps'][k] = []\n",
        "\n",
        "    for _ in input_document['sentences'][k]:\n",
        "      input_document['token_maps'][k].append(tid)\n",
        "      input_document['speakers'].append((tid, '_'))\n",
        "      input_document['genres'].append('wi')\n",
        "      tid += 1\n",
        "  return input_document"
      ],
      "metadata": {
        "cellView": "form",
        "id": "cDBDlw0eOaj3"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title function def to create batches\n",
        "def create_next_batch(states_dict, batche_size=1, num_batches=1):\n",
        "  batches = [[]]\n",
        "  states = []\n",
        "  for key, state in states_dict.items():\n",
        "    if state.extend_done():\n",
        "      continue\n",
        "\n",
        "    states.append(state)\n",
        "    if len(states) >= (batche_size * num_batches):\n",
        "      break\n",
        "  for state in states:\n",
        "    batches[-1].append(state.get_input_annotation())\n",
        "    if len(batches[-1]) >= batche_size:\n",
        "      if len(batches) >= num_batches:\n",
        "        break\n",
        "      batches.append([])\n",
        "  return states, batches\n"
      ],
      "metadata": {
        "id": "cg836m1dT0qF"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title sample document - from wikipedia\n",
        "d1_title = \"Eiffel Tower Wiki\"\n",
        "d1 = \"\"\"The Eiffel Tower (French: tour Eiffel) is a wrought-iron lattice tower on the Champ de Mars in Paris, France. It is named after the engineer Gustave Eiffel, whose company designed and built the tower.\n",
        "Locally nicknamed \"La dame de fer\" (French for \"Iron Lady\"), it was constructed from 1887 to 1889 as the centerpiece of the 1889 World's Fair.\n",
        "Although initially criticised by some of France's leading artists and intellectuals for its design, it has since become a global cultural icon of France and one of the most recognisable structures in the world.\n",
        "The Eiffel Tower is the most visited monument with an entrance fee in the world: 6.91 million people ascended it in 2015.\n",
        "It was designated a monument historique in 1964, and was named part of a UNESCO World Heritage Site (\"Paris, Banks of the Seine\") in 1991.\"\"\"\n",
        "\n",
        "input_document = create_document(d1, d1_title)\n",
        "states_dict = {d1_title: State(input_document)}"
      ],
      "metadata": {
        "id": "_QZw0SI1ONa9"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "emulate_predictions = True # @param\n",
        "\n",
        "if emulate_predictions:\n",
        "  predictioned_results = [\n",
        "      [\n",
        "          'It ## is named after -> The Eiffel Tower ( French : tour Eiffel ) ##'\n",
        "          ' is a wrought ;; the tower ## . ** _ -> It ## is named after ;;'\n",
        "      ],\n",
        "      ['it ## was constructed from -> [1 ;;'],\n",
        "      [\n",
        "          'its ## design , it -> [1 ;; it ## has since become -> its ## design'\n",
        "          \" , it ;; France ' s ## leading artists and -> France ## . [1 It ;;\"\n",
        "          \" France ## and one of -> France ' s ## leading artists and ;;\"\n",
        "      ],\n",
        "      [\n",
        "          'The Eiffel Tower ## is the most -> [1 ;; it ## in 2015 . -> The'\n",
        "          ' Eiffel Tower ## is the most ;; the world ## : 6 . -> the world ## .'\n",
        "          ' | _ ;;'\n",
        "      ],\n",
        "      [\n",
        "          'It ## was designated a -> [1 ;; Paris ## , Banks of -> Paris , [2'\n",
        "          ' France ## ] . [1 ;;'\n",
        "      ],\n",
        "  ]\n",
        "else:\n",
        "  predictioned_results = []"
      ],
      "metadata": {
        "id": "Swb2TFG3oKgE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "expand_only = False\n",
        "total_time = time.time()\n",
        "total_results = 0\n",
        "\n",
        "debug = True\n",
        "\n",
        "for step in range(100):  # while states\n",
        "  t = time.time()\n",
        "  states, batches = create_next_batch(states_dict)\n",
        "\n",
        "  if not states:\n",
        "    break\n",
        "\n",
        "  documents_processing = set([x.input_document['doc_key'] for x in states])\n",
        "\n",
        "  print(f'Processing documents: {documents_processing}')\n",
        "\n",
        "  if predictioned_results:\n",
        "    results = predictioned_results[step]\n",
        "  else:\n",
        "    predictions = predict_coreferences(batches, len(batches))\n",
        "    results = extract_result_string(predictions)\n",
        "\n",
        "  for state, result, batch in zip(states, results, batches):\n",
        "    state.extend(result)\n",
        "\n",
        "    if debug:\n",
        "      print('input batch[0]: ', batch)\n",
        "      print('mt5 output:     ', results)\n",
        "\n",
        "  total_results += len(results)\n",
        "  print(\n",
        "      f'time { time.time()-t}, round time/seq : {(time.time()-t)/len(results)}'\n",
        "      f' total time/seq: {(time.time()-total_time)/total_results}'\n",
        "  )\n",
        "  print()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "DTFYSNjdoDQl",
        "outputId": "967ace0d-2411-4c05-b248-a7066f30f5d6"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Processing documents: {'Eiffel Tower Wiki'}\n",
            "input batch[0]:  ['coref: w | _ The Eiffel Tower ( French : tour Eiffel ) is a wrought - iron lattice tower on the Champ de Mars in Paris , France . It is named after the engineer Gustave Eiffel , whose company designed and built the tower . ** _ Locally nicknamed \" La dame de fer \" ( French for \" Iron Lady \"), it was constructed from 1887 to 1889 as the centerpiece of the 1889 World \\' s Fair . _ Although initially criticised by some of France \\' s leading artists and intellectuals for its design , it has since become a global cultural icon of France and one of the most recognisable structures in the world . _ The Eiffel Tower is the most visited monument with an entrance fee in the world : 6 . 91 million people ascended it in 2015 . _ It was designated a monument historique in 1964 , and was named part of a UNESCO World Heritage Site (\" Paris , Banks of the Seine \") in 1991 .']\n",
            "mt5 output:      ['It ## is named after -> The Eiffel Tower ( French : tour Eiffel ) ## is a wrought ;; the tower ## . ** _ -> It ## is named after ;;']\n",
            "time 0.010899782180786133, round time/seq : 0.010912418365478516 total time/seq: 0.011353254318237305\n",
            "\n",
            "Processing documents: {'Eiffel Tower Wiki'}\n",
            "input batch[0]:  ['coref: w _ [1 The Eiffel Tower ( French : tour Eiffel ) ] is a wrought - iron lattice tower on the Champ de Mars in Paris , France . [1 It ] is named after the engineer Gustave Eiffel , whose company designed and built [1 the tower ] . | _ Locally nicknamed \" La dame de fer \" ( French for \" Iron Lady \"), it was constructed from 1887 to 1889 as the centerpiece of the 1889 World \\' s Fair . ** _ Although initially criticised by some of France \\' s leading artists and intellectuals for its design , it has since become a global cultural icon of France and one of the most recognisable structures in the world . _ The Eiffel Tower is the most visited monument with an entrance fee in the world : 6 . 91 million people ascended it in 2015 . _ It was designated a monument historique in 1964 , and was named part of a UNESCO World Heritage Site (\" Paris , Banks of the Seine \") in 1991 .']\n",
            "mt5 output:      ['it ## was constructed from -> [1 ;;']\n",
            "time 0.008285999298095703, round time/seq : 0.008303165435791016 total time/seq: 0.009858012199401855\n",
            "\n",
            "Processing documents: {'Eiffel Tower Wiki'}\n",
            "input batch[0]:  ['coref: w _ [1 The Eiffel Tower ( French : tour Eiffel ) ] is a wrought - iron lattice tower on the Champ de Mars in Paris , France . [1 It ] is named after the engineer Gustave Eiffel , whose company designed and built [1 the tower ] . _ Locally nicknamed \" La dame de fer \" ( French for \" Iron Lady \"), [1 it ] was constructed from 1887 to 1889 as the centerpiece of the 1889 World \\' s Fair . | _ Although initially criticised by some of France \\' s leading artists and intellectuals for its design , it has since become a global cultural icon of France and one of the most recognisable structures in the world . ** _ The Eiffel Tower is the most visited monument with an entrance fee in the world : 6 . 91 million people ascended it in 2015 . _ It was designated a monument historique in 1964 , and was named part of a UNESCO World Heritage Site (\" Paris , Banks of the Seine \") in 1991 .']\n",
            "mt5 output:      [\"its ## design , it -> [1 ;; it ## has since become -> its ## design , it ;; France ' s ## leading artists and -> France ## . [1 It ;; France ## and one of -> France ' s ## leading artists and ;;\"]\n",
            "time 0.008397102355957031, round time/seq : 0.008409738540649414 total time/seq: 0.00939615567525228\n",
            "\n",
            "Processing documents: {'Eiffel Tower Wiki'}\n",
            "input batch[0]:  ['coref: w _ [1 The Eiffel Tower ( French : tour Eiffel ) ] is a wrought - iron lattice tower on the Champ de Mars in Paris , [2 France ] . [1 It ] is named after the engineer Gustave Eiffel , whose company designed and built [1 the tower ] . _ Locally nicknamed \" La dame de fer \" ( French for \" Iron Lady \"), [1 it ] was constructed from 1887 to 1889 as the centerpiece of the 1889 World \\' s Fair . _ Although initially criticised by some of [2 France \\' s ] leading artists and intellectuals for [1 its ] design , [1 it ] has since become a global cultural icon of [2 France ] and one of the most recognisable structures in the world . | _ The Eiffel Tower is the most visited monument with an entrance fee in the world : 6 . 91 million people ascended it in 2015 . ** _ It was designated a monument historique in 1964 , and was named part of a UNESCO World Heritage Site (\" Paris , Banks of the Seine \") in 1991 .']\n",
            "mt5 output:      ['The Eiffel Tower ## is the most -> [1 ;; it ## in 2015 . -> The Eiffel Tower ## is the most ;; the world ## : 6 . -> the world ## . | _ ;;']\n",
            "time 0.0051686763763427734, round time/seq : 0.0051805973052978516 total time/seq: 0.008359313011169434\n",
            "\n",
            "Processing documents: {'Eiffel Tower Wiki'}\n",
            "input batch[0]:  ['coref: w _ [1 The Eiffel Tower ( French : tour Eiffel ) ] is a wrought - iron lattice tower on the Champ de Mars in Paris , [2 France ] . [1 It ] is named after the engineer Gustave Eiffel , whose company designed and built [1 the tower ] . _ Locally nicknamed \" La dame de fer \" ( French for \" Iron Lady \"), [1 it ] was constructed from 1887 to 1889 as the centerpiece of the 1889 World \\' s Fair . _ Although initially criticised by some of [2 France \\' s ] leading artists and intellectuals for [1 its ] design , [1 it ] has since become a global cultural icon of [2 France ] and one of the most recognisable structures in [3 the world ] . _ [1 The Eiffel Tower ] is the most visited monument with an entrance fee in [3 the world ] : 6 . 91 million people ascended [1 it ] in 2015 . | _ It was designated a monument historique in 1964 , and was named part of a UNESCO World Heritage Site (\" Paris , Banks of the Seine \") in 1991 . **']\n",
            "mt5 output:      ['It ## was designated a -> [1 ;; Paris ## , Banks of -> Paris , [2 France ## ] . [1 ;;']\n",
            "time 0.00478363037109375, round time/seq : 0.004794597625732422 total time/seq: 0.007658481597900391\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# @title get and print the output as annotated document\n",
        "\n",
        "for doc_name, s in states_dict.items():\n",
        "  pred_clusters = [cluster for name, cluster in s.cluster_name_to_cluster.items()]\n",
        "  print('predicted clusters with word indexes', pred_clusters)\n",
        "\n",
        "  text, text_map = [], []\n",
        "  for k, snt in states_dict[doc_name].input_document['sentences'].items():\n",
        "    m = states_dict[doc_name].input_document['token_maps'][k]\n",
        "    text += snt\n",
        "    text_map += m\n",
        "\n",
        "  cluster_annotations_start = []\n",
        "  cluster_annotations_end = []\n",
        "\n",
        "  # Cluster annotation per token\n",
        "  for tid in text_map:\n",
        "    cluster_annotations_start.append([])\n",
        "    cluster_annotations_end.append([])\n",
        "    for ci in pred_clusters:\n",
        "      for m in ci:\n",
        "\n",
        "        if tid == m[0]:\n",
        "          m_len = m[1] - m[0]\n",
        "          name = s.mention_index_to_cluster_name[str(m)]\n",
        "          cluster_annotations_start[-1].append((name, m_len))\n",
        "\n",
        "        if tid == m[1]:\n",
        "          cluster_annotations_end[-1].append(']')\n",
        "\n",
        "  # get the text with the coreference annotations\n",
        "  all_text = []\n",
        "  for tok, start, end in zip(text, cluster_annotations_start, cluster_annotations_end):\n",
        "\n",
        "    if start:\n",
        "      for x in sorted(start, key=lambda x : x[1], reverse=True):\n",
        "        all_text.append('['+str(x[0]))\n",
        "\n",
        "    all_text.append(tok)\n",
        "\n",
        "    if end:\n",
        "      all_text.append(''.join(end))\n",
        "\n",
        "  print(' '.join(all_text))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "cellView": "form",
        "id": "vhWCsybuSBqP",
        "outputId": "95f45cc5-7d01-4113-8ef4-3b00fd4aa89f"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "predicted clusters with word indexes [[(0, 8), (26, 26), (40, 41), (58, 58), (90, 90), (93, 93), (114, 116), (136, 136), (140, 140)], [(24, 24), (82, 84), (102, 102)], [(111, 112), (127, 128)], [(22, 24), (160, 160)]]\n",
            "[1 The Eiffel Tower ( French : tour Eiffel ) ] is a wrought - iron lattice tower on the Champ de Mars in [4 Paris , [2 France ]] . [1 It ] is named after the engineer Gustave Eiffel , whose company designed and built [1 the tower ] . Locally nicknamed \" La dame de fer \" ( French for \" Iron Lady \"), [1 it ] was constructed from 1887 to 1889 as the centerpiece of the 1889 World ' s Fair . Although initially criticised by some of [2 France ' s ] leading artists and intellectuals for [1 its ] design , [1 it ] has since become a global cultural icon of [2 France ] and one of the most recognisable structures in [3 the world ] . [1 The Eiffel Tower ] is the most visited monument with an entrance fee in [3 the world ] : 6 . 91 million people ascended [1 it ] in 2015 . [1 It ] was designated a monument historique in 1964 , and was named part of a UNESCO World Heritage Site (\" [4 Paris ] , Banks of the Seine \") in 1991 .\n"
          ]
        }
      ]
    }
  ]
}