{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "24429554-b6c7-4c40-9f77-d86465ab67c4",
   "metadata": {},
   "source": [
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/dbamman/anlp25/blob/main/11.nlp/Stanza%20Coref.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ead7214a-7abe-46d9-b947-3a8e5e135785",
   "metadata": {},
   "source": [
    "In this notebook, we'll explore coreference resolution using the Stanza library, which also does dependency parsing and other core NLP tasks. You can also visualize Stanza output at [http://stanza.run](http://stanza.run).\n",
    "\n",
    "The coreference resolution that Stanza uses is an implementation of [Conjunction-Aware Word-level Coreference Resolution (D'Oosterlinck et al, 2023)](https://arxiv.org/abs/2310.06165). You can find more information in [the documentation](https://stanfordnlp.github.io/stanza/coref.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "718233c1-166e-4c88-9127-ace1bd67186b",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install stanza==1.10.1\n",
    "!pip install peft"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f59c7098-8a2b-437e-b613-04064cd9fc94",
   "metadata": {},
   "outputs": [],
   "source": [
    "import stanza\n",
    "\n",
    "# load the stanza NLP pipeline with coref and dependency parsing\n",
    "pipe = stanza.Pipeline(\"en\", processors=\"tokenize,lemma,pos,depparse,coref\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02c81f3d-ccf6-4f86-8185-cd733f052adb",
   "metadata": {},
   "source": [
    "Now let's run a sentence through Stanza.\n",
    "\n",
    "> Although he was very busy with his work, Peter had had enough of it. He and his wife decided they needed a holiday. They travelled to Spain because they loved the country very much."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "492d8fbd-d6a5-4ee8-941c-1d9d11a32297",
   "metadata": {},
   "outputs": [],
   "source": [
    "doc = pipe(\"Although he was very busy with his work, Peter had had enough of it. He and his wife decided they needed a holiday. They travelled to Spain because they loved the country very much.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "755beab6-8d6b-4016-a2e4-9f9e06391e55",
   "metadata": {},
   "source": [
    "The coreference output is attached the the document as `.coref`.\n",
    "\n",
    "The output is a list of `CorefChain` objects. Each `CorefChain` object contains a list of `CorefMention` objects. You can view the properties that you can access through each of these classes by looking [in the source code](https://github.com/stanfordnlp/stanza/blob/main/stanza/models/coref/coref_chain.py).\n",
    "\n",
    "```py\n",
    "class CorefMention:\n",
    "    def __init__(self, sentence, start_word, end_word):\n",
    "        self.sentence = sentence\n",
    "        self.start_word = start_word\n",
    "        self.end_word = end_word\n",
    "\n",
    "class CorefChain:\n",
    "    def __init__(self, index, mentions, representative_text, representative_index):\n",
    "        self.index = index\n",
    "        self.mentions = mentions\n",
    "        self.representative_text = representative_text\n",
    "        self.representative_index = representative_index\n",
    "\n",
    "class CorefAttachment:\n",
    "    def __init__(self, chain, is_start, is_end, is_representative):\n",
    "        self.chain = chain\n",
    "        self.is_start = is_start\n",
    "        self.is_end = is_end\n",
    "        self.is_representative = is_representative\n",
    "\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f35f91a-ced4-4afa-b012-afa82e5bc638",
   "metadata": {},
   "source": [
    "Here, we will print out the spans for each of the coref chains."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21533ea2-ad48-4452-a6e6-1024f2b9d155",
   "metadata": {},
   "outputs": [],
   "source": [
    "for coref_chain in doc.coref:\n",
    "    print(f\"Representative Text: {coref_chain.representative_text}\")\n",
    "    for mention in coref_chain.mentions:\n",
    "        span = doc.sentences[mention.sentence].words[mention.start_word:mention.end_word]\n",
    "        span_text = \" \".join([word.text for word in span])\n",
    "        print(f\"{span_text}\\t{span[0].start_char}\\t{span[-1].end_char}\")\n",
    "\n",
    "    print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c0a7a47-f2c8-4da8-99ab-ff2a8166404d",
   "metadata": {},
   "source": [
    "Each coref mention consists of one or more words. We can try to get the root of the span by checking the dependency heads. In most cases, all the words except one in the coref span should point to other words inside the span; the word that depends on a word outside the span is the root.\n",
    "\n",
    "The syntactic relation of the entire mention to the rest of the sentence is best captured by this root."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f439f1cb-24c9-41f2-9460-afcb44840443",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_span_root(span: list[stanza.models.common.doc.Word]):\n",
    "    # if there's only one word, it is the root\n",
    "    if len(span) == 1:\n",
    "        return span[0]\n",
    "    # find the words whose heads that exceed the span\n",
    "    span_min = span[0].id\n",
    "    span_max = span[-1].id\n",
    "    roots = [word for word in span if (word.head) < span_min or (word.head) > span_max]\n",
    "    assert len(roots) > 0\n",
    "    # we just return the first one if there is more than one\n",
    "    return roots[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62bb192e-589a-4523-a158-c4176e51f050",
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_coref_chains(doc):\n",
    "    for coref_chain in doc.coref:\n",
    "        print(f\"Representative Text: {coref_chain.representative_text}\")\n",
    "        for mention in coref_chain.mentions:\n",
    "            span = doc.sentences[mention.sentence].words[mention.start_word:mention.end_word]\n",
    "            span_text = \" \".join([word.text for word in span])\n",
    "            root = get_span_root(span)\n",
    "            print(f\"{span_text}\\t{doc.sentences[mention.sentence].words[root.head - 1].text}\\t{root.deprel}\")\n",
    "        print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b5394287-79d5-4d2b-a0bb-a61ea32f4727",
   "metadata": {},
   "outputs": [],
   "source": [
    "doc2 = pipe(\"The trophy would not fit in the brown suitcase because it was too big\")\n",
    "print_coref_chains(doc2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e4d39dc-6aad-4ada-9ec7-328cd47f7dd0",
   "metadata": {},
   "outputs": [],
   "source": [
    "doc3 = pipe(\"The town councilors refused to give the demonstrators a permit because they feared violence.\")\n",
    "print_coref_chains(doc3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8b3384f-927a-4096-990c-bb7e8412b009",
   "metadata": {},
   "outputs": [],
   "source": [
    "doc4 = pipe(\"The town councilors refused to give the demonstrators a permit because they advocated violence.\")\n",
    "print_coref_chains(doc4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8f69c26-5f02-430c-8707-1d72a62d4969",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}