{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "fc867b65-d7bf-42a0-b6c8-f1d88b9bcb61",
   "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/HW11_LLM_Coref.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06488105-1f63-471f-b9c6-4236b511f181",
   "metadata": {},
   "source": [
    "# HW11: Coreference with LLMs\n",
    "\n",
    "In this homework, you will experiment with using LLMs for zero-shot or few-shot coreference resolution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05d10d53-04fc-470f-8cc4-d991434465dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "from tqdm import tqdm\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a00a39dc-845b-40a5-9a34-bdc107f585cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install stanza==1.10.1 peft"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38585bdd-89b6-44ef-a18e-4cbeef078739",
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget https://raw.githubusercontent.com/dbamman/anlp25/refs/heads/main/data/1342_pride_and_prejudice_brat.conll -O 1342_pride_and_prejudice_brat.conll\n",
    "!wget https://raw.githubusercontent.com/dbamman/anlp25/refs/heads/main/data/1342_pride_and_prejudice_sample.txt -O 1342_pride_and_prejudice_sample.txt\n",
    "!wget https://raw.githubusercontent.com/dbamman/anlp25/refs/heads/main/11.nlp/coref_utils.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7da7ee66-2a3c-4cf4-b740-dee6849d7cc7",
   "metadata": {},
   "source": [
    "## The coreference resolution task\n",
    "\n",
    "We formulate the coreference resolution task as follows: given an input text, output a sequence of coref chains $(C_1, \\ldots, C_i)$, each of which contains a sequence of coref mentions $C_i = (m_{i1}, \\ldots, m_{ij})$ ordered by start index. Each coref mention is a tuple of the start and end indices $m_{ij} = (\\text{start\\_index}, \\text{end\\_index})$ denoting the span of the mention in the text.\n",
    "\n",
    "To formalize this in the code, we set up the following classes:\n",
    "\n",
    "```python\n",
    "@dataclass\n",
    "class CorefMention:\n",
    "    start_idx: int\n",
    "    end_idx: int\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class CorefChain:\n",
    "    mentions: list[CorefMention] = field(default_factory=list)\n",
    "\n",
    "\n",
    "CorefOutput = list[CorefChain]\n",
    "```\n",
    "\n",
    "To avoid cluttering up the notebook, we provide a utility file with these types and other useful functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5b2e6d9-f825-4dfc-a766-ddafcfcfe451",
   "metadata": {},
   "outputs": [],
   "source": [
    "from coref_utils import CorefMention, CorefChain, CorefOutput"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a2a70a8-303f-47b1-9ad8-14bc9de03c55",
   "metadata": {},
   "source": [
    "## Evaluation data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34a1ca35-243c-4b56-bbe3-b0b1cf4a4053",
   "metadata": {},
   "source": [
    "We will use data from [LitBank](https://github.com/dbamman/litbank), which contains coreference annotations for novels in the public domain. In particular, we will be evaluate the coreference output on approximately 2,000 words of _Bleak House_ by Charles Dickens.\n",
    "\n",
    "Because most systems can't handle coreference on such long texts (and it would also strain the memory usage of the LLM), we do inference on chunks of sentences instead. The `load_conll_data` function takes care of this for us."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6909289-f30b-4705-ae68-de92ae8c92cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "from coref_utils import load_conll_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47fab26b-2706-45dd-8a37-0db5771054f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "text, gold_coref_chains = load_conll_data(\"./1342_pride_and_prejudice_sample.txt\", \"./1342_pride_and_prejudice_brat.conll\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ea19b4f-4f3c-4185-b036-29572653ba8d",
   "metadata": {},
   "source": [
    "Let's examine the first text chunk as an example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6be1f905-3962-409c-a816-943a47323abd",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(text[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6ff30bc-a542-43cd-b37a-00c3981eaac9",
   "metadata": {},
   "outputs": [],
   "source": [
    "for chain in gold_coref_chains[0][:10]:\n",
    "    print(\"===\")\n",
    "    for mention in chain.mentions:\n",
    "        print(f'\"{text[0][mention.start_idx:mention.end_idx]}\"')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0002528b-72b8-4b69-bee2-b3a596c377c5",
   "metadata": {},
   "source": [
    "## Baseline with `stanza`\n",
    "\n",
    "Here, we will evaluate a baseline with the `stanza` coreference resolution system.\n",
    "\n",
    "The $B^3$ precision and recall metrics are defined at the entity mention level. We follow previous works in evaluating mention detection and coreference separately:\n",
    "- We calculate span F1 to measure the performance of mention detection\n",
    "- We calculate the $B^3$ metrics on only the mentions that are shared between the gold and system outputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7210bea6-0b75-45b4-8727-c2926da622ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "from coref_utils import evaluate\n",
    "\n",
    "# Usage:\n",
    "# evaluate(gold, pred)\n",
    "# Returns an EvaluationOutput object that contains B3 precision and recall, as well as span precision/recall/f1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8eb2572-7765-46fc-acaf-8019a87f43e9",
   "metadata": {},
   "source": [
    "Now let's implement and run the Stanza baseline."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be721b93-5cde-4146-ba91-27a89cd290a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "## STANZA BASELINE\n",
    "import stanza\n",
    "\n",
    "def stanza_baseline(text_chunks: list[str]) -> CorefOutput:\n",
    "    \"\"\"\n",
    "    Run the stanza baseline on the text chunks.\n",
    "    \"\"\"\n",
    "    pipe = stanza.Pipeline(\"en\", processors=\"tokenize,lemma,pos,depparse,coref\")\n",
    "    results = []\n",
    "    for text_chunk in tqdm(text_chunks):\n",
    "        chains = []\n",
    "        doc = pipe(text_chunk)\n",
    "        for coref_chain in doc.coref:\n",
    "            mentions = []\n",
    "            for mention in coref_chain.mentions:\n",
    "                span = doc.sentences[mention.sentence].words[mention.start_word:mention.end_word]\n",
    "                mentions.append(CorefMention(start_idx=span[0].start_char, end_idx=span[-1].end_char))\n",
    "            chains.append(CorefChain(mentions=mentions))\n",
    "        results.append(chains)\n",
    "    return results\n",
    "\n",
    "baseline = stanza_baseline(text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd864363-168e-4660-9d0b-ea1baed060d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "evaluate(gold_coref_chains, baseline)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ca57f42-f519-4990-b33f-eebad4122050",
   "metadata": {},
   "source": [
    "## LLM prompting\n",
    "\n",
    "How can we prompt an LLM to do this task? How can we post-process the output into the structure that we desire?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "419a33c3-ec3b-4266-b041-c8be6906d7ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "# use the 4B model\n",
    "\n",
    "model = AutoModelForCausalLM.from_pretrained(\"Qwen/Qwen3-4B\", device_map=\"cuda\", dtype=\"auto\")\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"Qwen/Qwen3-4B\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de9a7f5a-bb4f-47a5-bdaf-a182a7782d40",
   "metadata": {},
   "source": [
    "We use the same code to call the LLM as in previous assignments. Feel free to modify this if you wish."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4423e169-081d-4ed7-ac9d-e1c223d4d9f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def call_llm(prompt, system_prompt=\"You are a helpful assistant.\", generation_config=None):  \n",
    "    if generation_config is None:\n",
    "        generation_config = {\n",
    "            \"max_new_tokens\": 10,\n",
    "            \"temperature\": 0.01\n",
    "        }\n",
    "    messages = [\n",
    "        {\"role\": \"system\", \"content\": system_prompt},\n",
    "        {\"role\": \"user\", \"content\": prompt}\n",
    "    ]\n",
    "    text = tokenizer.apply_chat_template(\n",
    "        messages,\n",
    "        tokenize=False,\n",
    "        add_generation_prompt=True,\n",
    "        enable_thinking=False\n",
    "    )\n",
    "\n",
    "    model_inputs = tokenizer([text], return_tensors=\"pt\").to(model.device)\n",
    "    \n",
    "    # conduct text completion\n",
    "    generated = model.generate(\n",
    "        **model_inputs,\n",
    "        **generation_config\n",
    "    )\n",
    "\n",
    "    # let's break this down:\n",
    "    #                      | we take the element of the batch (our batch size is 1)\n",
    "    #                      |  |-----------------------------| skip our original input\n",
    "    output_ids = generated[0][len(model_inputs.input_ids[0]):].tolist()\n",
    "\n",
    "    # decode into token space\n",
    "    return tokenizer.decode(output_ids, skip_special_tokens=True).strip(\"\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80aa6e3b-2d7c-4a85-b9ce-302b40ef080f",
   "metadata": {},
   "source": [
    "### Question 1\n",
    "\n",
    "What is the simplest solution you can think of to do this task with an LLM? **In a few sentences**, explain your approach. Then, **implement this method** and evaluate its performance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "503163f9-e50c-4168-af49-1cb27807a749",
   "metadata": {},
   "outputs": [],
   "source": [
    "def llm_baseline(text_chunks) -> CorefOutput:\n",
    "    \"\"\"For you to implement! Return a list of CorefChain objects\"\"\"\n",
    "    return [[CorefChain([])] for _ in text_chunks]\n",
    "\n",
    "output = llm_baseline(text)\n",
    "evaluate(gold_coref_chains, output)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ad59194-8020-403f-ad5b-bdde28db7a07",
   "metadata": {},
   "source": [
    "### Question 2\n",
    "\n",
    "What are some mistakes that you notice the system making? What are some improvements you can make? Name at least **two** improvements you want to test. Then, **implement these** and report how they affect the performance of the system. Feel free to experiment with more!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f01054e3-9867-4354-a8a5-4a9223ce0298",
   "metadata": {},
   "outputs": [],
   "source": [
    "def llm_improved(text_chunks) -> CorefOutput:\n",
    "    \"\"\"For you to implement! Return a list of CorefChain objects\"\"\"\n",
    "    return [[CorefChain([])] for _ in text_chunks]\n",
    "\n",
    "output = llm_improved(text)\n",
    "evaluate(gold_coref_chains, output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "942bb911-8e02-4250-a11c-04099ca9ba98",
   "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
}