modular_addition_3.3.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21b02306",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from matplotlib.patches import Ellipse, Circle\n",
    "\n",
    "\n",
    "seed = 1\n",
    "np.random.seed(seed)\n",
    "torch.manual_seed(seed)\n",
    "\n",
    "torch.set_default_tensor_type(torch.DoubleTensor)\n",
    "\n",
    "\n",
    "class BioLinear(nn.Module):\n",
    "\n",
    "    def __init__(self, in_dim, out_dim, in_fold=1, out_fold=1):\n",
    "        super(BioLinear, self).__init__()\n",
    "        self.in_dim = in_dim\n",
    "        self.out_dim = out_dim\n",
    "        self.linear = nn.Linear(in_dim, out_dim)\n",
    "        self.in_fold = in_fold\n",
    "        self.out_fold = out_fold\n",
    "        assert in_dim % in_fold == 0\n",
    "        assert out_dim % out_fold == 0\n",
    "        #compute in_cor, shape: (in_dim)\n",
    "        in_dim_fold = int(in_dim/in_fold)\n",
    "        out_dim_fold = int(out_dim/out_fold)\n",
    "        self.in_coordinates = torch.tensor(list(np.linspace(1/(2*in_dim_fold), 1-1/(2*in_dim_fold), num=in_dim_fold))*in_fold, dtype=torch.float)\n",
    "        self.out_coordinates = torch.tensor(list(np.linspace(1/(2*out_dim_fold), 1-1/(2*out_dim_fold), num=out_dim_fold))*out_fold, dtype=torch.float)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        return self.linear(x)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8fcb02d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "class BioMLP(nn.Module):\n",
    "    def __init__(self, in_dim=2, out_dim=2, w=2, depth=2, shp=None, token_embedding=False, embedding_size=None):\n",
    "        super(BioMLP, self).__init__()\n",
    "        if shp == None:\n",
    "            shp = [in_dim] + [w]*(depth-1) + [out_dim]\n",
    "            self.in_dim = in_dim\n",
    "            self.out_dim = out_dim\n",
    "            self.depth = depth\n",
    "                 \n",
    "        else:\n",
    "            self.in_dim = shp[0]\n",
    "            self.out_dim = shp[-1]\n",
    "            self.depth = len(shp) - 1\n",
    "        linear_list = []\n",
    "        for i in range(self.depth):\n",
    "            if i == 0:\n",
    "                # for modular addition\n",
    "                linear_list.append(BioLinear(shp[i], shp[i+1], in_fold=2))\n",
    "                # for regression\n",
    "                #linear_list.append(BioLinear(shp[i], shp[i+1], in_fold=1))\n",
    "            else:\n",
    "                linear_list.append(BioLinear(shp[i], shp[i+1]))\n",
    "        self.linears = nn.ModuleList(linear_list)\n",
    "        \n",
    "        \n",
    "        if token_embedding == True:\n",
    "            # embedding size: number of tokens * embedding dimension\n",
    "            self.embedding = torch.nn.Parameter(torch.normal(0,1,size=embedding_size))\n",
    "        \n",
    "        self.shp = shp\n",
    "        # parameters for the bio-inspired trick\n",
    "        self.l0 = 0.5 # distance between two nearby layers\n",
    "        self.in_perm = nn.Parameter(torch.tensor(np.arange(int(self.in_dim/self.linears[0].in_fold)), dtype=torch.float))\n",
    "        self.out_perm = nn.Parameter(torch.tensor(np.arange(int(self.out_dim/self.linears[-1].out_fold)), dtype=torch.float))\n",
    "        self.top_k = 30\n",
    "        self.token_embedding = token_embedding\n",
    "\n",
    "    def forward(self, x):\n",
    "        shp = x.shape\n",
    "        in_fold = self.linears[0].in_fold\n",
    "        x = x.reshape(shp[0], in_fold, int(shp[1]/in_fold))\n",
    "        x = x[:,:,self.in_perm.long()]\n",
    "        x = x.reshape(shp[0], shp[1])\n",
    "        f = torch.nn.SiLU()\n",
    "        for i in range(self.depth-1):\n",
    "            x = f(self.linears[i](x))\n",
    "        x = self.linears[-1](x)\n",
    "        \n",
    "        out_perm_inv = torch.zeros(self.out_dim, dtype=torch.long)\n",
    "        out_perm_inv[self.out_perm.long()] = torch.arange(self.out_dim)\n",
    "        x = x[:,out_perm_inv]\n",
    "        #x = x[:,self.out_perm]\n",
    "        \n",
    "        return x\n",
    "    \n",
    "    def get_linear_layers(self):\n",
    "        return self.linears\n",
    "    \n",
    "    def get_cc(self, weight_factor=2.0, bias_penalize=True, no_penalize_last=False):\n",
    "        # compute connection cost\n",
    "        cc = 0\n",
    "        num_linear = len(self.linears)\n",
    "        for i in range(num_linear):\n",
    "            if i == num_linear - 1 and no_penalize_last:\n",
    "                weight_factor = 0.\n",
    "            biolinear = self.linears[i]\n",
    "            dist = torch.abs(biolinear.out_coordinates.unsqueeze(dim=1) - biolinear.in_coordinates.unsqueeze(dim=0))\n",
    "            cc += torch.mean(torch.abs(biolinear.linear.weight)*(weight_factor*dist+self.l0))\n",
    "            if bias_penalize == True:\n",
    "                cc += torch.mean(torch.abs(biolinear.linear.bias)*(self.l0))\n",
    "        if self.token_embedding:\n",
    "            cc += torch.mean(torch.abs(self.embedding)*(self.l0))\n",
    "            #pass\n",
    "        return cc\n",
    "    \n",
    "    def swap_weight(self, weights, j, k, swap_type=\"out\"):\n",
    "        with torch.no_grad():  \n",
    "            if swap_type == \"in\":\n",
    "                temp = weights[:,j].clone()\n",
    "                weights[:,j] = weights[:,k].clone()\n",
    "                weights[:,k] = temp\n",
    "            elif swap_type == \"out\":\n",
    "                temp = weights[j].clone()\n",
    "                weights[j] = weights[k].clone()\n",
    "                weights[k] = temp\n",
    "            else:\n",
    "                raise Exception(\"Swap type {} is not recognized!\".format(swap_type))\n",
    "            \n",
    "    def swap_bias(self, biases, j, k):\n",
    "        with torch.no_grad():  \n",
    "            temp = biases[j].clone()\n",
    "            biases[j] = biases[k].clone()\n",
    "            biases[k] = temp\n",
    "    \n",
    "    def swap(self, i, j, k):\n",
    "        # in the ith layer (of neurons), swap the jth and the kth neuron. \n",
    "        # Note: n layers of weights means n+1 layers of neurons.\n",
    "        # (incoming, outgoing) * weights + biases are swapped. \n",
    "        linears = self.get_linear_layers()\n",
    "        num_linear = len(linears)\n",
    "        if i == 0:\n",
    "            # input layer, only has outgoing weights; update in_perm\n",
    "            weights = linears[i].linear.weight\n",
    "            infold = linears[i].in_fold\n",
    "            fold_dim = int(weights.shape[1]/infold)\n",
    "            for l in range(infold):\n",
    "                self.swap_weight(weights, j+fold_dim*l, k+fold_dim*l, swap_type=\"in\")\n",
    "            # change input_perm\n",
    "            self.swap_bias(self.in_perm, j, k)\n",
    "        elif i == num_linear:\n",
    "            # output layer, only has incoming weights and biases; update out_perm\n",
    "            weights = linears[i-1].linear.weight\n",
    "            biases = linears[i-1].linear.bias\n",
    "            self.swap_weight(weights, j, k, swap_type=\"out\")\n",
    "            self.swap_bias(biases, j, k)\n",
    "            # change output_perm\n",
    "            self.swap_bias(self.out_perm, j, k)\n",
    "        else:\n",
    "            # middle layer : (incoming, outgoing) * weights, and biases\n",
    "            weights_in = linears[i-1].linear.weight\n",
    "            weights_out = linears[i].linear.weight\n",
    "            biases = linears[i-1].linear.bias\n",
    "            self.swap_weight(weights_in, j, k, swap_type=\"out\")\n",
    "            self.swap_weight(weights_out, j, k, swap_type=\"in\")\n",
    "            self.swap_bias(biases, j, k)\n",
    "\n",
    "    def get_top_id(self, i, top_k=20):\n",
    "        linears = self.get_linear_layers()\n",
    "        num_linear = len(linears)\n",
    "        if i == 0:\n",
    "            # input layer\n",
    "            weights = linears[i].linear.weight\n",
    "            score = torch.sum(torch.abs(weights), dim=0)\n",
    "            in_fold = linears[0].in_fold\n",
    "            #print(score.shape)\n",
    "            score = torch.sum(score.reshape(in_fold, int(score.shape[0]/in_fold)), dim=0)\n",
    "        elif i == num_linear:\n",
    "            # output layer\n",
    "            weights = linears[i-1].linear.weight\n",
    "            score = torch.sum(torch.abs(weights), dim=1)\n",
    "        else:\n",
    "            weights_in = linears[i-1].linear.weight\n",
    "            weights_out = linears[i].linear.weight\n",
    "            score = torch.sum(torch.abs(weights_out), dim=0) + torch.sum(torch.abs(weights_in), dim=1)\n",
    "        #print(score.shape)\n",
    "        top_index = torch.flip(torch.argsort(score),[0])[:top_k]\n",
    "        return top_index\n",
    "    \n",
    "    def relocate_ij(self, i, j):\n",
    "        # In the ith layer (of neurons), relocate the jth neuron\n",
    "        linears = self.get_linear_layers()\n",
    "        num_linear = len(linears)\n",
    "        if i < num_linear:\n",
    "            num_neuron = int(linears[i].linear.weight.shape[1]/linears[i].in_fold)\n",
    "        else:\n",
    "            num_neuron = linears[i-1].linear.weight.shape[0]\n",
    "        ccs = []\n",
    "        for k in range(num_neuron):\n",
    "            self.swap(i,j,k)\n",
    "            ccs.append(self.get_cc())\n",
    "            self.swap(i,j,k)\n",
    "        k = torch.argmin(torch.stack(ccs))\n",
    "        self.swap(i,j,k)\n",
    "            \n",
    "    def relocate_i(self, i):\n",
    "        # Relocate neurons in the ith layer\n",
    "        top_id = self.get_top_id(i, top_k=self.top_k)\n",
    "        for j in top_id:\n",
    "            self.relocate_ij(i,j)\n",
    "            \n",
    "    def relocate(self):\n",
    "        # Relocate neurons in the whole model\n",
    "        linears = self.get_linear_layers()\n",
    "        num_linear = len(linears)\n",
    "        for i in range(num_linear+1):\n",
    "            self.relocate_i(i)\n",
    "            \n",
    "    def plot(self):\n",
    "        fig, ax = plt.subplots(figsize=(6,6))\n",
    "        shp = self.shp\n",
    "        s = 1/(2*max(shp))\n",
    "        for j in range(len(shp)):\n",
    "            N = shp[j]\n",
    "            if j == 0:\n",
    "                in_fold = self.linears[j].in_fold\n",
    "                N = int(N/in_fold)\n",
    "            for i in range(N):\n",
    "                if j == 0:\n",
    "                    for fold in range(in_fold):\n",
    "                        circle = Ellipse((1/(2*N)+i/N, 0.1*j+0.02*fold-0.01), s, s/10*((len(shp)-1)+0.4), color='black')\n",
    "                        ax.add_patch(circle)\n",
    "                else:\n",
    "                    for fold in range(in_fold):\n",
    "                        circle = Ellipse((1/(2*N)+i/N, 0.1*j), s, s/10*((len(shp)-1)+0.4), color='black')\n",
    "                        ax.add_patch(circle)\n",
    "\n",
    "\n",
    "        plt.ylim(-0.02,0.1*(len(shp)-1)+0.02)\n",
    "        plt.xlim(-0.02,1.02)\n",
    "\n",
    "        linears = self.linears\n",
    "        for ii in range(len(linears)):\n",
    "            biolinear = linears[ii]\n",
    "            p = biolinear.linear.weight.clone()\n",
    "            p_shp = p.shape\n",
    "            \n",
    "            p = p/torch.abs(p).max()\n",
    "            in_fold = biolinear.in_fold\n",
    "            fold_num = int(p_shp[1]/in_fold)\n",
    "            for i in range(p_shp[0]):\n",
    "                if ii == 0:\n",
    "                    for fold in range(in_fold):\n",
    "                        for j in range(fold_num):\n",
    "                            plt.plot([1/(2*p_shp[0])+i/p_shp[0], 1/(2*fold_num)+j/fold_num], [0.1*(ii+1),0.1*ii+0.02*fold-0.01], lw=1*np.abs(p[i,j].detach().numpy()), color=\"blue\" if p[i,j]>0 else \"red\")\n",
    "                else:\n",
    "                    for j in range(fold_num):\n",
    "                        plt.plot([1/(2*p_shp[0])+i/p_shp[0], 1/(2*fold_num)+j/fold_num], [0.1*(ii+1),0.1*ii], lw=0.5*np.abs(p[i,j].detach().numpy()), color=\"blue\" if p[i,j]>0 else \"red\")\n",
    "                    \n",
    "        ax.axis('off')\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d021439",
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "\n",
    "seed = 4\n",
    "np.random.seed(seed)\n",
    "torch.manual_seed(seed)\n",
    "\n",
    "### create model ###\n",
    "p = 59\n",
    "d = 32\n",
    "in_dim = 2*d\n",
    "out_dim = p\n",
    "\n",
    "shp = [in_dim, 50, 50, out_dim]\n",
    "model = BioMLP(shp=shp, token_embedding=True, embedding_size=(p, d))\n",
    "\n",
    "### create dataset ###\n",
    "\n",
    "x = np.arange(p)\n",
    "y = np.arange(p)\n",
    "XX, YY = np.meshgrid(x, y)\n",
    "data_id = np.transpose([XX.reshape(-1,), YY.reshape(-1,)])\n",
    "labels = (data_id[:,0] + data_id[:,1]) % p\n",
    "labels = torch.tensor(labels, dtype=torch.long)\n",
    "print(labels.shape)\n",
    "fraction = 0.8\n",
    "train_num = int(p**2*fraction)\n",
    "test_num = p**2 - train_num\n",
    "\n",
    "train_id = np.random.choice(p**2,train_num,replace=False)\n",
    "test_id = np.array(list(set(np.arange(p**2)) - set(train_id)))\n",
    "\n",
    "def get_data(id_):\n",
    "    global labels\n",
    "    inputs = torch.cat([model.embedding[data_id[id_][:,0]], model.embedding[data_id[id_][:,1]]], dim=1)\n",
    "    return inputs, labels[id_]\n",
    "\n",
    "\n",
    "### train ###\n",
    "optimizer = torch.optim.AdamW(model.parameters(), lr=0.001, weight_decay=0.0)\n",
    "steps = 20000\n",
    "log = 200\n",
    "lamb = 0.1\n",
    "swap_log = 200\n",
    "plot_log = 200\n",
    "\n",
    "for step in range(steps):\n",
    "    \n",
    "    if step == int(steps*1/4):\n",
    "        lamb *= 10\n",
    "        \n",
    "    if step == int(steps*3/4):\n",
    "        lamb *= 0.1\n",
    "    \n",
    "    \n",
    "    CEL = nn.CrossEntropyLoss()\n",
    "    \n",
    "    optimizer.zero_grad()\n",
    "    \n",
    "    \n",
    "    inputs_train, labels_train = get_data(train_id)\n",
    "    pred  = model(inputs_train)\n",
    "    loss = CEL(pred, labels_train)\n",
    "    acc = torch.mean((torch.argmax(pred, dim=1) == labels_train).float())\n",
    "    \n",
    "    inputs_test, labels_test = get_data(test_id)\n",
    "    pred_test  = model(inputs_test)\n",
    "    loss_test = CEL(pred_test, labels_test)\n",
    "    acc_test = torch.mean((torch.argmax(pred_test, dim=1) == labels_test).float())\n",
    "    \n",
    "    cc = model.get_cc(weight_factor=2.0, no_penalize_last=False)\n",
    "    #cc = model.get_cc(weight_factor=2.0, no_penalize_last=True) # this gives 3+2 Tree.\n",
    "    total_loss = loss + lamb*cc\n",
    "    total_loss.backward()\n",
    "    optimizer.step()\n",
    "    \n",
    "    if step % log == 0:\n",
    "        print(model.linears[1].linear.weight.grad[0])\n",
    "    \n",
    "    \n",
    "    if step % log == 0:\n",
    "        print(\"step = %d | total loss: %.2e | train loss: %.2e | test loss %.2e | cc: %.2e | train acc: %.2e | test acc: %.2e \"%(step, total_loss.detach().numpy(), loss.detach().numpy(), loss_test.detach().numpy(), cc.detach().numpy(), acc.detach().numpy(), acc_test.detach().numpy()))\n",
    "    if (step+1) % swap_log == 0:\n",
    "        model.relocate()\n",
    "        \n",
    "    if step % plot_log == 0:\n",
    "        pass\n",
    "        #print(\"plot!!!\")\n",
    "        model.plot()\n",
    "        plt.title(\"%.2e | %.2e\" % (loss.detach().numpy(), loss_test.detach().numpy()))\n",
    "        plt.show()\n",
    "        #plt.savefig(\"./video_figs/toy_swap/{0:05d}.png\".format(step))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6deeaa0",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3c7e6c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.plot()\n",
    "d = 32\n",
    "for i in range(d):\n",
    "    plt.text(i/d+0.007, -0.02, i+1, fontsize=6)\n",
    "    \n",
    "plt.text(-0.15,-0.005, \"input\", fontsize=15)\n",
    "plt.text(-0.1,0.095, \"L1\", fontsize=15)\n",
    "plt.text(-0.1,0.195, \"L2\", fontsize=15)\n",
    "plt.text(-0.17,0.295, \"output\", fontsize=15)\n",
    "plt.title(\"Modular addition\", y=1.05, fontsize=15)\n",
    "\n",
    "plt.text(0.318,-0.025,\"{\",rotation=90,fontsize=15)\n",
    "plt.text(0.327,-0.045,\"A\",fontsize=15)\n",
    "\n",
    "plt.text(0.502,-0.025,\"{\",rotation=90,fontsize=15)\n",
    "plt.text(0.52,-0.045,\"B\",fontsize=15)\n",
    "\n",
    "plt.text(0.632,-0.025,\"{\",rotation=90,fontsize=23)\n",
    "plt.text(0.655,-0.045,\"C\",fontsize=15)\n",
    "\n",
    "plt.savefig('./fig/modadd_graph.png', bbox_inches=\"tight\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bbf740f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.decomposition import PCA\n",
    "\n",
    "plt.figure(figsize=(5,17))\n",
    "\n",
    "plt.subplot(3,1,1)\n",
    "data = model.embedding[:,model.in_perm.long()[[10,11]]].detach().numpy()\n",
    "pca = PCA(n_components=2)\n",
    "data = pca.fit_transform(data)\n",
    "\n",
    "j = 0; k = 1\n",
    "for i in range(p):\n",
    "    plt.scatter(data[i,j], data[i,k], color=\"blue\",alpha=0.5)\n",
    "    plt.text(data[i,j], data[i,k], i, fontsize=10)\n",
    "    \n",
    "#plt.xlim(-0.9,1.3)\n",
    "#plt.ylim(-0.9,1.3)\n",
    "plt.xticks([])\n",
    "plt.yticks([])\n",
    "plt.axis('off')\n",
    "plt.arrow(0,0,0.15,0.0,head_length=0.02,head_width=0.015, color=\"black\")\n",
    "plt.arrow(0,0,0.0,0.15,head_length=0.02,head_width=0.015, color=\"black\")\n",
    "plt.title(\"A\", fontsize=45)\n",
    "\n",
    "\n",
    "    \n",
    "    \n",
    "plt.subplot(3,1,2)\n",
    "data = model.embedding[:,model.in_perm.long()[[16,17]]].detach().numpy()\n",
    "pca = PCA(n_components=2)\n",
    "data = pca.fit_transform(data)\n",
    "\n",
    "j = 0; k = 1\n",
    "for i in range(p):\n",
    "    plt.scatter(data[i,j], data[i,k], color=\"blue\",alpha=0.5)\n",
    "    plt.text(data[i,j], data[i,k], i, fontsize=10)\n",
    "    \n",
    "#plt.xlim(-0.9,1.3)\n",
    "#plt.ylim(-0.9,1.3)\n",
    "plt.xticks([])\n",
    "plt.yticks([])\n",
    "plt.axis('off')\n",
    "plt.arrow(0,0,0.15,0.0,head_length=0.02,head_width=0.015, color=\"black\")\n",
    "plt.arrow(0,0,0.0,0.15,head_length=0.02,head_width=0.015, color=\"black\")\n",
    "plt.title(\"B\", fontsize=45)\n",
    "\n",
    "\n",
    "ax = plt.subplot(313, projection = '3d')\n",
    "\n",
    "data = model.embedding[:,model.in_perm.long()[[20,21,22]]].detach().numpy()\n",
    "pca = PCA(n_components=3)\n",
    "data = pca.fit_transform(data)\n",
    "\n",
    "\n",
    "for i in range(p):\n",
    "    ax.scatter(data[i,0], data[i,1], data[i,2], color=\"blue\", alpha=0.5)\n",
    "    ax.text(data[i,0], data[i,1], data[i,2], i)\n",
    "    \n",
    "ax.view_init(-140, 40)\n",
    "ax.set_xticks([])\n",
    "ax.set_yticks([])\n",
    "ax.set_zticks([])\n",
    "lim = 0.4\n",
    "ax.set_xlim(-lim,lim)\n",
    "ax.set_ylim(-lim,lim)\n",
    "ax.set_zlim(-lim,lim)\n",
    "plt.axis('off') \n",
    "\n",
    "ax.quiver([0],[0],[0],[0.2],[0],[0], color=\"black\")\n",
    "ax.quiver([0],[0],[0],[0],[0.2],[0], color=\"black\")\n",
    "ax.quiver([0],[0],[0],[0],[0],[0.2], color=\"black\")\n",
    "\n",
    "plt.title(\"C\", fontsize=45)\n",
    "    \n",
    "\n",
    "plt.savefig('./fig/modadd_rings.pdf', bbox_inches=\"tight\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc9eb1cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(projection='3d')\n",
    "\n",
    "for i in range(p):\n",
    "    ax.scatter(data[i,0], data[i,1], data[i,2], color=\"blue\")\n",
    "    ax.text(data[i,0], data[i,1], data[i,2], i)\n",
    "    \n",
    "ax.view_init(-140, 20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64ed5a7f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# set some of embedding dimension to be zero (knockout experiments)\n",
    "\n",
    "\n",
    "embedding_original = model.embedding.data.clone()\n",
    "\n",
    "with torch.no_grad():\n",
    "    model.embedding.data[:,model.in_perm.long()[[0,1,2,3,4,5,6,7,8,9,12,13,14,15,18,19,23,24,25,26,27,28,29,30,31]]] = 0.\n",
    "    \n",
    "inputs, labels = get_data(id_)\n",
    "pred  = model(inputs)\n",
    "acc = torch.mean((torch.argmax(pred, dim=1) == labels).float())\n",
    "print(acc)\n",
    "\n",
    "\n",
    "model.embedding.data = embedding_original\n",
    "\n",
    "inputs, labels = get_data(id_)\n",
    "pred  = model(inputs)\n",
    "acc = torch.mean((torch.argmax(pred, dim=1) == labels).float())\n",
    "print(acc)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61e89b93",
   "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}