addressed reviewers comments

AlignmentResearch · Aug 30, 2023 · 99076b6 · 99076b6
1 parent e6e396b
commit 99076b6
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/tests/test_lenses.py b/tests/test_lenses.py
@@ -144,6 +144,8 @@ def test_tuned_lens_generate_smoke(random_small_model: trf.PreTrainedModel):
     )
     assert tokens.shape[-1] <= 11
     assert tokens.shape[-1] > 1
+    assert input_ids == tokens[:, :1]
+    assert input_ids == th.tensor([bos_token_id]), "Don't mutate input_ids!"
 
     tokens = tuned_lens.generate(
         model=random_small_model,

diff --git a/tuned_lens/nn/lenses.py b/tuned_lens/nn/lenses.py
@@ -343,7 +343,7 @@ def generate(
         """
         eos_token = model.generation_config.eos_token_id
 
-        tokens = input_ids.clone()
+        tokens = input_ids
         if tokens.ndim == 1:
             tokens = tokens.unsqueeze(0)
         batch, prompt_len = tokens.shape
@@ -364,8 +364,8 @@ def generate(
             new_logits = self.forward(new_hidden, layer)
             if do_sample:
                 new_logits = new_logits / temp
-                new_logits = th.nn.functional.log_softmax(new_logits, dim=-1)
-                new_tokens = th.multinomial(new_logits.exp(), num_samples=1)
+                probs = new_logits.softmax(dim=-1)
+                new_tokens = th.multinomial(probs, num_samples=1)
             else:
                 new_tokens = new_logits.argmax(dim=-1, keepdim=True)