Additional notebook tests

EleutherAI · Apr 22, 2023 · d61375c · d61375c
1 parent d94132d
commit d61375c
Showing 1 changed file with 21 additions and 10 deletions.
diff --git a/elk/rwkv_lm/test_rwkv.ipynb b/elk/rwkv_lm/test_rwkv.ipynb
@@ -518,7 +518,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 57,
+   "execution_count": 58,
    "metadata": {},
    "outputs": [
     {
@@ -531,42 +531,53 @@
     {
      "data": {
       "text/plain": [
-       "2048"
+       "torch.Size([243, 2048])"
       ]
      },
-     "execution_count": 57,
+     "execution_count": 58,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
     "tokens = tokenizer.encode(bad_review_text, return_tensors=\"pt\")\n",
     "print(f\"Input IDs shape = {tokens.shape}\")\n",
+    "\n",
     "states = model(tokens)\n",
     "tensor_final_state = states[-1]\n",
-    "tensor_final_state.shape\n",
-    "len(states[0][0])"
+    "tensor_final_state.shape"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 53,
+   "execution_count": 59,
    "metadata": {},
    "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Input IDs shape = torch.Size([1, 243])\n"
+     ]
+    },
     {
      "data": {
       "text/plain": [
-       "tensor([ 0.3164, -0.3691,  0.0165,  ..., -0.0081,  0.0175, -0.0771],\n",
-       "       device='cuda:0', dtype=torch.bfloat16)"
+       "torch.Size([2048])"
       ]
      },
-     "execution_count": 53,
+     "execution_count": 59,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
-    "tensor_final_state[-1]"
+    "tokens = tokenizer.encode(bad_review_text, return_tensors=\"pt\")\n",
+    "print(f\"Input IDs shape = {tokens.shape}\")\n",
+    "\n",
+    "states = model(tokens[0])\n",
+    "tensor_final_state = states[-1]\n",
+    "tensor_final_state.shape"
    ]
   }
  ],