allow whitespace before or after token choices

this helps in case of weird formatting. especially with bad prompts, or models that have become less coherent due to interventions

src/datasets/scores.py

@@ -47,30 +47,35 @@ def choice2id(tokenizer, c: str, whitespace_first=False) -> List[int]:
     # HACK: this whole function is messy, and specific to the llama tokenizer :(. I don't want it to fail silently, so I'm adding a few asserts. It's better to find out before 4 hours of data collection
     
     # Note some tokenizers differentiate between "yes", "\nyes" and " yes", and ideally we want all! 
-    ids2 = [
-        tokenizer(f' {c}', add_special_tokens=False)["input_ids"][-1],
-        tokenizer(f'\n{c}', add_special_tokens=False)["input_ids"][-1],
-        tokenizer(f'{c}', add_special_tokens=False)["input_ids"][0],
-    ]
+    ids2 = []
+    ids2 += tokenizer(f' {c}', add_special_tokens=False)["input_ids"]
+    ids2 += tokenizer(f'\n{c}', add_special_tokens=False)["input_ids"]
+    ids2 += tokenizer(f'{c}', add_special_tokens=False)["input_ids"]
     ids = list(set(ids2))
     
+    print(ids2)
+    print(ids)
+    print([f'`{t}`' for t in tokenizer.batch_decode(ids, skip_special_tokens=True)])
+    print([c.strip().startswith(tokenizer.decode(i)) for i in ids])
+
+
     # only include ones that decode to our original
-    ids = [i for i in ids2 if c.startswith(tokenizer.decode(i)) and len(tokenizer.decode(i))]
+    ids = [i for i in ids if c.strip().startswith(tokenizer.decode(i).strip()) and len(tokenizer.decode(i).strip())]
     assert len(ids)
     
     # QC: they should all decode to the same token
     decoded_ids = tokenizer.batch_decode(ids)
     shortest = sorted(decoded_ids, key=lambda s:len(s))[0]
     assert len(shortest)
-    assert all([decoded_ids[i].startswith(shortest) for i in range(len(decoded_ids))]), f"decoded_ids={decoded_ids}"
+    assert all([decoded_ids[i].strip().startswith(shortest) for i in range(len(decoded_ids))]), f"decoded_ids={decoded_ids}"
     
     # check that we can decode it
     c3 = tokenizer.batch_decode(ids)
     for c2 in c3:
-        if not c.startswith(c2) and len(c2):
+        if not c.strip().startswith(c2.strip()) and len(c2):
             print(c, c2, c3)
             ids = tokenizer(c, add_special_tokens=False)["input_ids"]
-            decoded_ids = [tokenizer.decode(i) for i in ids]
+            decoded_ids = [tokenizer.decode(i).strip() for i in ids]
             print(f"{c}=>{ids}=>{decoded_ids}")
             raise AssertionError(f'We should be able to encode and decode the choices, but it failed: tokenizer.decode(tokenizer(`{c}`))==`{c2}`!=`{c}`')
     return ids