fixes

src/datasets/batch.py

@@ -51,10 +51,10 @@ def batch_hidden_states(model, tokenizer, data: Dataset, n=100, batch_size=2, mc
             
             yield dict(
                 hs0=hs0['hidden_states'][j],
-                scores1=hs0["scores"][j],
+                scores0=hs0["scores"][j],
                 
                 hs1=hs1['hidden_states'][j],
-                scores2=hs1["scores"][j],                    
+                scores1=hs1["scores"][j],                    
                 
                 true=true_labels[j].item(),
                 index=index[j],

src/datasets/hs.py

@@ -24,10 +24,21 @@ from tqdm.auto import tqdm
 from torch.utils.data import DataLoader
 from datasets import Dataset
 import numpy as np
+import torch.nn.functional as F
 
 default_class2choices = {False: ['No', 'Negative', 'no', 'false', 'wrong', 'False'], True: ['Yes', 'Positive', 'yes', 'true', 'correct', 'right', 'True']}
 
-def scores2choice_probs(row, class2_ids, keys=["scores1", "scores2"] ):
+def scores2choice_probs(row, class2_ids, keys=["scores0", "scores1"] ):
+    """ Given next_token scores (logits) we take only the subset the corresponds to our
+    - negative tokens (e.g. False, no, ...) 
+    - and positive tokens (e.g. Yes, yes, affirmative, ...).
+    
+    example output:
+    {'choice_probs1': array([0.39, 0.31 ], dtype=float32),
+    'ans1': 0.44,
+    'choice_probs2': array([0.44, 0.45], dtype=float32),
+    'ans2': 0.502,}
+    """
     eps = 1e-5
     out = {}
     for key in keys:
@@ -39,7 +50,7 @@ def scores2choice_probs(row, class2_ids, keys=["scores1", "scores2"] ):
         out[key.replace("scores", "choice_probs")] = probs_c
         out[key.replace("scores", "ans")] = probs_c[1] / (np.sum(probs_c) + eps)
 
-        # # balance of logits (much more exagerated)
+        # # balance of logits (much more exaggerated)
         # scores_c = [scores[class2_ids[c]].sum() for c in class2_ids]
         # out[key.replace("scores", "ansb")] = torch.tensor(scores_c).softmax(-1)[1].item()
     return out

src/datasets/load.py

@@ -26,9 +26,9 @@ def ds2df(ds, cols=None):
     df = pd.DataFrame([rows_item(r) for r in df])
     
     # derived
-    df['dir_true'] = df['ans2'] - df['ans1']
-    df['conf'] = (df['ans1']-df['ans2']).abs()  
-    df['llm_prob'] = (df['ans1']+df['ans2'])/2
+    df['dir_true'] = df['ans1'] - df['ans0']
+    df['conf'] = (df['ans0']-df['ans1']).abs()  
+    df['llm_prob'] = (df['ans0']+df['ans1'])/2
     df['llm_ans'] = df['llm_prob']>0.5
     df['desired_ans'] = df.label ^ df.lie
     return df