Add links to existing colab notebook and more comments

src/callbacks.py

@@ -194,6 +194,9 @@ class GradientVarianceCallback(EvaluationCallbackBase):
     """Calculates gradient variance and distance between definitions and corresponding questions.
     Requires a tokenized eval dataset with keys: [<d1 definitions dataset>, '<d2 definitions dataset>', '<d1 questions>', '<d2 questions>'].
     Example: ['train_defs_d1consis', 'train_defs_d2consis', 'd1consis', 'd2consis']
+    
+    You can play with existing tensor logs in this colab notebook:
+    https://colab.research.google.com/drive/1K-bWitUMffNlB1cIl8ELq6jtnyG6_J5b?usp=sharing
     """
     def __init__(self, eval_dataset_tokenized,
                  keys,

src/lm_training_utils.py

@@ -55,6 +55,7 @@ def create_tokenizer(add_tokens_for_var_names=True, num_letters_per_var=3, max_x
     return tokenizer
 
 
+# !!! THIS FUNCTION IS NOT USED FOR THE EXPERIMENTS IN THE PAPER !!!
 def linear_probe(hugginface_model, eval_dataset_d1, eval_dataset_d2, save_path='logit_results.txt', device='cuda'):
     try:
         import statsmodels.api as sm

src/train_lm.py

@@ -400,7 +400,7 @@ def compute_objective(metrics: Dict[str, float]) -> float:
             trainer.log_metrics(f"eval_{k}", metrics)
             trainer.save_metrics(f"eval_{k}", metrics)
 
-        # linear probes
+        # !!! THIS BLOCK WITH LINEAR PROBES IS NOT USED FOR THE EXPERIMENTS IN THE PAPER !!!
         if training_args.do_lin_probe:
             from src.lm_training_utils import linear_probe
             logger.info('Starting linear probe')