{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "# Neural Bag-of-Words Classifiers\n\nLecture 2 | CMU ANLP Spring 2026 | Instructor: Sean Welleck"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "This is a notebook for [CMU CS11-711 Advanced NLP](https://cmu-l3.github.io/anlp-spring2026/) that trains neural network classifiers. Specifically, each model uses a bag-of-words variant to encode an input sequence into a continuous vector that is mapped to a probability distribution over the output classes. The model is trained to minimize cross-entropy loss using backpropagation.\n\n*Ackowledgements*: adapted from Graham Neubig's ANLP Fall 2024 [code](https://github.com/neubig/anlp-code/tree/main/02-textclass)"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Tweet classification\n",
    "\n",
    "We use the [`mteb/tweet_sentiment_extraction`](https://huggingface.co/datasets/mteb/tweet_sentiment_extraction) dataset, which consists of classifying an input tweet as positive, neutral, or negative sentiment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\"id\":\"cb774db0d1\",\"text\":\" I`d have responded, if I were going\",\"label\":1,\"label_text\":\"neutral\"}\n",
      "{\"id\":\"549e992a42\",\"text\":\" Sooo SAD I will miss you here in San Diego!!!\",\"label\":0,\"label_text\":\"negative\"}\n",
      "{\"id\":\"088c60f138\",\"text\":\"my boss is bullying me...\",\"label\":0,\"label_text\":\"negative\"}\n",
      "{\"id\":\"9642c003ef\",\"text\":\" what interview! leave me alone\",\"label\":0,\"label_text\":\"negative\"}\n"
     ]
    }
   ],
   "source": [
    "!head -n 4 train.jsonl"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Train a tokenizer\n",
    "\n",
    "Based on the examples above, splitting on whitespace isn't a great idea. Let's learn a BPE vocabulary using `sentencepiece`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sentencepiece as spm\n",
    "import json\n",
    "\n",
    "with open(\"bow_tokenizer_txt.txt\", \"w\", encoding=\"utf-8\") as f:\n",
    "    with open('train.jsonl', \"r\") as f2:\n",
    "        for line in f2:\n",
    "            j = json.loads(line)\n",
    "            words = j['text']\n",
    "            f.write(words + \"\\n\")\n",
    "\n",
    "import os\n",
    "\n",
    "options = dict(\n",
    "  input=\"bow_tokenizer_txt.txt\",\n",
    "  input_format=\"text\",\n",
    "  model_prefix=\"bow_tok\", \n",
    "  model_type=\"bpe\",\n",
    "  vocab_size=2048,\n",
    "  byte_fallback=True,\n",
    "  num_threads=os.cpu_count()\n",
    ")\n",
    "\n",
    "spm.SentencePieceTrainer.train(**options);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[['get', 1000],\n",
       " ['▁gl', 1001],\n",
       " ['▁away', 1002],\n",
       " ['eeee', 1003],\n",
       " ['▁left', 1004],\n",
       " ['▁mothers', 1005],\n",
       " ['?!', 1006],\n",
       " ['ily', 1007],\n",
       " ['oke', 1008],\n",
       " ['url', 1009],\n",
       " ['▁late', 1010],\n",
       " ['ire', 1011],\n",
       " ['hes', 1012],\n",
       " ['ner', 1013],\n",
       " ['▁Hope', 1014],\n",
       " ['▁Twitter', 1015],\n",
       " ['▁sha', 1016],\n",
       " ['▁bu', 1017],\n",
       " ['▁em', 1018],\n",
       " ['inking', 1019]]"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sp = spm.SentencePieceProcessor()\n",
    "sp.load('bow_tok.model')\n",
    "vocab = [[sp.id_to_piece(idx), idx] for idx in range(sp.get_piece_size())]\n",
    "vocab[1000:1020]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Data loading\n",
    "\n",
    "Read in the data, tokenize it, and split it into a training and dev set. There is a separate test set on [HuggingFace](https://huggingface.co/datasets/mteb/tweet_sentiment_extraction)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import defaultdict\n",
    "import json\n",
    "import random\n",
    "\n",
    "random.seed(123)\n",
    "\n",
    "label_to_text = {}\n",
    "def read_dataset(filename):\n",
    "    with open(filename, \"r\") as f:\n",
    "        for line in f:\n",
    "            j = json.loads(line)\n",
    "            words = j['text']\n",
    "            label = j['label']\n",
    "            label_to_text[label] = j['label_text']\n",
    "            tokens = sp.encode(words)\n",
    "            yield (tokens, label)\n",
    "\n",
    "# Read in the data\n",
    "ds = list(read_dataset(\"train.jsonl\"))\n",
    "random.shuffle(ds)\n",
    "train = ds[:-1000]\n",
    "dev = ds[1000:]\n",
    "\n",
    "nwords = len(sp)\n",
    "ntags = 3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model 1: Bag-of-Embeddings\n",
    "\n",
    "Our simplest model simply sums together 3-dimensional word embeddings (3 dimensions since we have three classes).\n",
    "\n",
    "First, for understanding purposes let's implement our own embedding layer.\n",
    "\n",
    "To do so, we multiply a one-hot vector representation of a token with a suitably-sized weight matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[402, 510, 953, 428, 413]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[0, 0, 0,  ..., 0, 0, 0],\n",
       "        [0, 0, 0,  ..., 0, 0, 0],\n",
       "        [0, 0, 0,  ..., 0, 0, 0],\n",
       "        [0, 0, 0,  ..., 0, 0, 0],\n",
       "        [0, 0, 0,  ..., 0, 0, 0]])"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "print(train[1][0][:5])\n",
    "\n",
    "torch.nn.functional.one_hot(torch.tensor(train[0][0]), num_classes=nwords)[:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Parameter containing:\n",
       "tensor([[-1.4509e+00,  9.2547e-01,  8.0536e-02,  ..., -3.2355e-01,\n",
       "         -1.0078e-01,  1.9657e+00],\n",
       "        [-1.1709e+00,  1.7886e+00,  2.0609e-01,  ..., -1.8114e+00,\n",
       "         -1.1534e+00, -1.8620e-01],\n",
       "        [ 1.7818e-01,  2.0482e+00,  2.8740e-01,  ..., -1.0381e-01,\n",
       "         -9.6792e-01, -1.8163e-01],\n",
       "        ...,\n",
       "        [-2.8673e-01,  2.4939e-01, -5.8121e-01,  ..., -1.1078e-01,\n",
       "         -8.1260e-02,  2.9797e-01],\n",
       "        [ 1.9415e-03, -3.5610e+00, -7.1223e-01,  ..., -2.9136e-01,\n",
       "          8.4358e-01, -1.5612e-01],\n",
       "        [-1.9175e+00,  1.5752e+00,  2.1777e-01,  ..., -1.4831e+00,\n",
       "          1.4779e+00,  1.3776e+00]], requires_grad=True)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch.nn as nn\n",
    "\n",
    "weight = nn.Parameter(torch.randn(nwords, 64))\n",
    "weight"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([5, 64])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "xs = torch.nn.functional.one_hot(torch.tensor(train[1][0]), num_classes=nwords)[:5]\n",
    "\n",
    "torch.matmul(xs.float(), weight).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Embedding(nn.Module):\n",
    "    def __init__(self, vocab_size, emb_size):\n",
    "        super(Embedding, self).__init__()\n",
    "        self.weight = nn.Parameter(torch.randn(vocab_size, emb_size))\n",
    "        self.vocab_size = vocab_size\n",
    "\n",
    "        nn.init.xavier_uniform_(self.weight)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        xs = torch.nn.functional.one_hot(x, num_classes=self.vocab_size).float()\n",
    "        return torch.matmul(xs, self.weight)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now here is our simple bag-of-words model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "class BoW(torch.nn.Module):\n",
    "    def __init__(self, vocab_size, num_labels):\n",
    "        super(BoW, self).__init__()\n",
    "        self.embedding = Embedding(vocab_size, num_labels)\n",
    "        nn.init.xavier_uniform_(self.embedding.weight)\n",
    "\n",
    "    def forward(self, tokens):\n",
    "        emb = self.embedding(tokens)\n",
    "        out = torch.sum(emb, dim=0) \n",
    "        logits = out.view(1, -1) \n",
    "        return logits"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's also implement cross-entropy loss ourselves this time:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ce_loss(logits, target):\n",
    "    log_probs = torch.nn.functional.log_softmax(logits, dim=1)\n",
    "    loss = -log_probs[:, target]\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's a training loop.\n",
    "\n",
    "We simply do \"single batch\" training here, i.e. loop over each training example one at a time and perform an update. We'll implement batching later on.\n",
    "\n",
    "You can use the SGD (Stochastic Gradient Descent) optimizer that was introduced in class, or this typically better optimizer Adam (we'll see it in a later class)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "iter 0: train loss/sent=0.9043, time=5.95s\n",
      "iter 0: valid acc=0.6844\n",
      "iter 1: train loss/sent=0.7722, time=6.14s\n",
      "iter 1: valid acc=0.7026\n",
      "iter 2: train loss/sent=0.7376, time=6.22s\n",
      "iter 2: valid acc=0.7099\n",
      "iter 3: train loss/sent=0.7222, time=5.74s\n",
      "iter 3: valid acc=0.7148\n",
      "iter 4: train loss/sent=0.7142, time=6.09s\n",
      "iter 4: valid acc=0.7153\n"
     ]
    }
   ],
   "source": [
    "import random\n",
    "import time\n",
    "\n",
    "# initialize the model\n",
    "model = BoW(nwords, ntags)\n",
    "criterion = ce_loss\n",
    "# optimizer = torch.optim.SGD(model.parameters(), lr=5e-4)\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=5e-4)\n",
    "\n",
    "for ITER in range(5):\n",
    "    # Perform training\n",
    "    random.shuffle(train)\n",
    "    train_loss = 0.0\n",
    "    start = time.time()\n",
    "    for x, y in train:\n",
    "        x = torch.tensor(x, dtype=torch.long)\n",
    "        y = torch.tensor([y])\n",
    "        logits = model(x)\n",
    "        loss = criterion(logits, y)\n",
    "        train_loss += loss.item()\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "    print(\"iter %r: train loss/sent=%.4f, time=%.2fs\" % (\n",
    "                ITER, train_loss/len(train), time.time()-start))\n",
    "    # Perform validation\n",
    "    test_correct = 0.0\n",
    "    for x, y in dev:\n",
    "        x = torch.tensor(x, dtype=torch.long)\n",
    "        logits = model(x)[0].detach()\n",
    "        predict = logits.argmax().item()\n",
    "        if predict == y:\n",
    "            test_correct += 1\n",
    "    print(\"iter %r: valid acc=%.4f\" % (ITER, test_correct/len(dev)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model 2: Bag-of-embeddings + output layer\n",
    "\n",
    "This is what we called `CBoW` in the lecture. Take a look at the code to see how it differs from the previous model.\n",
    "\n",
    "Also, it turns out to be important to initialize the weights well. We'll discuss this in a later class. Try removing the `nn.init` lines and see the performance change."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CBoW(torch.nn.Module):\n",
    "    def __init__(self, vocab_size, num_labels, emb_size):\n",
    "        super(CBoW, self).__init__()\n",
    "        self.embedding = nn.Embedding(vocab_size, emb_size)\n",
    "        self.output_layer = nn.Linear(emb_size, num_labels)\n",
    "\n",
    "        nn.init.xavier_uniform_(self.embedding.weight)\n",
    "        nn.init.xavier_uniform_(self.output_layer.weight)\n",
    "\n",
    "    def forward(self, tokens):\n",
    "        emb = self.embedding(tokens)    # [len(tokens) x emb_size]\n",
    "        emb_sum = torch.sum(emb, dim=0) # [emb_size]\n",
    "        h = emb_sum.view(1, -1)         # [1 x emb_size]\n",
    "        logits = self.output_layer(h)   # [1 x num_labels]\n",
    "        return logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "iter 0: train loss/sent=0.8398, time=10.39s\n",
      "iter 0: dev acc=0.7103\n",
      "iter 1: train loss/sent=0.7399, time=9.62s\n",
      "iter 1: dev acc=0.7204\n",
      "iter 2: train loss/sent=0.7158, time=10.31s\n",
      "iter 2: dev acc=0.7286\n",
      "iter 3: train loss/sent=0.7048, time=10.92s\n",
      "iter 3: dev acc=0.7341\n",
      "iter 4: train loss/sent=0.6967, time=10.54s\n",
      "iter 4: dev acc=0.7237\n"
     ]
    }
   ],
   "source": [
    "EMB_SIZE=32\n",
    "model = CBoW(nwords, ntags, EMB_SIZE)\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=5e-4)\n",
    "\n",
    "for ITER in range(5):\n",
    "    random.shuffle(train)\n",
    "    train_loss = 0.0\n",
    "    start = time.time()\n",
    "    model.train()\n",
    "    for x, y in train:\n",
    "        x = torch.tensor(x, dtype=torch.long)\n",
    "        y = torch.tensor([y])\n",
    "        logits = model(x)\n",
    "        loss = criterion(logits, y)\n",
    "        train_loss += loss.item()\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "    print(\"iter %r: train loss/sent=%.4f, time=%.2fs\" % (\n",
    "                ITER, train_loss/len(train), time.time()-start))\n",
    "    model.eval()\n",
    "    # Perform testing\n",
    "    test_correct = 0.0\n",
    "    for x, y in dev:\n",
    "        x = torch.tensor(x, dtype=torch.long)\n",
    "        logits = model(x)[0].detach()\n",
    "        predict = logits.argmax().item()\n",
    "        if predict == y:\n",
    "            test_correct += 1\n",
    "    print(\"iter %r: dev acc=%.4f\" % (ITER, test_correct/len(dev)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model 3: Deep CBoW\n",
    "\n",
    "Now we introduce a nonlinear layer involving a tanh activation. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DeepCBoW(torch.nn.Module):\n",
    "    def __init__(self, vocab_size, num_labels, emb_size, hid_size):\n",
    "        super(DeepCBoW, self).__init__()\n",
    "        self.embedding = nn.Embedding(vocab_size, emb_size)\n",
    "        self.linear1 = nn.Linear(emb_size, hid_size)    \n",
    "        self.output_layer = nn.Linear(hid_size, num_labels)\n",
    "\n",
    "        nn.init.xavier_uniform_(self.embedding.weight)\n",
    "        nn.init.xavier_uniform_(self.linear1.weight)     \n",
    "        nn.init.xavier_uniform_(self.output_layer.weight)\n",
    "\n",
    "    def forward(self, tokens):\n",
    "        emb = self.embedding(tokens)\n",
    "        emb_sum = torch.sum(emb, dim=0) \n",
    "        h = emb_sum.view(1, -1) \n",
    "        h = torch.tanh(self.linear1(h))  \n",
    "        logits = self.output_layer(h)\n",
    "        return logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 0: train loss/sent=0.8199, time=11.63s\n",
      "iter 0: dev acc=0.7249\n",
      "epoch 1: train loss/sent=0.6858, time=10.44s\n",
      "iter 1: dev acc=0.7553\n",
      "epoch 2: train loss/sent=0.6270, time=10.71s\n",
      "iter 2: dev acc=0.7830\n",
      "epoch 3: train loss/sent=0.5716, time=13.66s\n",
      "iter 3: dev acc=0.8120\n",
      "epoch 4: train loss/sent=0.5143, time=12.05s\n",
      "iter 4: dev acc=0.8405\n",
      "epoch 5: train loss/sent=0.4565, time=12.04s\n",
      "iter 5: dev acc=0.8695\n",
      "epoch 6: train loss/sent=0.3966, time=12.47s\n",
      "iter 6: dev acc=0.8929\n",
      "epoch 7: train loss/sent=0.3438, time=11.31s\n",
      "iter 7: dev acc=0.9133\n",
      "epoch 8: train loss/sent=0.2917, time=12.60s\n",
      "iter 8: dev acc=0.9288\n",
      "epoch 9: train loss/sent=0.2439, time=12.54s\n",
      "iter 9: dev acc=0.9400\n"
     ]
    }
   ],
   "source": [
    "EMB_SIZE=32\n",
    "model = DeepCBoW(nwords, ntags, EMB_SIZE, 32)\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=5e-4)\n",
    "\n",
    "for EPOCH in range(10):\n",
    "    random.shuffle(train)\n",
    "    train_loss = 0.0\n",
    "    start = time.time()\n",
    "    model.train()\n",
    "    for x, y in train:\n",
    "        x = torch.tensor(x, dtype=torch.long)\n",
    "        y = torch.tensor([y])\n",
    "        logits = model(x)\n",
    "        loss = criterion(logits, y)\n",
    "        train_loss += loss.item()\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "    print(\"epoch %r: train loss/sent=%.4f, time=%.2fs\" % (\n",
    "                EPOCH, train_loss/len(train), time.time()-start))\n",
    "    model.eval()\n",
    "    # Perform testing\n",
    "    test_correct = 0.0\n",
    "    for x, y in dev:\n",
    "        x = torch.tensor(x, dtype=torch.long)\n",
    "        logits = model(x)[0].detach()\n",
    "        predict = logits.argmax().item()\n",
    "        if predict == y:\n",
    "            test_correct += 1\n",
    "    print(\"iter %r: dev acc=%.4f\" % (EPOCH, test_correct/len(dev)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Go deep learning!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Classify an example with our trained model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'positive'"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tweet = \"I'm learning so much in advanced NLP!\"\n",
    "tokens = torch.tensor(sp.encode(tweet), dtype=torch.long)\n",
    "logits = model(tokens)[0].detach()\n",
    "predict = logits.argmax().item()\n",
    "label_to_text[predict]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Suggested exercises\n",
    "\n",
    "- Try changing the initialization of weights. Does the loss and/or dev accuracy change?\n",
    "- Generalize the `DeepCBoW` implementation to take in a `num_layers` parameter. How does performance change as the number of layers is increased?\n",
    "- Try different hyperparameters (e.g., learning rate, embedding size, hidden size, number of epochs). Can you identify any consistent trends?\n",
    "- Try out different qualitative examples. Can you find patterns in how the model succeeds / fails?\n",
    "- Implement batching by introducing a new `[PAD]` token. Make sure to mask out vectors for pad tokens in the model forward pass."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "anlp",
   "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.10.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}