update

evaluate_auc.py

@@ -435,6 +435,7 @@ def main():
     parser.add_argument("--model_name", type=str, default="n_embd_256_n_layer_16_n_head_16", help="Model checkpoint name")
     parser.add_argument("--dataset_subset_size", type=int, default=-1, help="Dataset subset size for evaluation")
     parser.add_argument("--n_bootstrap", type=int, default=1, help="Number of bootstrap samples")
+    parser.add_argument("--offset", type=float, default=365.25, help="Offset in days for prediction")
     # Optional filtering/chunking parameters:
     parser.add_argument("--filter_min_total", type=int, default=100, help="Minimum total count to filter tokens")
     parser.add_argument("--disease_chunk_size", type=int, default=200, help="Chunk size for processing diseases")
@@ -444,6 +445,7 @@ def main():
     model_name = args.model_name
     output_path = f'auc_evaluation_{model_name}'
     dataset_subset_size = args.dataset_subset_size
+    offset = args.offset
 
     # Create output folder if it doesn't exist.
     Path(output_path).mkdir(exist_ok=True, parents=True)
@@ -487,6 +489,7 @@ def main():
         disease_chunk_size=args.disease_chunk_size,
         device=device,
         seed=seed,
+        offset=offset,
         n_bootstrap=args.n_bootstrap,
         n_jobs=args.n_jobs,
     )