{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "9e544464",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import os\n",
    "sys.path.append(os.path.abspath('..'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "785c1981",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from tensor_transformers.nv_bert.modeling_tensor_cola import TensorBertConfig, TensorBertForSequenceClassification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "13cd16dc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['0', '1']\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "2"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from processors.glue import PROCESSORS\n",
    "# print(PROCESSORS.keys())\n",
    "# GLU dataset processor\n",
    "processor = PROCESSORS['wnli']()\n",
    "print(processor.get_labels())\n",
    "num_labels = len(processor.get_labels())\n",
    "num_labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "84a23ed6",
   "metadata": {},
   "outputs": [],
   "source": [
    "chkpoint_path = \"cola_bert_llama_mlp_large_0.7_only_lr-001\"\n",
    "# \"cola_bert_llama_mlp_large_0.7-004\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "e7f8d4d0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total number of parameters in the model: 247.30M\n"
     ]
    }
   ],
   "source": [
    "config = TensorBertConfig.from_json_file(\n",
    "    chkpoint_path + \"/config.json\" # chkpoint path/config.json\n",
    ")\n",
    "if config.vocab_size % 8 != 0:\n",
    "    config.vocab_size += 8 - (config.vocab_size % 8)\n",
    "\n",
    "model = TensorBertForSequenceClassification(\n",
    "        config,\n",
    "        num_labels=num_labels,\n",
    "    )\n",
    "\n",
    "total_params = sum(p.numel() for p in model.parameters())\n",
    "print(f\"Total number of parameters in the model: {total_params/1e6:.2f}M\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "dfdc35b0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "device(type='cpu')"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "next(model.parameters()).device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "7a5aa755",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "getattr(\n",
    "        config.config_forward, \"prune_mask\", False\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "79e65a32",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "odict_keys(['bert.embeddings.word_embeddings.weight', 'bert.embeddings.position_embeddings.weight', 'bert.embeddings.token_type_embeddings.weight', 'bert.embeddings.LayerNorm.weight', 'bert.embeddings.LayerNorm.bias', 'bert.encoder.layer.0.attention.self.query.bias', 'bert.encoder.layer.0.attention.self.query.tensor.factors.0', 'bert.encoder.layer.0.attention.self.query.tensor.factors.1', 'bert.encoder.layer.0.attention.self.key.bias', 'bert.encoder.layer.0.attention.self.key.tensor.factors.0', 'bert.encoder.layer.0.attention.self.key.tensor.factors.1', 'bert.encoder.layer.0.attention.self.value.bias', 'bert.encoder.layer.0.attention.self.value.tensor.factors.0', 'bert.encoder.layer.0.attention.self.value.tensor.factors.1', 'bert.encoder.layer.0.attention.output.dense.bias', 'bert.encoder.layer.0.attention.output.dense.tensor.factors.0', 'bert.encoder.layer.0.attention.output.dense.tensor.factors.1', 'bert.encoder.layer.0.attention.output.LayerNorm.weight', 'bert.encoder.layer.0.attention.output.LayerNorm.bias', 'bert.encoder.layer.0.intermediate.dense_act.gate_proj.bias', 'bert.encoder.layer.0.intermediate.dense_act.gate_proj.tensor.factors.0', 'bert.encoder.layer.0.intermediate.dense_act.gate_proj.tensor.factors.1', 'bert.encoder.layer.0.intermediate.dense_act.up_proj.bias', 'bert.encoder.layer.0.intermediate.dense_act.up_proj.tensor.factors.0', 'bert.encoder.layer.0.intermediate.dense_act.up_proj.tensor.factors.1', 'bert.encoder.layer.0.output.dense.bias', 'bert.encoder.layer.0.output.dense.tensor.factors.0', 'bert.encoder.layer.0.output.dense.tensor.factors.1', 'bert.encoder.layer.0.output.LayerNorm.weight', 'bert.encoder.layer.0.output.LayerNorm.bias', 'bert.encoder.layer.1.attention.self.query.bias', 'bert.encoder.layer.1.attention.self.query.tensor.factors.0', 'bert.encoder.layer.1.attention.self.query.tensor.factors.1', 'bert.encoder.layer.1.attention.self.key.bias', 'bert.encoder.layer.1.attention.self.key.tensor.factors.0', 'bert.encoder.layer.1.attention.self.key.tensor.factors.1', 'bert.encoder.layer.1.attention.self.value.bias', 'bert.encoder.layer.1.attention.self.value.tensor.factors.0', 'bert.encoder.layer.1.attention.self.value.tensor.factors.1', 'bert.encoder.layer.1.attention.output.dense.bias', 'bert.encoder.layer.1.attention.output.dense.tensor.factors.0', 'bert.encoder.layer.1.attention.output.dense.tensor.factors.1', 'bert.encoder.layer.1.attention.output.LayerNorm.weight', 'bert.encoder.layer.1.attention.output.LayerNorm.bias', 'bert.encoder.layer.1.intermediate.dense_act.gate_proj.bias', 'bert.encoder.layer.1.intermediate.dense_act.gate_proj.tensor.factors.0', 'bert.encoder.layer.1.intermediate.dense_act.gate_proj.tensor.factors.1', 'bert.encoder.layer.1.intermediate.dense_act.up_proj.bias', 'bert.encoder.layer.1.intermediate.dense_act.up_proj.tensor.factors.0', 'bert.encoder.layer.1.intermediate.dense_act.up_proj.tensor.factors.1', 'bert.encoder.layer.1.output.dense.bias', 'bert.encoder.layer.1.output.dense.tensor.factors.0', 'bert.encoder.layer.1.output.dense.tensor.factors.1', 'bert.encoder.layer.1.output.LayerNorm.weight', 'bert.encoder.layer.1.output.LayerNorm.bias', 'bert.encoder.layer.2.attention.self.query.bias', 'bert.encoder.layer.2.attention.self.query.tensor.factors.0', 'bert.encoder.layer.2.attention.self.query.tensor.factors.1', 'bert.encoder.layer.2.attention.self.key.bias', 'bert.encoder.layer.2.attention.self.key.tensor.factors.0', 'bert.encoder.layer.2.attention.self.key.tensor.factors.1', 'bert.encoder.layer.2.attention.self.value.bias', 'bert.encoder.layer.2.attention.self.value.tensor.factors.0', 'bert.encoder.layer.2.attention.self.value.tensor.factors.1', 'bert.encoder.layer.2.attention.output.dense.bias', 'bert.encoder.layer.2.attention.output.dense.tensor.factors.0', 'bert.encoder.layer.2.attention.output.dense.tensor.factors.1', 'bert.encoder.layer.2.attention.output.LayerNorm.weight', 'bert.encoder.layer.2.attention.output.LayerNorm.bias', 'bert.encoder.layer.2.intermediate.dense_act.gate_proj.bias', 'bert.encoder.layer.2.intermediate.dense_act.gate_proj.tensor.factors.0', 'bert.encoder.layer.2.intermediate.dense_act.gate_proj.tensor.factors.1', 'bert.encoder.layer.2.intermediate.dense_act.up_proj.bias', 'bert.encoder.layer.2.intermediate.dense_act.up_proj.tensor.factors.0', 'bert.encoder.layer.2.intermediate.dense_act.up_proj.tensor.factors.1', 'bert.encoder.layer.2.output.dense.bias', 'bert.encoder.layer.2.output.dense.tensor.factors.0', 'bert.encoder.layer.2.output.dense.tensor.factors.1', 'bert.encoder.layer.2.output.LayerNorm.weight', 'bert.encoder.layer.2.output.LayerNorm.bias', 'bert.encoder.layer.3.attention.self.query.bias', 'bert.encoder.layer.3.attention.self.query.tensor.factors.0', 'bert.encoder.layer.3.attention.self.query.tensor.factors.1', 'bert.encoder.layer.3.attention.self.key.bias', 'bert.encoder.layer.3.attention.self.key.tensor.factors.0', 'bert.encoder.layer.3.attention.self.key.tensor.factors.1', 'bert.encoder.layer.3.attention.self.value.bias', 'bert.encoder.layer.3.attention.self.value.tensor.factors.0', 'bert.encoder.layer.3.attention.self.value.tensor.factors.1', 'bert.encoder.layer.3.attention.output.dense.bias', 'bert.encoder.layer.3.attention.output.dense.tensor.factors.0', 'bert.encoder.layer.3.attention.output.dense.tensor.factors.1', 'bert.encoder.layer.3.attention.output.LayerNorm.weight', 'bert.encoder.layer.3.attention.output.LayerNorm.bias', 'bert.encoder.layer.3.intermediate.dense_act.gate_proj.bias', 'bert.encoder.layer.3.intermediate.dense_act.gate_proj.tensor.factors.0', 'bert.encoder.layer.3.intermediate.dense_act.gate_proj.tensor.factors.1', 'bert.encoder.layer.3.intermediate.dense_act.up_proj.bias', 'bert.encoder.layer.3.intermediate.dense_act.up_proj.tensor.factors.0', 'bert.encoder.layer.3.intermediate.dense_act.up_proj.tensor.factors.1', 'bert.encoder.layer.3.output.dense.bias', 'bert.encoder.layer.3.output.dense.tensor.factors.0', 'bert.encoder.layer.3.output.dense.tensor.factors.1', 'bert.encoder.layer.3.output.LayerNorm.weight', 'bert.encoder.layer.3.output.LayerNorm.bias', 'bert.encoder.layer.4.attention.self.query.bias', 'bert.encoder.layer.4.attention.self.query.tensor.factors.0', 'bert.encoder.layer.4.attention.self.query.tensor.factors.1', 'bert.encoder.layer.4.attention.self.key.bias', 'bert.encoder.layer.4.attention.self.key.tensor.factors.0', 'bert.encoder.layer.4.attention.self.key.tensor.factors.1', 'bert.encoder.layer.4.attention.self.value.bias', 'bert.encoder.layer.4.attention.self.value.tensor.factors.0', 'bert.encoder.layer.4.attention.self.value.tensor.factors.1', 'bert.encoder.layer.4.attention.output.dense.bias', 'bert.encoder.layer.4.attention.output.dense.tensor.factors.0', 'bert.encoder.layer.4.attention.output.dense.tensor.factors.1', 'bert.encoder.layer.4.attention.output.LayerNorm.weight', 'bert.encoder.layer.4.attention.output.LayerNorm.bias', 'bert.encoder.layer.4.intermediate.dense_act.gate_proj.bias', 'bert.encoder.layer.4.intermediate.dense_act.gate_proj.tensor.factors.0', 'bert.encoder.layer.4.intermediate.dense_act.gate_proj.tensor.factors.1', 'bert.encoder.layer.4.intermediate.dense_act.up_proj.bias', 'bert.encoder.layer.4.intermediate.dense_act.up_proj.tensor.factors.0', 'bert.encoder.layer.4.intermediate.dense_act.up_proj.tensor.factors.1', 'bert.encoder.layer.4.output.dense.bias', 'bert.encoder.layer.4.output.dense.tensor.factors.0', 'bert.encoder.layer.4.output.dense.tensor.factors.1', 'bert.encoder.layer.4.output.LayerNorm.weight', 'bert.encoder.layer.4.output.LayerNorm.bias', 'bert.encoder.layer.5.attention.self.query.bias', 'bert.encoder.layer.5.attention.self.query.tensor.factors.0', 'bert.encoder.layer.5.attention.self.query.tensor.factors.1', 'bert.encoder.layer.5.attention.self.key.bias', 'bert.encoder.layer.5.attention.self.key.tensor.factors.0', 'bert.encoder.layer.5.attention.self.key.tensor.factors.1', 'bert.encoder.layer.5.attention.self.value.bias', 'bert.encoder.layer.5.attention.self.value.tensor.factors.0', 'bert.encoder.layer.5.attention.self.value.tensor.factors.1', 'bert.encoder.layer.5.attention.output.dense.bias', 'bert.encoder.layer.5.attention.output.dense.tensor.factors.0', 'bert.encoder.layer.5.attention.output.dense.tensor.factors.1', 'bert.encoder.layer.5.attention.output.LayerNorm.weight', 'bert.encoder.layer.5.attention.output.LayerNorm.bias', 'bert.encoder.layer.5.intermediate.dense_act.gate_proj.bias', 'bert.encoder.layer.5.intermediate.dense_act.gate_proj.tensor.factors.0', 'bert.encoder.layer.5.intermediate.dense_act.gate_proj.tensor.factors.1', 'bert.encoder.layer.5.intermediate.dense_act.up_proj.bias', 'bert.encoder.layer.5.intermediate.dense_act.up_proj.tensor.factors.0', 'bert.encoder.layer.5.intermediate.dense_act.up_proj.tensor.factors.1', 'bert.encoder.layer.5.output.dense.bias', 'bert.encoder.layer.5.output.dense.tensor.factors.0', 'bert.encoder.layer.5.output.dense.tensor.factors.1', 'bert.encoder.layer.5.output.LayerNorm.weight', 'bert.encoder.layer.5.output.LayerNorm.bias', 'bert.encoder.layer.6.attention.self.query.bias', 'bert.encoder.layer.6.attention.self.query.tensor.factors.0', 'bert.encoder.layer.6.attention.self.query.tensor.factors.1', 'bert.encoder.layer.6.attention.self.key.bias', 'bert.encoder.layer.6.attention.self.key.tensor.factors.0', 'bert.encoder.layer.6.attention.self.key.tensor.factors.1', 'bert.encoder.layer.6.attention.self.value.bias', 'bert.encoder.layer.6.attention.self.value.tensor.factors.0', 'bert.encoder.layer.6.attention.self.value.tensor.factors.1', 'bert.encoder.layer.6.attention.output.dense.bias', 'bert.encoder.layer.6.attention.output.dense.tensor.factors.0', 'bert.encoder.layer.6.attention.output.dense.tensor.factors.1', 'bert.encoder.layer.6.attention.output.LayerNorm.weight', 'bert.encoder.layer.6.attention.output.LayerNorm.bias', 'bert.encoder.layer.6.intermediate.dense_act.gate_proj.bias', 'bert.encoder.layer.6.intermediate.dense_act.gate_proj.tensor.factors.0', 'bert.encoder.layer.6.intermediate.dense_act.gate_proj.tensor.factors.1', 'bert.encoder.layer.6.intermediate.dense_act.up_proj.bias', 'bert.encoder.layer.6.intermediate.dense_act.up_proj.tensor.factors.0', 'bert.encoder.layer.6.intermediate.dense_act.up_proj.tensor.factors.1', 'bert.encoder.layer.6.output.dense.bias', 'bert.encoder.layer.6.output.dense.tensor.factors.0', 'bert.encoder.layer.6.output.dense.tensor.factors.1', 'bert.encoder.layer.6.output.LayerNorm.weight', 'bert.encoder.layer.6.output.LayerNorm.bias', 'bert.encoder.layer.7.attention.self.query.bias', 'bert.encoder.layer.7.attention.self.query.tensor.factors.0', 'bert.encoder.layer.7.attention.self.query.tensor.factors.1', 'bert.encoder.layer.7.attention.self.key.bias', 'bert.encoder.layer.7.attention.self.key.tensor.factors.0', 'bert.encoder.layer.7.attention.self.key.tensor.factors.1', 'bert.encoder.layer.7.attention.self.value.bias', 'bert.encoder.layer.7.attention.self.value.tensor.factors.0', 'bert.encoder.layer.7.attention.self.value.tensor.factors.1', 'bert.encoder.layer.7.attention.output.dense.bias', 'bert.encoder.layer.7.attention.output.dense.tensor.factors.0', 'bert.encoder.layer.7.attention.output.dense.tensor.factors.1', 'bert.encoder.layer.7.attention.output.LayerNorm.weight', 'bert.encoder.layer.7.attention.output.LayerNorm.bias', 'bert.encoder.layer.7.intermediate.dense_act.gate_proj.bias', 'bert.encoder.layer.7.intermediate.dense_act.gate_proj.tensor.factors.0', 'bert.encoder.layer.7.intermediate.dense_act.gate_proj.tensor.factors.1', 'bert.encoder.layer.7.intermediate.dense_act.up_proj.bias', 'bert.encoder.layer.7.intermediate.dense_act.up_proj.tensor.factors.0', 'bert.encoder.layer.7.intermediate.dense_act.up_proj.tensor.factors.1', 'bert.encoder.layer.7.output.dense.bias', 'bert.encoder.layer.7.output.dense.tensor.factors.0', 'bert.encoder.layer.7.output.dense.tensor.factors.1', 'bert.encoder.layer.7.output.LayerNorm.weight', 'bert.encoder.layer.7.output.LayerNorm.bias', 'bert.encoder.layer.8.attention.self.query.bias', 'bert.encoder.layer.8.attention.self.query.tensor.factors.0', 'bert.encoder.layer.8.attention.self.query.tensor.factors.1', 'bert.encoder.layer.8.attention.self.key.bias', 'bert.encoder.layer.8.attention.self.key.tensor.factors.0', 'bert.encoder.layer.8.attention.self.key.tensor.factors.1', 'bert.encoder.layer.8.attention.self.value.bias', 'bert.encoder.layer.8.attention.self.value.tensor.factors.0', 'bert.encoder.layer.8.attention.self.value.tensor.factors.1', 'bert.encoder.layer.8.attention.output.dense.bias', 'bert.encoder.layer.8.attention.output.dense.tensor.factors.0', 'bert.encoder.layer.8.attention.output.dense.tensor.factors.1', 'bert.encoder.layer.8.attention.output.LayerNorm.weight', 'bert.encoder.layer.8.attention.output.LayerNorm.bias', 'bert.encoder.layer.8.intermediate.dense_act.gate_proj.bias', 'bert.encoder.layer.8.intermediate.dense_act.gate_proj.tensor.factors.0', 'bert.encoder.layer.8.intermediate.dense_act.gate_proj.tensor.factors.1', 'bert.encoder.layer.8.intermediate.dense_act.up_proj.bias', 'bert.encoder.layer.8.intermediate.dense_act.up_proj.tensor.factors.0', 'bert.encoder.layer.8.intermediate.dense_act.up_proj.tensor.factors.1', 'bert.encoder.layer.8.output.dense.bias', 'bert.encoder.layer.8.output.dense.tensor.factors.0', 'bert.encoder.layer.8.output.dense.tensor.factors.1', 'bert.encoder.layer.8.output.LayerNorm.weight', 'bert.encoder.layer.8.output.LayerNorm.bias', 'bert.encoder.layer.9.attention.self.query.bias', 'bert.encoder.layer.9.attention.self.query.tensor.factors.0', 'bert.encoder.layer.9.attention.self.query.tensor.factors.1', 'bert.encoder.layer.9.attention.self.key.bias', 'bert.encoder.layer.9.attention.self.key.tensor.factors.0', 'bert.encoder.layer.9.attention.self.key.tensor.factors.1', 'bert.encoder.layer.9.attention.self.value.bias', 'bert.encoder.layer.9.attention.self.value.tensor.factors.0', 'bert.encoder.layer.9.attention.self.value.tensor.factors.1', 'bert.encoder.layer.9.attention.output.dense.bias', 'bert.encoder.layer.9.attention.output.dense.tensor.factors.0', 'bert.encoder.layer.9.attention.output.dense.tensor.factors.1', 'bert.encoder.layer.9.attention.output.LayerNorm.weight', 'bert.encoder.layer.9.attention.output.LayerNorm.bias', 'bert.encoder.layer.9.intermediate.dense_act.gate_proj.bias', 'bert.encoder.layer.9.intermediate.dense_act.gate_proj.tensor.factors.0', 'bert.encoder.layer.9.intermediate.dense_act.gate_proj.tensor.factors.1', 'bert.encoder.layer.9.intermediate.dense_act.up_proj.bias', 'bert.encoder.layer.9.intermediate.dense_act.up_proj.tensor.factors.0', 'bert.encoder.layer.9.intermediate.dense_act.up_proj.tensor.factors.1', 'bert.encoder.layer.9.output.dense.bias', 'bert.encoder.layer.9.output.dense.tensor.factors.0', 'bert.encoder.layer.9.output.dense.tensor.factors.1', 'bert.encoder.layer.9.output.LayerNorm.weight', 'bert.encoder.layer.9.output.LayerNorm.bias', 'bert.encoder.layer.10.attention.self.query.bias', 'bert.encoder.layer.10.attention.self.query.tensor.factors.0', 'bert.encoder.layer.10.attention.self.query.tensor.factors.1', 'bert.encoder.layer.10.attention.self.key.bias', 'bert.encoder.layer.10.attention.self.key.tensor.factors.0', 'bert.encoder.layer.10.attention.self.key.tensor.factors.1', 'bert.encoder.layer.10.attention.self.value.bias', 'bert.encoder.layer.10.attention.self.value.tensor.factors.0', 'bert.encoder.layer.10.attention.self.value.tensor.factors.1', 'bert.encoder.layer.10.attention.output.dense.bias', 'bert.encoder.layer.10.attention.output.dense.tensor.factors.0', 'bert.encoder.layer.10.attention.output.dense.tensor.factors.1', 'bert.encoder.layer.10.attention.output.LayerNorm.weight', 'bert.encoder.layer.10.attention.output.LayerNorm.bias', 'bert.encoder.layer.10.intermediate.dense_act.gate_proj.bias', 'bert.encoder.layer.10.intermediate.dense_act.gate_proj.tensor.factors.0', 'bert.encoder.layer.10.intermediate.dense_act.gate_proj.tensor.factors.1', 'bert.encoder.layer.10.intermediate.dense_act.up_proj.bias', 'bert.encoder.layer.10.intermediate.dense_act.up_proj.tensor.factors.0', 'bert.encoder.layer.10.intermediate.dense_act.up_proj.tensor.factors.1', 'bert.encoder.layer.10.output.dense.bias', 'bert.encoder.layer.10.output.dense.tensor.factors.0', 'bert.encoder.layer.10.output.dense.tensor.factors.1', 'bert.encoder.layer.10.output.LayerNorm.weight', 'bert.encoder.layer.10.output.LayerNorm.bias', 'bert.encoder.layer.11.attention.self.query.bias', 'bert.encoder.layer.11.attention.self.query.tensor.factors.0', 'bert.encoder.layer.11.attention.self.query.tensor.factors.1', 'bert.encoder.layer.11.attention.self.key.bias', 'bert.encoder.layer.11.attention.self.key.tensor.factors.0', 'bert.encoder.layer.11.attention.self.key.tensor.factors.1', 'bert.encoder.layer.11.attention.self.value.bias', 'bert.encoder.layer.11.attention.self.value.tensor.factors.0', 'bert.encoder.layer.11.attention.self.value.tensor.factors.1', 'bert.encoder.layer.11.attention.output.dense.bias', 'bert.encoder.layer.11.attention.output.dense.tensor.factors.0', 'bert.encoder.layer.11.attention.output.dense.tensor.factors.1', 'bert.encoder.layer.11.attention.output.LayerNorm.weight', 'bert.encoder.layer.11.attention.output.LayerNorm.bias', 'bert.encoder.layer.11.intermediate.dense_act.gate_proj.bias', 'bert.encoder.layer.11.intermediate.dense_act.gate_proj.tensor.factors.0', 'bert.encoder.layer.11.intermediate.dense_act.gate_proj.tensor.factors.1', 'bert.encoder.layer.11.intermediate.dense_act.up_proj.bias', 'bert.encoder.layer.11.intermediate.dense_act.up_proj.tensor.factors.0', 'bert.encoder.layer.11.intermediate.dense_act.up_proj.tensor.factors.1', 'bert.encoder.layer.11.output.dense.bias', 'bert.encoder.layer.11.output.dense.tensor.factors.0', 'bert.encoder.layer.11.output.dense.tensor.factors.1', 'bert.encoder.layer.11.output.LayerNorm.weight', 'bert.encoder.layer.11.output.LayerNorm.bias', 'bert.encoder.layer.12.attention.self.query.bias', 'bert.encoder.layer.12.attention.self.query.tensor.factors.0', 'bert.encoder.layer.12.attention.self.query.tensor.factors.1', 'bert.encoder.layer.12.attention.self.key.bias', 'bert.encoder.layer.12.attention.self.key.tensor.factors.0', 'bert.encoder.layer.12.attention.self.key.tensor.factors.1', 'bert.encoder.layer.12.attention.self.value.bias', 'bert.encoder.layer.12.attention.self.value.tensor.factors.0', 'bert.encoder.layer.12.attention.self.value.tensor.factors.1', 'bert.encoder.layer.12.attention.output.dense.bias', 'bert.encoder.layer.12.attention.output.dense.tensor.factors.0', 'bert.encoder.layer.12.attention.output.dense.tensor.factors.1', 'bert.encoder.layer.12.attention.output.LayerNorm.weight', 'bert.encoder.layer.12.attention.output.LayerNorm.bias', 'bert.encoder.layer.12.intermediate.dense_act.gate_proj.bias', 'bert.encoder.layer.12.intermediate.dense_act.gate_proj.tensor.factors.0', 'bert.encoder.layer.12.intermediate.dense_act.gate_proj.tensor.factors.1', 'bert.encoder.layer.12.intermediate.dense_act.up_proj.bias', 'bert.encoder.layer.12.intermediate.dense_act.up_proj.tensor.factors.0', 'bert.encoder.layer.12.intermediate.dense_act.up_proj.tensor.factors.1', 'bert.encoder.layer.12.output.dense.bias', 'bert.encoder.layer.12.output.dense.tensor.factors.0', 'bert.encoder.layer.12.output.dense.tensor.factors.1', 'bert.encoder.layer.12.output.LayerNorm.weight', 'bert.encoder.layer.12.output.LayerNorm.bias', 'bert.encoder.layer.13.attention.self.query.bias', 'bert.encoder.layer.13.attention.self.query.tensor.factors.0', 'bert.encoder.layer.13.attention.self.query.tensor.factors.1', 'bert.encoder.layer.13.attention.self.key.bias', 'bert.encoder.layer.13.attention.self.key.tensor.factors.0', 'bert.encoder.layer.13.attention.self.key.tensor.factors.1', 'bert.encoder.layer.13.attention.self.value.bias', 'bert.encoder.layer.13.attention.self.value.tensor.factors.0', 'bert.encoder.layer.13.attention.self.value.tensor.factors.1', 'bert.encoder.layer.13.attention.output.dense.bias', 'bert.encoder.layer.13.attention.output.dense.tensor.factors.0', 'bert.encoder.layer.13.attention.output.dense.tensor.factors.1', 'bert.encoder.layer.13.attention.output.LayerNorm.weight', 'bert.encoder.layer.13.attention.output.LayerNorm.bias', 'bert.encoder.layer.13.intermediate.dense_act.gate_proj.bias', 'bert.encoder.layer.13.intermediate.dense_act.gate_proj.tensor.factors.0', 'bert.encoder.layer.13.intermediate.dense_act.gate_proj.tensor.factors.1', 'bert.encoder.layer.13.intermediate.dense_act.up_proj.bias', 'bert.encoder.layer.13.intermediate.dense_act.up_proj.tensor.factors.0', 'bert.encoder.layer.13.intermediate.dense_act.up_proj.tensor.factors.1', 'bert.encoder.layer.13.output.dense.bias', 'bert.encoder.layer.13.output.dense.tensor.factors.0', 'bert.encoder.layer.13.output.dense.tensor.factors.1', 'bert.encoder.layer.13.output.LayerNorm.weight', 'bert.encoder.layer.13.output.LayerNorm.bias', 'bert.encoder.layer.14.attention.self.query.bias', 'bert.encoder.layer.14.attention.self.query.tensor.factors.0', 'bert.encoder.layer.14.attention.self.query.tensor.factors.1', 'bert.encoder.layer.14.attention.self.key.bias', 'bert.encoder.layer.14.attention.self.key.tensor.factors.0', 'bert.encoder.layer.14.attention.self.key.tensor.factors.1', 'bert.encoder.layer.14.attention.self.value.bias', 'bert.encoder.layer.14.attention.self.value.tensor.factors.0', 'bert.encoder.layer.14.attention.self.value.tensor.factors.1', 'bert.encoder.layer.14.attention.output.dense.bias', 'bert.encoder.layer.14.attention.output.dense.tensor.factors.0', 'bert.encoder.layer.14.attention.output.dense.tensor.factors.1', 'bert.encoder.layer.14.attention.output.LayerNorm.weight', 'bert.encoder.layer.14.attention.output.LayerNorm.bias', 'bert.encoder.layer.14.intermediate.dense_act.gate_proj.bias', 'bert.encoder.layer.14.intermediate.dense_act.gate_proj.tensor.factors.0', 'bert.encoder.layer.14.intermediate.dense_act.gate_proj.tensor.factors.1', 'bert.encoder.layer.14.intermediate.dense_act.up_proj.bias', 'bert.encoder.layer.14.intermediate.dense_act.up_proj.tensor.factors.0', 'bert.encoder.layer.14.intermediate.dense_act.up_proj.tensor.factors.1', 'bert.encoder.layer.14.output.dense.bias', 'bert.encoder.layer.14.output.dense.tensor.factors.0', 'bert.encoder.layer.14.output.dense.tensor.factors.1', 'bert.encoder.layer.14.output.LayerNorm.weight', 'bert.encoder.layer.14.output.LayerNorm.bias', 'bert.encoder.layer.15.attention.self.query.bias', 'bert.encoder.layer.15.attention.self.query.tensor.factors.0', 'bert.encoder.layer.15.attention.self.query.tensor.factors.1', 'bert.encoder.layer.15.attention.self.key.bias', 'bert.encoder.layer.15.attention.self.key.tensor.factors.0', 'bert.encoder.layer.15.attention.self.key.tensor.factors.1', 'bert.encoder.layer.15.attention.self.value.bias', 'bert.encoder.layer.15.attention.self.value.tensor.factors.0', 'bert.encoder.layer.15.attention.self.value.tensor.factors.1', 'bert.encoder.layer.15.attention.output.dense.bias', 'bert.encoder.layer.15.attention.output.dense.tensor.factors.0', 'bert.encoder.layer.15.attention.output.dense.tensor.factors.1', 'bert.encoder.layer.15.attention.output.LayerNorm.weight', 'bert.encoder.layer.15.attention.output.LayerNorm.bias', 'bert.encoder.layer.15.intermediate.dense_act.gate_proj.bias', 'bert.encoder.layer.15.intermediate.dense_act.gate_proj.tensor.factors.0', 'bert.encoder.layer.15.intermediate.dense_act.gate_proj.tensor.factors.1', 'bert.encoder.layer.15.intermediate.dense_act.up_proj.bias', 'bert.encoder.layer.15.intermediate.dense_act.up_proj.tensor.factors.0', 'bert.encoder.layer.15.intermediate.dense_act.up_proj.tensor.factors.1', 'bert.encoder.layer.15.output.dense.bias', 'bert.encoder.layer.15.output.dense.tensor.factors.0', 'bert.encoder.layer.15.output.dense.tensor.factors.1', 'bert.encoder.layer.15.output.LayerNorm.weight', 'bert.encoder.layer.15.output.LayerNorm.bias', 'bert.encoder.layer.16.attention.self.query.bias', 'bert.encoder.layer.16.attention.self.query.tensor.factors.0', 'bert.encoder.layer.16.attention.self.query.tensor.factors.1', 'bert.encoder.layer.16.attention.self.key.bias', 'bert.encoder.layer.16.attention.self.key.tensor.factors.0', 'bert.encoder.layer.16.attention.self.key.tensor.factors.1', 'bert.encoder.layer.16.attention.self.value.bias', 'bert.encoder.layer.16.attention.self.value.tensor.factors.0', 'bert.encoder.layer.16.attention.self.value.tensor.factors.1', 'bert.encoder.layer.16.attention.output.dense.bias', 'bert.encoder.layer.16.attention.output.dense.tensor.factors.0', 'bert.encoder.layer.16.attention.output.dense.tensor.factors.1', 'bert.encoder.layer.16.attention.output.LayerNorm.weight', 'bert.encoder.layer.16.attention.output.LayerNorm.bias', 'bert.encoder.layer.16.intermediate.dense_act.gate_proj.bias', 'bert.encoder.layer.16.intermediate.dense_act.gate_proj.tensor.factors.0', 'bert.encoder.layer.16.intermediate.dense_act.gate_proj.tensor.factors.1', 'bert.encoder.layer.16.intermediate.dense_act.up_proj.bias', 'bert.encoder.layer.16.intermediate.dense_act.up_proj.tensor.factors.0', 'bert.encoder.layer.16.intermediate.dense_act.up_proj.tensor.factors.1', 'bert.encoder.layer.16.output.dense.bias', 'bert.encoder.layer.16.output.dense.tensor.factors.0', 'bert.encoder.layer.16.output.dense.tensor.factors.1', 'bert.encoder.layer.16.output.LayerNorm.weight', 'bert.encoder.layer.16.output.LayerNorm.bias', 'bert.encoder.layer.17.attention.self.query.bias', 'bert.encoder.layer.17.attention.self.query.tensor.factors.0', 'bert.encoder.layer.17.attention.self.query.tensor.factors.1', 'bert.encoder.layer.17.attention.self.key.bias', 'bert.encoder.layer.17.attention.self.key.tensor.factors.0', 'bert.encoder.layer.17.attention.self.key.tensor.factors.1', 'bert.encoder.layer.17.attention.self.value.bias', 'bert.encoder.layer.17.attention.self.value.tensor.factors.0', 'bert.encoder.layer.17.attention.self.value.tensor.factors.1', 'bert.encoder.layer.17.attention.output.dense.bias', 'bert.encoder.layer.17.attention.output.dense.tensor.factors.0', 'bert.encoder.layer.17.attention.output.dense.tensor.factors.1', 'bert.encoder.layer.17.attention.output.LayerNorm.weight', 'bert.encoder.layer.17.attention.output.LayerNorm.bias', 'bert.encoder.layer.17.intermediate.dense_act.gate_proj.bias', 'bert.encoder.layer.17.intermediate.dense_act.gate_proj.tensor.factors.0', 'bert.encoder.layer.17.intermediate.dense_act.gate_proj.tensor.factors.1', 'bert.encoder.layer.17.intermediate.dense_act.up_proj.bias', 'bert.encoder.layer.17.intermediate.dense_act.up_proj.tensor.factors.0', 'bert.encoder.layer.17.intermediate.dense_act.up_proj.tensor.factors.1', 'bert.encoder.layer.17.output.dense.bias', 'bert.encoder.layer.17.output.dense.tensor.factors.0', 'bert.encoder.layer.17.output.dense.tensor.factors.1', 'bert.encoder.layer.17.output.LayerNorm.weight', 'bert.encoder.layer.17.output.LayerNorm.bias', 'bert.encoder.layer.18.attention.self.query.bias', 'bert.encoder.layer.18.attention.self.query.tensor.factors.0', 'bert.encoder.layer.18.attention.self.query.tensor.factors.1', 'bert.encoder.layer.18.attention.self.key.bias', 'bert.encoder.layer.18.attention.self.key.tensor.factors.0', 'bert.encoder.layer.18.attention.self.key.tensor.factors.1', 'bert.encoder.layer.18.attention.self.value.bias', 'bert.encoder.layer.18.attention.self.value.tensor.factors.0', 'bert.encoder.layer.18.attention.self.value.tensor.factors.1', 'bert.encoder.layer.18.attention.output.dense.bias', 'bert.encoder.layer.18.attention.output.dense.tensor.factors.0', 'bert.encoder.layer.18.attention.output.dense.tensor.factors.1', 'bert.encoder.layer.18.attention.output.LayerNorm.weight', 'bert.encoder.layer.18.attention.output.LayerNorm.bias', 'bert.encoder.layer.18.intermediate.dense_act.gate_proj.bias', 'bert.encoder.layer.18.intermediate.dense_act.gate_proj.tensor.factors.0', 'bert.encoder.layer.18.intermediate.dense_act.gate_proj.tensor.factors.1', 'bert.encoder.layer.18.intermediate.dense_act.up_proj.bias', 'bert.encoder.layer.18.intermediate.dense_act.up_proj.tensor.factors.0', 'bert.encoder.layer.18.intermediate.dense_act.up_proj.tensor.factors.1', 'bert.encoder.layer.18.output.dense.bias', 'bert.encoder.layer.18.output.dense.tensor.factors.0', 'bert.encoder.layer.18.output.dense.tensor.factors.1', 'bert.encoder.layer.18.output.LayerNorm.weight', 'bert.encoder.layer.18.output.LayerNorm.bias', 'bert.encoder.layer.19.attention.self.query.bias', 'bert.encoder.layer.19.attention.self.query.tensor.factors.0', 'bert.encoder.layer.19.attention.self.query.tensor.factors.1', 'bert.encoder.layer.19.attention.self.key.bias', 'bert.encoder.layer.19.attention.self.key.tensor.factors.0', 'bert.encoder.layer.19.attention.self.key.tensor.factors.1', 'bert.encoder.layer.19.attention.self.value.bias', 'bert.encoder.layer.19.attention.self.value.tensor.factors.0', 'bert.encoder.layer.19.attention.self.value.tensor.factors.1', 'bert.encoder.layer.19.attention.output.dense.bias', 'bert.encoder.layer.19.attention.output.dense.tensor.factors.0', 'bert.encoder.layer.19.attention.output.dense.tensor.factors.1', 'bert.encoder.layer.19.attention.output.LayerNorm.weight', 'bert.encoder.layer.19.attention.output.LayerNorm.bias', 'bert.encoder.layer.19.intermediate.dense_act.gate_proj.bias', 'bert.encoder.layer.19.intermediate.dense_act.gate_proj.tensor.factors.0', 'bert.encoder.layer.19.intermediate.dense_act.gate_proj.tensor.factors.1', 'bert.encoder.layer.19.intermediate.dense_act.up_proj.bias', 'bert.encoder.layer.19.intermediate.dense_act.up_proj.tensor.factors.0', 'bert.encoder.layer.19.intermediate.dense_act.up_proj.tensor.factors.1', 'bert.encoder.layer.19.output.dense.bias', 'bert.encoder.layer.19.output.dense.tensor.factors.0', 'bert.encoder.layer.19.output.dense.tensor.factors.1', 'bert.encoder.layer.19.output.LayerNorm.weight', 'bert.encoder.layer.19.output.LayerNorm.bias', 'bert.encoder.layer.20.attention.self.query.bias', 'bert.encoder.layer.20.attention.self.query.tensor.factors.0', 'bert.encoder.layer.20.attention.self.query.tensor.factors.1', 'bert.encoder.layer.20.attention.self.key.bias', 'bert.encoder.layer.20.attention.self.key.tensor.factors.0', 'bert.encoder.layer.20.attention.self.key.tensor.factors.1', 'bert.encoder.layer.20.attention.self.value.bias', 'bert.encoder.layer.20.attention.self.value.tensor.factors.0', 'bert.encoder.layer.20.attention.self.value.tensor.factors.1', 'bert.encoder.layer.20.attention.output.dense.bias', 'bert.encoder.layer.20.attention.output.dense.tensor.factors.0', 'bert.encoder.layer.20.attention.output.dense.tensor.factors.1', 'bert.encoder.layer.20.attention.output.LayerNorm.weight', 'bert.encoder.layer.20.attention.output.LayerNorm.bias', 'bert.encoder.layer.20.intermediate.dense_act.gate_proj.bias', 'bert.encoder.layer.20.intermediate.dense_act.gate_proj.tensor.factors.0', 'bert.encoder.layer.20.intermediate.dense_act.gate_proj.tensor.factors.1', 'bert.encoder.layer.20.intermediate.dense_act.up_proj.bias', 'bert.encoder.layer.20.intermediate.dense_act.up_proj.tensor.factors.0', 'bert.encoder.layer.20.intermediate.dense_act.up_proj.tensor.factors.1', 'bert.encoder.layer.20.output.dense.bias', 'bert.encoder.layer.20.output.dense.tensor.factors.0', 'bert.encoder.layer.20.output.dense.tensor.factors.1', 'bert.encoder.layer.20.output.LayerNorm.weight', 'bert.encoder.layer.20.output.LayerNorm.bias', 'bert.encoder.layer.21.attention.self.query.bias', 'bert.encoder.layer.21.attention.self.query.tensor.factors.0', 'bert.encoder.layer.21.attention.self.query.tensor.factors.1', 'bert.encoder.layer.21.attention.self.key.bias', 'bert.encoder.layer.21.attention.self.key.tensor.factors.0', 'bert.encoder.layer.21.attention.self.key.tensor.factors.1', 'bert.encoder.layer.21.attention.self.value.bias', 'bert.encoder.layer.21.attention.self.value.tensor.factors.0', 'bert.encoder.layer.21.attention.self.value.tensor.factors.1', 'bert.encoder.layer.21.attention.output.dense.bias', 'bert.encoder.layer.21.attention.output.dense.tensor.factors.0', 'bert.encoder.layer.21.attention.output.dense.tensor.factors.1', 'bert.encoder.layer.21.attention.output.LayerNorm.weight', 'bert.encoder.layer.21.attention.output.LayerNorm.bias', 'bert.encoder.layer.21.intermediate.dense_act.gate_proj.bias', 'bert.encoder.layer.21.intermediate.dense_act.gate_proj.tensor.factors.0', 'bert.encoder.layer.21.intermediate.dense_act.gate_proj.tensor.factors.1', 'bert.encoder.layer.21.intermediate.dense_act.up_proj.bias', 'bert.encoder.layer.21.intermediate.dense_act.up_proj.tensor.factors.0', 'bert.encoder.layer.21.intermediate.dense_act.up_proj.tensor.factors.1', 'bert.encoder.layer.21.output.dense.bias', 'bert.encoder.layer.21.output.dense.tensor.factors.0', 'bert.encoder.layer.21.output.dense.tensor.factors.1', 'bert.encoder.layer.21.output.LayerNorm.weight', 'bert.encoder.layer.21.output.LayerNorm.bias', 'bert.encoder.layer.22.attention.self.query.bias', 'bert.encoder.layer.22.attention.self.query.tensor.factors.0', 'bert.encoder.layer.22.attention.self.query.tensor.factors.1', 'bert.encoder.layer.22.attention.self.key.bias', 'bert.encoder.layer.22.attention.self.key.tensor.factors.0', 'bert.encoder.layer.22.attention.self.key.tensor.factors.1', 'bert.encoder.layer.22.attention.self.value.bias', 'bert.encoder.layer.22.attention.self.value.tensor.factors.0', 'bert.encoder.layer.22.attention.self.value.tensor.factors.1', 'bert.encoder.layer.22.attention.output.dense.bias', 'bert.encoder.layer.22.attention.output.dense.tensor.factors.0', 'bert.encoder.layer.22.attention.output.dense.tensor.factors.1', 'bert.encoder.layer.22.attention.output.LayerNorm.weight', 'bert.encoder.layer.22.attention.output.LayerNorm.bias', 'bert.encoder.layer.22.intermediate.dense_act.gate_proj.bias', 'bert.encoder.layer.22.intermediate.dense_act.gate_proj.tensor.factors.0', 'bert.encoder.layer.22.intermediate.dense_act.gate_proj.tensor.factors.1', 'bert.encoder.layer.22.intermediate.dense_act.up_proj.bias', 'bert.encoder.layer.22.intermediate.dense_act.up_proj.tensor.factors.0', 'bert.encoder.layer.22.intermediate.dense_act.up_proj.tensor.factors.1', 'bert.encoder.layer.22.output.dense.bias', 'bert.encoder.layer.22.output.dense.tensor.factors.0', 'bert.encoder.layer.22.output.dense.tensor.factors.1', 'bert.encoder.layer.22.output.LayerNorm.weight', 'bert.encoder.layer.22.output.LayerNorm.bias', 'bert.encoder.layer.23.attention.self.query.bias', 'bert.encoder.layer.23.attention.self.query.tensor.factors.0', 'bert.encoder.layer.23.attention.self.query.tensor.factors.1', 'bert.encoder.layer.23.attention.self.key.bias', 'bert.encoder.layer.23.attention.self.key.tensor.factors.0', 'bert.encoder.layer.23.attention.self.key.tensor.factors.1', 'bert.encoder.layer.23.attention.self.value.bias', 'bert.encoder.layer.23.attention.self.value.tensor.factors.0', 'bert.encoder.layer.23.attention.self.value.tensor.factors.1', 'bert.encoder.layer.23.attention.output.dense.bias', 'bert.encoder.layer.23.attention.output.dense.tensor.factors.0', 'bert.encoder.layer.23.attention.output.dense.tensor.factors.1', 'bert.encoder.layer.23.attention.output.LayerNorm.weight', 'bert.encoder.layer.23.attention.output.LayerNorm.bias', 'bert.encoder.layer.23.intermediate.dense_act.gate_proj.bias', 'bert.encoder.layer.23.intermediate.dense_act.gate_proj.tensor.factors.0', 'bert.encoder.layer.23.intermediate.dense_act.gate_proj.tensor.factors.1', 'bert.encoder.layer.23.intermediate.dense_act.up_proj.bias', 'bert.encoder.layer.23.intermediate.dense_act.up_proj.tensor.factors.0', 'bert.encoder.layer.23.intermediate.dense_act.up_proj.tensor.factors.1', 'bert.encoder.layer.23.output.dense.bias', 'bert.encoder.layer.23.output.dense.tensor.factors.0', 'bert.encoder.layer.23.output.dense.tensor.factors.1', 'bert.encoder.layer.23.output.LayerNorm.weight', 'bert.encoder.layer.23.output.LayerNorm.bias', 'bert.pooler.dense_act.linear.weight', 'bert.pooler.dense_act.linear.bias', 'cls.predictions.bias', 'cls.predictions.transform.dense_act.linear.weight', 'cls.predictions.transform.dense_act.linear.bias', 'cls.predictions.transform.LayerNorm.weight', 'cls.predictions.transform.LayerNorm.bias', 'cls.predictions.decoder.weight', 'cls.seq_relationship.weight', 'cls.seq_relationship.bias'])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "_IncompatibleKeys(missing_keys=['classifier.weight', 'classifier.bias'], unexpected_keys=['cls.predictions.bias', 'cls.predictions.transform.dense_act.linear.weight', 'cls.predictions.transform.dense_act.linear.bias', 'cls.predictions.transform.LayerNorm.weight', 'cls.predictions.transform.LayerNorm.bias', 'cls.predictions.decoder.weight', 'cls.seq_relationship.weight', 'cls.seq_relationship.bias'])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "checkpoint = torch.load(chkpoint_path + \"/model.pt\")\n",
    "checkpoint = checkpoint[\"model\"] if \"model\" in checkpoint.keys() else checkpoint\n",
    "print(checkpoint.keys())\n",
    "\n",
    "model.load_state_dict(checkpoint, strict=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "25e110ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tokenization import BertTokenizer\n",
    "vocab_file = \"../vocab\"\n",
    "tokenizer = BertTokenizer(\n",
    "        vocab_file = vocab_file,\n",
    "        do_lower_case=True,\n",
    "        max_len=512,\n",
    "    ) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "3cb94bc0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# from transformers import BertTokenizer\n",
    "# tok2 = BertTokenizer.from_pretrained(\"bert-large-uncased\", do_lower_case=True, max_len=512)\n",
    "# tok2.vocab_size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "32e2c801",
   "metadata": {},
   "outputs": [],
   "source": [
    "from processors.glue import PROCESSORS, convert_examples_to_features\n",
    "import pickle\n",
    "import logging\n",
    "from utils import is_main_process, mkdir_by_main_process, format_step, get_world_size\n",
    "from torch.utils.data import DataLoader, RandomSampler, SequentialSampler, TensorDataset\n",
    "\n",
    "\n",
    "logger = logging.getLogger(__name__)\n",
    "\n",
    "def get_train_features(\n",
    "    data_dir,\n",
    "    bert_model,\n",
    "    max_seq_length,\n",
    "    do_lower_case,\n",
    "    local_rank,\n",
    "    train_batch_size,\n",
    "    gradient_accumulation_steps,\n",
    "    num_train_epochs,\n",
    "    tokenizer,\n",
    "    processor,\n",
    "):\n",
    "    cached_train_features_file = os.path.join(\n",
    "        data_dir,\n",
    "        \"{0}_{1}_{2}\".format(\n",
    "            list(filter(None, bert_model.split(\"/\"))).pop(),\n",
    "            str(max_seq_length),\n",
    "            str(do_lower_case),\n",
    "        ),\n",
    "    )\n",
    "    train_features = None\n",
    "    try:\n",
    "        with open(cached_train_features_file, \"rb\") as reader:\n",
    "            train_features = pickle.load(reader)\n",
    "        logger.info(\n",
    "            \"Loaded pre-processed features from {}\".format(cached_train_features_file)\n",
    "        )\n",
    "    except:\n",
    "        logger.info(\n",
    "            \"Did not find pre-processed features from {}\".format(\n",
    "                cached_train_features_file\n",
    "            )\n",
    "        )\n",
    "        train_examples = processor.get_train_examples(data_dir)\n",
    "        train_features, _ = convert_examples_to_features(\n",
    "            train_examples,\n",
    "            processor.get_labels(),\n",
    "            max_seq_length,\n",
    "            tokenizer,\n",
    "        )\n",
    "        if is_main_process():\n",
    "            logger.info(\n",
    "                \"  Saving train features into cached file %s\",\n",
    "                cached_train_features_file,\n",
    "            )\n",
    "            with open(cached_train_features_file, \"wb\") as writer:\n",
    "                pickle.dump(train_features, writer)\n",
    "    return train_features\n",
    "\n",
    "\n",
    "\n",
    "def gen_tensor_dataset(features):\n",
    "    all_input_ids = torch.tensor(\n",
    "        [f.input_ids for f in features],\n",
    "        dtype=torch.long,\n",
    "    )\n",
    "    all_input_mask = torch.tensor(\n",
    "        [f.input_mask for f in features],\n",
    "        dtype=torch.long,\n",
    "    )\n",
    "    all_segment_ids = torch.tensor(\n",
    "        [f.segment_ids for f in features],\n",
    "        dtype=torch.long,\n",
    "    )\n",
    "    all_label_ids = torch.tensor(\n",
    "        [f.label_id for f in features],\n",
    "        dtype=torch.long,\n",
    "    )\n",
    "    return TensorDataset(\n",
    "        all_input_ids,\n",
    "        all_input_mask,\n",
    "        all_segment_ids,\n",
    "        all_label_ids,\n",
    "    )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "743b50bd",
   "metadata": {},
   "outputs": [
    {
     "ename": "FileNotFoundError",
     "evalue": "[Errno 2] No such file or directory: '../glue_data/WNLI/train.tsv'",
     "output_type": "error",
     "traceback": [
      "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
      "\u001b[31mFileNotFoundError\u001b[39m                         Traceback (most recent call last)",
      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[10]\u001b[39m\u001b[32m, line 32\u001b[39m, in \u001b[36mget_train_features\u001b[39m\u001b[34m(data_dir, bert_model, max_seq_length, do_lower_case, local_rank, train_batch_size, gradient_accumulation_steps, num_train_epochs, tokenizer, processor)\u001b[39m\n\u001b[32m     31\u001b[39m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[32m---> \u001b[39m\u001b[32m32\u001b[39m     \u001b[38;5;28;01mwith\u001b[39;00m \u001b[38;5;28;43mopen\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mcached_train_features_file\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[33;43m\"\u001b[39;49m\u001b[33;43mrb\u001b[39;49m\u001b[33;43m\"\u001b[39;49m\u001b[43m)\u001b[49m \u001b[38;5;28;01mas\u001b[39;00m reader:\n\u001b[32m     33\u001b[39m         train_features = pickle.load(reader)\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/miniforge3/envs/colaBert/lib/python3.11/site-packages/IPython/core/interactiveshell.py:343\u001b[39m, in \u001b[36m_modified_open\u001b[39m\u001b[34m(file, *args, **kwargs)\u001b[39m\n\u001b[32m    337\u001b[39m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\n\u001b[32m    338\u001b[39m         \u001b[33mf\u001b[39m\u001b[33m\"\u001b[39m\u001b[33mIPython won\u001b[39m\u001b[33m'\u001b[39m\u001b[33mt let you open fd=\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mfile\u001b[38;5;132;01m}\u001b[39;00m\u001b[33m by default \u001b[39m\u001b[33m\"\u001b[39m\n\u001b[32m    339\u001b[39m         \u001b[33m\"\u001b[39m\u001b[33mas it is likely to crash IPython. If you know what you are doing, \u001b[39m\u001b[33m\"\u001b[39m\n\u001b[32m    340\u001b[39m         \u001b[33m\"\u001b[39m\u001b[33myou can use builtins\u001b[39m\u001b[33m'\u001b[39m\u001b[33m open.\u001b[39m\u001b[33m\"\u001b[39m\n\u001b[32m    341\u001b[39m     )\n\u001b[32m--> \u001b[39m\u001b[32m343\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mio_open\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfile\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m*\u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m*\u001b[49m\u001b[43m*\u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[31mFileNotFoundError\u001b[39m: [Errno 2] No such file or directory: '../glue_data/WNLI/tensor-bert_512_True'",
      "\nDuring handling of the above exception, another exception occurred:\n",
      "\u001b[31mFileNotFoundError\u001b[39m                         Traceback (most recent call last)",
      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[11]\u001b[39m\u001b[32m, line 13\u001b[39m\n\u001b[32m     10\u001b[39m tokenizer = tokenizer\n\u001b[32m     11\u001b[39m processor = processor\n\u001b[32m---> \u001b[39m\u001b[32m13\u001b[39m train_features = \u001b[43mget_train_features\u001b[49m\u001b[43m(\u001b[49m\n\u001b[32m     14\u001b[39m \u001b[43m            \u001b[49m\u001b[43mdata_dir\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m     15\u001b[39m \u001b[43m            \u001b[49m\u001b[43mbert_model\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m     16\u001b[39m \u001b[43m            \u001b[49m\u001b[43mmax_seq_length\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m     17\u001b[39m \u001b[43m            \u001b[49m\u001b[43mdo_lower_case\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m     18\u001b[39m \u001b[43m            \u001b[49m\u001b[43mlocal_rank\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m     19\u001b[39m \u001b[43m            \u001b[49m\u001b[43mtrain_batch_size\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m     20\u001b[39m \u001b[43m            \u001b[49m\u001b[43mgradient_accumulation_steps\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m     21\u001b[39m \u001b[43m            \u001b[49m\u001b[43mnum_train_epochs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m     22\u001b[39m \u001b[43m            \u001b[49m\u001b[43mtokenizer\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m     23\u001b[39m \u001b[43m            \u001b[49m\u001b[43mprocessor\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m     24\u001b[39m \u001b[43m        \u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m     25\u001b[39m train_data = gen_tensor_dataset(train_features)\n\u001b[32m     26\u001b[39m train_sampler = RandomSampler(train_data)\n",
      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[10]\u001b[39m\u001b[32m, line 43\u001b[39m, in \u001b[36mget_train_features\u001b[39m\u001b[34m(data_dir, bert_model, max_seq_length, do_lower_case, local_rank, train_batch_size, gradient_accumulation_steps, num_train_epochs, tokenizer, processor)\u001b[39m\n\u001b[32m     37\u001b[39m \u001b[38;5;28;01mexcept\u001b[39;00m:\n\u001b[32m     38\u001b[39m     logger.info(\n\u001b[32m     39\u001b[39m         \u001b[33m\"\u001b[39m\u001b[33mDid not find pre-processed features from \u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[33m\"\u001b[39m.format(\n\u001b[32m     40\u001b[39m             cached_train_features_file\n\u001b[32m     41\u001b[39m         )\n\u001b[32m     42\u001b[39m     )\n\u001b[32m---> \u001b[39m\u001b[32m43\u001b[39m     train_examples = \u001b[43mprocessor\u001b[49m\u001b[43m.\u001b[49m\u001b[43mget_train_examples\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdata_dir\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m     44\u001b[39m     train_features, _ = convert_examples_to_features(\n\u001b[32m     45\u001b[39m         train_examples,\n\u001b[32m     46\u001b[39m         processor.get_labels(),\n\u001b[32m     47\u001b[39m         max_seq_length,\n\u001b[32m     48\u001b[39m         tokenizer,\n\u001b[32m     49\u001b[39m     )\n\u001b[32m     50\u001b[39m     \u001b[38;5;28;01mif\u001b[39;00m is_main_process():\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/Documents/TensorBERT-FT/processors/glue.py:386\u001b[39m, in \u001b[36mWnliProcessor.get_train_examples\u001b[39m\u001b[34m(self, data_dir)\u001b[39m\n\u001b[32m    383\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34mget_train_examples\u001b[39m(\u001b[38;5;28mself\u001b[39m, data_dir):\n\u001b[32m    384\u001b[39m \u001b[38;5;250m    \u001b[39m\u001b[33;03m\"\"\"See base class.\"\"\"\u001b[39;00m\n\u001b[32m    385\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m._create_examples(\n\u001b[32m--> \u001b[39m\u001b[32m386\u001b[39m         \u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43m_read_tsv\u001b[49m\u001b[43m(\u001b[49m\u001b[43mos\u001b[49m\u001b[43m.\u001b[49m\u001b[43mpath\u001b[49m\u001b[43m.\u001b[49m\u001b[43mjoin\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdata_dir\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[33;43m\"\u001b[39;49m\u001b[33;43mtrain.tsv\u001b[39;49m\u001b[33;43m\"\u001b[39;49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m,\n\u001b[32m    387\u001b[39m         \u001b[33m\"\u001b[39m\u001b[33mtrain\u001b[39m\u001b[33m\"\u001b[39m,\n\u001b[32m    388\u001b[39m     )\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/Documents/TensorBERT-FT/processors/glue.py:72\u001b[39m, in \u001b[36mDataProcessor._read_tsv\u001b[39m\u001b[34m(cls, input_file, quotechar)\u001b[39m\n\u001b[32m     69\u001b[39m \u001b[38;5;129m@classmethod\u001b[39m\n\u001b[32m     70\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34m_read_tsv\u001b[39m(\u001b[38;5;28mcls\u001b[39m, input_file, quotechar=\u001b[38;5;28;01mNone\u001b[39;00m):\n\u001b[32m     71\u001b[39m \u001b[38;5;250m    \u001b[39m\u001b[33;03m\"\"\"Reads a tab separated value file.\"\"\"\u001b[39;00m\n\u001b[32m---> \u001b[39m\u001b[32m72\u001b[39m     \u001b[38;5;28;01mwith\u001b[39;00m \u001b[38;5;28;43mopen\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43minput_file\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[33;43m\"\u001b[39;49m\u001b[33;43mr\u001b[39;49m\u001b[33;43m\"\u001b[39;49m\u001b[43m)\u001b[49m \u001b[38;5;28;01mas\u001b[39;00m f:\n\u001b[32m     73\u001b[39m         reader = csv.reader(f, delimiter=\u001b[33m\"\u001b[39m\u001b[38;5;130;01m\\t\u001b[39;00m\u001b[33m\"\u001b[39m, quotechar=quotechar)\n\u001b[32m     74\u001b[39m         lines = []\n",
      "\u001b[31mFileNotFoundError\u001b[39m: [Errno 2] No such file or directory: '../glue_data/WNLI/train.tsv'"
     ]
    }
   ],
   "source": [
    "data_dir = \"../glue_data/WNLI\"\n",
    "bert_model = \"tensor-bert\"\n",
    "max_seq_length = max_pos_embeddings = 512\n",
    "do_lower_case = True\n",
    "local_rank = 0\n",
    "train_batch_size = 8\n",
    "gradient_accumulation_steps = 1\n",
    "num_train_epochs = 3\n",
    "\n",
    "tokenizer = tokenizer\n",
    "processor = processor\n",
    "\n",
    "train_features = get_train_features(\n",
    "            data_dir,\n",
    "            bert_model,\n",
    "            max_seq_length,\n",
    "            do_lower_case,\n",
    "            local_rank,\n",
    "            train_batch_size,\n",
    "            gradient_accumulation_steps,\n",
    "            num_train_epochs,\n",
    "            tokenizer,\n",
    "            processor,\n",
    "        )\n",
    "train_data = gen_tensor_dataset(train_features)\n",
    "train_sampler = RandomSampler(train_data)\n",
    "train_dataloader = DataLoader(\n",
    "            train_data,\n",
    "            sampler=train_sampler,\n",
    "            batch_size=train_batch_size,\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b97b356d",
   "metadata": {},
   "source": [
    "# Explore Model a bit\n",
    "- attention\n",
    "> - q, k, v\n",
    "> - output \n",
    "- intermediate \n",
    "> - gate proj\n",
    "> - up proj\n",
    "- out proj"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "ca2a2363",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn.functional as F\n",
    "from transformers.activations import gelu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "42117cd5",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.eval();\n",
    "encoderLayers = model.bert.encoder.layer\n",
    "block = encoderLayers[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "b4ee6130",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TensorBertLayer(\n",
       "  (attention): TensorBertAttention(\n",
       "    (self): TensorBertSelfAttention(\n",
       "      (query): TensorizedLinear(\n",
       "        (lr_act): GELU(approximate='none')\n",
       "        (tensor): TensorTrain(\n",
       "          (factors): ParameterList(\n",
       "              (0): Parameter containing: [torch.float32 of size 1x1024x384]\n",
       "              (1): Parameter containing: [torch.float32 of size 384x1024x1]\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (key): TensorizedLinear(\n",
       "        (lr_act): GELU(approximate='none')\n",
       "        (tensor): TensorTrain(\n",
       "          (factors): ParameterList(\n",
       "              (0): Parameter containing: [torch.float32 of size 1x1024x384]\n",
       "              (1): Parameter containing: [torch.float32 of size 384x1024x1]\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (value): TensorizedLinear(\n",
       "        (lr_act): GELU(approximate='none')\n",
       "        (tensor): TensorTrain(\n",
       "          (factors): ParameterList(\n",
       "              (0): Parameter containing: [torch.float32 of size 1x1024x384]\n",
       "              (1): Parameter containing: [torch.float32 of size 384x1024x1]\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (dropout): Dropout(p=0.1, inplace=False)\n",
       "    )\n",
       "    (output): TensorBertSelfOutput(\n",
       "      (dense): TensorizedLinear(\n",
       "        (lr_act): GELU(approximate='none')\n",
       "        (tensor): TensorTrain(\n",
       "          (factors): ParameterList(\n",
       "              (0): Parameter containing: [torch.float32 of size 1x1024x384]\n",
       "              (1): Parameter containing: [torch.float32 of size 384x1024x1]\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (LayerNorm): LayerNorm((1024,), eps=1e-12, elementwise_affine=True)\n",
       "      (dropout): Dropout(p=0.1, inplace=False)\n",
       "    )\n",
       "  )\n",
       "  (intermediate): TensorBertIntermediate(\n",
       "    (dense_act): TensorLinearActivation(\n",
       "      (gate_proj): TensorizedLinear(\n",
       "        (lr_act): GELU(approximate='none')\n",
       "        (tensor): TensorTrain(\n",
       "          (factors): ParameterList(\n",
       "              (0): Parameter containing: [torch.float32 of size 1x1024x512]\n",
       "              (1): Parameter containing: [torch.float32 of size 512x2736x1]\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (up_proj): TensorizedLinear(\n",
       "        (lr_act): GELU(approximate='none')\n",
       "        (tensor): TensorTrain(\n",
       "          (factors): ParameterList(\n",
       "              (0): Parameter containing: [torch.float32 of size 1x1024x512]\n",
       "              (1): Parameter containing: [torch.float32 of size 512x2736x1]\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (output): TensorBertOutput(\n",
       "    (dense): TensorizedLinear(\n",
       "      (lr_act): GELU(approximate='none')\n",
       "      (tensor): TensorTrain(\n",
       "        (factors): ParameterList(\n",
       "            (0): Parameter containing: [torch.float32 of size 1x2736x512]\n",
       "            (1): Parameter containing: [torch.float32 of size 512x1024x1]\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (LayerNorm): LayerNorm((1024,), eps=1e-12, elementwise_affine=True)\n",
       "    (dropout): Dropout(p=0.1, inplace=False)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "block"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5b48085",
   "metadata": {},
   "source": [
    "### Attention Block\n",
    "- inputs are seq, bsc, hidden"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "3372b9b9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TensorBertSelfAttention(\n",
       "  (query): TensorizedLinear(\n",
       "    (lr_act): GELU(approximate='none')\n",
       "    (tensor): TensorTrain(\n",
       "      (factors): ParameterList(\n",
       "          (0): Parameter containing: [torch.float32 of size 1x1024x384]\n",
       "          (1): Parameter containing: [torch.float32 of size 384x1024x1]\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (key): TensorizedLinear(\n",
       "    (lr_act): GELU(approximate='none')\n",
       "    (tensor): TensorTrain(\n",
       "      (factors): ParameterList(\n",
       "          (0): Parameter containing: [torch.float32 of size 1x1024x384]\n",
       "          (1): Parameter containing: [torch.float32 of size 384x1024x1]\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (value): TensorizedLinear(\n",
       "    (lr_act): GELU(approximate='none')\n",
       "    (tensor): TensorTrain(\n",
       "      (factors): ParameterList(\n",
       "          (0): Parameter containing: [torch.float32 of size 1x1024x384]\n",
       "          (1): Parameter containing: [torch.float32 of size 384x1024x1]\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (dropout): Dropout(p=0.1, inplace=False)\n",
       ")"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "attn = block.attention.self\n",
    "attn.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "5e5fa7e3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([16, 1, 1024])"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.manual_seed(0)\n",
    "exIn = torch.randn((16,1,1024))  # seq len, bsz, dim\n",
    "attn_mask = torch.zeros((16,1,1)) \n",
    "attnOut = attn(exIn, attn_mask)\n",
    "attnOut.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52417685",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 16 1024\n",
      "torch.Size([1, 16, 16, 64]) torch.Size([1, 16, 16, 64]) torch.Size([1, 16, 16, 64])\n"
     ]
    }
   ],
   "source": [
    "# Simulate normal attention where inputs are bsz, seq len, dim\n",
    "\n",
    "exIn2 = exIn.transpose(0,1) # bsz, seq len, dim\n",
    "bsz, tgt_len = exIn2.shape[0],exIn2.shape[1]\n",
    "print(bsz, tgt_len, exIn2.shape[2])\n",
    "\n",
    "num_heads = attn.num_attention_heads\n",
    "hsz = attn.attention_head_size\n",
    "\n",
    "query = attn.query(exIn2).reshape(bsz,-1,num_heads, hsz).transpose(1, 2)\n",
    "key = attn.key(exIn2).reshape(bsz,-1,num_heads, hsz).transpose(1, 2)\n",
    "value = attn.value(exIn2).reshape(bsz,-1,num_heads, hsz).transpose(1, 2)\n",
    "print(query.shape, key.shape, value.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "84c38b5c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "attn_scores: torch.Size([1, 16, 16, 16])\n",
      "attn_probs: torch.Size([1, 16, 16, 16])\n",
      "attn_output (before merge): torch.Size([1, 16, 16, 64])\n",
      "attn_output (final): torch.Size([1, 16, 1024])\n"
     ]
    }
   ],
   "source": [
    "# Step 1: Compute raw attention scores\n",
    "# Q @ K^T\n",
    "# Shape: [bsz, num_heads, seq_len, seq_len]\n",
    "attn_scores = torch.matmul(query, key.transpose(-2, -1))\n",
    "print(\"attn_scores:\", attn_scores.shape)\n",
    "# Step 2: Scale by sqrt(head_dim)\n",
    "attn_scores = attn_scores / (hsz ** 0.5)\n",
    "\n",
    "# Step 3: Apply softmax to get attention probabilities\n",
    "attn_probs = F.softmax(attn_scores, dim=-1)\n",
    "print(\"attn_probs:\", attn_probs.shape)\n",
    "\n",
    "\n",
    "# Step 4: Multiply attention probabilities by values\n",
    "# Shape: [bsz, num_heads, seq_len, head_dim]\n",
    "attn_output = torch.matmul(attn_probs, value)\n",
    "print(\"attn_output (before merge):\", attn_output.shape)\n",
    "\n",
    "# Step 5: Concatenate heads and reshape to original embedding size\n",
    "# Currently: [bsz, num_heads, seq_len, head_dim]\n",
    "# Desired: [bsz, seq_len, num_heads*head_dim]\n",
    "attn_output = attn_output.transpose(1, 2).reshape(bsz, tgt_len, -1)\n",
    "print(\"attn_output (final):\", attn_output.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "c779809f",
   "metadata": {},
   "outputs": [],
   "source": [
    "faOut = F.scaled_dot_product_attention(query, key, value,\n",
    "    attn_mask=attn_mask, dropout_p=0.0, \n",
    "    is_causal=False)\n",
    "faOut = faOut.transpose(1,2)\n",
    "faOut = faOut.reshape(bsz, tgt_len, -1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "253a2478",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor(0., grad_fn=<MaxBackward1>),\n",
       " tensor(8.9407e-08, grad_fn=<MaxBackward1>))"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(attn_output.transpose(0,1) -  attnOut).abs().max(), (faOut.transpose(0,1) -  attnOut).abs().max()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ce6b261",
   "metadata": {},
   "source": [
    "### Output and Intermediate Block"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "id": "ce30a4a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "exIn1, exIn2 = torch.randn((16,1,1024)), torch.randn((16,1,1024))\n",
    "block.attention.output.eval()\n",
    "blockOut = block.attention.output(exIn1.transpose(0,1), exIn2.transpose(0,1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "9be18bac",
   "metadata": {},
   "outputs": [],
   "source": [
    "blockOut2 = block.attention.output.LayerNorm(block.attention.output.dense(exIn1) + exIn2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "95c33af9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0., grad_fn=<MaxBackward1>)"
      ]
     },
     "execution_count": 54,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(blockOut - blockOut2.transpose(0,1)).abs().max()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "id": "09702e91",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TensorBertIntermediate(\n",
       "  (dense_act): TensorLinearActivation(\n",
       "    (gate_proj): TensorizedLinear(\n",
       "      (lr_act): GELU(approximate='none')\n",
       "      (tensor): TensorTrain(\n",
       "        (factors): ParameterList(\n",
       "            (0): Parameter containing: [torch.float32 of size 1x1024x512]\n",
       "            (1): Parameter containing: [torch.float32 of size 512x2736x1]\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (up_proj): TensorizedLinear(\n",
       "      (lr_act): GELU(approximate='none')\n",
       "      (tensor): TensorTrain(\n",
       "        (factors): ParameterList(\n",
       "            (0): Parameter containing: [torch.float32 of size 1x1024x512]\n",
       "            (1): Parameter containing: [torch.float32 of size 512x2736x1]\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 87,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "block.intermediate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "id": "1711f6ce",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(Parameter containing:\n",
       " tensor([-0.0136,  0.0021, -0.0034,  ..., -0.0092, -0.0159,  0.0171],\n",
       "        requires_grad=True),\n",
       " Parameter containing:\n",
       " tensor([ 0.0033, -0.0035,  0.0143,  ...,  0.0114,  0.0028, -0.0012],\n",
       "        requires_grad=True))"
      ]
     },
     "execution_count": 90,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "block.intermediate.dense_act.gate_proj.bias, block.intermediate.dense_act.up_proj.bias"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "id": "f8012804",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n",
      "torch.Size([16, 3, 2736])\n"
     ]
    }
   ],
   "source": [
    "exIn = torch.randn((16,3,1024)) # seq len, bsz, hidden dim\n",
    "mlpLayer = block.intermediate.dense_act\n",
    "mlpLayer.eval()\n",
    "print(mlpLayer.only_lr_act) # only low rank activation\n",
    "mlpOut = mlpLayer(exIn)\n",
    "print(mlpOut.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "id": "9c6360ab",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([3, 16, 2736])\n"
     ]
    }
   ],
   "source": [
    "mlpOut2 = mlpLayer.gate_proj(exIn.transpose(0,1)) * mlpLayer.up_proj(exIn.transpose(0,1))\n",
    "print(mlpOut2.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "id": "b22a6f72",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0., grad_fn=<MaxBackward1>)"
      ]
     },
     "execution_count": 73,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(mlpOut2 - mlpOut.transpose(0,1)).abs().max()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe57c0dc",
   "metadata": {},
   "source": [
    "## Check Tensorized Linear Module"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "id": "98292093",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TensorizedLinear(\n",
       "  (lr_act): GELU(approximate='none')\n",
       "  (tensor): TensorTrain(\n",
       "    (factors): ParameterList(\n",
       "        (0): Parameter containing: [torch.float32 of size 1x2736x512]\n",
       "        (1): Parameter containing: [torch.float32 of size 512x1024x1]\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 74,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tensoredizedLinear = block.output.dense\n",
    "tensoredizedLinear"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "id": "98ca9cf5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(2736, 1024, [2736, 1024])"
      ]
     },
     "execution_count": 75,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tensoredizedLinear.in_features, tensoredizedLinear.out_features, tensoredizedLinear.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "id": "d96edead",
   "metadata": {},
   "outputs": [],
   "source": [
    "exIn = torch.randn((5, 2736))\n",
    "exOut = block.output.dense(exIn)\n",
    "exOut\n",
    "\n",
    "lin1Weight = tensoredizedLinear.tensor.factors[0].squeeze()\n",
    "lin2Weight = tensoredizedLinear.tensor.factors[1].squeeze()\n",
    "lin1Weight.shape, lin2Weight.shape\n",
    "\n",
    "manOut = (tensoredizedLinear.lr_act(exIn@lin1Weight))@lin2Weight + tensoredizedLinear.bias"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "id": "e1917724",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Parameter containing:\n",
       "tensor([-0.0255, -0.0019,  0.0155,  ..., -0.0167, -0.0054, -0.0067],\n",
       "       requires_grad=True)"
      ]
     },
     "execution_count": 84,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tensoredizedLinear.bias"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "id": "ca010f30",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0., grad_fn=<MaxBackward1>)"
      ]
     },
     "execution_count": 83,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(manOut - exOut).abs().max()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "673b12ed",
   "metadata": {},
   "source": [
    "# Check Backward pass works"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "fc558bbd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([8, 512]),\n",
       " torch.Size([8, 512]),\n",
       " torch.Size([8, 512]),\n",
       " torch.Size([8]))"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "batch = next(iter(train_dataloader))\n",
    "input_ids, input_mask, segment_ids, label_ids = batch\n",
    "input_ids.shape, input_mask.shape, segment_ids.shape, label_ids.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "48026114",
   "metadata": {},
   "outputs": [],
   "source": [
    "with torch.no_grad():\n",
    "    out = model(input_ids, segment_ids, input_mask)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "b5d6cf09",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = 'cuda'\n",
    "model = model.to(device)\n",
    "input_ids, input_mask, segment_ids, label_ids = [x.to(device) for x in batch]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "b83b5c32",
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_fct = torch.nn.CrossEntropyLoss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "25f1764f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "bert.embeddings.word_embeddings.weight torch.Size([30528, 1024])\n",
      "bert.embeddings.position_embeddings.weight torch.Size([512, 1024])\n",
      "bert.embeddings.token_type_embeddings.weight torch.Size([2, 1024])\n",
      "bert.embeddings.LayerNorm.weight torch.Size([1024])\n",
      "bert.embeddings.LayerNorm.bias torch.Size([1024])\n",
      "bert.encoder.layer.0.attention.self.query.bias torch.Size([1024])\n",
      "bert.encoder.layer.0.attention.self.query.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.0.attention.self.query.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.0.attention.self.key.bias torch.Size([1024])\n",
      "bert.encoder.layer.0.attention.self.key.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.0.attention.self.key.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.0.attention.self.value.bias torch.Size([1024])\n",
      "bert.encoder.layer.0.attention.self.value.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.0.attention.self.value.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.0.attention.output.dense.bias torch.Size([1024])\n",
      "bert.encoder.layer.0.attention.output.dense.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.0.attention.output.dense.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.0.attention.output.LayerNorm.weight torch.Size([1024])\n",
      "bert.encoder.layer.0.attention.output.LayerNorm.bias torch.Size([1024])\n",
      "bert.encoder.layer.0.intermediate.dense_act.gate_proj.bias torch.Size([2736])\n",
      "bert.encoder.layer.0.intermediate.dense_act.gate_proj.tensor.factors.0 torch.Size([1, 1024, 512])\n",
      "bert.encoder.layer.0.intermediate.dense_act.gate_proj.tensor.factors.1 torch.Size([512, 2736, 1])\n",
      "bert.encoder.layer.0.intermediate.dense_act.up_proj.bias torch.Size([2736])\n",
      "bert.encoder.layer.0.intermediate.dense_act.up_proj.tensor.factors.0 torch.Size([1, 1024, 512])\n",
      "bert.encoder.layer.0.intermediate.dense_act.up_proj.tensor.factors.1 torch.Size([512, 2736, 1])\n",
      "bert.encoder.layer.0.output.dense.bias torch.Size([1024])\n",
      "bert.encoder.layer.0.output.dense.tensor.factors.0 torch.Size([1, 2736, 512])\n",
      "bert.encoder.layer.0.output.dense.tensor.factors.1 torch.Size([512, 1024, 1])\n",
      "bert.encoder.layer.0.output.LayerNorm.weight torch.Size([1024])\n",
      "bert.encoder.layer.0.output.LayerNorm.bias torch.Size([1024])\n",
      "bert.encoder.layer.1.attention.self.query.bias torch.Size([1024])\n",
      "bert.encoder.layer.1.attention.self.query.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.1.attention.self.query.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.1.attention.self.key.bias torch.Size([1024])\n",
      "bert.encoder.layer.1.attention.self.key.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.1.attention.self.key.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.1.attention.self.value.bias torch.Size([1024])\n",
      "bert.encoder.layer.1.attention.self.value.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.1.attention.self.value.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.1.attention.output.dense.bias torch.Size([1024])\n",
      "bert.encoder.layer.1.attention.output.dense.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.1.attention.output.dense.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.1.attention.output.LayerNorm.weight torch.Size([1024])\n",
      "bert.encoder.layer.1.attention.output.LayerNorm.bias torch.Size([1024])\n",
      "bert.encoder.layer.1.intermediate.dense_act.gate_proj.bias torch.Size([2736])\n",
      "bert.encoder.layer.1.intermediate.dense_act.gate_proj.tensor.factors.0 torch.Size([1, 1024, 512])\n",
      "bert.encoder.layer.1.intermediate.dense_act.gate_proj.tensor.factors.1 torch.Size([512, 2736, 1])\n",
      "bert.encoder.layer.1.intermediate.dense_act.up_proj.bias torch.Size([2736])\n",
      "bert.encoder.layer.1.intermediate.dense_act.up_proj.tensor.factors.0 torch.Size([1, 1024, 512])\n",
      "bert.encoder.layer.1.intermediate.dense_act.up_proj.tensor.factors.1 torch.Size([512, 2736, 1])\n",
      "bert.encoder.layer.1.output.dense.bias torch.Size([1024])\n",
      "bert.encoder.layer.1.output.dense.tensor.factors.0 torch.Size([1, 2736, 512])\n",
      "bert.encoder.layer.1.output.dense.tensor.factors.1 torch.Size([512, 1024, 1])\n",
      "bert.encoder.layer.1.output.LayerNorm.weight torch.Size([1024])\n",
      "bert.encoder.layer.1.output.LayerNorm.bias torch.Size([1024])\n",
      "bert.encoder.layer.2.attention.self.query.bias torch.Size([1024])\n",
      "bert.encoder.layer.2.attention.self.query.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.2.attention.self.query.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.2.attention.self.key.bias torch.Size([1024])\n",
      "bert.encoder.layer.2.attention.self.key.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.2.attention.self.key.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.2.attention.self.value.bias torch.Size([1024])\n",
      "bert.encoder.layer.2.attention.self.value.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.2.attention.self.value.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.2.attention.output.dense.bias torch.Size([1024])\n",
      "bert.encoder.layer.2.attention.output.dense.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.2.attention.output.dense.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.2.attention.output.LayerNorm.weight torch.Size([1024])\n",
      "bert.encoder.layer.2.attention.output.LayerNorm.bias torch.Size([1024])\n",
      "bert.encoder.layer.2.intermediate.dense_act.gate_proj.bias torch.Size([2736])\n",
      "bert.encoder.layer.2.intermediate.dense_act.gate_proj.tensor.factors.0 torch.Size([1, 1024, 512])\n",
      "bert.encoder.layer.2.intermediate.dense_act.gate_proj.tensor.factors.1 torch.Size([512, 2736, 1])\n",
      "bert.encoder.layer.2.intermediate.dense_act.up_proj.bias torch.Size([2736])\n",
      "bert.encoder.layer.2.intermediate.dense_act.up_proj.tensor.factors.0 torch.Size([1, 1024, 512])\n",
      "bert.encoder.layer.2.intermediate.dense_act.up_proj.tensor.factors.1 torch.Size([512, 2736, 1])\n",
      "bert.encoder.layer.2.output.dense.bias torch.Size([1024])\n",
      "bert.encoder.layer.2.output.dense.tensor.factors.0 torch.Size([1, 2736, 512])\n",
      "bert.encoder.layer.2.output.dense.tensor.factors.1 torch.Size([512, 1024, 1])\n",
      "bert.encoder.layer.2.output.LayerNorm.weight torch.Size([1024])\n",
      "bert.encoder.layer.2.output.LayerNorm.bias torch.Size([1024])\n",
      "bert.encoder.layer.3.attention.self.query.bias torch.Size([1024])\n",
      "bert.encoder.layer.3.attention.self.query.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.3.attention.self.query.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.3.attention.self.key.bias torch.Size([1024])\n",
      "bert.encoder.layer.3.attention.self.key.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.3.attention.self.key.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.3.attention.self.value.bias torch.Size([1024])\n",
      "bert.encoder.layer.3.attention.self.value.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.3.attention.self.value.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.3.attention.output.dense.bias torch.Size([1024])\n",
      "bert.encoder.layer.3.attention.output.dense.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.3.attention.output.dense.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.3.attention.output.LayerNorm.weight torch.Size([1024])\n",
      "bert.encoder.layer.3.attention.output.LayerNorm.bias torch.Size([1024])\n",
      "bert.encoder.layer.3.intermediate.dense_act.gate_proj.bias torch.Size([2736])\n",
      "bert.encoder.layer.3.intermediate.dense_act.gate_proj.tensor.factors.0 torch.Size([1, 1024, 512])\n",
      "bert.encoder.layer.3.intermediate.dense_act.gate_proj.tensor.factors.1 torch.Size([512, 2736, 1])\n",
      "bert.encoder.layer.3.intermediate.dense_act.up_proj.bias torch.Size([2736])\n",
      "bert.encoder.layer.3.intermediate.dense_act.up_proj.tensor.factors.0 torch.Size([1, 1024, 512])\n",
      "bert.encoder.layer.3.intermediate.dense_act.up_proj.tensor.factors.1 torch.Size([512, 2736, 1])\n",
      "bert.encoder.layer.3.output.dense.bias torch.Size([1024])\n",
      "bert.encoder.layer.3.output.dense.tensor.factors.0 torch.Size([1, 2736, 512])\n",
      "bert.encoder.layer.3.output.dense.tensor.factors.1 torch.Size([512, 1024, 1])\n",
      "bert.encoder.layer.3.output.LayerNorm.weight torch.Size([1024])\n",
      "bert.encoder.layer.3.output.LayerNorm.bias torch.Size([1024])\n",
      "bert.encoder.layer.4.attention.self.query.bias torch.Size([1024])\n",
      "bert.encoder.layer.4.attention.self.query.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.4.attention.self.query.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.4.attention.self.key.bias torch.Size([1024])\n",
      "bert.encoder.layer.4.attention.self.key.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.4.attention.self.key.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.4.attention.self.value.bias torch.Size([1024])\n",
      "bert.encoder.layer.4.attention.self.value.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.4.attention.self.value.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.4.attention.output.dense.bias torch.Size([1024])\n",
      "bert.encoder.layer.4.attention.output.dense.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.4.attention.output.dense.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.4.attention.output.LayerNorm.weight torch.Size([1024])\n",
      "bert.encoder.layer.4.attention.output.LayerNorm.bias torch.Size([1024])\n",
      "bert.encoder.layer.4.intermediate.dense_act.gate_proj.bias torch.Size([2736])\n",
      "bert.encoder.layer.4.intermediate.dense_act.gate_proj.tensor.factors.0 torch.Size([1, 1024, 512])\n",
      "bert.encoder.layer.4.intermediate.dense_act.gate_proj.tensor.factors.1 torch.Size([512, 2736, 1])\n",
      "bert.encoder.layer.4.intermediate.dense_act.up_proj.bias torch.Size([2736])\n",
      "bert.encoder.layer.4.intermediate.dense_act.up_proj.tensor.factors.0 torch.Size([1, 1024, 512])\n",
      "bert.encoder.layer.4.intermediate.dense_act.up_proj.tensor.factors.1 torch.Size([512, 2736, 1])\n",
      "bert.encoder.layer.4.output.dense.bias torch.Size([1024])\n",
      "bert.encoder.layer.4.output.dense.tensor.factors.0 torch.Size([1, 2736, 512])\n",
      "bert.encoder.layer.4.output.dense.tensor.factors.1 torch.Size([512, 1024, 1])\n",
      "bert.encoder.layer.4.output.LayerNorm.weight torch.Size([1024])\n",
      "bert.encoder.layer.4.output.LayerNorm.bias torch.Size([1024])\n",
      "bert.encoder.layer.5.attention.self.query.bias torch.Size([1024])\n",
      "bert.encoder.layer.5.attention.self.query.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.5.attention.self.query.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.5.attention.self.key.bias torch.Size([1024])\n",
      "bert.encoder.layer.5.attention.self.key.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.5.attention.self.key.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.5.attention.self.value.bias torch.Size([1024])\n",
      "bert.encoder.layer.5.attention.self.value.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.5.attention.self.value.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.5.attention.output.dense.bias torch.Size([1024])\n",
      "bert.encoder.layer.5.attention.output.dense.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.5.attention.output.dense.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.5.attention.output.LayerNorm.weight torch.Size([1024])\n",
      "bert.encoder.layer.5.attention.output.LayerNorm.bias torch.Size([1024])\n",
      "bert.encoder.layer.5.intermediate.dense_act.gate_proj.bias torch.Size([2736])\n",
      "bert.encoder.layer.5.intermediate.dense_act.gate_proj.tensor.factors.0 torch.Size([1, 1024, 512])\n",
      "bert.encoder.layer.5.intermediate.dense_act.gate_proj.tensor.factors.1 torch.Size([512, 2736, 1])\n",
      "bert.encoder.layer.5.intermediate.dense_act.up_proj.bias torch.Size([2736])\n",
      "bert.encoder.layer.5.intermediate.dense_act.up_proj.tensor.factors.0 torch.Size([1, 1024, 512])\n",
      "bert.encoder.layer.5.intermediate.dense_act.up_proj.tensor.factors.1 torch.Size([512, 2736, 1])\n",
      "bert.encoder.layer.5.output.dense.bias torch.Size([1024])\n",
      "bert.encoder.layer.5.output.dense.tensor.factors.0 torch.Size([1, 2736, 512])\n",
      "bert.encoder.layer.5.output.dense.tensor.factors.1 torch.Size([512, 1024, 1])\n",
      "bert.encoder.layer.5.output.LayerNorm.weight torch.Size([1024])\n",
      "bert.encoder.layer.5.output.LayerNorm.bias torch.Size([1024])\n",
      "bert.encoder.layer.6.attention.self.query.bias torch.Size([1024])\n",
      "bert.encoder.layer.6.attention.self.query.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.6.attention.self.query.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.6.attention.self.key.bias torch.Size([1024])\n",
      "bert.encoder.layer.6.attention.self.key.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.6.attention.self.key.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.6.attention.self.value.bias torch.Size([1024])\n",
      "bert.encoder.layer.6.attention.self.value.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.6.attention.self.value.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.6.attention.output.dense.bias torch.Size([1024])\n",
      "bert.encoder.layer.6.attention.output.dense.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.6.attention.output.dense.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.6.attention.output.LayerNorm.weight torch.Size([1024])\n",
      "bert.encoder.layer.6.attention.output.LayerNorm.bias torch.Size([1024])\n",
      "bert.encoder.layer.6.intermediate.dense_act.gate_proj.bias torch.Size([2736])\n",
      "bert.encoder.layer.6.intermediate.dense_act.gate_proj.tensor.factors.0 torch.Size([1, 1024, 512])\n",
      "bert.encoder.layer.6.intermediate.dense_act.gate_proj.tensor.factors.1 torch.Size([512, 2736, 1])\n",
      "bert.encoder.layer.6.intermediate.dense_act.up_proj.bias torch.Size([2736])\n",
      "bert.encoder.layer.6.intermediate.dense_act.up_proj.tensor.factors.0 torch.Size([1, 1024, 512])\n",
      "bert.encoder.layer.6.intermediate.dense_act.up_proj.tensor.factors.1 torch.Size([512, 2736, 1])\n",
      "bert.encoder.layer.6.output.dense.bias torch.Size([1024])\n",
      "bert.encoder.layer.6.output.dense.tensor.factors.0 torch.Size([1, 2736, 512])\n",
      "bert.encoder.layer.6.output.dense.tensor.factors.1 torch.Size([512, 1024, 1])\n",
      "bert.encoder.layer.6.output.LayerNorm.weight torch.Size([1024])\n",
      "bert.encoder.layer.6.output.LayerNorm.bias torch.Size([1024])\n",
      "bert.encoder.layer.7.attention.self.query.bias torch.Size([1024])\n",
      "bert.encoder.layer.7.attention.self.query.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.7.attention.self.query.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.7.attention.self.key.bias torch.Size([1024])\n",
      "bert.encoder.layer.7.attention.self.key.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.7.attention.self.key.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.7.attention.self.value.bias torch.Size([1024])\n",
      "bert.encoder.layer.7.attention.self.value.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.7.attention.self.value.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.7.attention.output.dense.bias torch.Size([1024])\n",
      "bert.encoder.layer.7.attention.output.dense.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.7.attention.output.dense.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.7.attention.output.LayerNorm.weight torch.Size([1024])\n",
      "bert.encoder.layer.7.attention.output.LayerNorm.bias torch.Size([1024])\n",
      "bert.encoder.layer.7.intermediate.dense_act.gate_proj.bias torch.Size([2736])\n",
      "bert.encoder.layer.7.intermediate.dense_act.gate_proj.tensor.factors.0 torch.Size([1, 1024, 512])\n",
      "bert.encoder.layer.7.intermediate.dense_act.gate_proj.tensor.factors.1 torch.Size([512, 2736, 1])\n",
      "bert.encoder.layer.7.intermediate.dense_act.up_proj.bias torch.Size([2736])\n",
      "bert.encoder.layer.7.intermediate.dense_act.up_proj.tensor.factors.0 torch.Size([1, 1024, 512])\n",
      "bert.encoder.layer.7.intermediate.dense_act.up_proj.tensor.factors.1 torch.Size([512, 2736, 1])\n",
      "bert.encoder.layer.7.output.dense.bias torch.Size([1024])\n",
      "bert.encoder.layer.7.output.dense.tensor.factors.0 torch.Size([1, 2736, 512])\n",
      "bert.encoder.layer.7.output.dense.tensor.factors.1 torch.Size([512, 1024, 1])\n",
      "bert.encoder.layer.7.output.LayerNorm.weight torch.Size([1024])\n",
      "bert.encoder.layer.7.output.LayerNorm.bias torch.Size([1024])\n",
      "bert.encoder.layer.8.attention.self.query.bias torch.Size([1024])\n",
      "bert.encoder.layer.8.attention.self.query.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.8.attention.self.query.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.8.attention.self.key.bias torch.Size([1024])\n",
      "bert.encoder.layer.8.attention.self.key.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.8.attention.self.key.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.8.attention.self.value.bias torch.Size([1024])\n",
      "bert.encoder.layer.8.attention.self.value.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.8.attention.self.value.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.8.attention.output.dense.bias torch.Size([1024])\n",
      "bert.encoder.layer.8.attention.output.dense.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.8.attention.output.dense.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.8.attention.output.LayerNorm.weight torch.Size([1024])\n",
      "bert.encoder.layer.8.attention.output.LayerNorm.bias torch.Size([1024])\n",
      "bert.encoder.layer.8.intermediate.dense_act.gate_proj.bias torch.Size([2736])\n",
      "bert.encoder.layer.8.intermediate.dense_act.gate_proj.tensor.factors.0 torch.Size([1, 1024, 512])\n",
      "bert.encoder.layer.8.intermediate.dense_act.gate_proj.tensor.factors.1 torch.Size([512, 2736, 1])\n",
      "bert.encoder.layer.8.intermediate.dense_act.up_proj.bias torch.Size([2736])\n",
      "bert.encoder.layer.8.intermediate.dense_act.up_proj.tensor.factors.0 torch.Size([1, 1024, 512])\n",
      "bert.encoder.layer.8.intermediate.dense_act.up_proj.tensor.factors.1 torch.Size([512, 2736, 1])\n",
      "bert.encoder.layer.8.output.dense.bias torch.Size([1024])\n",
      "bert.encoder.layer.8.output.dense.tensor.factors.0 torch.Size([1, 2736, 512])\n",
      "bert.encoder.layer.8.output.dense.tensor.factors.1 torch.Size([512, 1024, 1])\n",
      "bert.encoder.layer.8.output.LayerNorm.weight torch.Size([1024])\n",
      "bert.encoder.layer.8.output.LayerNorm.bias torch.Size([1024])\n",
      "bert.encoder.layer.9.attention.self.query.bias torch.Size([1024])\n",
      "bert.encoder.layer.9.attention.self.query.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.9.attention.self.query.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.9.attention.self.key.bias torch.Size([1024])\n",
      "bert.encoder.layer.9.attention.self.key.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.9.attention.self.key.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.9.attention.self.value.bias torch.Size([1024])\n",
      "bert.encoder.layer.9.attention.self.value.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.9.attention.self.value.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.9.attention.output.dense.bias torch.Size([1024])\n",
      "bert.encoder.layer.9.attention.output.dense.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.9.attention.output.dense.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.9.attention.output.LayerNorm.weight torch.Size([1024])\n",
      "bert.encoder.layer.9.attention.output.LayerNorm.bias torch.Size([1024])\n",
      "bert.encoder.layer.9.intermediate.dense_act.gate_proj.bias torch.Size([2736])\n",
      "bert.encoder.layer.9.intermediate.dense_act.gate_proj.tensor.factors.0 torch.Size([1, 1024, 512])\n",
      "bert.encoder.layer.9.intermediate.dense_act.gate_proj.tensor.factors.1 torch.Size([512, 2736, 1])\n",
      "bert.encoder.layer.9.intermediate.dense_act.up_proj.bias torch.Size([2736])\n",
      "bert.encoder.layer.9.intermediate.dense_act.up_proj.tensor.factors.0 torch.Size([1, 1024, 512])\n",
      "bert.encoder.layer.9.intermediate.dense_act.up_proj.tensor.factors.1 torch.Size([512, 2736, 1])\n",
      "bert.encoder.layer.9.output.dense.bias torch.Size([1024])\n",
      "bert.encoder.layer.9.output.dense.tensor.factors.0 torch.Size([1, 2736, 512])\n",
      "bert.encoder.layer.9.output.dense.tensor.factors.1 torch.Size([512, 1024, 1])\n",
      "bert.encoder.layer.9.output.LayerNorm.weight torch.Size([1024])\n",
      "bert.encoder.layer.9.output.LayerNorm.bias torch.Size([1024])\n",
      "bert.encoder.layer.10.attention.self.query.bias torch.Size([1024])\n",
      "bert.encoder.layer.10.attention.self.query.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.10.attention.self.query.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.10.attention.self.key.bias torch.Size([1024])\n",
      "bert.encoder.layer.10.attention.self.key.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.10.attention.self.key.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.10.attention.self.value.bias torch.Size([1024])\n",
      "bert.encoder.layer.10.attention.self.value.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.10.attention.self.value.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.10.attention.output.dense.bias torch.Size([1024])\n",
      "bert.encoder.layer.10.attention.output.dense.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.10.attention.output.dense.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.10.attention.output.LayerNorm.weight torch.Size([1024])\n",
      "bert.encoder.layer.10.attention.output.LayerNorm.bias torch.Size([1024])\n",
      "bert.encoder.layer.10.intermediate.dense_act.gate_proj.bias torch.Size([2736])\n",
      "bert.encoder.layer.10.intermediate.dense_act.gate_proj.tensor.factors.0 torch.Size([1, 1024, 512])\n",
      "bert.encoder.layer.10.intermediate.dense_act.gate_proj.tensor.factors.1 torch.Size([512, 2736, 1])\n",
      "bert.encoder.layer.10.intermediate.dense_act.up_proj.bias torch.Size([2736])\n",
      "bert.encoder.layer.10.intermediate.dense_act.up_proj.tensor.factors.0 torch.Size([1, 1024, 512])\n",
      "bert.encoder.layer.10.intermediate.dense_act.up_proj.tensor.factors.1 torch.Size([512, 2736, 1])\n",
      "bert.encoder.layer.10.output.dense.bias torch.Size([1024])\n",
      "bert.encoder.layer.10.output.dense.tensor.factors.0 torch.Size([1, 2736, 512])\n",
      "bert.encoder.layer.10.output.dense.tensor.factors.1 torch.Size([512, 1024, 1])\n",
      "bert.encoder.layer.10.output.LayerNorm.weight torch.Size([1024])\n",
      "bert.encoder.layer.10.output.LayerNorm.bias torch.Size([1024])\n",
      "bert.encoder.layer.11.attention.self.query.bias torch.Size([1024])\n",
      "bert.encoder.layer.11.attention.self.query.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.11.attention.self.query.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.11.attention.self.key.bias torch.Size([1024])\n",
      "bert.encoder.layer.11.attention.self.key.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.11.attention.self.key.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.11.attention.self.value.bias torch.Size([1024])\n",
      "bert.encoder.layer.11.attention.self.value.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.11.attention.self.value.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.11.attention.output.dense.bias torch.Size([1024])\n",
      "bert.encoder.layer.11.attention.output.dense.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.11.attention.output.dense.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.11.attention.output.LayerNorm.weight torch.Size([1024])\n",
      "bert.encoder.layer.11.attention.output.LayerNorm.bias torch.Size([1024])\n",
      "bert.encoder.layer.11.intermediate.dense_act.gate_proj.bias torch.Size([2736])\n",
      "bert.encoder.layer.11.intermediate.dense_act.gate_proj.tensor.factors.0 torch.Size([1, 1024, 512])\n",
      "bert.encoder.layer.11.intermediate.dense_act.gate_proj.tensor.factors.1 torch.Size([512, 2736, 1])\n",
      "bert.encoder.layer.11.intermediate.dense_act.up_proj.bias torch.Size([2736])\n",
      "bert.encoder.layer.11.intermediate.dense_act.up_proj.tensor.factors.0 torch.Size([1, 1024, 512])\n",
      "bert.encoder.layer.11.intermediate.dense_act.up_proj.tensor.factors.1 torch.Size([512, 2736, 1])\n",
      "bert.encoder.layer.11.output.dense.bias torch.Size([1024])\n",
      "bert.encoder.layer.11.output.dense.tensor.factors.0 torch.Size([1, 2736, 512])\n",
      "bert.encoder.layer.11.output.dense.tensor.factors.1 torch.Size([512, 1024, 1])\n",
      "bert.encoder.layer.11.output.LayerNorm.weight torch.Size([1024])\n",
      "bert.encoder.layer.11.output.LayerNorm.bias torch.Size([1024])\n",
      "bert.encoder.layer.12.attention.self.query.bias torch.Size([1024])\n",
      "bert.encoder.layer.12.attention.self.query.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.12.attention.self.query.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.12.attention.self.key.bias torch.Size([1024])\n",
      "bert.encoder.layer.12.attention.self.key.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.12.attention.self.key.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.12.attention.self.value.bias torch.Size([1024])\n",
      "bert.encoder.layer.12.attention.self.value.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.12.attention.self.value.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.12.attention.output.dense.bias torch.Size([1024])\n",
      "bert.encoder.layer.12.attention.output.dense.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.12.attention.output.dense.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.12.attention.output.LayerNorm.weight torch.Size([1024])\n",
      "bert.encoder.layer.12.attention.output.LayerNorm.bias torch.Size([1024])\n",
      "bert.encoder.layer.12.intermediate.dense_act.gate_proj.bias torch.Size([2736])\n",
      "bert.encoder.layer.12.intermediate.dense_act.gate_proj.tensor.factors.0 torch.Size([1, 1024, 512])\n",
      "bert.encoder.layer.12.intermediate.dense_act.gate_proj.tensor.factors.1 torch.Size([512, 2736, 1])\n",
      "bert.encoder.layer.12.intermediate.dense_act.up_proj.bias torch.Size([2736])\n",
      "bert.encoder.layer.12.intermediate.dense_act.up_proj.tensor.factors.0 torch.Size([1, 1024, 512])\n",
      "bert.encoder.layer.12.intermediate.dense_act.up_proj.tensor.factors.1 torch.Size([512, 2736, 1])\n",
      "bert.encoder.layer.12.output.dense.bias torch.Size([1024])\n",
      "bert.encoder.layer.12.output.dense.tensor.factors.0 torch.Size([1, 2736, 512])\n",
      "bert.encoder.layer.12.output.dense.tensor.factors.1 torch.Size([512, 1024, 1])\n",
      "bert.encoder.layer.12.output.LayerNorm.weight torch.Size([1024])\n",
      "bert.encoder.layer.12.output.LayerNorm.bias torch.Size([1024])\n",
      "bert.encoder.layer.13.attention.self.query.bias torch.Size([1024])\n",
      "bert.encoder.layer.13.attention.self.query.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.13.attention.self.query.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.13.attention.self.key.bias torch.Size([1024])\n",
      "bert.encoder.layer.13.attention.self.key.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.13.attention.self.key.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.13.attention.self.value.bias torch.Size([1024])\n",
      "bert.encoder.layer.13.attention.self.value.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.13.attention.self.value.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.13.attention.output.dense.bias torch.Size([1024])\n",
      "bert.encoder.layer.13.attention.output.dense.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.13.attention.output.dense.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.13.attention.output.LayerNorm.weight torch.Size([1024])\n",
      "bert.encoder.layer.13.attention.output.LayerNorm.bias torch.Size([1024])\n",
      "bert.encoder.layer.13.intermediate.dense_act.gate_proj.bias torch.Size([2736])\n",
      "bert.encoder.layer.13.intermediate.dense_act.gate_proj.tensor.factors.0 torch.Size([1, 1024, 512])\n",
      "bert.encoder.layer.13.intermediate.dense_act.gate_proj.tensor.factors.1 torch.Size([512, 2736, 1])\n",
      "bert.encoder.layer.13.intermediate.dense_act.up_proj.bias torch.Size([2736])\n",
      "bert.encoder.layer.13.intermediate.dense_act.up_proj.tensor.factors.0 torch.Size([1, 1024, 512])\n",
      "bert.encoder.layer.13.intermediate.dense_act.up_proj.tensor.factors.1 torch.Size([512, 2736, 1])\n",
      "bert.encoder.layer.13.output.dense.bias torch.Size([1024])\n",
      "bert.encoder.layer.13.output.dense.tensor.factors.0 torch.Size([1, 2736, 512])\n",
      "bert.encoder.layer.13.output.dense.tensor.factors.1 torch.Size([512, 1024, 1])\n",
      "bert.encoder.layer.13.output.LayerNorm.weight torch.Size([1024])\n",
      "bert.encoder.layer.13.output.LayerNorm.bias torch.Size([1024])\n",
      "bert.encoder.layer.14.attention.self.query.bias torch.Size([1024])\n",
      "bert.encoder.layer.14.attention.self.query.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.14.attention.self.query.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.14.attention.self.key.bias torch.Size([1024])\n",
      "bert.encoder.layer.14.attention.self.key.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.14.attention.self.key.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.14.attention.self.value.bias torch.Size([1024])\n",
      "bert.encoder.layer.14.attention.self.value.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.14.attention.self.value.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.14.attention.output.dense.bias torch.Size([1024])\n",
      "bert.encoder.layer.14.attention.output.dense.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.14.attention.output.dense.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.14.attention.output.LayerNorm.weight torch.Size([1024])\n",
      "bert.encoder.layer.14.attention.output.LayerNorm.bias torch.Size([1024])\n",
      "bert.encoder.layer.14.intermediate.dense_act.gate_proj.bias torch.Size([2736])\n",
      "bert.encoder.layer.14.intermediate.dense_act.gate_proj.tensor.factors.0 torch.Size([1, 1024, 512])\n",
      "bert.encoder.layer.14.intermediate.dense_act.gate_proj.tensor.factors.1 torch.Size([512, 2736, 1])\n",
      "bert.encoder.layer.14.intermediate.dense_act.up_proj.bias torch.Size([2736])\n",
      "bert.encoder.layer.14.intermediate.dense_act.up_proj.tensor.factors.0 torch.Size([1, 1024, 512])\n",
      "bert.encoder.layer.14.intermediate.dense_act.up_proj.tensor.factors.1 torch.Size([512, 2736, 1])\n",
      "bert.encoder.layer.14.output.dense.bias torch.Size([1024])\n",
      "bert.encoder.layer.14.output.dense.tensor.factors.0 torch.Size([1, 2736, 512])\n",
      "bert.encoder.layer.14.output.dense.tensor.factors.1 torch.Size([512, 1024, 1])\n",
      "bert.encoder.layer.14.output.LayerNorm.weight torch.Size([1024])\n",
      "bert.encoder.layer.14.output.LayerNorm.bias torch.Size([1024])\n",
      "bert.encoder.layer.15.attention.self.query.bias torch.Size([1024])\n",
      "bert.encoder.layer.15.attention.self.query.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.15.attention.self.query.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.15.attention.self.key.bias torch.Size([1024])\n",
      "bert.encoder.layer.15.attention.self.key.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.15.attention.self.key.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.15.attention.self.value.bias torch.Size([1024])\n",
      "bert.encoder.layer.15.attention.self.value.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.15.attention.self.value.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.15.attention.output.dense.bias torch.Size([1024])\n",
      "bert.encoder.layer.15.attention.output.dense.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.15.attention.output.dense.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.15.attention.output.LayerNorm.weight torch.Size([1024])\n",
      "bert.encoder.layer.15.attention.output.LayerNorm.bias torch.Size([1024])\n",
      "bert.encoder.layer.15.intermediate.dense_act.gate_proj.bias torch.Size([2736])\n",
      "bert.encoder.layer.15.intermediate.dense_act.gate_proj.tensor.factors.0 torch.Size([1, 1024, 512])\n",
      "bert.encoder.layer.15.intermediate.dense_act.gate_proj.tensor.factors.1 torch.Size([512, 2736, 1])\n",
      "bert.encoder.layer.15.intermediate.dense_act.up_proj.bias torch.Size([2736])\n",
      "bert.encoder.layer.15.intermediate.dense_act.up_proj.tensor.factors.0 torch.Size([1, 1024, 512])\n",
      "bert.encoder.layer.15.intermediate.dense_act.up_proj.tensor.factors.1 torch.Size([512, 2736, 1])\n",
      "bert.encoder.layer.15.output.dense.bias torch.Size([1024])\n",
      "bert.encoder.layer.15.output.dense.tensor.factors.0 torch.Size([1, 2736, 512])\n",
      "bert.encoder.layer.15.output.dense.tensor.factors.1 torch.Size([512, 1024, 1])\n",
      "bert.encoder.layer.15.output.LayerNorm.weight torch.Size([1024])\n",
      "bert.encoder.layer.15.output.LayerNorm.bias torch.Size([1024])\n",
      "bert.encoder.layer.16.attention.self.query.bias torch.Size([1024])\n",
      "bert.encoder.layer.16.attention.self.query.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.16.attention.self.query.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.16.attention.self.key.bias torch.Size([1024])\n",
      "bert.encoder.layer.16.attention.self.key.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.16.attention.self.key.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.16.attention.self.value.bias torch.Size([1024])\n",
      "bert.encoder.layer.16.attention.self.value.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.16.attention.self.value.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.16.attention.output.dense.bias torch.Size([1024])\n",
      "bert.encoder.layer.16.attention.output.dense.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.16.attention.output.dense.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.16.attention.output.LayerNorm.weight torch.Size([1024])\n",
      "bert.encoder.layer.16.attention.output.LayerNorm.bias torch.Size([1024])\n",
      "bert.encoder.layer.16.intermediate.dense_act.gate_proj.bias torch.Size([2736])\n",
      "bert.encoder.layer.16.intermediate.dense_act.gate_proj.tensor.factors.0 torch.Size([1, 1024, 512])\n",
      "bert.encoder.layer.16.intermediate.dense_act.gate_proj.tensor.factors.1 torch.Size([512, 2736, 1])\n",
      "bert.encoder.layer.16.intermediate.dense_act.up_proj.bias torch.Size([2736])\n",
      "bert.encoder.layer.16.intermediate.dense_act.up_proj.tensor.factors.0 torch.Size([1, 1024, 512])\n",
      "bert.encoder.layer.16.intermediate.dense_act.up_proj.tensor.factors.1 torch.Size([512, 2736, 1])\n",
      "bert.encoder.layer.16.output.dense.bias torch.Size([1024])\n",
      "bert.encoder.layer.16.output.dense.tensor.factors.0 torch.Size([1, 2736, 512])\n",
      "bert.encoder.layer.16.output.dense.tensor.factors.1 torch.Size([512, 1024, 1])\n",
      "bert.encoder.layer.16.output.LayerNorm.weight torch.Size([1024])\n",
      "bert.encoder.layer.16.output.LayerNorm.bias torch.Size([1024])\n",
      "bert.encoder.layer.17.attention.self.query.bias torch.Size([1024])\n",
      "bert.encoder.layer.17.attention.self.query.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.17.attention.self.query.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.17.attention.self.key.bias torch.Size([1024])\n",
      "bert.encoder.layer.17.attention.self.key.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.17.attention.self.key.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.17.attention.self.value.bias torch.Size([1024])\n",
      "bert.encoder.layer.17.attention.self.value.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.17.attention.self.value.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.17.attention.output.dense.bias torch.Size([1024])\n",
      "bert.encoder.layer.17.attention.output.dense.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.17.attention.output.dense.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.17.attention.output.LayerNorm.weight torch.Size([1024])\n",
      "bert.encoder.layer.17.attention.output.LayerNorm.bias torch.Size([1024])\n",
      "bert.encoder.layer.17.intermediate.dense_act.gate_proj.bias torch.Size([2736])\n",
      "bert.encoder.layer.17.intermediate.dense_act.gate_proj.tensor.factors.0 torch.Size([1, 1024, 512])\n",
      "bert.encoder.layer.17.intermediate.dense_act.gate_proj.tensor.factors.1 torch.Size([512, 2736, 1])\n",
      "bert.encoder.layer.17.intermediate.dense_act.up_proj.bias torch.Size([2736])\n",
      "bert.encoder.layer.17.intermediate.dense_act.up_proj.tensor.factors.0 torch.Size([1, 1024, 512])\n",
      "bert.encoder.layer.17.intermediate.dense_act.up_proj.tensor.factors.1 torch.Size([512, 2736, 1])\n",
      "bert.encoder.layer.17.output.dense.bias torch.Size([1024])\n",
      "bert.encoder.layer.17.output.dense.tensor.factors.0 torch.Size([1, 2736, 512])\n",
      "bert.encoder.layer.17.output.dense.tensor.factors.1 torch.Size([512, 1024, 1])\n",
      "bert.encoder.layer.17.output.LayerNorm.weight torch.Size([1024])\n",
      "bert.encoder.layer.17.output.LayerNorm.bias torch.Size([1024])\n",
      "bert.encoder.layer.18.attention.self.query.bias torch.Size([1024])\n",
      "bert.encoder.layer.18.attention.self.query.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.18.attention.self.query.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.18.attention.self.key.bias torch.Size([1024])\n",
      "bert.encoder.layer.18.attention.self.key.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.18.attention.self.key.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.18.attention.self.value.bias torch.Size([1024])\n",
      "bert.encoder.layer.18.attention.self.value.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.18.attention.self.value.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.18.attention.output.dense.bias torch.Size([1024])\n",
      "bert.encoder.layer.18.attention.output.dense.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.18.attention.output.dense.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.18.attention.output.LayerNorm.weight torch.Size([1024])\n",
      "bert.encoder.layer.18.attention.output.LayerNorm.bias torch.Size([1024])\n",
      "bert.encoder.layer.18.intermediate.dense_act.gate_proj.bias torch.Size([2736])\n",
      "bert.encoder.layer.18.intermediate.dense_act.gate_proj.tensor.factors.0 torch.Size([1, 1024, 512])\n",
      "bert.encoder.layer.18.intermediate.dense_act.gate_proj.tensor.factors.1 torch.Size([512, 2736, 1])\n",
      "bert.encoder.layer.18.intermediate.dense_act.up_proj.bias torch.Size([2736])\n",
      "bert.encoder.layer.18.intermediate.dense_act.up_proj.tensor.factors.0 torch.Size([1, 1024, 512])\n",
      "bert.encoder.layer.18.intermediate.dense_act.up_proj.tensor.factors.1 torch.Size([512, 2736, 1])\n",
      "bert.encoder.layer.18.output.dense.bias torch.Size([1024])\n",
      "bert.encoder.layer.18.output.dense.tensor.factors.0 torch.Size([1, 2736, 512])\n",
      "bert.encoder.layer.18.output.dense.tensor.factors.1 torch.Size([512, 1024, 1])\n",
      "bert.encoder.layer.18.output.LayerNorm.weight torch.Size([1024])\n",
      "bert.encoder.layer.18.output.LayerNorm.bias torch.Size([1024])\n",
      "bert.encoder.layer.19.attention.self.query.bias torch.Size([1024])\n",
      "bert.encoder.layer.19.attention.self.query.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.19.attention.self.query.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.19.attention.self.key.bias torch.Size([1024])\n",
      "bert.encoder.layer.19.attention.self.key.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.19.attention.self.key.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.19.attention.self.value.bias torch.Size([1024])\n",
      "bert.encoder.layer.19.attention.self.value.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.19.attention.self.value.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.19.attention.output.dense.bias torch.Size([1024])\n",
      "bert.encoder.layer.19.attention.output.dense.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.19.attention.output.dense.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.19.attention.output.LayerNorm.weight torch.Size([1024])\n",
      "bert.encoder.layer.19.attention.output.LayerNorm.bias torch.Size([1024])\n",
      "bert.encoder.layer.19.intermediate.dense_act.gate_proj.bias torch.Size([2736])\n",
      "bert.encoder.layer.19.intermediate.dense_act.gate_proj.tensor.factors.0 torch.Size([1, 1024, 512])\n",
      "bert.encoder.layer.19.intermediate.dense_act.gate_proj.tensor.factors.1 torch.Size([512, 2736, 1])\n",
      "bert.encoder.layer.19.intermediate.dense_act.up_proj.bias torch.Size([2736])\n",
      "bert.encoder.layer.19.intermediate.dense_act.up_proj.tensor.factors.0 torch.Size([1, 1024, 512])\n",
      "bert.encoder.layer.19.intermediate.dense_act.up_proj.tensor.factors.1 torch.Size([512, 2736, 1])\n",
      "bert.encoder.layer.19.output.dense.bias torch.Size([1024])\n",
      "bert.encoder.layer.19.output.dense.tensor.factors.0 torch.Size([1, 2736, 512])\n",
      "bert.encoder.layer.19.output.dense.tensor.factors.1 torch.Size([512, 1024, 1])\n",
      "bert.encoder.layer.19.output.LayerNorm.weight torch.Size([1024])\n",
      "bert.encoder.layer.19.output.LayerNorm.bias torch.Size([1024])\n",
      "bert.encoder.layer.20.attention.self.query.bias torch.Size([1024])\n",
      "bert.encoder.layer.20.attention.self.query.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.20.attention.self.query.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.20.attention.self.key.bias torch.Size([1024])\n",
      "bert.encoder.layer.20.attention.self.key.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.20.attention.self.key.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.20.attention.self.value.bias torch.Size([1024])\n",
      "bert.encoder.layer.20.attention.self.value.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.20.attention.self.value.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.20.attention.output.dense.bias torch.Size([1024])\n",
      "bert.encoder.layer.20.attention.output.dense.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.20.attention.output.dense.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.20.attention.output.LayerNorm.weight torch.Size([1024])\n",
      "bert.encoder.layer.20.attention.output.LayerNorm.bias torch.Size([1024])\n",
      "bert.encoder.layer.20.intermediate.dense_act.gate_proj.bias torch.Size([2736])\n",
      "bert.encoder.layer.20.intermediate.dense_act.gate_proj.tensor.factors.0 torch.Size([1, 1024, 512])\n",
      "bert.encoder.layer.20.intermediate.dense_act.gate_proj.tensor.factors.1 torch.Size([512, 2736, 1])\n",
      "bert.encoder.layer.20.intermediate.dense_act.up_proj.bias torch.Size([2736])\n",
      "bert.encoder.layer.20.intermediate.dense_act.up_proj.tensor.factors.0 torch.Size([1, 1024, 512])\n",
      "bert.encoder.layer.20.intermediate.dense_act.up_proj.tensor.factors.1 torch.Size([512, 2736, 1])\n",
      "bert.encoder.layer.20.output.dense.bias torch.Size([1024])\n",
      "bert.encoder.layer.20.output.dense.tensor.factors.0 torch.Size([1, 2736, 512])\n",
      "bert.encoder.layer.20.output.dense.tensor.factors.1 torch.Size([512, 1024, 1])\n",
      "bert.encoder.layer.20.output.LayerNorm.weight torch.Size([1024])\n",
      "bert.encoder.layer.20.output.LayerNorm.bias torch.Size([1024])\n",
      "bert.encoder.layer.21.attention.self.query.bias torch.Size([1024])\n",
      "bert.encoder.layer.21.attention.self.query.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.21.attention.self.query.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.21.attention.self.key.bias torch.Size([1024])\n",
      "bert.encoder.layer.21.attention.self.key.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.21.attention.self.key.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.21.attention.self.value.bias torch.Size([1024])\n",
      "bert.encoder.layer.21.attention.self.value.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.21.attention.self.value.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.21.attention.output.dense.bias torch.Size([1024])\n",
      "bert.encoder.layer.21.attention.output.dense.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.21.attention.output.dense.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.21.attention.output.LayerNorm.weight torch.Size([1024])\n",
      "bert.encoder.layer.21.attention.output.LayerNorm.bias torch.Size([1024])\n",
      "bert.encoder.layer.21.intermediate.dense_act.gate_proj.bias torch.Size([2736])\n",
      "bert.encoder.layer.21.intermediate.dense_act.gate_proj.tensor.factors.0 torch.Size([1, 1024, 512])\n",
      "bert.encoder.layer.21.intermediate.dense_act.gate_proj.tensor.factors.1 torch.Size([512, 2736, 1])\n",
      "bert.encoder.layer.21.intermediate.dense_act.up_proj.bias torch.Size([2736])\n",
      "bert.encoder.layer.21.intermediate.dense_act.up_proj.tensor.factors.0 torch.Size([1, 1024, 512])\n",
      "bert.encoder.layer.21.intermediate.dense_act.up_proj.tensor.factors.1 torch.Size([512, 2736, 1])\n",
      "bert.encoder.layer.21.output.dense.bias torch.Size([1024])\n",
      "bert.encoder.layer.21.output.dense.tensor.factors.0 torch.Size([1, 2736, 512])\n",
      "bert.encoder.layer.21.output.dense.tensor.factors.1 torch.Size([512, 1024, 1])\n",
      "bert.encoder.layer.21.output.LayerNorm.weight torch.Size([1024])\n",
      "bert.encoder.layer.21.output.LayerNorm.bias torch.Size([1024])\n",
      "bert.encoder.layer.22.attention.self.query.bias torch.Size([1024])\n",
      "bert.encoder.layer.22.attention.self.query.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.22.attention.self.query.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.22.attention.self.key.bias torch.Size([1024])\n",
      "bert.encoder.layer.22.attention.self.key.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.22.attention.self.key.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.22.attention.self.value.bias torch.Size([1024])\n",
      "bert.encoder.layer.22.attention.self.value.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.22.attention.self.value.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.22.attention.output.dense.bias torch.Size([1024])\n",
      "bert.encoder.layer.22.attention.output.dense.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.22.attention.output.dense.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.22.attention.output.LayerNorm.weight torch.Size([1024])\n",
      "bert.encoder.layer.22.attention.output.LayerNorm.bias torch.Size([1024])\n",
      "bert.encoder.layer.22.intermediate.dense_act.gate_proj.bias torch.Size([2736])\n",
      "bert.encoder.layer.22.intermediate.dense_act.gate_proj.tensor.factors.0 torch.Size([1, 1024, 512])\n",
      "bert.encoder.layer.22.intermediate.dense_act.gate_proj.tensor.factors.1 torch.Size([512, 2736, 1])\n",
      "bert.encoder.layer.22.intermediate.dense_act.up_proj.bias torch.Size([2736])\n",
      "bert.encoder.layer.22.intermediate.dense_act.up_proj.tensor.factors.0 torch.Size([1, 1024, 512])\n",
      "bert.encoder.layer.22.intermediate.dense_act.up_proj.tensor.factors.1 torch.Size([512, 2736, 1])\n",
      "bert.encoder.layer.22.output.dense.bias torch.Size([1024])\n",
      "bert.encoder.layer.22.output.dense.tensor.factors.0 torch.Size([1, 2736, 512])\n",
      "bert.encoder.layer.22.output.dense.tensor.factors.1 torch.Size([512, 1024, 1])\n",
      "bert.encoder.layer.22.output.LayerNorm.weight torch.Size([1024])\n",
      "bert.encoder.layer.22.output.LayerNorm.bias torch.Size([1024])\n",
      "bert.encoder.layer.23.attention.self.query.bias torch.Size([1024])\n",
      "bert.encoder.layer.23.attention.self.query.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.23.attention.self.query.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.23.attention.self.key.bias torch.Size([1024])\n",
      "bert.encoder.layer.23.attention.self.key.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.23.attention.self.key.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.23.attention.self.value.bias torch.Size([1024])\n",
      "bert.encoder.layer.23.attention.self.value.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.23.attention.self.value.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.23.attention.output.dense.bias torch.Size([1024])\n",
      "bert.encoder.layer.23.attention.output.dense.tensor.factors.0 torch.Size([1, 1024, 384])\n",
      "bert.encoder.layer.23.attention.output.dense.tensor.factors.1 torch.Size([384, 1024, 1])\n",
      "bert.encoder.layer.23.attention.output.LayerNorm.weight torch.Size([1024])\n",
      "bert.encoder.layer.23.attention.output.LayerNorm.bias torch.Size([1024])\n",
      "bert.encoder.layer.23.intermediate.dense_act.gate_proj.bias torch.Size([2736])\n",
      "bert.encoder.layer.23.intermediate.dense_act.gate_proj.tensor.factors.0 torch.Size([1, 1024, 512])\n",
      "bert.encoder.layer.23.intermediate.dense_act.gate_proj.tensor.factors.1 torch.Size([512, 2736, 1])\n",
      "bert.encoder.layer.23.intermediate.dense_act.up_proj.bias torch.Size([2736])\n",
      "bert.encoder.layer.23.intermediate.dense_act.up_proj.tensor.factors.0 torch.Size([1, 1024, 512])\n",
      "bert.encoder.layer.23.intermediate.dense_act.up_proj.tensor.factors.1 torch.Size([512, 2736, 1])\n",
      "bert.encoder.layer.23.output.dense.bias torch.Size([1024])\n",
      "bert.encoder.layer.23.output.dense.tensor.factors.0 torch.Size([1, 2736, 512])\n",
      "bert.encoder.layer.23.output.dense.tensor.factors.1 torch.Size([512, 1024, 1])\n",
      "bert.encoder.layer.23.output.LayerNorm.weight torch.Size([1024])\n",
      "bert.encoder.layer.23.output.LayerNorm.bias torch.Size([1024])\n",
      "bert.pooler.dense_act.linear.weight torch.Size([1024, 1024])\n",
      "bert.pooler.dense_act.linear.bias torch.Size([1024])\n",
      "classifier.weight torch.Size([2, 1024])\n",
      "classifier.bias torch.Size([2])\n"
     ]
    }
   ],
   "source": [
    "model.train()\n",
    "out = model(input_ids, segment_ids, input_mask)\n",
    "loss = loss_fct(out.view(-1, num_labels), label_ids.view(-1))\n",
    "loss.backward()\n",
    "\n",
    "for name, param in model.named_parameters():\n",
    "    print(name,param.shape)\n",
    "    # if param.grad is not None:\n",
    "    #     print(f\"{name}: {param.grad.shape}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f340e458",
   "metadata": {},
   "outputs": [],
   "source": [
    "# import json\n",
    "\n",
    "# def save_state_dict_structure(state_dict, file_path):\n",
    "#     def serialize_state_dict(state_dict):\n",
    "#         serialized = {}\n",
    "#         for key, value in state_dict.items():\n",
    "#             if isinstance(value, torch.Tensor):\n",
    "#                 serialized[key] = f\"Tensor with shape {list(value.shape)}\"\n",
    "#             elif isinstance(value, (int, float, str)):\n",
    "#                 serialized[key] = value\n",
    "#             elif isinstance(value, dict):\n",
    "#                 serialized[key] = serialize_state_dict(value)\n",
    "#             else:\n",
    "#                 serialized[key] = str(type(value))\n",
    "#         return serialized\n",
    "\n",
    "#     with open(file_path, 'w') as f:\n",
    "#         json.dump(serialize_state_dict(state_dict), f, indent=4)\n",
    "\n",
    "# save_state_dict_structure(stateDict, \"state_dict_structure.json\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4a54a77",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b6c70fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "\t"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "colaBert",
   "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
}