{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "32b50d1b-aba7-4080-a6c6-9b65cb57ce6a",
   "metadata": {},
   "source": [
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/dbamman/anlp25/blob/main/7.lm//HW7_BERT_Parts_of_Speech.ipynb)\n",
    "\n",
    "# HW7: Part-of-speech tagging with BERT\n",
    "\n",
    "In this homework, we will be using the contextual embeddings from a large masked language model (BERT) to do token classification with the Huggingface `transformers` library."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd9c1390-5303-40a1-bb08-1b8e5c972aa4",
   "metadata": {},
   "source": [
    "Part-of-speech (POS) tagging is a core task in NLP that involves identifying the part of speech for each token in a sentence. We will use the Universal Dependencies POS tags, as well as a manually annotated English UD corpus, to build a high-performing POS tagger.\n",
    "\n",
    "To read more about the Universal Dependencies project, [see this website](https://universaldependencies.org/introduction.html). To read more about the universal POS tags that we will be relying on, [see this page](https://universaldependencies.org/u/pos/index.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88c2090b-847b-4ffc-adbf-b7caed6d6b09",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer, BertModel\n",
    "import torch\n",
    "from torch import nn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f58b672-c32c-4e63-b5a5-d49a686e439c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Make sure the GPU is available\n",
    "torch.cuda.is_available()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19bed838-ee36-48d6-b9b3-a77d7145153f",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(\"bert-base-multilingual-cased\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82daced0-4a59-464b-aabd-d7e79de3602c",
   "metadata": {},
   "source": [
    "## Familiarizing ourselves with POS"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "107153aa-d4d8-41de-b0e6-b9cab98d0ba4",
   "metadata": {},
   "source": [
    "Download the datasets from the UD Github."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a787435-08a8-410f-8fff-ac229a5d28bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget https://github.com/UniversalDependencies/UD_English-EWT/raw/refs/heads/master/en_ewt-ud-train.conllu\n",
    "!wget https://github.com/UniversalDependencies/UD_English-EWT/raw/refs/heads/master/en_ewt-ud-dev.conllu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5dcd9167-cc46-4aee-8a6d-f4807cb565bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# from https://universaldependencies.org/u/pos/index.html\n",
    "idx2label = [\n",
    "    \"ADJ\",\n",
    "    \"ADP\",\n",
    "    \"ADV\",\n",
    "    \"AUX\",\n",
    "    \"CCONJ\",\n",
    "    \"DET\",\n",
    "    \"INTJ\",\n",
    "    \"NOUN\",\n",
    "    \"NUM\",\n",
    "    \"PART\",\n",
    "    \"PRON\",\n",
    "    \"PROPN\",\n",
    "    \"PUNCT\",\n",
    "    \"SCONJ\",\n",
    "    \"SYM\",\n",
    "    \"VERB\",\n",
    "    \"X\",\n",
    "]\n",
    "\n",
    "label2idx = {\n",
    "    lab: idx for idx, lab in enumerate(idx2label)\n",
    "}\n",
    "\n",
    "# the conllu files follow the format described here:\n",
    "# https://universaldependencies.org/format.html\n",
    "def parse_conllu(file, n=100):\n",
    "    # read up to n sentences from the provided file\n",
    "    dataset = []\n",
    "    this_sentence = []\n",
    "    this_labels = []\n",
    "    for line in open(file, \"r\"):\n",
    "        if line[0] == \"#\":\n",
    "            continue\n",
    "        line = line.strip()\n",
    "        if len(line) == 0:\n",
    "            dataset.append({\n",
    "                \"tokens\": this_sentence,\n",
    "                \"upos\": this_labels\n",
    "            })\n",
    "            this_sentence = []\n",
    "            this_labels = []\n",
    "\n",
    "            if len(dataset) >= n:\n",
    "                return dataset\n",
    "            \n",
    "            continue\n",
    "        idx, word, _, upos, _, _, _, _, _, _ = line.split(\"\\t\")\n",
    "        if \"-\" in idx:\n",
    "            # skip multiword tokens because each individual one is also included\n",
    "            continue\n",
    "        if upos not in label2idx:\n",
    "            print(line)\n",
    "        this_sentence.append(word)\n",
    "        this_labels.append(label2idx[upos])\n",
    "\n",
    "    return dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a516b5f2-9e7b-46b8-a772-77b076516ef5",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = parse_conllu(\"en_ewt-ud-train.conllu\", n=20_000)\n",
    "dev_dataset = parse_conllu(\"en_ewt-ud-dev.conllu\", n=5_000)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5342706-4fc6-49ad-98aa-691b9cbc5342",
   "metadata": {},
   "source": [
    "### Question 1\n",
    "\n",
    "Take a look through the dataset to get a feel for this task. Then, manually annotate the following sentence:\n",
    "\n",
    "```\n",
    "Emily took a flight to Paris.\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c73a672-c1d5-44e7-818a-21d1719413a0",
   "metadata": {},
   "source": [
    "_Fill me in:_\n",
    "\n",
    "```\n",
    "Emily  -\n",
    "took   -\n",
    "a      -\n",
    "flight -\n",
    "to     -\n",
    "Paris  -\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2ac37aa5-b8e2-40ea-b980-3a728312993e",
   "metadata": {},
   "source": [
    "## Setting up data helpers\n",
    "\n",
    "The annotations are provided at the word level, but BERT tokenizes into subword tokens, so we need to make sure the labels still align.\n",
    "\n",
    "Because we usually process inputs in parallel batches for efficiency, we also define a `collate` function that prepares a batch of examples for training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3028c833-b2cc-43a1-b5b7-55dadc7ae58a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def tokenize(batch):\n",
    "    # tokenize the words\n",
    "    output = tokenizer(batch[\"tokens\"], is_split_into_words=True)\n",
    "\n",
    "    # labels are provided on the word level, but we might tokenize into subword tokens\n",
    "    # we also inject special tokens (like [CLS] and [SEP])\n",
    "    # to account for this, we give null (-100) labels to any subword token that isn't the\n",
    "    # first token of the word.\n",
    "    # we configure the CrossEntropyLoss to ignore the -100 value\n",
    "    # (see `ignore_index` in https://docs.pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html)\n",
    "    all_new_labels = []\n",
    "    for i, labels in enumerate(batch[\"upos\"]):\n",
    "        new_labels = []\n",
    "        prev_id = None\n",
    "        for word_id in output.word_ids(batch_index=i):\n",
    "            if word_id is None or word_id == prev_id:\n",
    "                new_labels.append(-100)\n",
    "            else:\n",
    "                new_labels.append(labels[word_id])\n",
    "            prev_id = word_id\n",
    "        all_new_labels.append(new_labels)\n",
    "\n",
    "    output[\"labels\"] = all_new_labels\n",
    "    return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "beb42c14-fb41-414e-8a9d-adae92bb628a",
   "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 != \"labels\" else -100))\n",
    "        for k, v in outputs.items()\n",
    "    }\n",
    "    return outputs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba6a8ba7-8adb-41f1-9ed5-d0b261a8f713",
   "metadata": {},
   "source": [
    "## Define model\n",
    "\n",
    "Our model architecture is as follows:\n",
    "\n",
    "1. We encode the tokens with our BERT model.\n",
    "2. We then use a linear classification layer on the last layer of attention outputs for _each token_, giving us a classification for each token.\n",
    "\n",
    "### Question 2\n",
    "\n",
    "When working with neural networks, a useful debugging strategy is making sure the shapes of your tensors are correct.\n",
    "**Fill in the correct values for the shape assertions.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16efb270-6148-489c-af8e-7b39a450e04e",
   "metadata": {},
   "outputs": [],
   "source": [
    "class POSBert(nn.Module):\n",
    "    def __init__(self, bert_model, num_labels=2):\n",
    "        super().__init__()\n",
    "        self.num_labels = num_labels\n",
    "        self.hidden_size = bert_model.config.hidden_size\n",
    "\n",
    "        self.bert_model = bert_model\n",
    "        self.classification = nn.Linear(self.hidden_size, self.num_labels)\n",
    "\n",
    "    def forward(self, batch):\n",
    "        device = next(self.parameters()).device\n",
    "\n",
    "        # batch[\"input_ids\"] has shape (B, L) where B is the batch size and L is the padded length\n",
    "        B, L = batch[\"input_ids\"].shape\n",
    "\n",
    "        model_output = self.bert_model(\n",
    "            batch[\"input_ids\"].to(device),\n",
    "            attention_mask=batch[\"attention_mask\"].to(device),\n",
    "            token_type_ids=batch[\"token_type_ids\"].to(device)\n",
    "        )\n",
    "        hidden = model_output.last_hidden_state\n",
    "        assert hidden.shape == (0, 0, 0, 0)  # CHANGE ME\n",
    "        \n",
    "        logits = self.classification(hidden)\n",
    "        assert logits.shape == (0, 0, 0, 0)  # CHANGE ME\n",
    "        return logits"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "007b390c-d21b-4e35-b745-b3f16a3bef83",
   "metadata": {},
   "source": [
    "## Train model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58be713e-c03c-4811-9780-58c5a0de04b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "from torch.utils.data import DataLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a209fff7-1fab-451a-8489-1325b9a58784",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics import f1_score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b91bfb6d-dd31-429c-b713-5fb9177e541d",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 64\n",
    "learning_rate = 3e-5\n",
    "num_epochs = 1\n",
    "\n",
    "def filter_preds(predicted, target):\n",
    "    f_predicted = []\n",
    "    f_target = []\n",
    "    for pred, tar in zip(predicted, target):\n",
    "        if tar == -100:\n",
    "            continue\n",
    "        f_predicted.append(pred)\n",
    "        f_target.append(tar)\n",
    "    return np.array(f_predicted), np.array(f_target)\n",
    "\n",
    "\n",
    "def train(model, loss_fn, optimizer, loader, num_epochs=1):\n",
    "    model.train()\n",
    "    step = 0\n",
    "    for epoch in range(num_epochs):\n",
    "        for batch in tqdm(loader, desc=\"train\"):\n",
    "            optimizer.zero_grad()\n",
    "    \n",
    "            output = model(batch)\n",
    "            loss = loss_fn(output.flatten(0, 1), batch[\"labels\"].flatten(0, 1).to(\"cuda\"))\n",
    "    \n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "            # We `yield` at the end of each step, allowing us to iterate\n",
    "            # over the training steps in the `training_loop` function.\n",
    "            # This is called a `generator` in Python\n",
    "\n",
    "            # See this page to learn more about generators:\n",
    "            # https://wiki.python.org/moin/Generators\n",
    "            yield {\"epoch\": epoch, \"step\": step}\n",
    "\n",
    "            step += 1\n",
    "\n",
    "def validate(model, loader):\n",
    "    with torch.no_grad():\n",
    "        model.eval()\n",
    "        predicted = []\n",
    "        target = []\n",
    "        for batch in tqdm(loader, desc=\"dev\", disable=True):\n",
    "            output = model(batch)\n",
    "            predicted.append(output.flatten(0, 1).argmax(-1).cpu().numpy())\n",
    "            target.append(batch[\"labels\"].flatten(0, 1).numpy())\n",
    "\n",
    "        predicted = np.concatenate(predicted)\n",
    "        target = np.concatenate(target)\n",
    "\n",
    "        predicted, target = filter_preds(predicted, target)\n",
    "        \n",
    "        accuracy = (predicted == target).sum() / len(predicted)\n",
    "        f1 = f1_score(target, predicted, average=\"macro\")\n",
    "        return {\"accuracy\": accuracy, \"f1\": f1}\n",
    "\n",
    "\n",
    "def training_loop(model, loss, optimizer, train_loader, val_loader):\n",
    "    # Train loop\n",
    "    dev_scores = []\n",
    "\n",
    "    dev_scores.append(validate(model, val_loader))\n",
    "    for i, step in enumerate(train(model, loss, optimizer, train_loader)):\n",
    "        if i % 10 == 0:\n",
    "            # Validation loop every 10 train steps\n",
    "            dev_scores.append(validate(model, val_loader))\n",
    "\n",
    "    return dev_scores"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3eca9f78-f4ac-4b5a-91bb-0dee451b35c5",
   "metadata": {},
   "source": [
    "### Question 3:\n",
    "\n",
    "Fill in the code to train the POS model with cross entropy loss and the Adam optimizer. Use the learning rate that we defined above.\n",
    "\n",
    "We also use [Pytorch data loaders](https://docs.pytorch.org/tutorials/beginner/basics/data_tutorial.html#preparing-your-data-for-training-with-dataloaders) to help batch our data efficiently. We want to shuffle batches on each epoch for the train set for more stable training, but keep the dev batches unshuffled for consistency in our evaluation.\n",
    "\n",
    "The training might take a few minutes, but if it takes longer than 10 minutes, make sure that you have moved the model to the GPU before training!\n",
    "\n",
    "The lab material from Tuesday might be a useful reference here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "765b0a1b-7115-468f-a376-9946161b3955",
   "metadata": {},
   "outputs": [],
   "source": [
    "bert_model = BertModel.from_pretrained(\"bert-base-multilingual-cased\")\n",
    "pos_model = # FILL ME IN\n",
    "loss_fn = # FILL ME IN\n",
    "optimizer = # FILL ME IN\n",
    "\n",
    "train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, collate_fn=collate)\n",
    "dev_loader = # FILL ME IN\n",
    "\n",
    "scores = training_loop(pos_model, loss_fn, optimizer, train_loader, dev_loader)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e9fa94a-c37d-41ba-af78-8b5208536ae0",
   "metadata": {},
   "source": [
    "## Inspect output\n",
    "\n",
    "We'll first plot the accuracy and F1 validation curves. You should have a performance with high 90s in accuracy, and high-80s or low-90s in F1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "756f7aa4-64fc-4f3c-9415-1e626085a6ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "import seaborn as sns\n",
    "sns.set()\n",
    "\n",
    "sns.lineplot(x=list(range(len(scores))), y=[score[\"accuracy\"] for score in scores])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d09b35d-8922-4b6e-a512-38f99a024633",
   "metadata": {},
   "outputs": [],
   "source": [
    "sns.lineplot(x=list(range(len(scores))), y=[score[\"f1\"] for score in scores])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c65669c-c093-4a57-8785-a9ed4cb627b8",
   "metadata": {},
   "source": [
    "### Question 4\n",
    "\n",
    "Use the `visualize_results` function to test the model on multiple sentences.\n",
    "\n",
    "1. Find a sentence where the same word occurs twice with two different parts of speech.\n",
    "2. Find a pair of sentences where the same token has different output labels.\n",
    "\n",
    "Write a paragraph reflecting on why the architecture of BERT might be useful for POS tagging. Compare to some other methods we have learned in class, like static word embeddings or logistic regression on bags of words. What would be challenging about this task using those other methods? Why are contextual representations useful?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f90b93a-3cd5-4611-a379-24adc91ed57f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def visualize_results(sentence, model):\n",
    "    tokens = tokenizer(sentence, return_tensors=\"pt\")\n",
    "    preds = model(tokens).argmax(-1)\n",
    "    for tok, pred in zip(tokens.input_ids[0][1:-1], preds[0][1:-1]):\n",
    "        print(tokenizer.decode(tok), idx2label[pred.item()])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9ec8631-3010-4456-bd8e-d28f80915123",
   "metadata": {},
   "outputs": [],
   "source": [
    "visualize_results(\"Naitian is a PhD student at the UC Berkeley School of Information in California.\", pos_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4607d120-5bf6-4c42-9493-31490087ebaa",
   "metadata": {},
   "source": [
    "### Question 5\n",
    "\n",
    "You may have noticed that we are using a pre-trained multilingual BERT model. Pick another language that has UD annotations from the [list on this page](https://universaldependencies.org/). Download the `conllu` file from the linked Github repository for that language and train a new model on POS tagging for that language.\n",
    "\n",
    "Compare the performance between the English results from above and the language you chose. If you are familiar with the language, feel free to also inspect the model outputs on some of your own sentences.\n",
    "\n",
    "If you find yourself running out of memory on the GPU, you may want to restart the kernel to clear the previous model before doing this last part."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb74ff36-f924-4775-a5d1-66d2f0415530",
   "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
}