[Minor] Allow other field passing in generate besides input_ids

stanfordnlp · Mar 5, 2024 · e32bc0a · e32bc0a
1 parent fcf4870
commit e32bc0a
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 7 deletions.
diff --git a/pyvene/models/intervenable_base.py b/pyvene/models/intervenable_base.py
@@ -1496,7 +1496,7 @@ def generate(
         base_outputs = None
         if output_original_output:
             # returning un-intervened output
-            base_outputs = self.model.generate(inputs=base["input_ids"], **kwargs)
+            base_outputs = self.model.generate(**base, **kwargs)
 
         set_handlers_to_remove = None
         try:
@@ -1522,7 +1522,7 @@ def generate(
 
             # run intervened generate
             counterfactual_outputs = self.model.generate(
-                inputs=base["input_ids"], **kwargs
+                **base, **kwargs
             )
 
             collected_activations = []

diff --git a/pyvene_101.ipynb b/pyvene_101.ipynb
@@ -139,6 +139,7 @@
     }
    ],
    "source": [
+    "import torch\n",
     "import pyvene as pv\n",
     "\n",
     "_, tokenizer, gpt2 = pv.create_gpt2()\n",
@@ -222,6 +223,7 @@
     }
    ],
    "source": [
+    "import torch\n",
     "import pyvene as pv\n",
     "\n",
     "_, tokenizer, gpt2 = pv.create_gpt2()\n",
@@ -262,6 +264,7 @@
     }
    ],
    "source": [
+    "import torch\n",
     "import copy\n",
     "import pyvene as pv\n",
     "\n",
@@ -1068,17 +1071,14 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 4,
    "id": "f718e2d6",
    "metadata": {},
    "outputs": [
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "/sailhome/wuzhengx/.local/lib/python3.8/site-packages/torch/_utils.py:831: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.  To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage()\n",
-      "  return self.fget.__get__(instance, owner)()\n",
-      "The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n",
       "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
      ]
     },
@@ -2647,7 +2647,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.12"
+   "version": "3.10.13"
   },
   "toc-autonumbering": true,
   "toc-showcode": false,