Score choices after the answer newline

Co-Authored-By: PI[openai-codex] <288921227+claudypoo@users.noreply.github.com>

antipasto/eval.py

@@ -18,6 +18,23 @@ def get_choice_ids(tokenizer, positive_word="yes", negative_word="no") -> List[L
     return [list(negative_choices.values()), list(positive_choices.values())]
 
 
+def append_choice_newline(input_ids, tokenizer, attention_mask=None):
+    """Append the newline between ``My choice:`` and the scored answer."""
+    newline_ids = tokenizer("\n", add_special_tokens=False)["input_ids"]
+    suffix = torch.tensor(newline_ids, dtype=input_ids.dtype, device=input_ids.device)
+    suffix = suffix.view(*([1] * (input_ids.ndim - 1)), -1).expand(
+        *input_ids.shape[:-1], -1
+    )
+    input_ids = torch.cat([input_ids, suffix], dim=-1)
+    if attention_mask is None:
+        return input_ids
+    suffix_mask = torch.ones(
+        suffix.shape, dtype=attention_mask.dtype, device=attention_mask.device
+    )
+    attention_mask = torch.cat([attention_mask, suffix_mask], dim=-1)
+    return input_ids, attention_mask
+
+
 def calc_nll(input_ids, logits, attention_mask):
     """Calculate per-sequence NLL from input_ids and logits.
     
@@ -128,8 +145,6 @@ def gen_with_choices(model, tokenizer, input_ids, attention_mask, choice_ids, co
                 top_tokens = [(tokenizer.decode([tid.item()]), p.exp().item()) 
                               for tid, p in zip(topk_ids[i], topk_probs[i])]
                 top_tokens_str = ", ".join([f"{tok!r} ({prob:.2%})" for tok, prob in top_tokens])
-                question_s = tokenizer.batch_decode(input_ids)[i]
-                
                 logger.warning(
                     f"Low choice prob mass: {pmass[i].item():.2%} < 10% of max ({maxp[i].item():.2%}). "
                     # f"Your choices might not match the model's tokenization. OR the model might be refusing or incoherent"

antipasto/train/daily_dilemas.py

@@ -11,7 +11,7 @@ from tabulate import tabulate
 from torch.utils.data import DataLoader
 from transformers import DataCollatorWithPadding
 from datasets import concatenate_datasets
-from antipasto.eval import gen_with_choices
+from antipasto.eval import append_choice_newline, gen_with_choices
 from antipasto.transfer_analysis import VALUE_CLUSTERS
 from antipasto.metrics import (
     compute_centered_regression,
@@ -279,7 +279,7 @@ def format_messages(
             f"Input truncated to max_size={max_size} tokens for dilemma_idx={row['dilemma_idx']}, idx={row['idx']}. Consider increasing max_size."
         )
 
-    return {"input_ids": inputs_ids.squeeze(0)}
+    return {"input_ids": append_choice_newline(inputs_ids.squeeze(0), tokenizer)}
 
 
 def load_and_process_daily_dilemmas_eval_dataset(

antipasto/train/train_adapter.py

@@ -33,7 +33,7 @@ from transformers import DataCollatorWithPadding
 
 from antipasto import ControlVector
 from antipasto.config import TrainingConfig, proj_root
-from antipasto.eval import gen_with_choices, get_choice_ids
+from antipasto.eval import append_choice_newline, gen_with_choices, get_choice_ids
 from antipasto.peft_utils.adapter_scaling import ScaleAdapter, get_scale_adapter_fn
 from antipasto.peft_utils.antipasto_adapter import register_antipasto_peft
 from antipasto.peft_utils.layer_selection import (
@@ -1538,8 +1538,9 @@ Action: Tell a white lie"""
         return_dict=True,
         return_attention_mask=True,
     ).to(model.device)
-    input_ids = batch["input_ids"]
-    attn_mask = batch["attention_mask"]
+    input_ids, attn_mask = append_choice_newline(
+        batch["input_ids"], tokenizer, batch["attention_mask"]
+    )
 
     model.eval()
 

tests/test_eval.py

@@ -0,0 +1,22 @@
+import torch
+
+from antipasto.eval import append_choice_newline
+
+
+class NewlineTokenizer:
+    def __call__(self, text, add_special_tokens):
+        assert text == "\n"
+        assert not add_special_tokens
+        return {"input_ids": [107]}
+
+
+def test_append_choice_newline_appends_before_scoring():
+    input_ids = torch.tensor([[1, 2], [3, 4]])
+    attention_mask = torch.tensor([[1, 1], [1, 0]])
+
+    input_ids, attention_mask = append_choice_newline(
+        input_ids, NewlineTokenizer(), attention_mask
+    )
+
+    assert input_ids.tolist() == [[1, 2, 107], [3, 4, 107]]
+    assert attention_mask.tolist() == [[1, 1, 1], [1, 0, 1]]