{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "75f8568c-9762-49a6-b8d1-bc0d7f7937ab",
   "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/HW12_Relation_Extraction.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a48bb6f-c5b8-440b-acac-69c243de17c0",
   "metadata": {},
   "source": [
    "# Relation extraction with LLMs\n",
    "\n",
    "In this homework, we will explore the challenges and affordances of using LLMs for relation extraction, and how we can evaluate LLM RE systems."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f4a14b5-7c22-43f3-b450-9416369930e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "\n",
    "from tqdm import tqdm\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18f12f96-fad6-444a-952a-3a606c53a08b",
   "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": "code",
   "execution_count": null,
   "id": "d9ec965b-389b-4d34-a73b-89fd69c7745e",
   "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\": 500,\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": "55f729c9-f1ab-4e97-a26a-bc2b100da989",
   "metadata": {},
   "source": [
    "## Load data\n",
    "\n",
    "We will be using the relationship triples you extracted during the in-class activity on Tuesday. These have been preprocessed to match each triple to a paragraph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3486d34-5b2b-43c1-b94c-61dc840cefac",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e68a724b-1ed6-46b0-9e2f-f2239880c008",
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget https://raw.githubusercontent.com/dbamman/anlp25/refs/heads/main/11.nlp/movie_relations.json -O movie_relations.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2d1abd9-71cd-4fbc-8e7c-4f5fa1b94414",
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_data(path):\n",
    "    df = pd.read_json(path)\n",
    "    df = df.sample(50, random_state=42)\n",
    "    texts = df.paragraph_text.to_list()\n",
    "    labels = df.triples.to_list()\n",
    "    return texts, labels\n",
    "\n",
    "texts, triples = read_data(\"./movie_relations.json\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0aa4f86b-9c2d-470e-9451-60b6b2ca1904",
   "metadata": {},
   "source": [
    "## Setting up the LLM"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d726198-d1cd-4bf6-a222-4187f9782ebb",
   "metadata": {},
   "source": [
    "**Question 1:** Come up with **at least two different prompts or prompting methods** to perform relationship extraction using LLMs based on the relation categories we defined in the lab activity on Tuesday. Your output should be relationship triples. To enforce this, we create a `RelationTriple` wrapper class for your output.\n",
    "\n",
    "Here's an example from the dataset:\n",
    "\n",
    "Input:\n",
    "```\n",
    "Aboard the space station, Peiqiang discovers that MOSS, the station's computer commander, has decided to abandon Earth and repurpose the station as an interstellar ark to seed a new planet with Earth's biosphere. Breaking out of forced hibernation, he is joined by fellow Russian cosmonaut Maxim Makarov, whom MOSS awakens to stop Liu. While spacewalking, Makarov is killed by the spacecraft's automated security measures. Liu enters the control room, but his attempts to override the evacuation procedures are revoked. Qi's group arrives at the Sulawesi Supply Depot to find that, while most engines around the planet have been restored, the combined thrust is insufficient to divert Earth's trajectory as it approaches Jupiter's Roche limit. MOSS broadcasts a final message to the world, but Peiqiang refuses to follow the computer's instructions.\n",
    "```\n",
    "\n",
    "Output (you will want to return a list of `RelationTriple`s):\n",
    "```\n",
    "<Liu Peiqiang,business,Maxim Makarov>\n",
    "<Liu Peiqiang,nemeses,MOSS>\n",
    "<Maxim Makarov,nemeses,MOSS>\n",
    "```\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8e7155f-cbcb-4121-8a43-55a29f81c966",
   "metadata": {},
   "outputs": [],
   "source": [
    "relations = [\n",
    "    \"family\",\n",
    "    \"nemeses\",\n",
    "    \"romantic\",\n",
    "    \"friends\",\n",
    "    \"business\"\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83098025-f2ba-4b7f-aa18-2f3b7a71ee51",
   "metadata": {},
   "outputs": [],
   "source": [
    "class RelationTriple():\n",
    "    def __init__(self, head, tail, relation):\n",
    "        self.head = head\n",
    "        self.tail = tail\n",
    "        self.relation = relation\n",
    "\n",
    "    @classmethod\n",
    "    def from_triple(cls, triple: str):\n",
    "        parts = triple.strip(\"<>\").split(\",\")\n",
    "        parts = [part.strip() for part in parts]\n",
    "        if len(parts) != 3:\n",
    "            raise ValueError(f\"triple {triple} is malformed\")\n",
    "        head, relation, tail = parts\n",
    "        if relation not in relations:\n",
    "            raise ValueError(f\"triple {triple} has unsupported relation {relation}\")\n",
    "        return cls(head, tail, relation)\n",
    "\n",
    "    @classmethod\n",
    "    def validate(cls, triple: str):\n",
    "        parts = triple.strip(\"<>\").split(\",\")\n",
    "        parts = [part.strip() for part in parts]\n",
    "        if len(parts) != 3:\n",
    "            return False\n",
    "        head, relation, tail = parts\n",
    "        if relation not in relations:\n",
    "            return False\n",
    "        return True\n",
    "\n",
    "    def __str__(self):\n",
    "        return f\"<{self.head},{self.relation},{self.tail}>\"\n",
    "\n",
    "    def __repr__(self):\n",
    "        return f\"<{self.head},{self.relation},{self.tail}>\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3fee361-196b-4aaf-95b8-a7837d8bdc18",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_relations_one(text: str) -> list[RelationTriple]:\n",
    "    # TODO: fill me in!\n",
    "    pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "513dee85-4f85-4281-815d-64da251451a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_relations_two(text: str) -> list[RelationTriple]:\n",
    "    # TODO: fill me in!\n",
    "    pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c5076ae-5bc2-4845-965d-10b26b187171",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_on_data(fn, texts) -> list[list[RelationTriple]]:\n",
    "    return [\n",
    "        fn(text) for text in tqdm(texts)\n",
    "    ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8df9f244-6f9a-4ab4-909a-b03be5025c65",
   "metadata": {},
   "outputs": [],
   "source": [
    "first_outputs = run_on_data(generate_relations_one, texts)\n",
    "second_outputs = run_on_data(generate_relations_two, texts)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec461765-7ca8-4abe-a3d2-22f77c055853",
   "metadata": {},
   "source": [
    "## Evaluating output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fdb7b33c-edd4-4bad-b966-46465d4e2aa0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_gold_labels(labels: list[list[str]]):\n",
    "    return [\n",
    "        [RelationTriple.from_triple(triple) for triple in paragraph if RelationTriple.validate(triple)] for paragraph in labels\n",
    "    ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd02c702-ac62-4f59-b9de-0f88d4f9238d",
   "metadata": {},
   "outputs": [],
   "source": [
    "gold_labels = get_gold_labels(triples)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bdbd0af3-4d9f-431b-8644-eb4d5b9c5a3b",
   "metadata": {},
   "source": [
    "### Strict matching\n",
    "\n",
    "**Question 2:** **Implement the following functions** in order to calculate the precision / recall / F1 of your model output on both prompts.\n",
    "\n",
    "`get_confusion_matrix` should return a `ConfusionMatrix` containing the number of false/true positives/negatives calculated for the gold and predicted labels for one paragraph. It should use `correct_fn` to compute whether two triples match.\n",
    "\n",
    "You shoudl calculate `precision`, `recall`, and `f1` over the gold and predicted labels for the entire list of paragraphs by adding up the confusion matrices for each one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "271847f2-3860-4b16-9ab7-167500b88c39",
   "metadata": {},
   "outputs": [],
   "source": [
    "def strict_correct_fn(gold: RelationTriple, pred: RelationTriple) -> bool:\n",
    "    return gold.head == pred.head and gold.relation == pred.relation and gold.tail == pred.tail"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d89c5c79-f9b8-4fca-b10d-e21cdc0aa4ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ConfusionMatrix():\n",
    "    def __init__(self, tp=0, fp=0, tn=0, fn=0):\n",
    "        self.tp = tp\n",
    "        self.fp = fp\n",
    "        self.tn = tn\n",
    "        self.fn = fn\n",
    "\n",
    "    def __add__(self, other):\n",
    "        return ConfusionMatrix(\n",
    "            self.tp + other.tp,\n",
    "            self.fp + other.fp,\n",
    "            self.tn + other.tn,\n",
    "            self.fn + other.fn,\n",
    "        )\n",
    "\n",
    "    def to_numpy(self):\n",
    "        return np.array([self.tp, self.fp, self.fn, self.tn])\n",
    "\n",
    "def get_confusion_matrix(gold: list[RelationTriple], pred: list[RelationTriple], correct_fn) -> ConfusionMatrix:\n",
    "    pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90a3ab9c-4cd5-4b4f-a5af-8059bb143767",
   "metadata": {},
   "outputs": [],
   "source": [
    "def precision(confusion_matrix: ConfusionMatrix) -> float:\n",
    "    pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dbfdc94c-4e66-40a7-a7e6-00654332b2f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def recall(confusion_matrix: ConfusionMatrix) -> float:\n",
    "    pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "274446f8-3c00-48b3-aad7-bede12af67c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def f1(confusion_matrix: ConfusionMatrix) -> float:\n",
    "    pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b16719a-358f-47af-85ee-493609e238bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def summed_confusion_matrix(gold: list[list[RelationTriple]], pred: list[list[RelationTriple]], correct_fn) -> ConfusionMatrix:\n",
    "    return sum([\n",
    "        get_confusion_matrix(g, p, correct_fn) for g, p in zip(gold, pred)\n",
    "    ])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26d7cabb-e0e1-4b63-a270-9496aa74b49c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "c01da0b7-547c-4ac3-8227-df4391970745",
   "metadata": {},
   "source": [
    "### LLM-as-judge\n",
    "\n",
    "**Question 3:** Use the LLM to adjudicate the output by **implementing the `llm_correct_fn`**, then computing new precision, recall, and F1 scores."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76e51bc9-ffa0-42c6-8e9c-226e425b6f96",
   "metadata": {},
   "outputs": [],
   "source": [
    "def llm_correct_fn(gold: RelationTriple, pred: RelationTriple) -> bool:\n",
    "    pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "853b19ed-bbb6-4e7f-aa4c-f01ae8cac709",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "fe68bd1c-ea80-4d85-bee7-cad7e27dca44",
   "metadata": {},
   "source": [
    "### Evaluating the evaluation\n",
    "\n",
    "**Question 4:** For each of the evaluation methods (strict matching and LLM-as-judge), sample 10 false positives and 10 false negatives. What proportion of these are incorrectly evaluated? **In a few sentences,** compare the evaluation methods and reflect on the challenges and potential methods for evaluating relationship extraction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53c2c8d3-a226-417b-b01a-e6b0ed261527",
   "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
}
