tests/test_skip_chain_crf.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import random\n",
    "import torch\n",
    "from tqdm import tqdm\n",
    "from torch.utils.data import DataLoader, IterableDataset\n",
    "from pseudo_data import SkipChainDataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "    This is a dataset simulating a general graph.\n",
      "    An item looks like:\n",
      "        _7____________5___E2__A5_____7__5____22___D7_@@@@@\n",
      "         |                           |             |\n",
      "         D                           D             D\n",
      "    where:\n",
      "    - 7 is labeled as \"D\" since one of 7s is a successor of a \"D\"\n",
      "    - similarly, 5 is labeled as \"A\" and 2 is labeled as \"E\"\n",
      "    - 3 numbers (7, 5, 2) occurs in an item, and each occurs 3 times,\n",
      "        \n",
      "    For graphical models, we can build one ternary factor or 3 binary factors:\n",
      "            7                   7\n",
      "           / \\         or       |\n",
      "          7---7              7_/ \\_7\n",
      "    \n",
      "=== 0 ===\n",
      "_______6_2F6__4__C44_A2____2_6_@@@@@@@@@@@@@@@@@@@\n",
      "_______F_A_F__C___CC__A____A_F_@@@@@@@@@@@@@@@@@@@\n",
      "=== 1 ===\n",
      "______C3_7___3_5E5_____5___3_F77______@@@@@@@@@@@@\n",
      "_______C_F___C_E_E_____E___C__FF______@@@@@@@@@@@@\n",
      "=== 2 ===\n",
      "_7__6_E6__B3___3___6G7______7_______3__@@@@@@@@@@@\n",
      "_G__E__E___B___B___E_G______G_______B__@@@@@@@@@@@\n",
      "=== 3 ===\n",
      "_____2_3__2_7__B27_______E3G7__3@@@@@@@@@@@@@@@@@@\n",
      "_____B_E__B_G___BG________E_G__E@@@@@@@@@@@@@@@@@@\n"
     ]
    }
   ],
   "source": [
    "dataset = SkipChainDataset(50, True)\n",
    "print(dataset.__doc__)\n",
    "train_data = DataLoader(dataset, batch_size=20)\n",
    "test_data = DataLoader(dataset, batch_size=100)\n",
    "\n",
    "for i, batch in enumerate(train_data):\n",
    "    nodes, masks, bin_edges, bin_masks, ter_edges, ter_masks, label = batch\n",
    "    print(f\"=== {i} ===\")\n",
    "    print(dataset.to_sent(nodes))\n",
    "    print(dataset.to_label(label))\n",
    "    if i == 3:\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "loss: 1.8042, accu: 1.0000: : 400it [00:09, 42.75it/s]\n"
     ]
    }
   ],
   "source": [
    "from torch_random_fields.models import GeneralCRF\n",
    "import einops\n",
    "from einops.layers.torch import Rearrange\n",
    "from torch_random_fields.models.constants import Inference, Learning\n",
    "\n",
    "\n",
    "class Model(torch.nn.Module):\n",
    "    def __init__(self) -> None:\n",
    "        super().__init__()\n",
    "        self.embed = torch.nn.Embedding(num_embeddings=len(dataset.word2idx), embedding_dim=10)\n",
    "        self.pred = torch.nn.Sequential(\n",
    "            Rearrange(\"B T C -> B C T\"),\n",
    "            torch.nn.Conv1d(10, len(dataset.label2idx), 3, padding=\"same\"),\n",
    "            Rearrange(\"B C T -> B T C\"),\n",
    "        )\n",
    "        # self.pred = torch.nn.Linear(10, len(dataset.label2idx))\n",
    "        self.crf = GeneralCRF(\n",
    "            num_states=len(dataset.label2idx),\n",
    "            feature_size=10,\n",
    "            beam_size=64,\n",
    "            low_rank=10,\n",
    "            learning=Learning.PIECEWISE,\n",
    "            inference=Inference.BELIEF_PROPAGATION,\n",
    "            support_ternary=True,\n",
    "        )\n",
    "\n",
    "    def forward(self, nodes, masks, bin_edges, bin_masks, ter_edges, ter_masks, targets):\n",
    "        feats = self.embed(nodes)\n",
    "        unaries = self.pred(feats)\n",
    "        loss = self.crf(\n",
    "            unaries=unaries,\n",
    "            masks=masks,\n",
    "            binary_edges=bin_edges,\n",
    "            binary_masks=bin_masks,\n",
    "            ternary_edges=ter_edges,\n",
    "            ternary_masks=ter_masks,\n",
    "            targets=targets,\n",
    "            node_features=feats,\n",
    "        )\n",
    "        return loss\n",
    "\n",
    "    def decode(self, nodes, masks, bin_edges, bin_masks, ter_edges, ter_masks):\n",
    "        return self(nodes, masks, bin_edges, bin_masks, ter_edges, ter_masks, None)[1]\n",
    "\n",
    "    def evaulate(self, nodes, masks, bin_edges, bin_masks, ter_edges, ter_masks, label):\n",
    "        for b in test_data:\n",
    "            pred = self(nodes, masks, bin_edges, bin_masks, ter_edges, ter_masks, None)[1]\n",
    "            break\n",
    "        corr = pred == label\n",
    "        accu = corr[masks].sum() / masks.sum()\n",
    "        return accu\n",
    "\n",
    "\n",
    "pbar = tqdm(train_data)\n",
    "model = Model()\n",
    "opt = torch.optim.AdamW(model.parameters(), .01, weight_decay=0.01)\n",
    "for i, batch in enumerate(pbar):\n",
    "    nodes, masks, bin_edges, bin_masks, ter_edges, ter_masks, label = batch\n",
    "    loss = model(nodes, masks, bin_edges, bin_masks, ter_edges, ter_masks, label)\n",
    "    opt.zero_grad()\n",
    "    loss.backward()\n",
    "    opt.step()\n",
    "\n",
    "    if i % 100 == 0:\n",
    "        accu = model.evaulate(nodes, masks, bin_edges, bin_masks, ter_edges, ter_masks, label)\n",
    "        pbar.set_description(\"loss: {:.4f}, accu: {:.4f}\".format(loss.item(), accu.item()))\n",
    "\n",
    "    if i == 400:\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "_5___73___5___A7____________E3_B5_____3____7______\n",
      "_B___AE___B____A_____________E__B_____E____A______\n",
      "_B___AE___B____A_____________E__B_____E____A______\n"
     ]
    }
   ],
   "source": [
    "bid = 1\n",
    "node_len = masks[bid].sum()\n",
    "print(dataset.to_sent(nodes[bid, :node_len]))\n",
    "print(dataset.to_label(label[bid, :node_len]))\n",
    "\n",
    "pred = model.decode(nodes, masks, bin_edges, bin_masks, ter_edges, ter_masks)\n",
    "print(dataset.to_label(pred[bid, :node_len]))\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.7.11 ('pystruct3.7')",
   "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.7.11"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "26e7100929c5974aafe5d65382914e883abf933defba9fccc362b643f50b42c7"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}