{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6cd7ac33-83ff-4630-89ed-14bab402ec09",
   "metadata": {},
   "source": [
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/dbamman/anlp25/blob/main/8.clustering/HW8_Word_Sense_Induction.ipynb)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c2a21ba-3b6e-4f8f-bc82-46cb9a5f9b31",
   "metadata": {},
   "source": [
    "# HW8: Unsupervised Word Sense Induction\n",
    "\n",
    "The same _word type_ can have different _senses_, or meanings. For example, \"class\" could refer to a category (\"I have never flown first class.\") or a course (\"ANLP is my favorite class.\") Indeed, with how frequently language changes, a new word sense can come into common usage before dictionaries can be updated with this new information. In this setting, we might be interested in _inducing_ word senses: how can we surface different senses of a word without knowing the definition _a priori_?\n",
    "\n",
    "In this homework, you will be working on this classic NLP task by clustering BERT token embeddings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bacde848-1a2c-4c70-abb1-40ca8bc86b29",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer, BertModel\n",
    "import pandas as pd\n",
    "import torch\n",
    "from torch import nn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90e44f3f-5570-4985-a765-7dd99fc1e96d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Make sure the GPU is available\n",
    "torch.cuda.is_available()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "197f373a-c613-489a-ac1e-22bfb40b8bb6",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(\"bert-base-uncased\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da683e4c-2dc6-48aa-befa-d27579c9daa7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import xml.etree.ElementTree as ET\n",
    "\n",
    "def parse_semeval_file(filepath):\n",
    "    data = ET.parse(filepath)\n",
    "    root = data.getroot()\n",
    "    tag = root.tag\n",
    "    word = tag.strip().split(\".\")[0]\n",
    "\n",
    "    data = []\n",
    "\n",
    "    for sentence in root:\n",
    "        split = sentence.text.lower().split(\" \")\n",
    "        if word not in split:\n",
    "            continue\n",
    "        data.append({\"word\": word, \"sentence\": split})\n",
    "\n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34cc50e2-a4e3-4ba3-a993-d2b2db90f0b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# download data file\n",
    "!wget https://raw.githubusercontent.com/dbamman/anlp25/main/data/wsi/address.n.xml"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55995f1c-6c25-4a61-bd31-ebecaef67049",
   "metadata": {},
   "source": [
    "## Exploring the data\n",
    "\n",
    "### Question 1\n",
    "\n",
    "Load the data for \"address\" (`address.n.xml`). Look through the examples and **identify two sentences that use \"address\" with different senses**. Report each sentence and what definition of \"address\" is being used."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e80d7f5f-aaec-47bb-9338-de99735ddd42",
   "metadata": {},
   "outputs": [],
   "source": [
    "data = parse_semeval_file(\"address.n.xml\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d5098e9-3f21-417d-9ebb-377ce54d897a",
   "metadata": {},
   "source": [
    "## Setting up data helpers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa9c2b7f-0678-41a6-95f1-1cf9e711efba",
   "metadata": {},
   "outputs": [],
   "source": [
    "def tokenize(batch):\n",
    "    # tokenize the words\n",
    "    output = tokenizer(batch[\"sentence\"], is_split_into_words=True)\n",
    "\n",
    "    # find index of first subword token belonging to word of interest\n",
    "    token_indices = []\n",
    "    for i, (sentence, word) in enumerate(zip(batch[\"sentence\"], batch[\"word\"])):\n",
    "        target_id = sentence.index(word)\n",
    "        token_index = None\n",
    "        for token_id, word_id in enumerate(output.word_ids(batch_index=i)):\n",
    "            if word_id == target_id:\n",
    "                token_index = token_id\n",
    "                break\n",
    "        token_indices.append(token_index)\n",
    "\n",
    "    assert not any(x is None for x in token_indices), \"Target token not found in sentence!\"\n",
    "    assert len(token_indices) == len(batch[\"sentence\"]), \"Token indices is the wrong length!\"\n",
    "\n",
    "    output[\"token_indices\"] = token_indices\n",
    "    return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "946b0e12-11a6-4976-88c5-d25277dba347",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.nn.utils.rnn import pad_sequence\n",
    "\n",
    "def collate(items):\n",
    "    # Converts from a list of dicts to a dict of lists\n",
    "    batch = {\n",
    "        k: [item[k] for item in items] for k in items[0]\n",
    "    }\n",
    "    # Tokenizes and pads each batch\n",
    "    outputs = tokenize(batch)\n",
    "    outputs = {\n",
    "        k: pad_sequence([torch.tensor(l) for l in v], batch_first=True, padding_value=0) if k != \"token_indices\" else torch.tensor(v)\n",
    "        for k, v in outputs.items()\n",
    "    }\n",
    "    return outputs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13a9e79f-7805-4e38-be51-fcabff882905",
   "metadata": {},
   "source": [
    "## Running BERT and extracting token representations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a48bc527-b92c-468d-bde4-b2efa115effa",
   "metadata": {},
   "outputs": [],
   "source": [
    "bert_model = BertModel.from_pretrained(\"bert-base-uncased\").to(\"cuda\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30fcb99c-5f9b-4dac-840f-e003e2521444",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import DataLoader\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2925abbe-2403-48d8-966d-733c1eacc25c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_token_embeddings(data, model, batch_size=128):\n",
    "    embeddings = []\n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        inference_loader = DataLoader(data, batch_size=batch_size, shuffle=False, collate_fn=collate)\n",
    "        for batch in tqdm(inference_loader):\n",
    "            output = bert_model(\n",
    "                batch[\"input_ids\"].to(\"cuda\"),\n",
    "                batch[\"attention_mask\"].to(\"cuda\"),\n",
    "                batch[\"token_type_ids\"].to(\"cuda\")\n",
    "            )\n",
    "            # extract the token representation from the last hidden state\n",
    "            batch_reps = output.last_hidden_state[range(len(batch[\"token_indices\"])), batch[\"token_indices\"], :]\n",
    "            embeddings.append(batch_reps.detach().cpu())\n",
    "    embeddings = torch.concat(embeddings, dim=0)\n",
    "    return embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb7e31e8-3303-4c7a-ade7-c0a00b3e0184",
   "metadata": {},
   "outputs": [],
   "source": [
    "embeddings = get_token_embeddings(data, bert_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e44590b-19e6-4299-b541-68c0428286d8",
   "metadata": {},
   "source": [
    "## Clustering with K-Means\n",
    "\n",
    "### Question 2\n",
    "We will begin by clustering with K-Means. **Write the code to cluster the embeddings with $k=5$.** Use `random_state=0` to ensure consistency. Use `diagnose_clustering` to examine the cluster outputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f86b4fde-db99-4c46-8da0-331c17e2ae47",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "def cosine_similarity(one, two):\n",
    "  return np.dot(one,two) / (np.sqrt(np.dot(one,one)) * np.sqrt(np.dot(two,two)))\n",
    "\n",
    "\n",
    "def diagnose_clustering(clustering):\n",
    "    # For each cluster, print out the n documents closest to the cluster center\n",
    "    # To support agglomerative clustering, we calculate the cluster center post-hoc\n",
    "    clusters = {}\n",
    "    for idx, label in enumerate(clustering.labels_):\n",
    "        if label not in clusters:\n",
    "            clusters[label] = []\n",
    "        clusters[label].append(idx)\n",
    "    \n",
    "    for label in clusters:\n",
    "        sims = {}\n",
    "        cluster_vecs = embeddings[clusters[label]]\n",
    "        normalized = cluster_vecs / torch.linalg.norm(cluster_vecs, dim=1, keepdims=True)\n",
    "        cluster_center = normalized.mean(dim=0)\n",
    "        for idx in clusters[label]:\n",
    "            sim = cosine_similarity(cluster_center, embeddings[idx])\n",
    "            sims[idx] = sim\n",
    "        for k, v in sorted(sims.items(), key=lambda item: item[1], reverse=True)[:5]:\n",
    "            print(k,\"%.3f\" % v, \" \".join(data[k][\"sentence\"]))\n",
    "        \n",
    "        print()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81155685-654c-45ce-86a8-8eac536d41de",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.cluster import KMeans, AgglomerativeClustering\n",
    "\n",
    "kmeans_clusters =  # FILL ME IN. Should be an instance of sklearn.cluster.KMeans\n",
    "diagnose_clustering(kmeans_clusters)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a01dfbb0-02a1-4f09-99ff-8fbe15787197",
   "metadata": {},
   "source": [
    "## Clustering with Agglomerative\n",
    "\n",
    "### Question 3\n",
    "\n",
    "K-Means operates in the Euclidean metric space, but we generally use cosine similarity when assessing the similarity of token embeddings. **Use agglomerative clustering with the `cosine` metric, `average` linkage, and $k=5$ clusters.** Again, examine the outputs with `diagnose_clustering`.\n",
    "\n",
    "**In a few sentences,** make some qualitative comparisons between how well these methods work at inducing word senses. Do the clusters surface coherent senses? Are they distinct? Do you find overlaps in senses between clusters?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "917200eb-bf31-44f2-b1f7-b2843f4a78f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "agglom = # FILL ME IN; this will take longer to run than K-Means"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24d1d917-42c5-4c49-b8e8-fbf1edd8f85e",
   "metadata": {},
   "outputs": [],
   "source": [
    "diagnose_clustering(agglom)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95638d6f-7dc5-48df-96a9-ca4bbe92eec1",
   "metadata": {},
   "source": [
    "## Evaluating\n",
    "\n",
    "In the absence of any ground truth, we will use the silhouette score to perform intrinsic evaluation of clusters. You can read more about the silhouette score in the [scikit-learn documentation](https://scikit-learn.org/stable/modules/generated/sklearn.metrics.silhouette_score.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ea15136-d9da-4415-b46e-94e28885e74a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics import silhouette_score"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7112005-4c71-4539-887e-3f48fb04b3f6",
   "metadata": {},
   "source": [
    "### Question 4\n",
    "\n",
    "**Compute the silhouette scores** for the agglomerative clustering output and the K-Means output from before, with the `cosine` metric. Compare them; do they align with your qualitative judgment?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7e0ffee-e60a-43a0-82a4-f2f02d1d74e3",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "c967d845-f72e-43bd-82e9-c41d6579e18a",
   "metadata": {},
   "source": [
    "### Question 5\n",
    "\n",
    "For both K-Means and agglomerative clustering, **plot the silhouette scores** for a range of cluster numbers $k = {2, \\ldots, 10}$. **In a few sentences,** do these plots align with your expectations? What optimal number of clusters do they suggest, and does that align with your existing understanding of the word \"address\"? Finally, what might be some reasons they do or don't align with your expectations, and what does that tell us about evaluating unsupervised models more generally?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5cbaa40b-e759-4492-a39e-cdc18ab521d2",
   "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}