quickstart_demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 🚞 Zero-shot RE Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # if you're running this in a colab notebook, you can run this cell to install the necessary dependencies\n",
    "# pip install glirel\n",
    "# !python -m spacy download en_core_web_sm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "config.json not found in /home/jackboylan/tmp/GLiREL/logs/redocred/redocred-2024-09-17__17-53-16/model_42900\n",
      "/home/jackboylan/miniconda3/envs/glirel/lib/python3.10/site-packages/transformers/convert_slow_tokenizer.py:562: UserWarning: The sentencepiece tokenizer that you are converting to a fast tokenizer uses the byte fallback option which is not implemented in the fast tokenizers. In practice this means that the fast version of the tokenizer can produce unknown tokens whereas the sentencepiece version would have converted these unknown tokens into a sequence of byte tokens matching the original piece of text.\n",
      "  warnings.warn(\n",
      "/home/jackboylan/tmp/GLiREL/glirel/model.py:580: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
      "  state_dict = torch.load(model_file, map_location=torch.device(map_location))\n"
     ]
    }
   ],
   "source": [
    "from glirel import GLiREL\n",
    "\n",
    "save_path = 'logs/redocred/redocred-2024-09-17__17-53-16/model_42900'\n",
    "model = GLiREL.from_pretrained(save_path, map_location='cpu')\n",
    "# model = GLiREL.from_pretrained('jackboyla/glirel_beta')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Inference\n",
    "\n",
    "To infer, the model needs `tokens`, `NER`, and `zero shot labels`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Eval data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['Derren', 'Nesbitt', 'had', 'a', 'history', 'of', 'being', 'cast', 'in', '\"', 'Doctor', 'Who', '\"', ',', 'having', 'played', 'villainous', 'warlord', 'Tegana', 'in', 'the', '1964', 'First', 'Doctor', 'serial', '\"', 'Marco', 'Polo', '\"', '.']\n",
      "\n",
      "[[26, 27, 'Q2989881', 'Marco Polo'], [22, 23, 'Q2989412', 'First Doctor']]\n",
      "['characters']\n"
     ]
    }
   ],
   "source": [
    "import json\n",
    "with open('./data/few_rel_all.jsonl', 'r') as f:\n",
    "    data = [json.loads(line) for line in f]\n",
    "\n",
    "i = 0\n",
    "\n",
    "tokens = data[i]['tokenized_text']\n",
    "ner = data[i]['ner']\n",
    "labels = list(set([r['relation_text'] for r in data[i]['relations']]))\n",
    "print(tokens)\n",
    "print()\n",
    "print(ner)\n",
    "print(labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['country of origin', 'licensed to broadcast to', 'father', 'followed by', 'characters']\n"
     ]
    }
   ],
   "source": [
    "labels = ['country of origin', 'licensed to broadcast to', 'father', 'followed by'] + labels\n",
    "print(labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of relations: 10\n",
      "\n",
      "Descending Order by Score:\n",
      "{'head_pos': [26, 28], 'tail_pos': [22, 24], 'head_text': ['Marco', 'Polo'], 'tail_text': ['First', 'Doctor'], 'label': 'followed by', 'score': 0.011230885982513428}\n",
      "{'head_pos': [22, 24], 'tail_pos': [26, 28], 'head_text': ['First', 'Doctor'], 'tail_text': ['Marco', 'Polo'], 'label': 'characters', 'score': 0.010938132181763649}\n",
      "{'head_pos': [22, 24], 'tail_pos': [26, 28], 'head_text': ['First', 'Doctor'], 'tail_text': ['Marco', 'Polo'], 'label': 'followed by', 'score': 0.010783118195831776}\n",
      "{'head_pos': [26, 28], 'tail_pos': [22, 24], 'head_text': ['Marco', 'Polo'], 'tail_text': ['First', 'Doctor'], 'label': 'characters', 'score': 0.01027392502874136}\n",
      "{'head_pos': [26, 28], 'tail_pos': [22, 24], 'head_text': ['Marco', 'Polo'], 'tail_text': ['First', 'Doctor'], 'label': 'father', 'score': 0.0004802523762919009}\n",
      "{'head_pos': [22, 24], 'tail_pos': [26, 28], 'head_text': ['First', 'Doctor'], 'tail_text': ['Marco', 'Polo'], 'label': 'father', 'score': 0.00039520618156529963}\n",
      "{'head_pos': [26, 28], 'tail_pos': [22, 24], 'head_text': ['Marco', 'Polo'], 'tail_text': ['First', 'Doctor'], 'label': 'licensed to broadcast to', 'score': 0.0002560171706136316}\n",
      "{'head_pos': [22, 24], 'tail_pos': [26, 28], 'head_text': ['First', 'Doctor'], 'tail_text': ['Marco', 'Polo'], 'label': 'licensed to broadcast to', 'score': 0.00021301423839759082}\n",
      "{'head_pos': [26, 28], 'tail_pos': [22, 24], 'head_text': ['Marco', 'Polo'], 'tail_text': ['First', 'Doctor'], 'label': 'country of origin', 'score': 0.00020898687944281846}\n",
      "{'head_pos': [22, 24], 'tail_pos': [26, 28], 'head_text': ['First', 'Doctor'], 'tail_text': ['Marco', 'Polo'], 'label': 'country of origin', 'score': 0.0001598332601133734}\n"
     ]
    }
   ],
   "source": [
    "relations = model.predict_relations(tokens, labels, threshold=0.0, ner=ner)\n",
    "\n",
    "print('Number of relations:', len(relations))  # num entity pairs (both directions) * num classes.... provided they're over the threshold\n",
    "\n",
    "sorted_data_desc = sorted(relations, key=lambda x: x['score'], reverse=True)\n",
    "print(\"\\nDescending Order by Score:\")\n",
    "for item in sorted_data_desc:\n",
    "    print(item)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Real-world example\n",
    "\n",
    "Constrain the entity types that can associated with a relationship.\n",
    "e.g:\n",
    "\n",
    "`co-founder` can only have a head `PERSON` entity and a tail `ORG` entity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Text: Apple Inc. was founded by Steve Jobs, Steve Wozniak, and Ronald Wayne in April 1976. The company is headquartered in Cupertino, California.\n",
      "Entities detected: [[0, 1, 'ORG', 'Apple Inc.'], [5, 6, 'PERSON', 'Steve Jobs'], [8, 9, 'PERSON', 'Steve Wozniak'], [12, 13, 'PERSON', 'Ronald Wayne'], [15, 16, 'DATE', 'April 1976'], [23, 23, 'GPE', 'Cupertino'], [25, 25, 'GPE', 'California']]\n",
      "Constraining relations by entity type\n",
      "Number of relations: 6\n",
      "\n",
      "Descending Order by Score:\n",
      "['Apple', 'Inc.'] --> headquartered in --> ['Cupertino'] | score: 0.9072741866111755\n",
      "['Apple', 'Inc.'] --> headquartered in --> ['California'] | score: 0.8888104557991028\n",
      "['Apple', 'Inc.'] --> founded on date --> ['April', '1976'] | score: 0.8402661681175232\n",
      "['Steve', 'Jobs'] --> founder --> ['Apple', 'Inc.'] | score: 0.8150324821472168\n",
      "['Steve', 'Wozniak'] --> founder --> ['Apple', 'Inc.'] | score: 0.8128281831741333\n",
      "['Ronald', 'Wayne'] --> founder --> ['Apple', 'Inc.'] | score: 0.7810325026512146\n"
     ]
    }
   ],
   "source": [
    "# Real-world example\n",
    "import spacy\n",
    "from glirel.modules.utils import constrain_relations_by_entity_type\n",
    "\n",
    "nlp = spacy.load('en_core_web_sm')\n",
    "\n",
    "\n",
    "text = \"Apple Inc. was founded by Steve Jobs, Steve Wozniak, and Ronald Wayne in April 1976. The company is headquartered in Cupertino, California.\"\n",
    "\n",
    "# text = \"Jack Dorsey's father, Tim Dorsey, is a licensed pilot. Jack met his wife Sarah Paulson in New York in 2003. They have one son, Edward.\"\n",
    "\n",
    "labels = {\"glirel_labels\": {\n",
    "    'co-founder': {\"allowed_head\": [\"PERSON\"], \"allowed_tail\": [\"ORG\"]}, \n",
    "    'country of origin': {\"allowed_head\": [\"PERSON\", \"ORG\"], \"allowed_tail\": [\"LOC\", \"GPE\"]}, \n",
    "    'licensed to broadcast to': {\"allowed_head\": [\"ORG\"]},  \n",
    "    'no relation': {},  \n",
    "    'parent': {\"allowed_head\": [\"PERSON\"], \"allowed_tail\": [\"PERSON\"]}, \n",
    "    'followed by': {\"allowed_head\": [\"PERSON\", \"ORG\"], \"allowed_tail\": [\"PERSON\", \"ORG\"]},  \n",
    "    'located in or next to body of water': {\"allowed_head\": [\"LOC\", \"GPE\", \"FAC\"], \"allowed_tail\": [\"LOC\", \"GPE\"]},  \n",
    "    'spouse': {\"allowed_head\": [\"PERSON\"], \"allowed_tail\": [\"PERSON\"]},  \n",
    "    'child': {\"allowed_head\": [\"PERSON\"], \"allowed_tail\": [\"PERSON\"]},  \n",
    "    'founder': {\"allowed_head\": [\"PERSON\"], \"allowed_tail\": [\"ORG\"]},  \n",
    "    'founded on date': {\"allowed_head\": [\"ORG\"], \"allowed_tail\": [\"DATE\"]},\n",
    "    'headquartered in': {\"allowed_head\": [\"ORG\"], \"allowed_tail\": [\"LOC\", \"GPE\", \"FAC\"]},  \n",
    "    'acquired by': {\"allowed_head\": [\"ORG\"], \"allowed_tail\": [\"ORG\", \"PERSON\"]},  \n",
    "    'subsidiary of': {\"allowed_head\": [\"ORG\"], \"allowed_tail\": [\"ORG\", \"PERSON\"]}, \n",
    "    }\n",
    "}\n",
    "\n",
    "\n",
    "def predict_and_show(text, labels):\n",
    "    doc = nlp(text)\n",
    "    print(f\"Text: {text}\")\n",
    "\n",
    "    tokens = [token.text for token in doc]\n",
    "\n",
    "    # NOTE: the end index should be inclusive\n",
    "    ner = [[ent.start, (ent.end - 1), ent.label_, ent.text] for ent in doc.ents]\n",
    "    print(f\"Entities detected: {ner}\")\n",
    "\n",
    "    labels_and_constraints = None\n",
    "    if isinstance(labels, dict):\n",
    "        labels = labels[\"glirel_labels\"]\n",
    "        labels_and_constraints = labels\n",
    "        labels = list(labels.keys())\n",
    "\n",
    "    relations = model.predict_relations(tokens, labels, threshold=0.0, ner=ner, top_k=1)\n",
    "\n",
    "    if isinstance(labels_and_constraints, dict):\n",
    "        print('Constraining relations by entity type')\n",
    "        relations = constrain_relations_by_entity_type(doc.ents, labels_and_constraints, relations)\n",
    "\n",
    "    print('Number of relations:', len(relations))\n",
    "\n",
    "    sorted_data_desc = sorted(relations, key=lambda x: x['score'], reverse=True)\n",
    "    print(\"\\nDescending Order by Score:\")\n",
    "    for item in sorted_data_desc:\n",
    "        print(f\"{item['head_text']} --> {item['label']} --> {item['tail_text']} | score: {item['score']}\")\n",
    "\n",
    "predict_and_show(text, labels)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A simple list of relation types can also be passed, although this generally results in noisier results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Text: Jack knows Gill. They live in the same house in London. They are not related.\n",
      "Entities detected: [[0, 0, 'PERSON', 'Jack'], [2, 2, 'PERSON', 'Gill'], [11, 11, 'GPE', 'London']]\n",
      "Number of relations: 6\n",
      "\n",
      "Descending Order by Score:\n",
      "['Jack'] --> lives in --> ['London'] | score: 0.9570847153663635\n",
      "['Gill'] --> lives in --> ['London'] | score: 0.9562698006629944\n",
      "['Jack'] --> knows --> ['Gill'] | score: 0.8528702259063721\n",
      "['Gill'] --> knows --> ['Jack'] | score: 0.8421204090118408\n",
      "['London'] --> lives in --> ['Gill'] | score: 0.6627970337867737\n",
      "['London'] --> lives in --> ['Jack'] | score: 0.6488385796546936\n"
     ]
    }
   ],
   "source": [
    "text = \"Jack knows Gill. They live in the same house in London. They are not related.\"\n",
    "labels = ['family relation', 'knows', 'lives with', 'loves', 'licensed to broadcast to', 'father', 'followed by', 'no relation', 'lives in',]\n",
    "predict_and_show(text, labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The token has not been saved to the git credentials helper. Pass `add_to_git_credential=True` in this function directly or `--add-to-git-credential` if using via `huggingface-cli` if you want to set the git credential as well.\n",
      "Token is valid (permission: write).\n",
      "Your token has been saved to /home/jackboylan/.cache/huggingface/token\n",
      "Login successful\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1bcf8fc4b83f4bfe82215ac0bd0dcfec",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "pytorch_model.bin:   0%|          | 0.00/1.87G [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "CommitInfo(commit_url='https://huggingface.co/jackboyla/glirel_beta/commit/9d864f1374760c9d5d9321a25d93bdf8895d0964', commit_message='Push model using huggingface_hub.', commit_description='', oid='9d864f1374760c9d5d9321a25d93bdf8895d0964', pr_url=None, pr_revision=None, pr_num=None)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# import huggingface_hub\n",
    "# import os\n",
    "\n",
    "# huggingface_hub.login(os.environ['HF_TOKEN'])\n",
    "\n",
    "# model.save_pretrained(\n",
    "#     './release_model/glirel_beta', \n",
    "#     push_to_hub=True, \n",
    "#     repo_id='jackboyla/glirel_beta'\n",
    "# )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Co-reference Resolution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['The', 'Loud', 'Tour', 'was', 'the', 'fourth', 'overall', 'and', 'third', 'world', 'concert', 'tour', 'by', 'Barbadian', 'recording', 'artist', 'Rihanna', '.', 'Performing', 'in', 'over', 'twenty', 'countries', 'in', 'the', 'Americas', 'and', 'Europe', ',', 'the', 'tour', 'was', 'launched', 'in', 'support', 'of', 'Rihanna', \"'s\", 'fifth', 'studio', 'album', 'Loud', '(', '2010', ')', '.', 'Critics', 'acclaimed', 'the', 'show', 'for', 'its', 'liveliness', 'and', 'higher', 'caliber', 'of', 'quality', 'when', 'compared', 'to', 'Rihanna', \"'s\", 'previous', 'tours', '.', 'The', 'Loud', 'Tour', 'was', 'a', 'large', 'commercial', 'success', ',', 'experiencing', 'demand', 'for', 'an', 'extension', 'of', 'shows', 'in', 'the', 'United', 'Kingdom', 'due', 'to', 'popularity', '.', 'In', 'London', ',', 'Rihanna', 'played', 'a', 'record', 'breaking', '10', 'dates', 'at', 'The', 'O2', 'Arena', '.', 'The', 'tour', 'ultimately', 'grossed', 'an', 'estimated', 'value', 'of', 'US$', '90', 'million', 'from', '98', 'reported', 'shows', 'and', 'a', 'total', 'audience', 'of', '1,200,800', '.', 'The', 'Loud', 'Tour', 'became', 'the', 'seventh', '-', 'highest', 'grossing', 'tour', 'of', '2011', '.']\n",
      "\n",
      "[[1, 3, 'MISC', 'Loud Tour'], [13, 14, 'LOC', 'Barbadian'], [16, 17, 'PER', 'Rihanna'], [21, 22, 'NUM', 'twenty'], [25, 26, 'LOC', 'Americas'], [27, 28, 'LOC', 'Europe'], [36, 37, 'PER', 'Rihanna'], [41, 42, 'MISC', 'Loud'], [43, 44, 'TIME', '2010'], [61, 62, 'PER', 'Rihanna'], [67, 69, 'MISC', 'Loud Tour'], [84, 86, 'LOC', 'United Kingdom'], [91, 92, 'LOC', 'London'], [93, 94, 'PER', 'Rihanna'], [98, 99, 'NUM', '10'], [101, 104, 'LOC', 'The O2 Arena'], [113, 116, 'NUM', 'US$90 million'], [117, 118, 'NUM', '98'], [125, 126, 'NUM', '1,200,800'], [128, 130, 'MISC', 'Loud Tour'], [138, 139, 'TIME', '2011']]\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "['continent',\n",
       " 'country',\n",
       " 'located in the administrative territorial entity',\n",
       " 'performer',\n",
       " 'publication date',\n",
       " 'notable work',\n",
       " 'country of citizenship',\n",
       " 'SELF']"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import json\n",
    "with open('data/redocred_test.jsonl', 'r') as f:\n",
    "    data = [json.loads(line) for line in f] \n",
    "\n",
    "examples = data[:5]\n",
    "batch_tokenized_text = [example['tokenized_text'] for example in examples]\n",
    "batch_ner = [example['ner'] for example in examples]\n",
    "print(batch_tokenized_text[0])\n",
    "print()\n",
    "print(batch_ner[0])\n",
    "print()\n",
    "\n",
    "batch_labels = [list(set([r['relation_text'] for r in example['relations']])) for example in examples]\n",
    "batch_labels[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'head_pos': [1, 3],\n",
       "  'tail_pos': [16, 17],\n",
       "  'head_text': ['Loud', 'Tour'],\n",
       "  'tail_text': ['Rihanna'],\n",
       "  'label': 'performer',\n",
       "  'score': 1.0},\n",
       " {'head_pos': [1, 3],\n",
       "  'tail_pos': [36, 37],\n",
       "  'head_text': ['Loud', 'Tour'],\n",
       "  'tail_text': ['Rihanna'],\n",
       "  'label': 'performer',\n",
       "  'score': 1.0},\n",
       " {'head_pos': [1, 3],\n",
       "  'tail_pos': [43, 44],\n",
       "  'head_text': ['Loud', 'Tour'],\n",
       "  'tail_text': ['2010'],\n",
       "  'label': 'publication date',\n",
       "  'score': 1.0},\n",
       " {'head_pos': [1, 3],\n",
       "  'tail_pos': [61, 62],\n",
       "  'head_text': ['Loud', 'Tour'],\n",
       "  'tail_text': ['Rihanna'],\n",
       "  'label': 'performer',\n",
       "  'score': 1.0},\n",
       " {'head_pos': [1, 3],\n",
       "  'tail_pos': [67, 69],\n",
       "  'head_text': ['Loud', 'Tour'],\n",
       "  'tail_text': ['Loud', 'Tour'],\n",
       "  'label': 'SELF',\n",
       "  'score': 1.0},\n",
       " {'head_pos': [1, 3],\n",
       "  'tail_pos': [93, 94],\n",
       "  'head_text': ['Loud', 'Tour'],\n",
       "  'tail_text': ['Rihanna'],\n",
       "  'label': 'performer',\n",
       "  'score': 1.0},\n",
       " {'head_pos': [1, 3],\n",
       "  'tail_pos': [128, 130],\n",
       "  'head_text': ['Loud', 'Tour'],\n",
       "  'tail_text': ['Loud', 'Tour'],\n",
       "  'label': 'SELF',\n",
       "  'score': 1.0},\n",
       " {'head_pos': [16, 17],\n",
       "  'tail_pos': [1, 3],\n",
       "  'head_text': ['Rihanna'],\n",
       "  'tail_text': ['Loud', 'Tour'],\n",
       "  'label': 'notable work',\n",
       "  'score': 1.0},\n",
       " {'head_pos': [16, 17],\n",
       "  'tail_pos': [13, 14],\n",
       "  'head_text': ['Rihanna'],\n",
       "  'tail_text': ['Barbadian'],\n",
       "  'label': 'country of citizenship',\n",
       "  'score': 1.0},\n",
       " {'head_pos': [16, 17],\n",
       "  'tail_pos': [36, 37],\n",
       "  'head_text': ['Rihanna'],\n",
       "  'tail_text': ['Rihanna'],\n",
       "  'label': 'SELF',\n",
       "  'score': 1.0},\n",
       " {'head_pos': [16, 17],\n",
       "  'tail_pos': [41, 42],\n",
       "  'head_text': ['Rihanna'],\n",
       "  'tail_text': ['Loud'],\n",
       "  'label': 'notable work',\n",
       "  'score': 1.0},\n",
       " {'head_pos': [16, 17],\n",
       "  'tail_pos': [61, 62],\n",
       "  'head_text': ['Rihanna'],\n",
       "  'tail_text': ['Rihanna'],\n",
       "  'label': 'SELF',\n",
       "  'score': 1.0},\n",
       " {'head_pos': [16, 17],\n",
       "  'tail_pos': [67, 69],\n",
       "  'head_text': ['Rihanna'],\n",
       "  'tail_text': ['Loud', 'Tour'],\n",
       "  'label': 'notable work',\n",
       "  'score': 1.0},\n",
       " {'head_pos': [16, 17],\n",
       "  'tail_pos': [93, 94],\n",
       "  'head_text': ['Rihanna'],\n",
       "  'tail_text': ['Rihanna'],\n",
       "  'label': 'SELF',\n",
       "  'score': 1.0},\n",
       " {'head_pos': [16, 17],\n",
       "  'tail_pos': [128, 130],\n",
       "  'head_text': ['Rihanna'],\n",
       "  'tail_text': ['Loud', 'Tour'],\n",
       "  'label': 'notable work',\n",
       "  'score': 1.0},\n",
       " {'head_pos': [36, 37],\n",
       "  'tail_pos': [1, 3],\n",
       "  'head_text': ['Rihanna'],\n",
       "  'tail_text': ['Loud', 'Tour'],\n",
       "  'label': 'notable work',\n",
       "  'score': 1.0},\n",
       " {'head_pos': [36, 37],\n",
       "  'tail_pos': [13, 14],\n",
       "  'head_text': ['Rihanna'],\n",
       "  'tail_text': ['Barbadian'],\n",
       "  'label': 'country of citizenship',\n",
       "  'score': 1.0},\n",
       " {'head_pos': [36, 37],\n",
       "  'tail_pos': [16, 17],\n",
       "  'head_text': ['Rihanna'],\n",
       "  'tail_text': ['Rihanna'],\n",
       "  'label': 'SELF',\n",
       "  'score': 1.0},\n",
       " {'head_pos': [36, 37],\n",
       "  'tail_pos': [41, 42],\n",
       "  'head_text': ['Rihanna'],\n",
       "  'tail_text': ['Loud'],\n",
       "  'label': 'notable work',\n",
       "  'score': 1.0},\n",
       " {'head_pos': [36, 37],\n",
       "  'tail_pos': [61, 62],\n",
       "  'head_text': ['Rihanna'],\n",
       "  'tail_text': ['Rihanna'],\n",
       "  'label': 'SELF',\n",
       "  'score': 1.0},\n",
       " {'head_pos': [36, 37],\n",
       "  'tail_pos': [67, 69],\n",
       "  'head_text': ['Rihanna'],\n",
       "  'tail_text': ['Loud', 'Tour'],\n",
       "  'label': 'notable work',\n",
       "  'score': 1.0},\n",
       " {'head_pos': [36, 37],\n",
       "  'tail_pos': [93, 94],\n",
       "  'head_text': ['Rihanna'],\n",
       "  'tail_text': ['Rihanna'],\n",
       "  'label': 'SELF',\n",
       "  'score': 1.0},\n",
       " {'head_pos': [36, 37],\n",
       "  'tail_pos': [128, 130],\n",
       "  'head_text': ['Rihanna'],\n",
       "  'tail_text': ['Loud', 'Tour'],\n",
       "  'label': 'notable work',\n",
       "  'score': 1.0},\n",
       " {'head_pos': [41, 42],\n",
       "  'tail_pos': [16, 17],\n",
       "  'head_text': ['Loud'],\n",
       "  'tail_text': ['Rihanna'],\n",
       "  'label': 'performer',\n",
       "  'score': 1.0},\n",
       " {'head_pos': [41, 42],\n",
       "  'tail_pos': [36, 37],\n",
       "  'head_text': ['Loud'],\n",
       "  'tail_text': ['Rihanna'],\n",
       "  'label': 'performer',\n",
       "  'score': 1.0},\n",
       " {'head_pos': [41, 42],\n",
       "  'tail_pos': [43, 44],\n",
       "  'head_text': ['Loud'],\n",
       "  'tail_text': ['2010'],\n",
       "  'label': 'publication date',\n",
       "  'score': 1.0},\n",
       " {'head_pos': [41, 42],\n",
       "  'tail_pos': [61, 62],\n",
       "  'head_text': ['Loud'],\n",
       "  'tail_text': ['Rihanna'],\n",
       "  'label': 'performer',\n",
       "  'score': 1.0},\n",
       " {'head_pos': [41, 42],\n",
       "  'tail_pos': [93, 94],\n",
       "  'head_text': ['Loud'],\n",
       "  'tail_text': ['Rihanna'],\n",
       "  'label': 'performer',\n",
       "  'score': 1.0},\n",
       " {'head_pos': [61, 62],\n",
       "  'tail_pos': [13, 14],\n",
       "  'head_text': ['Rihanna'],\n",
       "  'tail_text': ['Barbadian'],\n",
       "  'label': 'country of citizenship',\n",
       "  'score': 1.0},\n",
       " {'head_pos': [61, 62],\n",
       "  'tail_pos': [16, 17],\n",
       "  'head_text': ['Rihanna'],\n",
       "  'tail_text': ['Rihanna'],\n",
       "  'label': 'SELF',\n",
       "  'score': 1.0},\n",
       " {'head_pos': [61, 62],\n",
       "  'tail_pos': [36, 37],\n",
       "  'head_text': ['Rihanna'],\n",
       "  'tail_text': ['Rihanna'],\n",
       "  'label': 'SELF',\n",
       "  'score': 1.0},\n",
       " {'head_pos': [61, 62],\n",
       "  'tail_pos': [41, 42],\n",
       "  'head_text': ['Rihanna'],\n",
       "  'tail_text': ['Loud'],\n",
       "  'label': 'notable work',\n",
       "  'score': 1.0},\n",
       " {'head_pos': [61, 62],\n",
       "  'tail_pos': [67, 69],\n",
       "  'head_text': ['Rihanna'],\n",
       "  'tail_text': ['Loud', 'Tour'],\n",
       "  'label': 'notable work',\n",
       "  'score': 1.0},\n",
       " {'head_pos': [61, 62],\n",
       "  'tail_pos': [93, 94],\n",
       "  'head_text': ['Rihanna'],\n",
       "  'tail_text': ['Rihanna'],\n",
       "  'label': 'SELF',\n",
       "  'score': 1.0},\n",
       " {'head_pos': [61, 62],\n",
       "  'tail_pos': [128, 130],\n",
       "  'head_text': ['Rihanna'],\n",
       "  'tail_text': ['Loud', 'Tour'],\n",
       "  'label': 'notable work',\n",
       "  'score': 1.0},\n",
       " {'head_pos': [67, 69],\n",
       "  'tail_pos': [36, 37],\n",
       "  'head_text': ['Loud', 'Tour'],\n",
       "  'tail_text': ['Rihanna'],\n",
       "  'label': 'performer',\n",
       "  'score': 1.0},\n",
       " {'head_pos': [67, 69],\n",
       "  'tail_pos': [43, 44],\n",
       "  'head_text': ['Loud', 'Tour'],\n",
       "  'tail_text': ['2010'],\n",
       "  'label': 'publication date',\n",
       "  'score': 1.0},\n",
       " {'head_pos': [67, 69],\n",
       "  'tail_pos': [61, 62],\n",
       "  'head_text': ['Loud', 'Tour'],\n",
       "  'tail_text': ['Rihanna'],\n",
       "  'label': 'performer',\n",
       "  'score': 1.0},\n",
       " {'head_pos': [67, 69],\n",
       "  'tail_pos': [93, 94],\n",
       "  'head_text': ['Loud', 'Tour'],\n",
       "  'tail_text': ['Rihanna'],\n",
       "  'label': 'performer',\n",
       "  'score': 1.0},\n",
       " {'head_pos': [67, 69],\n",
       "  'tail_pos': [128, 130],\n",
       "  'head_text': ['Loud', 'Tour'],\n",
       "  'tail_text': ['Loud', 'Tour'],\n",
       "  'label': 'SELF',\n",
       "  'score': 1.0},\n",
       " {'head_pos': [91, 92],\n",
       "  'tail_pos': [84, 86],\n",
       "  'head_text': ['London'],\n",
       "  'tail_text': ['United', 'Kingdom'],\n",
       "  'label': 'located in the administrative territorial entity',\n",
       "  'score': 1.0},\n",
       " {'head_pos': [93, 94],\n",
       "  'tail_pos': [61, 62],\n",
       "  'head_text': ['Rihanna'],\n",
       "  'tail_text': ['Rihanna'],\n",
       "  'label': 'SELF',\n",
       "  'score': 1.0},\n",
       " {'head_pos': [93, 94],\n",
       "  'tail_pos': [67, 69],\n",
       "  'head_text': ['Rihanna'],\n",
       "  'tail_text': ['Loud', 'Tour'],\n",
       "  'label': 'notable work',\n",
       "  'score': 1.0},\n",
       " {'head_pos': [93, 94],\n",
       "  'tail_pos': [128, 130],\n",
       "  'head_text': ['Rihanna'],\n",
       "  'tail_text': ['Loud', 'Tour'],\n",
       "  'label': 'notable work',\n",
       "  'score': 1.0},\n",
       " {'head_pos': [101, 104],\n",
       "  'tail_pos': [84, 86],\n",
       "  'head_text': ['The', 'O2', 'Arena'],\n",
       "  'tail_text': ['United', 'Kingdom'],\n",
       "  'label': 'located in the administrative territorial entity',\n",
       "  'score': 1.0},\n",
       " {'head_pos': [128, 130],\n",
       "  'tail_pos': [93, 94],\n",
       "  'head_text': ['Loud', 'Tour'],\n",
       "  'tail_text': ['Rihanna'],\n",
       "  'label': 'performer',\n",
       "  'score': 1.0},\n",
       " {'head_pos': [101, 104],\n",
       "  'tail_pos': [91, 92],\n",
       "  'head_text': ['The', 'O2', 'Arena'],\n",
       "  'tail_text': ['London'],\n",
       "  'label': 'located in the administrative territorial entity',\n",
       "  'score': 0.9999502897262573}]"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.base_config.fixed_relation_types = False\n",
    "model = model.to('cpu')\n",
    "relations = model.batch_predict_relations(batch_tokenized_text, batch_labels, threshold=0.01, ner=batch_ner, top_k=1)\n",
    "relations[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Clusters: [[(128, 130), (67, 69), (1, 3)], [(13, 14)], [(61, 62), (93, 94), (36, 37), (16, 17)], [(41, 42)], [(43, 44)], [(84, 86)], [(91, 92)], [(101, 104)]]\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[{'h_idx': 0, 't_idx': 2, 'r': 'performer'},\n",
       " {'h_idx': 0, 't_idx': 3, 'r': 'follows'},\n",
       " {'h_idx': 0, 't_idx': 4, 'r': 'publication date'},\n",
       " {'h_idx': 2, 't_idx': 0, 'r': 'notable work'},\n",
       " {'h_idx': 2, 't_idx': 1, 'r': 'country of citizenship'},\n",
       " {'h_idx': 2, 't_idx': 3, 'r': 'notable work'},\n",
       " {'h_idx': 3, 't_idx': 0, 'r': 'followed by'},\n",
       " {'h_idx': 3, 't_idx': 2, 'r': 'performer'},\n",
       " {'h_idx': 3, 't_idx': 4, 'r': 'publication date'},\n",
       " {'h_idx': 6,\n",
       "  't_idx': 5,\n",
       "  'r': 'located in the administrative territorial entity'},\n",
       " {'h_idx': 7,\n",
       "  't_idx': 5,\n",
       "  'r': 'located in the administrative territorial entity'},\n",
       " {'h_idx': 7,\n",
       "  't_idx': 6,\n",
       "  'r': 'located in the administrative territorial entity'}]"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from glirel.modules.utils import get_coreference_clusters, aggregate_cluster_relations\n",
    "\n",
    "clusters_batch, entity_to_cluster_idx_batch = get_coreference_clusters(relations)\n",
    "print(\"Clusters:\", clusters_batch[0])\n",
    "print()\n",
    "cluster_relations_batch = aggregate_cluster_relations(entity_to_cluster_idx_batch, relations)\n",
    "cluster_relations_batch[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Head Cluster: [['Loud', 'Tour'], ['Loud', 'Tour'], ['Loud', 'Tour']]\n",
      "Relation: performer\n",
      "Tail Cluster: [['Rihanna'], ['Rihanna'], ['Rihanna'], ['Rihanna']]\n",
      "\n",
      "Head Cluster: [['Loud', 'Tour'], ['Loud', 'Tour'], ['Loud', 'Tour']]\n",
      "Relation: follows\n",
      "Tail Cluster: [['Loud']]\n",
      "\n",
      "Head Cluster: [['Loud', 'Tour'], ['Loud', 'Tour'], ['Loud', 'Tour']]\n",
      "Relation: publication date\n",
      "Tail Cluster: [['2010']]\n",
      "\n",
      "Head Cluster: [['Rihanna'], ['Rihanna'], ['Rihanna'], ['Rihanna']]\n",
      "Relation: notable work\n",
      "Tail Cluster: [['Loud', 'Tour'], ['Loud', 'Tour'], ['Loud', 'Tour']]\n",
      "\n",
      "Head Cluster: [['Rihanna'], ['Rihanna'], ['Rihanna'], ['Rihanna']]\n",
      "Relation: country of citizenship\n",
      "Tail Cluster: [['Barbadian']]\n",
      "\n",
      "Head Cluster: [['Rihanna'], ['Rihanna'], ['Rihanna'], ['Rihanna']]\n",
      "Relation: notable work\n",
      "Tail Cluster: [['Loud']]\n",
      "\n",
      "Head Cluster: [['Loud']]\n",
      "Relation: followed by\n",
      "Tail Cluster: [['Loud', 'Tour'], ['Loud', 'Tour'], ['Loud', 'Tour']]\n",
      "\n",
      "Head Cluster: [['Loud']]\n",
      "Relation: performer\n",
      "Tail Cluster: [['Rihanna'], ['Rihanna'], ['Rihanna'], ['Rihanna']]\n",
      "\n",
      "Head Cluster: [['Loud']]\n",
      "Relation: publication date\n",
      "Tail Cluster: [['2010']]\n",
      "\n",
      "Head Cluster: [['London']]\n",
      "Relation: located in the administrative territorial entity\n",
      "Tail Cluster: [['United', 'Kingdom']]\n",
      "\n",
      "Head Cluster: [['The', 'O2', 'Arena']]\n",
      "Relation: located in the administrative territorial entity\n",
      "Tail Cluster: [['United', 'Kingdom']]\n",
      "\n",
      "Head Cluster: [['The', 'O2', 'Arena']]\n",
      "Relation: located in the administrative territorial entity\n",
      "Tail Cluster: [['London']]\n",
      "\n",
      "Head Cluster: [['Orlov'], ['Vladimir', 'Mitrofanovich', 'Orlov'], ['Orlov']]\n",
      "Relation: date of birth\n",
      "Tail Cluster: [['July', '15', ',', '1895']]\n",
      "\n",
      "Head Cluster: [['Orlov'], ['Vladimir', 'Mitrofanovich', 'Orlov'], ['Orlov']]\n",
      "Relation: date of death\n",
      "Tail Cluster: [['28', 'July', '1938'], ['July', '28', ',', '1938']]\n",
      "\n",
      "Head Cluster: [['Orlov'], ['Vladimir', 'Mitrofanovich', 'Orlov'], ['Orlov']]\n",
      "Relation: country of citizenship\n",
      "Tail Cluster: [['Russian']]\n",
      "\n",
      "Head Cluster: [['Orlov'], ['Vladimir', 'Mitrofanovich', 'Orlov'], ['Orlov']]\n",
      "Relation: military branch\n",
      "Tail Cluster: [['Soviet', 'Naval', 'Forces']]\n",
      "\n",
      "Head Cluster: [['Orlov'], ['Vladimir', 'Mitrofanovich', 'Orlov'], ['Orlov']]\n",
      "Relation: place of birth\n",
      "Tail Cluster: [['Kherson']]\n",
      "\n",
      "Head Cluster: [['Orlov'], ['Vladimir', 'Mitrofanovich', 'Orlov'], ['Orlov']]\n",
      "Relation: educated at\n",
      "Tail Cluster: [['St', 'Petersburg', 'University']]\n",
      "\n",
      "Head Cluster: [['Orlov'], ['Vladimir', 'Mitrofanovich', 'Orlov'], ['Orlov']]\n",
      "Relation: military branch\n",
      "Tail Cluster: [['Baltic', 'Fleet'], ['Baltic', 'Fleet']]\n",
      "\n",
      "Head Cluster: [['Orlov'], ['Vladimir', 'Mitrofanovich', 'Orlov'], ['Orlov']]\n",
      "Relation: military branch\n",
      "Tail Cluster: [['Black', 'Sea', 'Fleet']]\n",
      "\n",
      "Head Cluster: [['Orlov'], ['Vladimir', 'Mitrofanovich', 'Orlov'], ['Orlov']]\n",
      "Relation: military branch\n",
      "Tail Cluster: [['Soviet', 'Navy']]\n",
      "\n",
      "Head Cluster: [['Soviet', 'Naval', 'Forces']]\n",
      "Relation: located in the administrative territorial entity\n",
      "Tail Cluster: [['Russian']]\n",
      "\n",
      "Head Cluster: [['Soviet', 'Naval', 'Forces']]\n",
      "Relation: has part\n",
      "Tail Cluster: [['Black', 'Sea', 'Fleet']]\n",
      "\n",
      "Head Cluster: [['Kherson']]\n",
      "Relation: located in the administrative territorial entity\n",
      "Tail Cluster: [['Russian']]\n",
      "\n",
      "Head Cluster: [['St', 'Petersburg', 'University']]\n",
      "Relation: located in the administrative territorial entity\n",
      "Tail Cluster: [['Russian']]\n",
      "\n",
      "Head Cluster: [['Baltic', 'Fleet'], ['Baltic', 'Fleet']]\n",
      "Relation: located in the administrative territorial entity\n",
      "Tail Cluster: [['Russian']]\n",
      "\n",
      "Head Cluster: [['Black', 'Sea', 'Fleet']]\n",
      "Relation: part of\n",
      "Tail Cluster: [['Soviet', 'Navy']]\n",
      "\n",
      "Head Cluster: [['Soviet', 'Navy']]\n",
      "Relation: has part\n",
      "Tail Cluster: [['Black', 'Sea', 'Fleet']]\n",
      "\n",
      "Head Cluster: [['Ecuadorian', 'Constituent', 'Assembly']]\n",
      "Relation: inception\n",
      "Tail Cluster: [['2007'], ['30', 'September', '2007'], ['November', '29', ',', '2007']]\n",
      "\n",
      "Head Cluster: [['Ecuadorian', 'Constituent', 'Assembly']]\n",
      "Relation: located in the administrative territorial entity\n",
      "Tail Cluster: [['Ecuador'], ['Ecuador'], ['Ecuador'], ['Ecuador'], ['Ecuadorian']]\n",
      "\n",
      "Head Cluster: [['Ecuadorian', 'Constituent', 'Assembly']]\n",
      "Relation: inception\n",
      "Tail Cluster: [['15', 'April', '2007']]\n",
      "\n",
      "Head Cluster: [['Ecuador'], ['Ecuador'], ['Ecuador'], ['Ecuador'], ['Ecuadorian']]\n",
      "Relation: legislative body\n",
      "Tail Cluster: [['Ecuadorian', 'Constituent', 'Assembly']]\n",
      "\n",
      "Head Cluster: [['Ecuador'], ['Ecuador'], ['Ecuador'], ['Ecuador'], ['Ecuadorian']]\n",
      "Relation: inception\n",
      "Tail Cluster: [['2007'], ['30', 'September', '2007'], ['November', '29', ',', '2007']]\n",
      "\n",
      "Head Cluster: [['Ecuador'], ['Ecuador'], ['Ecuador'], ['Ecuador'], ['Ecuadorian']]\n",
      "Relation: dissolved, abolished or demolished date\n",
      "Tail Cluster: [['2008'], ['2008'], ['July', ',', '2008'], ['2008'], ['September', '2008']]\n",
      "\n",
      "Head Cluster: [['Ecuador'], ['Ecuador'], ['Ecuador'], ['Ecuador'], ['Ecuadorian']]\n",
      "Relation: head of state\n",
      "Tail Cluster: [['Correa'], ['Rafael', 'Correa']]\n",
      "\n",
      "Head Cluster: [['Ecuador'], ['Ecuador'], ['Ecuador'], ['Ecuador'], ['Ecuadorian']]\n",
      "Relation: legislative body\n",
      "Tail Cluster: [['PAIS', 'Alliance']]\n",
      "\n",
      "Head Cluster: [['Constitution', 'of', 'Ecuador']]\n",
      "Relation: publication date\n",
      "Tail Cluster: [['2008'], ['2008'], ['July', ',', '2008'], ['2008'], ['September', '2008']]\n",
      "\n",
      "Head Cluster: [['Constitution', 'of', 'Ecuador']]\n",
      "Relation: country\n",
      "Tail Cluster: [['Ecuador'], ['Ecuador'], ['Ecuador'], ['Ecuador'], ['Ecuadorian']]\n",
      "\n",
      "Head Cluster: [['Correa'], ['Rafael', 'Correa']]\n",
      "Relation: applies to jurisdiction\n",
      "Tail Cluster: [['Ecuador'], ['Ecuador'], ['Ecuador'], ['Ecuador'], ['Ecuadorian']]\n",
      "\n",
      "Head Cluster: [['Correa'], ['Rafael', 'Correa']]\n",
      "Relation: member of political party\n",
      "Tail Cluster: [['PAIS', 'Alliance']]\n",
      "\n",
      "Head Cluster: [['PAIS', 'Alliance']]\n",
      "Relation: located in the administrative territorial entity\n",
      "Tail Cluster: [['Ecuador'], ['Ecuador'], ['Ecuador'], ['Ecuador'], ['Ecuadorian']]\n",
      "\n",
      "Head Cluster: [['PAIS', 'Alliance']]\n",
      "Relation: chairperson\n",
      "Tail Cluster: [['Correa'], ['Rafael', 'Correa']]\n",
      "\n",
      "Head Cluster: [['Montecristi']]\n",
      "Relation: located in the administrative territorial entity\n",
      "Tail Cluster: [['Ecuador'], ['Ecuador'], ['Ecuador'], ['Ecuador'], ['Ecuadorian']]\n",
      "\n",
      "Head Cluster: [['Bantu', 'homeland'], ['bantustan'], ['Bantustans'], ['Bantustan'], ['Bantustans'], ['Bantustans'], ['Bantustans']]\n",
      "Relation: country\n",
      "Tail Cluster: [['South', 'Africa'], ['South', 'African'], ['South', 'Africa'], ['South', 'Africa'], ['South', 'Africa'], ['South', 'African'], ['South', 'African'], ['South', 'Africa']]\n",
      "\n",
      "Head Cluster: [['Bantu', 'homeland'], ['bantustan'], ['Bantustans'], ['Bantustan'], ['Bantustans'], ['Bantustans'], ['Bantustans']]\n",
      "Relation: located in the administrative territorial entity\n",
      "Tail Cluster: [['South', 'Africa'], ['South', 'African'], ['South', 'Africa'], ['South', 'Africa'], ['South', 'Africa'], ['South', 'African'], ['South', 'African'], ['South', 'Africa']]\n",
      "\n",
      "Head Cluster: [['Bantu', 'homeland'], ['bantustan'], ['Bantustans'], ['Bantustan'], ['Bantustans'], ['Bantustans'], ['Bantustans']]\n",
      "Relation: inception\n",
      "Tail Cluster: [['1940s']]\n",
      "\n",
      "Head Cluster: [['Bantu', 'homeland'], ['bantustan'], ['Bantustans'], ['Bantustan'], ['Bantustans'], ['Bantustans'], ['Bantustans']]\n",
      "Relation: dissolved, abolished or demolished date\n",
      "Tail Cluster: [['1994']]\n",
      "\n",
      "Head Cluster: [['South', 'Africa'], ['South', 'African'], ['South', 'Africa'], ['South', 'Africa'], ['South', 'Africa'], ['South', 'African'], ['South', 'African'], ['South', 'Africa']]\n",
      "Relation: official language\n",
      "Tail Cluster: [['Bantu'], ['Bantu']]\n",
      "\n",
      "Head Cluster: [['South', 'Africa'], ['South', 'African'], ['South', 'Africa'], ['South', 'Africa'], ['South', 'Africa'], ['South', 'African'], ['South', 'African'], ['South', 'Africa']]\n",
      "Relation: contains administrative territorial entity\n",
      "Tail Cluster: [['Transkei']]\n",
      "\n",
      "Head Cluster: [['South', 'Africa'], ['South', 'African'], ['South', 'Africa'], ['South', 'Africa'], ['South', 'Africa'], ['South', 'African'], ['South', 'African'], ['South', 'Africa']]\n",
      "Relation: contains administrative territorial entity\n",
      "Tail Cluster: [['Bophuthatswana'], ['Ciskei'], ['TBVC']]\n",
      "\n",
      "Head Cluster: [['South', 'Africa'], ['South', 'African'], ['South', 'Africa'], ['South', 'Africa'], ['South', 'Africa'], ['South', 'African'], ['South', 'African'], ['South', 'Africa']]\n",
      "Relation: contains administrative territorial entity\n",
      "Tail Cluster: [['Venda']]\n",
      "\n",
      "Head Cluster: [['South', 'Africa'], ['South', 'African'], ['South', 'Africa'], ['South', 'Africa'], ['South', 'Africa'], ['South', 'African'], ['South', 'African'], ['South', 'Africa']]\n",
      "Relation: contains administrative territorial entity\n",
      "Tail Cluster: [['KwaZulu']]\n",
      "\n",
      "Head Cluster: [['South', 'Africa'], ['South', 'African'], ['South', 'Africa'], ['South', 'Africa'], ['South', 'Africa'], ['South', 'African'], ['South', 'African'], ['South', 'Africa']]\n",
      "Relation: contains administrative territorial entity\n",
      "Tail Cluster: [['Lebowa']]\n",
      "\n",
      "Head Cluster: [['South', 'Africa'], ['South', 'African'], ['South', 'Africa'], ['South', 'Africa'], ['South', 'Africa'], ['South', 'African'], ['South', 'African'], ['South', 'Africa']]\n",
      "Relation: dissolved, abolished or demolished date\n",
      "Tail Cluster: [['1994']]\n",
      "\n",
      "Head Cluster: [['South', 'West', 'Africa'], ['South', 'West', 'Africa'], ['South', 'West', 'Africa']]\n",
      "Relation: country\n",
      "Tail Cluster: [['South', 'Africa'], ['South', 'African'], ['South', 'Africa'], ['South', 'Africa'], ['South', 'Africa'], ['South', 'African'], ['South', 'African'], ['South', 'Africa']]\n",
      "\n",
      "Head Cluster: [['South', 'West', 'Africa'], ['South', 'West', 'Africa'], ['South', 'West', 'Africa']]\n",
      "Relation: country\n",
      "Tail Cluster: [['Namibia']]\n",
      "\n",
      "Head Cluster: [['Namibia']]\n",
      "Relation: official language\n",
      "Tail Cluster: [['South', 'West', 'Africa'], ['South', 'West', 'Africa'], ['South', 'West', 'Africa']]\n",
      "\n",
      "Head Cluster: [['Namibia']]\n",
      "Relation: continent\n",
      "Tail Cluster: [['Asia']]\n",
      "\n",
      "Head Cluster: [['Afrikaans']]\n",
      "Relation: country\n",
      "Tail Cluster: [['South', 'Africa'], ['South', 'African'], ['South', 'Africa'], ['South', 'Africa'], ['South', 'Africa'], ['South', 'African'], ['South', 'African'], ['South', 'Africa']]\n",
      "\n",
      "Head Cluster: [['Transkei']]\n",
      "Relation: country\n",
      "Tail Cluster: [['South', 'Africa'], ['South', 'African'], ['South', 'Africa'], ['South', 'Africa'], ['South', 'Africa'], ['South', 'African'], ['South', 'African'], ['South', 'Africa']]\n",
      "\n",
      "Head Cluster: [['Transkei']]\n",
      "Relation: located in the administrative territorial entity\n",
      "Tail Cluster: [['South', 'Africa'], ['South', 'African'], ['South', 'Africa'], ['South', 'Africa'], ['South', 'Africa'], ['South', 'African'], ['South', 'African'], ['South', 'Africa']]\n",
      "\n",
      "Head Cluster: [['Transkei']]\n",
      "Relation: follows\n",
      "Tail Cluster: [['Bophuthatswana'], ['Ciskei'], ['TBVC']]\n",
      "\n",
      "Head Cluster: [['Transkei']]\n",
      "Relation: replaced by\n",
      "Tail Cluster: [['Bophuthatswana'], ['Ciskei'], ['TBVC']]\n",
      "\n",
      "Head Cluster: [['Transkei']]\n",
      "Relation: follows\n",
      "Tail Cluster: [['Venda']]\n",
      "\n",
      "Head Cluster: [['Bophuthatswana'], ['Ciskei'], ['TBVC']]\n",
      "Relation: country\n",
      "Tail Cluster: [['South', 'Africa'], ['South', 'African'], ['South', 'Africa'], ['South', 'Africa'], ['South', 'Africa'], ['South', 'African'], ['South', 'African'], ['South', 'Africa']]\n",
      "\n",
      "Head Cluster: [['Bophuthatswana'], ['Ciskei'], ['TBVC']]\n",
      "Relation: located in the administrative territorial entity\n",
      "Tail Cluster: [['South', 'Africa'], ['South', 'African'], ['South', 'Africa'], ['South', 'Africa'], ['South', 'Africa'], ['South', 'African'], ['South', 'African'], ['South', 'Africa']]\n",
      "\n",
      "Head Cluster: [['Bophuthatswana'], ['Ciskei'], ['TBVC']]\n",
      "Relation: followed by\n",
      "Tail Cluster: [['Transkei']]\n",
      "\n",
      "Head Cluster: [['Bophuthatswana'], ['Ciskei'], ['TBVC']]\n",
      "Relation: replaced by\n",
      "Tail Cluster: [['Transkei']]\n",
      "\n",
      "Head Cluster: [['Bophuthatswana'], ['Ciskei'], ['TBVC']]\n",
      "Relation: follows\n",
      "Tail Cluster: [['Bophuthatswana'], ['Ciskei'], ['TBVC']]\n",
      "\n",
      "Head Cluster: [['Bophuthatswana'], ['Ciskei'], ['TBVC']]\n",
      "Relation: replaced by\n",
      "Tail Cluster: [['Bophuthatswana'], ['Ciskei'], ['TBVC']]\n",
      "\n",
      "Head Cluster: [['Bophuthatswana'], ['Ciskei'], ['TBVC']]\n",
      "Relation: follows\n",
      "Tail Cluster: [['Venda']]\n",
      "\n",
      "Head Cluster: [['Bophuthatswana'], ['Ciskei'], ['TBVC']]\n",
      "Relation: replaced by\n",
      "Tail Cluster: [['Venda']]\n",
      "\n",
      "Head Cluster: [['Venda']]\n",
      "Relation: located in the administrative territorial entity\n",
      "Tail Cluster: [['South', 'Africa'], ['South', 'African'], ['South', 'Africa'], ['South', 'Africa'], ['South', 'Africa'], ['South', 'African'], ['South', 'African'], ['South', 'Africa']]\n",
      "\n",
      "Head Cluster: [['Venda']]\n",
      "Relation: followed by\n",
      "Tail Cluster: [['Transkei']]\n",
      "\n",
      "Head Cluster: [['Venda']]\n",
      "Relation: followed by\n",
      "Tail Cluster: [['Bophuthatswana'], ['Ciskei'], ['TBVC']]\n",
      "\n",
      "Head Cluster: [['Venda']]\n",
      "Relation: replaces\n",
      "Tail Cluster: [['Bophuthatswana'], ['Ciskei'], ['TBVC']]\n",
      "\n",
      "Head Cluster: [['KwaZulu']]\n",
      "Relation: located in the administrative territorial entity\n",
      "Tail Cluster: [['South', 'Africa'], ['South', 'African'], ['South', 'Africa'], ['South', 'Africa'], ['South', 'Africa'], ['South', 'African'], ['South', 'African'], ['South', 'Africa']]\n",
      "\n",
      "Head Cluster: [['Lebowa']]\n",
      "Relation: located in the administrative territorial entity\n",
      "Tail Cluster: [['South', 'Africa'], ['South', 'African'], ['South', 'Africa'], ['South', 'Africa'], ['South', 'Africa'], ['South', 'African'], ['South', 'African'], ['South', 'Africa']]\n",
      "\n",
      "Head Cluster: [['QwaQwa']]\n",
      "Relation: located in the administrative territorial entity\n",
      "Tail Cluster: [['South', 'Africa'], ['South', 'African'], ['South', 'Africa'], ['South', 'Africa'], ['South', 'Africa'], ['South', 'African'], ['South', 'African'], ['South', 'Africa']]\n",
      "\n",
      "Head Cluster: [['Ovamboland']]\n",
      "Relation: located in the administrative territorial entity\n",
      "Tail Cluster: [['South', 'Africa'], ['South', 'African'], ['South', 'Africa'], ['South', 'Africa'], ['South', 'Africa'], ['South', 'African'], ['South', 'African'], ['South', 'Africa']]\n",
      "\n",
      "Head Cluster: [['Ovamboland']]\n",
      "Relation: located in the administrative territorial entity\n",
      "Tail Cluster: [['South', 'West', 'Africa'], ['South', 'West', 'Africa'], ['South', 'West', 'Africa']]\n",
      "\n",
      "Head Cluster: [['Kavangoland']]\n",
      "Relation: located in the administrative territorial entity\n",
      "Tail Cluster: [['South', 'Africa'], ['South', 'African'], ['South', 'Africa'], ['South', 'Africa'], ['South', 'Africa'], ['South', 'African'], ['South', 'African'], ['South', 'Africa']]\n",
      "\n",
      "Head Cluster: [['Kavangoland']]\n",
      "Relation: located in the administrative territorial entity\n",
      "Tail Cluster: [['South', 'West', 'Africa'], ['South', 'West', 'Africa'], ['South', 'West', 'Africa']]\n",
      "\n",
      "Head Cluster: [['East', 'Caprivi']]\n",
      "Relation: located in the administrative territorial entity\n",
      "Tail Cluster: [['South', 'Africa'], ['South', 'African'], ['South', 'Africa'], ['South', 'Africa'], ['South', 'Africa'], ['South', 'African'], ['South', 'African'], ['South', 'Africa']]\n",
      "\n",
      "Head Cluster: [['East', 'Caprivi']]\n",
      "Relation: located in the administrative territorial entity\n",
      "Tail Cluster: [['South', 'West', 'Africa'], ['South', 'West', 'Africa'], ['South', 'West', 'Africa']]\n",
      "\n",
      "Head Cluster: [['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell']]\n",
      "Relation: country of origin\n",
      "Tail Cluster: [['American']]\n",
      "\n",
      "Head Cluster: [['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell']]\n",
      "Relation: performer\n",
      "Tail Cluster: [['Manson'], ['Marilyn', 'Manson'], ['Marilyn', 'Manson'], ['Marilyn', 'Manson'], ['Manson'], ['Manson'], ['Manson']]\n",
      "\n",
      "Head Cluster: [['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell']]\n",
      "Relation: performer\n",
      "Tail Cluster: [['Sneaker', 'Pimps'], ['Sneaker', 'Pimps'], ['Sneaker', 'Pimps']]\n",
      "\n",
      "Head Cluster: [['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell']]\n",
      "Relation: producer\n",
      "Tail Cluster: [['Sneaker', 'Pimps'], ['Sneaker', 'Pimps'], ['Sneaker', 'Pimps']]\n",
      "\n",
      "Head Cluster: [['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell']]\n",
      "Relation: publication date\n",
      "Tail Cluster: [['1997'], ['1997']]\n",
      "\n",
      "Head Cluster: [['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell']]\n",
      "Relation: part of\n",
      "Tail Cluster: [['Spawn'], ['Spawn']]\n",
      "\n",
      "Head Cluster: [['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell']]\n",
      "Relation: lyrics by\n",
      "Tail Cluster: [['Twiggy', 'Ramirez']]\n",
      "\n",
      "Head Cluster: [['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell']]\n",
      "Relation: producer\n",
      "Tail Cluster: [['Sean', 'Beavan']]\n",
      "\n",
      "Head Cluster: [['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell']]\n",
      "Relation: follows\n",
      "Tail Cluster: [['Antichrist', 'Superstar']]\n",
      "\n",
      "Head Cluster: [['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell']]\n",
      "Relation: part of\n",
      "Tail Cluster: [['Mechanical', 'Animals']]\n",
      "\n",
      "Head Cluster: [['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell']]\n",
      "Relation: director\n",
      "Tail Cluster: [['Matthew', 'Rolston']]\n",
      "\n",
      "Head Cluster: [['Manson'], ['Marilyn', 'Manson'], ['Marilyn', 'Manson'], ['Marilyn', 'Manson'], ['Manson'], ['Manson'], ['Manson']]\n",
      "Relation: notable work\n",
      "Tail Cluster: [['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell']]\n",
      "\n",
      "Head Cluster: [['Manson'], ['Marilyn', 'Manson'], ['Marilyn', 'Manson'], ['Marilyn', 'Manson'], ['Manson'], ['Manson'], ['Manson']]\n",
      "Relation: country of citizenship\n",
      "Tail Cluster: [['American']]\n",
      "\n",
      "Head Cluster: [['Manson'], ['Marilyn', 'Manson'], ['Marilyn', 'Manson'], ['Marilyn', 'Manson'], ['Manson'], ['Manson'], ['Manson']]\n",
      "Relation: notable work\n",
      "Tail Cluster: [['Antichrist', 'Superstar']]\n",
      "\n",
      "Head Cluster: [['Manson'], ['Marilyn', 'Manson'], ['Marilyn', 'Manson'], ['Marilyn', 'Manson'], ['Manson'], ['Manson'], ['Manson']]\n",
      "Relation: notable work\n",
      "Tail Cluster: [['Mechanical', 'Animals']]\n",
      "\n",
      "Head Cluster: [['Sneaker', 'Pimps'], ['Sneaker', 'Pimps'], ['Sneaker', 'Pimps']]\n",
      "Relation: notable work\n",
      "Tail Cluster: [['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell']]\n",
      "\n",
      "Head Cluster: [['Sneaker', 'Pimps'], ['Sneaker', 'Pimps'], ['Sneaker', 'Pimps']]\n",
      "Relation: notable work\n",
      "Tail Cluster: [['Spawn'], ['Spawn']]\n",
      "\n",
      "Head Cluster: [['Sneaker', 'Pimps'], ['Sneaker', 'Pimps'], ['Sneaker', 'Pimps']]\n",
      "Relation: has part\n",
      "Tail Cluster: [['Kelli', 'Ali']]\n",
      "\n",
      "Head Cluster: [['Sneaker', 'Pimps'], ['Sneaker', 'Pimps'], ['Sneaker', 'Pimps']]\n",
      "Relation: notable work\n",
      "Tail Cluster: [['Mechanical', 'Animals']]\n",
      "\n",
      "Head Cluster: [['Spawn'], ['Spawn']]\n",
      "Relation: has part\n",
      "Tail Cluster: [['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell']]\n",
      "\n",
      "Head Cluster: [['Spawn'], ['Spawn']]\n",
      "Relation: publication date\n",
      "Tail Cluster: [['1997'], ['1997']]\n",
      "\n",
      "Head Cluster: [['Spawn'], ['Spawn']]\n",
      "Relation: has part\n",
      "Tail Cluster: [['Matthew', 'Rolston']]\n",
      "\n",
      "Head Cluster: [['Twiggy', 'Ramirez']]\n",
      "Relation: notable work\n",
      "Tail Cluster: [['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell']]\n",
      "\n",
      "Head Cluster: [['Sean', 'Beavan']]\n",
      "Relation: notable work\n",
      "Tail Cluster: [['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell']]\n",
      "\n",
      "Head Cluster: [['John', 'Milton']]\n",
      "Relation: notable work\n",
      "Tail Cluster: [['Paradise', 'Lost']]\n",
      "\n",
      "Head Cluster: [['John', 'Milton']]\n",
      "Relation: has part\n",
      "Tail Cluster: [['Radiohead']]\n",
      "\n",
      "Head Cluster: [['Paradise', 'Lost']]\n",
      "Relation: author\n",
      "Tail Cluster: [['John', 'Milton']]\n",
      "\n",
      "Head Cluster: [['Paradise', 'Lost']]\n",
      "Relation: publication date\n",
      "Tail Cluster: [['1667']]\n",
      "\n",
      "Head Cluster: [['Paradise', 'Lost']]\n",
      "Relation: has part\n",
      "Tail Cluster: [['Karma', 'Police']]\n",
      "\n",
      "Head Cluster: [['Kelli', 'Ali']]\n",
      "Relation: part of\n",
      "Tail Cluster: [['Sneaker', 'Pimps'], ['Sneaker', 'Pimps'], ['Sneaker', 'Pimps']]\n",
      "\n",
      "Head Cluster: [['Antichrist', 'Superstar']]\n",
      "Relation: followed by\n",
      "Tail Cluster: [['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell']]\n",
      "\n",
      "Head Cluster: [['Antichrist', 'Superstar']]\n",
      "Relation: performer\n",
      "Tail Cluster: [['Manson'], ['Marilyn', 'Manson'], ['Marilyn', 'Manson'], ['Marilyn', 'Manson'], ['Manson'], ['Manson'], ['Manson']]\n",
      "\n",
      "Head Cluster: [['Antichrist', 'Superstar']]\n",
      "Relation: publication date\n",
      "Tail Cluster: [['1996']]\n",
      "\n",
      "Head Cluster: [['Antichrist', 'Superstar']]\n",
      "Relation: followed by\n",
      "Tail Cluster: [['Mechanical', 'Animals']]\n",
      "\n",
      "Head Cluster: [['Mechanical', 'Animals']]\n",
      "Relation: has part\n",
      "Tail Cluster: [['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell']]\n",
      "\n",
      "Head Cluster: [['Mechanical', 'Animals']]\n",
      "Relation: performer\n",
      "Tail Cluster: [['Manson'], ['Marilyn', 'Manson'], ['Marilyn', 'Manson'], ['Marilyn', 'Manson'], ['Manson'], ['Manson'], ['Manson']]\n",
      "\n",
      "Head Cluster: [['Mechanical', 'Animals']]\n",
      "Relation: producer\n",
      "Tail Cluster: [['Manson'], ['Marilyn', 'Manson'], ['Marilyn', 'Manson'], ['Marilyn', 'Manson'], ['Manson'], ['Manson'], ['Manson']]\n",
      "\n",
      "Head Cluster: [['Mechanical', 'Animals']]\n",
      "Relation: follows\n",
      "Tail Cluster: [['Antichrist', 'Superstar']]\n",
      "\n",
      "Head Cluster: [['Mechanical', 'Animals']]\n",
      "Relation: publication date\n",
      "Tail Cluster: [['1998']]\n",
      "\n",
      "Head Cluster: [['Jonathan', 'Glazer'], ['Glazer']]\n",
      "Relation: notable work\n",
      "Tail Cluster: [['Karma', 'Police']]\n",
      "\n",
      "Head Cluster: [['Radiohead']]\n",
      "Relation: notable work\n",
      "Tail Cluster: [['Karma', 'Police']]\n",
      "\n",
      "Head Cluster: [['Karma', 'Police']]\n",
      "Relation: publication date\n",
      "Tail Cluster: [['1997'], ['1997']]\n",
      "\n",
      "Head Cluster: [['Karma', 'Police']]\n",
      "Relation: director\n",
      "Tail Cluster: [['Jonathan', 'Glazer'], ['Glazer']]\n",
      "\n",
      "Head Cluster: [['Karma', 'Police']]\n",
      "Relation: performer\n",
      "Tail Cluster: [['Radiohead']]\n",
      "\n",
      "Head Cluster: [['Matthew', 'Rolston']]\n",
      "Relation: notable work\n",
      "Tail Cluster: [['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell'], ['Long', 'Hard', 'Road', 'Out', 'of', 'Hell']]\n"
     ]
    }
   ],
   "source": [
    "for cluster, cluster_relations_list, tokenized_text in zip(clusters_batch, cluster_relations_batch, batch_tokenized_text):\n",
    "    for cluster_relations in cluster_relations_list:\n",
    "        print()\n",
    "        cluster_h_idx = cluster_relations['h_idx']\n",
    "        cluster_t_idx = cluster_relations['t_idx']\n",
    "        head_cluster = [tokenized_text[s:e] for s, e in cluster[cluster_h_idx]]\n",
    "        print(f\"Head Cluster: {head_cluster}\")\n",
    "        print(f\"Relation: {cluster_relations['r']}\")\n",
    "        tail_cluster = [tokenized_text[s:e] for s, e in cluster[cluster_t_idx]]\n",
    "        print(f\"Tail Cluster: {tail_cluster}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Coreference in the Wild\n",
    "Give a real news story"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "09/18/2024 12:39:32 - INFO - \t missing_keys: []\n",
      "09/18/2024 12:39:32 - INFO - \t unexpected_keys: []\n",
      "09/18/2024 12:39:32 - INFO - \t mismatched_keys: []\n",
      "09/18/2024 12:39:32 - INFO - \t error_msgs: []\n",
      "09/18/2024 12:39:32 - INFO - \t Model Parameters: 90.5M, Transformer: 82.1M, Coref head: 8.4M\n",
      "09/18/2024 12:39:32 - INFO - \t Number of eval relation types per instance: [20]\n",
      "09/18/2024 12:39:32 - WARNING - \t Token length 1178 is longer than max length 512. Truncating.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b93e4e4ca991494ab4dae9357012933d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/1 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "len of (cluster, cluster_relations_list, tokenized_text): (79, 59, 1178)\n",
      "\n",
      "Head Cluster: [['Nysa'], ['Nysa'], ['Nysa']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 0\n",
      "Tail Cluster: [['BBC', 'News'], ['BBC', 'iPlayer'], ['iPlayer']]\n",
      "\n",
      "Head Cluster: [['Budapest'], ['Warsaw'], ['Budapest'], ['Hungarian'], ['Budapest'], ['Hungarian'], ['Budapest'], ['Vienna']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 2\n",
      "Tail Cluster: [['Hungary'], ['Hungary']]\n",
      "\n",
      "Head Cluster: [['Budapest'], ['Warsaw'], ['Budapest'], ['Hungarian'], ['Budapest'], ['Hungarian'], ['Budapest'], ['Vienna']]\n",
      "Relation: mayor issues warning\n",
      "on cluster idx: 2\n",
      "Tail Cluster: [['Viktor', 'Orban']]\n",
      "\n",
      "Head Cluster: [['Budapest'], ['Warsaw'], ['Budapest'], ['Hungarian'], ['Budapest'], ['Hungarian'], ['Budapest'], ['Vienna']]\n",
      "Relation: prime minister declares disaster\n",
      "on cluster idx: 2\n",
      "Tail Cluster: [['Viktor', 'Orban']]\n",
      "\n",
      "Head Cluster: [['Kordian', 'Kolbiarz'], ['Kolbiarz']]\n",
      "Relation: mayor issues warning\n",
      "on cluster idx: 6\n",
      "Tail Cluster: [['Nysa'], ['Nysa'], ['Nysa']]\n",
      "\n",
      "Head Cluster: [['Danube'], ['the', 'Danube', 'River'], ['Danube']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 7\n",
      "Tail Cluster: [['Budapest'], ['Warsaw'], ['Budapest'], ['Hungarian'], ['Budapest'], ['Hungarian'], ['Budapest'], ['Vienna']]\n",
      "\n",
      "Head Cluster: [['Danube'], ['the', 'Danube', 'River'], ['Danube']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 7\n",
      "Tail Cluster: [['Hungary'], ['Hungary']]\n",
      "\n",
      "Head Cluster: [['Danube'], ['the', 'Danube', 'River'], ['Danube']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 7\n",
      "Tail Cluster: [['EU']]\n",
      "\n",
      "Head Cluster: [['Danube'], ['the', 'Danube', 'River'], ['Danube']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 7\n",
      "Tail Cluster: [['Slovakia']]\n",
      "\n",
      "Head Cluster: [['Danube'], ['the', 'Danube', 'River'], ['Danube']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 7\n",
      "Tail Cluster: [['Bratislava']]\n",
      "\n",
      "Head Cluster: [['BBC', 'News'], ['BBC', 'iPlayer'], ['iPlayer']]\n",
      "Relation: prime minister declares disaster\n",
      "on cluster idx: 9\n",
      "Tail Cluster: [['Viktor', 'Orban']]\n",
      "\n",
      "Head Cluster: [['Hungary'], ['Hungary']]\n",
      "Relation: prime minister declares disaster\n",
      "on cluster idx: 11\n",
      "Tail Cluster: [['Viktor', 'Orban']]\n",
      "\n",
      "Head Cluster: [['Donald', 'Tusk']]\n",
      "Relation: government allocates funds\n",
      "on cluster idx: 29\n",
      "Tail Cluster: [['Polish']]\n",
      "\n",
      "Head Cluster: [['Slovakia']]\n",
      "Relation: city experiences flood damage\n",
      "on cluster idx: 33\n",
      "Tail Cluster: [['Bratislava']]\n",
      "\n",
      "Head Cluster: [['Old', 'Town']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 34\n",
      "Tail Cluster: [['Slovakia']]\n",
      "\n",
      "Head Cluster: [['Old', 'Town']]\n",
      "Relation: city experiences flood damage\n",
      "on cluster idx: 34\n",
      "Tail Cluster: [['Bratislava']]\n",
      "\n",
      "Head Cluster: [['Bratislava']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 35\n",
      "Tail Cluster: [['Slovakia']]\n",
      "\n",
      "Head Cluster: [['Rudawa']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 39\n",
      "Tail Cluster: [['Poland'], ['Poland'], ['Poland'], ['Poland'], ['Poland'], ['Poland'], ['Poland'], ['Poland']]\n",
      "\n",
      "Head Cluster: [['Rudawa']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 39\n",
      "Tail Cluster: [['Europe'], ['Europe'], ['Europe'], ['Europe'], ['Europe']]\n",
      "\n",
      "Head Cluster: [['Rudawa']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 39\n",
      "Tail Cluster: [['Polish'], ['Polish']]\n",
      "\n",
      "Head Cluster: [['Rudawa']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 39\n",
      "Tail Cluster: [['Monday'], ['Monday'], ['1950'], ['Monday', 'morning'], ['Monday'], ['morning'], ['Monday']]\n",
      "\n",
      "Head Cluster: [['Rudawa']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 39\n",
      "Tail Cluster: [['Hungary'], ['Hungary']]\n",
      "\n",
      "Head Cluster: [['Rudawa']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 39\n",
      "Tail Cluster: [['the', 'Czech', 'Republic'], ['the', 'Czech', 'Republic'], ['Czech'], ['the', 'Czech', 'Republic']]\n",
      "\n",
      "Head Cluster: [['Rudawa']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 39\n",
      "Tail Cluster: [['the', 'coming', 'days'], ['four', 'days'], ['less', 'than', 'two', 'weeks'], ['days']]\n",
      "\n",
      "Head Cluster: [['Rudawa']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 39\n",
      "Tail Cluster: [['Austria'], ['Austria'], ['Austria']]\n",
      "\n",
      "Head Cluster: [['Rudawa']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 39\n",
      "Tail Cluster: [['Romania'], ['Romania']]\n",
      "\n",
      "Head Cluster: [['Rudawa']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 39\n",
      "Tail Cluster: [['later', 'this', 'week'], ['mid', '-', 'week']]\n",
      "\n",
      "Head Cluster: [['Rudawa']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 39\n",
      "Tail Cluster: [['Storm', 'Boris'], ['Storm', 'Boris'], ['Storm', 'Boris']]\n",
      "\n",
      "Head Cluster: [['Rudawa']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 39\n",
      "Tail Cluster: [['Jesenik']]\n",
      "\n",
      "Head Cluster: [['Rudawa']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 39\n",
      "Tail Cluster: [['473', 'mm']]\n",
      "\n",
      "Head Cluster: [['Rudawa']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 39\n",
      "Tail Cluster: [['19']]\n",
      "\n",
      "Head Cluster: [['Rudawa']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 39\n",
      "Tail Cluster: [['Thursday']]\n",
      "\n",
      "Head Cluster: [['Rudawa']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 39\n",
      "Tail Cluster: [['monthly']]\n",
      "\n",
      "Head Cluster: [['Rudawa']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 39\n",
      "Tail Cluster: [['Austrian']]\n",
      "\n",
      "Head Cluster: [['Rudawa']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 39\n",
      "Tail Cluster: [['St', 'Polten']]\n",
      "\n",
      "Head Cluster: [['Rudawa']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 39\n",
      "Tail Cluster: [['Karl', 'Nehammer']]\n",
      "\n",
      "Head Cluster: [['Rudawa']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 39\n",
      "Tail Cluster: [['Climate', 'Ministry']]\n",
      "\n",
      "Head Cluster: [['Rudawa']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 39\n",
      "Tail Cluster: [['300']]\n",
      "\n",
      "Head Cluster: [['Rudawa']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 39\n",
      "Tail Cluster: [['253']]\n",
      "\n",
      "Head Cluster: [['Rudawa']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 39\n",
      "Tail Cluster: [['29', 'September']]\n",
      "\n",
      "Head Cluster: [['Rudawa']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 39\n",
      "Tail Cluster: [['Emil', 'Dragomir']]\n",
      "\n",
      "Head Cluster: [['Rudawa']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 39\n",
      "Tail Cluster: [['Slobozia', 'Conachi']]\n",
      "\n",
      "Head Cluster: [['Rudawa']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 39\n",
      "Tail Cluster: [['Thousands']]\n",
      "\n",
      "Head Cluster: [['Rudawa']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 39\n",
      "Tail Cluster: [['Nysa']]\n",
      "\n",
      "Head Cluster: [['Rudawa']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 39\n",
      "Tail Cluster: [['Paczków']]\n",
      "\n",
      "Head Cluster: [['Rudawa']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 39\n",
      "Tail Cluster: [['Klodzko', 'city']]\n",
      "\n",
      "Head Cluster: [['Rudawa']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 39\n",
      "Tail Cluster: [['Michal', 'Piszko']]\n",
      "\n",
      "Head Cluster: [['Rudawa']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 39\n",
      "Tail Cluster: [['Polish']]\n",
      "\n",
      "Head Cluster: [['Rudawa']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 39\n",
      "Tail Cluster: [['Sunday']]\n",
      "\n",
      "Head Cluster: [['Rudawa']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 39\n",
      "Tail Cluster: [['Tuesday']]\n",
      "\n",
      "Head Cluster: [['Rudawa']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 39\n",
      "Tail Cluster: [['south', '-', 'east']]\n",
      "\n",
      "Head Cluster: [['Rudawa']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 39\n",
      "Tail Cluster: [['Germany']]\n",
      "\n",
      "Head Cluster: [['Rudawa']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 39\n",
      "Tail Cluster: [['100', 'mm']]\n",
      "\n",
      "Head Cluster: [['Rudawa']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 39\n",
      "Tail Cluster: [['Italy']]\n",
      "\n",
      "Head Cluster: [['Rudawa']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 39\n",
      "Tail Cluster: [['100', '-', '150']]\n",
      "\n",
      "Head Cluster: [['Rudawa']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 39\n",
      "Tail Cluster: [['Arctic']]\n",
      "\n",
      "Head Cluster: [['Rudawa']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 39\n",
      "Tail Cluster: [['Mediterranean']]\n",
      "\n",
      "Head Cluster: [['Rudawa']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 39\n",
      "Tail Cluster: [['one']]\n",
      "\n",
      "Head Cluster: [['Rudawa']]\n",
      "Relation: country affected by floods\n",
      "on cluster idx: 39\n",
      "Tail Cluster: [['about', '7', '%']]\n"
     ]
    }
   ],
   "source": [
    "# Real-world example\n",
    "import spacy\n",
    "from glirel.modules.utils import constrain_relations_by_entity_type\n",
    "from glirel.modules.utils import get_coreference_clusters, aggregate_cluster_relations\n",
    "\n",
    "from fastcoref import FCoref\n",
    "fcoref_model = FCoref(device='cpu', enable_progress_bar=False)\n",
    "\n",
    "nlp = spacy.load('en_core_web_sm')\n",
    "\n",
    "\n",
    "text = \"\"\"Polish city urged to evacuate as floods batter central Europe\n",
    "A flooded area in Nysa, Poland\n",
    "Image source,Reuters\n",
    "Image caption,\n",
    "A flooded area in Nysa, Poland\n",
    "\n",
    "Laura Gozzi, Nick Thorpe, Adam Easton and Rob Cameron\n",
    "BBC News\n",
    "Reporting from\n",
    "London, Budapest, Warsaw and Prague\n",
    "Published\n",
    "16 September 2024\n",
    "The mayor of a Polish city has asked all 44,000 residents to evacuate, as widespread flooding continues to batter central Europe.\n",
    "\n",
    "Nysa mayor Kordian Kolbiarz asked people to head for higher ground, citing the risk of an embankment breaching and releasing a cascade of water into the town from a nearby lake.\n",
    "\n",
    "The death toll from the floods that hit over the weekend rose to at least 16 on Monday, with seven confirmed fatalities in Romania. Casualties were also recorded in Austria, the Czech Republic and Poland.\n",
    "\n",
    "Budapest said it would close roads near the river Danube which runs through the Hungarian capital, citing the risk of flooding later this week.\n",
    "\n",
    "“Please evacuate your belongings, yourselves, your loved ones. It is worth getting to the top floor of the building immediately, because the wave may be several metres high. This means that the whole town will be flooded,” Nysa Mayor Kolbiarz wrote.\n",
    "\n",
    "Polish Prime Minister Donald Tusk said one billion zloty (£197m) would be allocated for flood victims in the country, adding that Poland would also apply for EU relief funds. His government has also declared a state of natural disaster.\n",
    "\n",
    "Although conditions have stabilised in some places, others are bracing themselves for more disruption and danger from the floods, sparked by Storm Boris.\n",
    "\n",
    "In Slovakia, the overflowing of the Danube River caused flooding in the Old Town area of the capital, Bratislava, with local media reporting that water levels exceeded 9m (30ft) and were expected to rise further.\n",
    "\n",
    "A dog being lifted by Polish rescuers\n",
    "Image source,Getty Images\n",
    "Image caption,\n",
    "Polish rescuers and soldiers evacuated local residents and their pets in the village of Rudawa, southern Poland\n",
    "\n",
    "Watch on BBC iPlayer (UK Only)\n",
    "From Above - Storm Boris\n",
    "Attribution\n",
    "iPlayer\n",
    "Hungary is bracing itself for floods in the coming days. Warnings are in force along 500km (310 miles) of the Danube.\n",
    "\n",
    "The river is rising by about a metre every 24 hours, with Budapest's mayor offering residents a million sandbags to protect against floodwaters.\n",
    "\n",
    "Some tram lines will not operate, while roads along the river will be closed in the Hungarian capital from Monday evening. Trains between Budapest and Vienna have also been cancelled.\n",
    "\n",
    "Prime Minister Viktor Orban said on X that he had postponed all his international obligations \"due to the extreme weather conditions and the ongoing floods in Hungary\".\n",
    "\n",
    "The highest rainfall totals have been in the Czech Republic. In the north-eastern town of Jesenik, 473mm (19in) of rain has fallen since Thursday morning - five times the average monthly rainfall.\n",
    "\n",
    "The Czech fire service delivered bottles of drinking water to stranded villages, where people were told not to drink water from their taps or their wells as it is likely heavily contaminated.\n",
    "\n",
    "In the Austrian town of St Polten, more rain has fallen in four days than in the whole of the wettest autumn on record, in 1950.\n",
    "\n",
    "Chancellor Karl Nehammer said the armed forces had been deployed to offer assistance to storm-hit regions. Austria's Climate Ministry said €300m (£253m) in recovery funds would be made available.\n",
    "\n",
    "Most parties paused campaigning for the federal elections due in less than two weeks, on 29 September.\n",
    "\n",
    "Villages and town were submerged in eastern Romania. Emil Dragomir, mayor of Slobozia Conachi, told media that the flooding had had a devastating impact.\n",
    "\n",
    "\"If you were here, you would cry instantly, because people are desperate, their whole lives' work is gone, there were people who were left with just the clothes they had on,\" he said.\n",
    "\n",
    "Thousands of people have been evacuted in Poland, including the personnel and patients of a hospital in the town of Nysa. Roads have been badly disrupted and train traffic was suspended in many parts of the country.\n",
    "\n",
    "On Monday morning, the mayor of Paczków in south-west Poland appealed to residents to evacuate after water began overflowing in a nearby reservoir, endangering the town.\n",
    "\n",
    "\n",
    "Media caption,\n",
    "Airborne rescues as Europe hit by floods\n",
    "\n",
    "In other parts of Poland, however, water levels are now falling, according to local officials.\n",
    "\n",
    "The mayor of Klodzko city, Michal Piszko, told Polish media the water had receded and the indications were the worst was now over.\n",
    "\n",
    "Video footage from Monday morning showed that city centre streets which were inundated on Sunday were now water-free, although the footage also revealed the extent of damage done to the buildings.\n",
    "\n",
    "Where will Storm Boris go next?\n",
    "More rain is expected throughout Monday and Tuesday in Austria, the Czech Republic and south-east Germany, where another 100mm could fall.\n",
    "\n",
    "While it may still take days for the flood waters to subside, the weather will improve in central Europe from mid-week with much drier conditions.\n",
    "\n",
    "Storm Boris will, however, now move further south into Italy, where it will reintensify and bring heavy rain. The Emilia-Romagna region is set to be worst hit, with 100-150mm of rain falling.\n",
    "\n",
    "The record rainfall seen in central Europe has been caused by a number of factors, including climate change.\n",
    "\n",
    "Different weather elements came together to create a “perfect storm” in which very cold air from the Arctic met warm air from the Mediterranean.\n",
    "\n",
    "A pattern of atmospheric pressure also meant that Storm Boris was stuck in one place for a long time.\n",
    "\n",
    "Scientists say that a warmer atmosphere holds more moisture, leading to more intense rainfall. Warmer oceans also lead to more evaporation, feeding storm systems.\n",
    "\n",
    "For every 1C rise in the global average temperature, the atmosphere is able to hold about 7% more moisture, external.\"\"\"\n",
    "\n",
    "\n",
    "labels = [\n",
    "    \"SELF\",\n",
    "    # Positive relation labels\n",
    "    \"mayor responsible for evacuation\",\n",
    "    \"flood causes evacuation\",\n",
    "    \"country affected by floods\",\n",
    "    \"storm responsible for floods\",\n",
    "    \"government allocates funds\",\n",
    "    \"river causes flood\",\n",
    "    \"mayor issues warning\",\n",
    "    \"rescue teams assist residents\",\n",
    "    \"prime minister declares disaster\",\n",
    "    \"city experiences flood damage\",\n",
    "    \"rescue of people and animals\",\n",
    "    \"flooding disrupts transport\",\n",
    "    \"climate change causes extreme weather\",\n",
    "    \"scientists explain weather patterns\",\n",
    "    \n",
    "    # Negative relation labels\n",
    "    \"mayor responsible for storm\",\n",
    "    \"train causes flood\",\n",
    "    \"flooding creates more jobs\",\n",
    "    \"rescue teams cause damage\",\n",
    "    \"election delayed by storm\"\n",
    "]\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "def predict_and_show(text, labels, coreference=False):\n",
    "    doc = nlp(text)\n",
    "    # print(f\"Text: {text}\")\n",
    "\n",
    "    tokens = [token.text for token in doc]\n",
    "\n",
    "    # NOTE: the end index should be inclusive\n",
    "    ner = [[ent.start, (ent.end - 1), ent.label_, ent.text] for ent in doc.ents]\n",
    "    # print(f\"Entities detected: {ner}\")\n",
    "\n",
    "    labels_and_constraints = None\n",
    "    if isinstance(labels, dict):\n",
    "        labels = labels[\"glirel_labels\"]\n",
    "        labels_and_constraints = labels\n",
    "        labels = list(labels.keys())\n",
    "\n",
    "    relations = model.batch_predict_relations([tokens], [labels], threshold=0.01, ner=[ner], top_k=1)\n",
    "\n",
    "    if isinstance(labels_and_constraints, dict):\n",
    "        print('Constraining relations by entity type')\n",
    "        relations = constrain_relations_by_entity_type(doc.ents, labels_and_constraints, relations)\n",
    "\n",
    "\n",
    "\n",
    "    mention_char_spans = [(m.start_char, m.end_char) for m in doc.ents]\n",
    "    if len(mention_char_spans) == 0:\n",
    "        raise ValueError(\"No entities detected in the text. Please provide entities for coreference resolution.\")\n",
    "    else:\n",
    "        output = fcoref_model.predict(\n",
    "            texts=doc.text, custom_mentions=mention_char_spans\n",
    "        )\n",
    "    char_idx2tok_idx = {(ent.start_char, ent.end_char): (ent.start, ent.end) for ent in doc.ents}\n",
    "\n",
    "    entity_to_cluster_idx_batch = {}\n",
    "    fcoref_output = output.get_clusters(as_strings=False)\n",
    "    # map positions to cluster id\n",
    "    for cluster_id, cluster in enumerate(fcoref_output):\n",
    "        for mention in cluster:\n",
    "            token_pos = char_idx2tok_idx[mention]\n",
    "            entity_to_cluster_idx_batch[(token_pos[0], token_pos[1])] = cluster_id\n",
    "    all_indices = [idx for cluster in fcoref_output for idx in cluster]\n",
    "    # add singletons\n",
    "    cluster_id = len(fcoref_output)\n",
    "    for k, token_pos in char_idx2tok_idx.items():\n",
    "        if k not in all_indices:\n",
    "            entity_to_cluster_idx_batch[(token_pos[0], token_pos[1])] = cluster_id\n",
    "            fcoref_output.append([k])\n",
    "            cluster_id += 1\n",
    "    # sort by start idx of keys\n",
    "    entity_to_cluster_idx_batch = dict(sorted(entity_to_cluster_idx_batch.items(), key=lambda x: x[0][0]))\n",
    "    \n",
    "    clusters_batch = [\n",
    "        [[(char_idx2tok_idx[(s, e)][0], char_idx2tok_idx[(s, e)][1]) for s, e in cluster] for cluster in fcoref_output]\n",
    "    ]\n",
    "\n",
    "\n",
    "    # clusters_batch, entity_to_cluster_idx_batch = get_coreference_clusters(relations)\n",
    "    print()\n",
    "    cluster_relations_batch = aggregate_cluster_relations(entity_to_cluster_idx_batch, relations)\n",
    "    cluster_relations_batch[0]\n",
    "\n",
    "    batch_tokenized_text = [tokens]\n",
    "    \n",
    "    for cluster, cluster_relations_list, tokenized_text in zip(clusters_batch, cluster_relations_batch, batch_tokenized_text):\n",
    "        print(f\"len of (cluster, cluster_relations_list, tokenized_text): {len(cluster), len(cluster_relations_list), len(tokenized_text)}\")\n",
    "        for cluster_relations in cluster_relations_list:\n",
    "            print()\n",
    "            cluster_h_idx = cluster_relations['h_idx']\n",
    "            cluster_t_idx = cluster_relations['t_idx']\n",
    "            head_cluster = [tokenized_text[s:e] for s, e in cluster[cluster_h_idx]]\n",
    "            print(f\"Head Cluster: {head_cluster}\")\n",
    "            print(f\"Relation: {cluster_relations['r']}\")\n",
    "            print(f\"on cluster idx: {cluster_h_idx}\")             \n",
    "            tail_cluster = [tokenized_text[s:e] for s, e in cluster[cluster_t_idx]]\n",
    "            print(f\"Tail Cluster: {tail_cluster}\")\n",
    "\n",
    "\n",
    "\n",
    "predict_and_show(text, labels, coreference=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "gliner_dev",
   "language": "python",
   "name": "python3"
  },
  "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}