From dcda29bae6c51268800a084ae3680458321b551f Mon Sep 17 00:00:00 2001
From: 51616 <rujikorn.ch@gmail.com>
Date: Mon, 2 Jun 2025 09:02:20 +0000
Subject: [PATCH] watcher + partial refactor eval + debug train w/ pwc and
 hotpot

---
 .gitignore                               |   4 +-
 README.md                                |  16 +
 configs/pwc_hotpot_qa.yaml               |  24 +-
 intx_sft.py                              |  73 +++--
 pyproject.toml                           |   5 +-
 run_eval.py                              |  84 +++++
 src/ctx_to_lora/data/collator.py         | 108 ++++++
 src/ctx_to_lora/data/definitions.py      |  20 ++
 src/ctx_to_lora/data/processing.py       |  32 +-
 src/ctx_to_lora/eval.py                  |   0
 eval.py => src/ctx_to_lora/eval_utils.py | 397 ++++-------------------
 src/ctx_to_lora/metrics.py               |  17 +-
 watcher.py                               |  73 ++---
 13 files changed, 409 insertions(+), 444 deletions(-)
 create mode 100644 run_eval.py
 delete mode 100644 src/ctx_to_lora/eval.py
 rename eval.py => src/ctx_to_lora/eval_utils.py (66%)

diff --git a/.gitignore b/.gitignore
index f1fe41d..03042d4 100644
--- a/.gitignore
+++ b/.gitignore
@@ -24,4 +24,6 @@ watcher_state.yaml
 ds_config.json
 eval_results/
 .github/
-*.code-workspace
\ No newline at end of file
+*.code-workspace
+.wandb/
+.ruff_cache/
diff --git a/README.md b/README.md
index 8b40768..939a315 100644
--- a/README.md
+++ b/README.md
@@ -33,6 +33,22 @@ print(model.decode(outputs))
 ```bash
 WANDB_MODE=disabled uv run python intx_sft.py configs/context_numbers_10.yaml --model_name_or_path=meta-llama/Llama-3.2-1B-Instruct --num_train_epochs=100 --per_device_train_batch_size=64 --per_device_eval_batch_size=64 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj --per_rank_gen=True --per_layer_processing=True --decoder_depth=2 --seed=1 --gen_lora_l1_reg_coef=0 --use_token_mixing=False --lora_r=8 --bf16=True --tf32=True --dataloader_num_workers=8 --dataloader_prefetch_factor=8
 ```
+### PwC + Hotpot training (for testing/debugging)
+```bash
+uv run intx_sft.py configs/pwc_hotpot_qa.yaml \
+--model_name_or_path=google/gemma-2-2b-it --num_train_epochs=1 --per_device_train_batch_size=32 \
+--gradient_accumulation_steps=1 --per_device_eval_batch_size=32 --exp_setup=hyper_lora --aggregator_type=perceiver \
+--target_modules=down_proj \
+--num_self_attends_per_block=4 --num_latent_factor=2 \
+--lora_r=8 \
+--eval_steps=1000 --save_steps=1000 --learning_rate=4e-5 --lora_dropout=0.0 \
+--neftune_noise_alpha=5 --use_light_weight_lora=False \
+--load_best_model_at_end=True --metric_for_best_model=pwc_loss --add_negative_prompt=False \
+--add_repeat_prompt=False \
+--use_sequence_packing=True --per_rank_gen=True \
+--per_layer_processing=True \
+--gen_lora_l1_reg_coef=0.1
+```
 
 <!-- ### HyperLoRA w/ context_numbers_128
 ```bash
diff --git a/configs/pwc_hotpot_qa.yaml b/configs/pwc_hotpot_qa.yaml
index d2cb994..26da9d3 100644
--- a/configs/pwc_hotpot_qa.yaml
+++ b/configs/pwc_hotpot_qa.yaml
@@ -1,6 +1,6 @@
 output_dir: "" # just a placeholder
 bf16: true
-model_name_or_path: meta-llama/Llama-3.2-1B-Instruct
+model_name_or_path: google/gemma-2-2b-it
 label_names: ["labels"]
 # eval_on_start: True
 # eval_strategy: "steps"
@@ -21,29 +21,31 @@ per_device_eval_batch_size: 8
 max_val_samples_per_ds: 1000
 # optim: schedule_free_adamw
 
-learning_rate: 0.00002
+learning_rate: 0.00004
 # lr_scheduler_type: "constant_with_warmup"
 neftune_noise_alpha: 1
 weight_decay: 0.01
 
+warmup_steps: 100
+
+dataloader_prefetch_factor: 8
+dataloader_num_workers: 8
 
 # LoRA
 lora_r: 8
-lora_dropout: 0.05
+lora_dropout: 0.02
 target_modules:
   - down_proj
-  - up_proj
-  - gate_proj
 
 # data
 train_ds_names:
-- pwc
-- hotpot_qa
+  - pwc
+  - hotpot_qa
 
 val_ds_names:
-- pwc
-- hotpot_qa
+  - pwc
+  - hotpot_qa
 
 test_ds_names:
-- pwc
-- hotpot_qa
\ No newline at end of file
+  - pwc
+  - hotpot_qa
diff --git a/intx_sft.py b/intx_sft.py
index 1279a01..30e4c2b 100755
--- a/intx_sft.py
+++ b/intx_sft.py
@@ -6,6 +6,7 @@ from math import ceil
 
 import numpy as np
 import torch
+import wandb
 from datasets import (
     disable_caching,
     interleave_datasets,
@@ -17,8 +18,6 @@ from transformers import (
 )
 from transformers.utils import is_liger_kernel_available
 
-import wandb
-
 from ctx_to_lora.configs import (
     AggregatorArguments,
     ArgumentParser,
@@ -32,7 +31,6 @@ from ctx_to_lora.configs import (
     TrainingArguments,
 )
 from ctx_to_lora.data.collator import train_collator, train_packed_collator
-from ctx_to_lora.data.definitions import DS_LEN
 from ctx_to_lora.data.processing import get_tokenized_dataset
 from ctx_to_lora.metrics import (
     Evaluator,
@@ -65,6 +63,21 @@ logger = logging.getLogger()
 LOCAL_RANK = int(os.getenv("LOCAL_RANK", "0"))
 
 
+def get_ds_prob(train_ds_len: list[int], total_len: int):
+    # if a dataset is smaller than 1%, make it 1%
+    probs = [0 for _ in train_ds_len]
+    for i, ds_len in enumerate(train_ds_len):
+        if ds_len / total_len <= 0.01:
+            probs[i] = 0.01
+    res_probs = 1 - sum(probs)
+    res_total_len = sum([l for l in train_ds_len if l / total_len > 0.01])
+    for i, ds_len in enumerate(train_ds_len):
+        if ds_len / total_len > 0.01:
+            probs[i] = ds_len / res_total_len * res_probs
+    assert sum(probs) == 1
+    return probs
+
+
 def main():
     ############ Argument parsing
     parser = ArgumentParser(
@@ -283,44 +296,44 @@ def main():
             val_ds[ds_name] = val_ds[ds_name].select(val_indices)
 
     # if data_args.streaming:
+    # max_steps = ceil(
+    #     sum(DS_LEN[ds] for ds in train_ds)
+    #     * training_args.num_train_epochs
+    #     / training_args.per_device_train_batch_size
+    #     / training_args.gradient_accumulation_steps
+    #     / training_args.world_size
+    # )
+    # training_args.max_steps = max_steps
+
+    # # interleaving streaming datasets
+    # # simplify the probs for smaller datasets
+    # # slightly upsample those datasets
+    # probs = [
+    #     0.01 if "fw_qa" not in ds_name else 1 + 0.01 - 0.01 * len(train_ds)
+    #     for ds_name in train_ds
+    # ]
+    # train_ds = interleave_datasets(
+    #     list(train_ds.values()),
+    #     probabilities=probs,
+    #     stopping_strategy="all_exhausted",
+    #     seed=training_args.seed,
+    # )
+    # else:
+    train_ds_len = [len(ds) for ds in train_ds.values()]
+    total_len = sum(train_ds_len)
     max_steps = ceil(
-        sum(DS_LEN[ds] for ds in train_ds)
+        total_len
         * training_args.num_train_epochs
         / training_args.per_device_train_batch_size
         / training_args.gradient_accumulation_steps
         / training_args.world_size
     )
     training_args.max_steps = max_steps
-
-    # interleaving streaming datasets
-    # simplify the probs for smaller datasets
-    # slightly upsample those datasets
-    probs = [
-        0.01 if "fw_qa" not in ds_name else 1 + 0.01 - 0.01 * len(train_ds)
-        for ds_name in train_ds
-    ]
     train_ds = interleave_datasets(
         list(train_ds.values()),
-        probabilities=probs,
-        stopping_strategy="all_exhausted",
+        probabilities=get_ds_prob(train_ds_len, total_len),
         seed=training_args.seed,
     )
-    # else:
-    #     train_ds_len = [len(ds) for ds in train_ds.values()]
-    #     total_len = sum(train_ds_len)
-    #     max_steps = ceil(
-    #         total_len
-    #         * training_args.num_train_epochs
-    #         / training_args.per_device_train_batch_size
-    #         / training_args.gradient_accumulation_steps
-    #         / training_args.world_size
-    #     )
-    #     training_args.max_steps = max_steps
-    #     train_ds = interleave_datasets(
-    #         list(train_ds.values()),
-    #         probabilities=[l / total_len for l in train_ds_len],
-    #         seed=training_args.seed,
-    #     )
 
     logger.info(f"train_ds: {train_ds}")
     logger.info(f"val_ds: {val_ds}")
diff --git a/pyproject.toml b/pyproject.toml
index 2d4700b..6f2e6c3 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -59,6 +59,9 @@ exclude = [
     "plots",
     "tmp",
     "wandb",
+    ".wandb",
+    ".ruff_cache",
+    "assets",
 ]
 typeCheckingMode = "off"
 
@@ -69,6 +72,4 @@ ignore = ["E", "F"]
 
 [tool.isort]
 profile = "black"
-known_first_party = ["wandb"]
 known_local_folder = ["ctx_to_lora"]
-known_third_party = ["wandb"]
diff --git a/run_eval.py b/run_eval.py
new file mode 100644
index 0000000..19cb6f8
--- /dev/null
+++ b/run_eval.py
@@ -0,0 +1,84 @@
+import logging
+
+from ctx_to_lora.eval_utils import (
+    run_eval,
+)
+
+logger = logging.getLogger()
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Evaluate a checkpoint")
+    parser.add_argument(
+        "--model_name_or_path",
+        type=str,
+        default=None,
+        help="Evaluate a base model from HuggingFace Hub, without loading checkpoint",
+    )
+    parser.add_argument(
+        "--checkpoint_path",
+        type=str,
+        default=None,
+        help="Path to the checkpoint to evaluate",
+    )
+    parser.add_argument(
+        "--split",
+        type=str,
+        choices=["validation", "test"],
+        default="validation",
+        help="Which split to evaluate on",
+    )
+    parser.add_argument(
+        "--datasets",
+        type=str,
+        nargs="+",
+        help=(
+            "Specific datasets to evaluate on."
+            "If not provided, uses default from args.yaml"
+        ),
+    )
+    parser.add_argument(
+        "--eval_batch_size",
+        type=int,
+        default=8,
+        help="Eval batch size for teacher forcing",
+    )
+    parser.add_argument(
+        "--eval_batch_size_gen",
+        type=int,
+        default=32,
+        help="Eval batch size for generation",
+    )
+    parser.add_argument(
+        "--remove_context",
+        action="store_true",
+        help="Remove context when evaluating the base model.",
+    )
+
+    cli_args = vars(parser.parse_args())
+    # setup_logging(output_dir, debug=os.getenv("DEBUG", False))
+
+    eval_batch_size_gen = cli_args.pop("eval_batch_size_gen")
+    eval_batch_size = cli_args.pop("eval_batch_size")
+    run_eval(
+        **cli_args,
+        # cli_args.checkpoint_path,
+        # cli_args.model_name_or_path,
+        # cli_args.eval_batch_size,
+        # args,
+        # split=cli_args.split,
+        eval_batch_size=eval_batch_size,
+        generative=False,
+    )
+    run_eval(
+        **cli_args,
+        # cli_args.checkpoint_path,
+        # cli_args.model_name_or_path,
+        # cli_args.eval_batch_size_gen,
+        # args,
+        # split=cli_args.split,
+        eval_batch_size=eval_batch_size_gen,
+        generative=True,
+    )
diff --git a/src/ctx_to_lora/data/collator.py b/src/ctx_to_lora/data/collator.py
index 2cf62d8..0a6a6ce 100644
--- a/src/ctx_to_lora/data/collator.py
+++ b/src/ctx_to_lora/data/collator.py
@@ -1,3 +1,4 @@
+import numpy as np
 import torch
 from transformers.data import DataCollatorWithFlattening
 
@@ -56,3 +57,110 @@ def train_collator(inp_list, tokenizer):
         out["ctx_attn_mask"] = ctx_attn_mask
 
     return out
+
+
+def eval_collator(inp_list, tokenizer):
+    # input is a list of tokenized sequences
+    padding_kwargs = dict(
+        padding=True,
+        padding_side="right",
+        pad_to_multiple_of=8,
+        return_tensors="pt",
+    )
+
+    ctx_ids = None
+    if "ctx_ids" in inp_list[0]:
+        # have to be manual since it has [ctx_len, features] shape
+        # pad to the longest ctx_len in the batch
+        # which can have a different length from the input_ids, attn_mask, labels
+
+        ctx_ids = [example.pop("ctx_ids") for example in inp_list]
+        ctx_ids = torch.nn.utils.rnn.pad_sequence(
+            ctx_ids,
+            batch_first=True,
+            padding_value=0,
+        )
+        # exotic keys won't be padded, so we need to pad them as well
+        ctx_attn_mask = [example.pop("ctx_attn_mask") for example in inp_list]
+        ctx_attn_mask = torch.nn.utils.rnn.pad_sequence(
+            ctx_attn_mask,
+            batch_first=True,
+            padding_value=0,
+        )
+
+    chat_ids = None
+    if "chat_ids" in inp_list[0]:
+        chat_ids = [x.pop("chat_ids") for x in inp_list]
+        chat_ids = torch.nn.utils.rnn.pad_sequence(
+            chat_ids,
+            batch_first=True,
+            padding_value=0,
+        )
+        chat_attn_mask = [x.pop("chat_attn_mask") for x in inp_list]
+        chat_attn_mask = torch.nn.utils.rnn.pad_sequence(
+            chat_attn_mask,
+            batch_first=True,
+            padding_value=0,
+        )
+        chat_labels = [x.pop("chat_labels") for x in inp_list]
+        chat_labels = torch.nn.utils.rnn.pad_sequence(
+            chat_labels,
+            batch_first=True,
+            padding_value=-100,
+        )
+        chat_labels = torch.where(chat_attn_mask == 0, -100, chat_labels)
+
+    labels = [x.pop("labels") for x in inp_list]
+    padded_seq = tokenizer.pad(inp_list, **padding_kwargs)
+
+    # hacky explicit padding since the labels are not padded by default
+    labels = tokenizer.pad({"input_ids": labels}, **padding_kwargs)["input_ids"]
+    labels = torch.where(padded_seq["attention_mask"] == 0, -100, labels)
+    out = {**padded_seq, "labels": labels}
+    if ctx_ids is not None:
+        out["ctx_ids"] = ctx_ids
+        out["ctx_attn_mask"] = ctx_attn_mask
+    if chat_ids is not None:
+        out["chat_ids"] = chat_ids
+        out["chat_attn_mask"] = chat_attn_mask
+        out["chat_labels"] = chat_labels
+    return out
+
+
+def generation_collator(inp_list, tokenizer):
+    padding_kwargs = dict(padding=True, padding_side="left", return_tensors="pt")
+    input_ids = [x.pop("input_ids") for x in inp_list]
+    attn_mask = [x.pop("attention_mask") for x in inp_list]
+    labels = [x.pop("labels") for x in inp_list]
+    for i, label in enumerate(labels):
+        # remove the response tokens
+        idx = np.argmax(label != -100)
+        input_ids[i] = input_ids[i][:idx]
+        attn_mask[i] = attn_mask[i][:idx]
+    out = tokenizer.pad(
+        {"input_ids": input_ids, "attention_mask": attn_mask}, **padding_kwargs
+    )
+
+    # we don't include the labels in the output
+    # since we don't want to compute the loss on the labels
+    # during generation
+    if "ctx_ids" in inp_list[0]:
+        # have to be manual since it has [ctx_len, features] shape
+        # pad to the longest ctx_len in the batch
+        # which can have a different length from the input_ids, attn_mask, labels
+        ctx_ids = [example.pop("ctx_ids") for example in inp_list]
+        ctx_ids = torch.nn.utils.rnn.pad_sequence(
+            ctx_ids,
+            batch_first=True,
+            padding_value=0,
+        )
+        # exotic keys won't be padded, so we need to pad them as well
+        ctx_attn_mask = [example.pop("ctx_attn_mask") for example in inp_list]
+        ctx_attn_mask = torch.nn.utils.rnn.pad_sequence(
+            ctx_attn_mask,
+            batch_first=True,
+            padding_value=0,
+        )
+        out["ctx_ids"] = ctx_ids
+        out["ctx_attn_mask"] = ctx_attn_mask
+    return out
diff --git a/src/ctx_to_lora/data/definitions.py b/src/ctx_to_lora/data/definitions.py
index a0de952..1357984 100644
--- a/src/ctx_to_lora/data/definitions.py
+++ b/src/ctx_to_lora/data/definitions.py
@@ -333,6 +333,26 @@ for ds_name in LONGBENCH_TASKS + LONGBENCH_E_TASKS:
         )
     )
 
+
+CLOSED_QA_DATASETS = {
+    "longbench/narrativeqa",
+    "longbench/qasper",
+    "longbench/multifieldqa_en",
+    "longbench/hotpotqa",
+    "longbench/2wikimqa",
+    "longbench/musique",
+    "hotpot_qa",
+    "squad",
+    "triviaqa_retrieved",
+    "negative_nq",
+}
+
+
+for ds_name in list(CLOSED_QA_DATASETS):
+    if ds_name.startswith("longbench/"):
+        CLOSED_QA_DATASETS.add(f"{ds_name}_e")
+
+
 # for training closed qa datasets, e.g., hotpot_qa, squad, etc.
 CLOSED_QA_INTX_TEMPLATES = [
     "Answer the question based on the given passages. Only give me the answer and do not output any other words.\n\nQuestion: {input}",
diff --git a/src/ctx_to_lora/data/processing.py b/src/ctx_to_lora/data/processing.py
index fb0f368..9104fe0 100644
--- a/src/ctx_to_lora/data/processing.py
+++ b/src/ctx_to_lora/data/processing.py
@@ -254,18 +254,18 @@ def get_ds_kwargs(ds_name: str, split: str) -> dict[str, Any]:
         # return DS_KWARGS[ds_name][split]
         kwargs = DS_KWARGS[ds_name][split]
 
-    # custom logic for slicing iterable datasets
-    take, skip = None, None
-    kw_split = kwargs["split"]
-    if ("[" in kw_split) and kw_split.endswith("]"):
-        kwargs["split"], slice = kw_split.split("[")
-        slice = slice.strip("]")
-        skip = slice.split(":")[0]
-        if skip:
-            kwargs["skip"] = int(skip)
-        take = slice.split(":")[1]
-        if take:
-            kwargs["take"] = int(take)
+    # # custom logic for slicing iterable datasets
+    # take, skip = None, None
+    # kw_split = kwargs["split"]
+    # if ("[" in kw_split) and kw_split.endswith("]"):
+    #     kwargs["split"], slice = kw_split.split("[")
+    #     slice = slice.strip("]")
+    #     skip = slice.split(":")[0]
+    #     if skip:
+    #         kwargs["skip"] = int(skip)
+    #     take = slice.split(":")[1]
+    #     if take:
+    #         kwargs["take"] = int(take)
     return kwargs
 
 
@@ -439,7 +439,9 @@ def _load_and_process_dataset(
         # TODO: drop samples or split into multiple contexts
         ds = ds.filter(filter_long_samples, batched=True, num_proc=num_proc)
         if add_negative_prompt:
-            # TODO: make explicit dataset for this
+            # TODO: add explicit dataset for this
+            # w/ qs that are answerable 0-shot
+            # here we add the context to the prompt
             ds = ds.map(
                 add_negative_prompt_fn,
                 batched=True,
@@ -484,8 +486,6 @@ def get_tokenized_dataset(
     load_and_process_kwargs = dict(
         ds_name=ds_name,
         split=split,
-        base_model_max_len=base_model_max_len,
-        ctx_model_max_len=ctx_model_max_len,
         add_negative_prompt=add_negative_prompt,
         add_repeat_prompt=add_repeat_prompt,
         repeat_prob=repeat_prob,
@@ -544,8 +544,10 @@ def construct_and_tokenize_ctx_qa(
     num_proc=None,
 ):
     kwargs = dict(
+        base_model_max_len=base_model_max_len,
         tokenizer=repr(tokenizer),
         tokenizer_kwargs=json.dumps(tokenizer_kwargs),
+        ctx_model_max_len=ctx_model_max_len,
         ctx_tokenizer=repr(ctx_tokenizer),
         ctx_tokenizer_kwargs=json.dumps(ctx_tokenizer_kwargs),
         add_ctx_to_chat=add_ctx_to_chat,
diff --git a/src/ctx_to_lora/eval.py b/src/ctx_to_lora/eval.py
deleted file mode 100644
index e69de29..0000000
diff --git a/eval.py b/src/ctx_to_lora/eval_utils.py
similarity index 66%
rename from eval.py
rename to src/ctx_to_lora/eval_utils.py
index 6f095e3..9616949 100644
--- a/eval.py
+++ b/src/ctx_to_lora/eval_utils.py
@@ -1,4 +1,3 @@
-import gc
 import json
 import logging
 import os
@@ -7,10 +6,8 @@ import string
 import sys
 from argparse import Namespace
 from collections import Counter, defaultdict
-from collections.abc import Callable
 from dataclasses import fields
 from functools import partial
-from typing import Any
 
 import numpy as np
 import pandas as pd
@@ -18,9 +15,7 @@ import torch
 import yaml
 from datasets import disable_caching
 from peft import PeftModel
-from rouge_score import rouge_scorer
 from transformers import (
-    EvalPrediction,
     GenerationConfig,
     PreTrainedModel,
     Seq2SeqTrainer,
@@ -30,20 +25,31 @@ from transformers import (
 )
 from transformers.utils import is_liger_kernel_available
 
-from ctx_to_lora.data.definitions import LONGBENCH_E_TASKS, LONGBENCH_TASKS
+from ctx_to_lora.data.collator import eval_collator, generation_collator
+from ctx_to_lora.data.definitions import (
+    CLOSED_QA_DATASETS,
+    LONGBENCH_E_TASKS,
+    LONGBENCH_TASKS,
+)
 from ctx_to_lora.data.processing import get_tokenized_dataset
+from ctx_to_lora.metrics import (
+    Evaluator,
+    compute_metrics,
+    compute_per_token_acc,
+    compute_perplexity,
+    compute_prefix_matching,
+    compute_rouge,
+)
 from ctx_to_lora.model_loading import get_model, get_tokenizer
 from ctx_to_lora.modeling import hypernet
 from ctx_to_lora.modeling.hypernet import ModulatedPretrainedModel
-
-# bandaid for loading old models (before restructure)
-sys.modules["ctx_to_lora.modeling_utils"] = hypernet
+from ctx_to_lora.utils import clear_gpu, setup_logging
 
 logger = logging.getLogger()
 
 
-# TODO: add negative samples eval (per-sample?)
-# TODO: add length info for more fine grain eval
+# bandaid for loading old models (before restructure)
+sys.modules["ctx_to_lora.modeling_utils"] = hypernet
 
 
 # longbench metrics
@@ -93,30 +99,6 @@ def compute_qa_f1_score(pred_texts: list[str], label_texts: list[str]):
     return dict(qa_f1=np.mean(res))
 
 
-CLOSED_QA_DATASETS = {
-    "longbench/narrativeqa",
-    "longbench/qasper",
-    "longbench/multifieldqa_en",
-    "longbench/hotpotqa",
-    "longbench/2wikimqa",
-    "longbench/musique",
-    "hotpot_qa",
-    "squad",
-    "triviaqa_retrieved",
-    "negative_nq",
-}
-
-for ds_name in list(CLOSED_QA_DATASETS):
-    CLOSED_QA_DATASETS.add(f"{ds_name}_e")
-
-
-def clear_gpu():
-    gc.collect()
-    torch.cuda.empty_cache()
-    torch.cuda.reset_max_memory_allocated()
-    torch.cuda.reset_max_memory_cached()
-
-
 def add_longbench_tasks(ds_names: list[str]):
     if "longbench" in ds_names:
         ds_names.remove("longbench")
@@ -126,106 +108,6 @@ def add_longbench_tasks(ds_names: list[str]):
         ds_names += LONGBENCH_E_TASKS
 
 
-def compute_rouge(pred_texts, label_texts):
-    out = defaultdict(list)
-    scorer = rouge_scorer.RougeScorer(["rougeL"], use_stemmer=True)
-    for pred_text, label_text in zip(pred_texts, label_texts):
-        scores = scorer.score(pred_text, label_text)
-        for k, v in scores.items():
-            out[f"{k}.f1"].append(v.fmeasure)
-    for k in out:
-        out[k] = np.mean(out[k])
-    return out
-
-
-def compute_per_token_acc(shift_logits, shift_labels, valid_masks):
-    indices = torch.where(valid_masks)
-    # acc = (shift_logits.argmax(-1) == shift_labels)[indices].float().mean().item()
-    # return {"per_token_acc": acc}
-    acc = (shift_logits.argmax(-1) == shift_labels)[indices].float()
-    return {
-        "per_token_accs": acc.flatten().tolist(),
-        "n_per_token_accs": valid_masks.sum().item(),
-    }
-
-
-def compute_prefix_matching(shift_logits, shift_labels, valid_masks):
-    lengths = valid_masks.sum(dim=1)
-
-    is_wrong = (shift_logits.argmax(-1) != shift_labels) * valid_masks
-    is_correct = (shift_logits.argmax(-1) == shift_labels) * valid_masks
-    # NOTE: not reliable for multi-turn conversations
-    # ie, all tokens in the following user's turn will be correct
-    # still monotonically correlate with perf though
-    wrong_pos = torch.argmax(is_wrong, dim=1) - torch.argmax(valid_masks, dim=1)
-    perf = wrong_pos / lengths
-
-    # if all tokens are correct, set to 1
-    perf = torch.where(is_correct.sum(dim=1) == lengths, 1, perf)
-    # return {"prefix_matching": perf.mean().item()}
-    return {
-        "prefix_matchings": perf.tolist(),
-        "n_prefix_matchings": valid_masks.shape[0],
-    }
-
-
-@torch.no_grad()
-def compute_perplexity(shift_logits, shift_labels, valid_masks):
-    loss_fct = torch.nn.CrossEntropyLoss(reduction="none")
-    loss = loss_fct(shift_logits.transpose(1, 2), shift_labels)
-    loss = (loss * valid_masks).sum(dim=1) / valid_masks.sum(dim=1)
-    # perplexity = torch.exp(loss).mean().item()
-    # return {"perplexity": perplexity}
-    preplexities = torch.exp(loss)
-    return {
-        "perplexities": preplexities.tolist(),
-        "n_perplexities": valid_masks.shape[0],
-    }
-
-
-class Evaluator:
-    def __init__(self, metric_fns: list[Callable]):
-        self.metric_fns = metric_fns
-        self.reset()
-
-    def reset(self):
-        self.accum_metrics = defaultdict(list)
-        self.count = defaultdict(list)
-
-    def update(self, shift_logits, shift_labels, valid_masks):
-        for metric_fn in self.metric_fns:
-            metric = metric_fn(shift_logits, shift_labels, valid_masks)
-            for k, v in metric.items():
-                if k.startswith("n_"):
-                    self.count[k[2:]].append(v)
-                else:
-                    self.accum_metrics[k] += v
-
-    def compute(self):
-        # Get result across entire eval set
-        result = {
-            k: np.sum(v) / np.sum(self.count[k]) for k, v in self.accum_metrics.items()
-        }
-        # Reset batch statistics
-        self.reset()
-        return result
-
-
-@torch.inference_mode()
-def compute_metrics(
-    eval_pred: EvalPrediction,
-    compute_result: bool,
-    evaluator: Evaluator,
-):
-    logits, labels = eval_pred.predictions, eval_pred.label_ids
-    shift_logits = logits[..., :-1, :]
-    shift_labels = labels[..., 1:]
-    valid_masks = torch.where(shift_labels != -100, 1, 0)
-    evaluator.update(shift_logits, shift_labels, valid_masks)
-    if compute_result:
-        return evaluator.compute()
-
-
 def save_generated_text(samples, output_dir, split):
     os.makedirs(output_dir, exist_ok=True)
     # Create any necessary subdirectories if split contains path separators
@@ -239,7 +121,7 @@ def save_generated_text(samples, output_dir, split):
 
 
 def create_metrics_csv(
-    metrics_dict: dict[str, dict[str, Any]],
+    metrics_dict: dict[str, dict[str, any]],
     output_dir: str,
     model_name: str,
     is_hypernet_model: bool = False,
@@ -261,12 +143,10 @@ def create_metrics_csv(
     # Collect all unique metric names, length groups, and tasks
     all_metrics = set()
     all_length_groups = set()
-    all_tasks = set()
+    all_splits = set()
 
-    # FIXME: fix name split (currently wrong)
     for split_name, metrics in metrics_dict.items():
-        task_name = split_name.split("_")[1] if "_" in split_name else split_name
-        all_tasks.add(task_name)
+        all_splits.add(split_name)
 
         for metric_key in metrics.keys():
             if not metric_key.startswith(split_name):
@@ -310,7 +190,7 @@ def create_metrics_csv(
             return (2, 0, 0)  # Put any malformed strings last
 
     all_length_groups = sorted(list(all_length_groups), key=sort_length_groups)
-    all_tasks = sorted(all_tasks)
+    all_splits = sorted(all_splits)
 
     # Create single row with hierarchical columns
     row_data = {
@@ -319,16 +199,9 @@ def create_metrics_csv(
     }
 
     # Create hierarchical column structure: task_metric_lengthgroup
-    for task in all_tasks:
+    for task in all_splits:
         # Find the corresponding split in metrics_dict
-        split_key = None
-        for split_name in metrics_dict.keys():
-            if task in split_name:
-                split_key = split_name
-                break
-
-        if split_key is None:
-            continue
+        split_key = split_name
 
         metrics = metrics_dict[split_key]
 
@@ -553,113 +426,6 @@ def eval_teacher_forcing(eval_trainer, datasets, split, remove_context):
     return out
 
 
-def train_collator(inp_list, tokenizer):
-    # input is a list of tokenized sequences
-    padding_kwargs = dict(
-        padding=True,
-        padding_side="right",
-        pad_to_multiple_of=8,
-        return_tensors="pt",
-    )
-
-    ctx_ids = None
-    if "ctx_ids" in inp_list[0]:
-        # have to be manual since it has [ctx_len, features] shape
-        # pad to the longest ctx_len in the batch
-        # which can have a different length from the input_ids, attn_mask, labels
-
-        ctx_ids = [example.pop("ctx_ids") for example in inp_list]
-        ctx_ids = torch.nn.utils.rnn.pad_sequence(
-            ctx_ids,
-            batch_first=True,
-            padding_value=0,
-        )
-        # exotic keys won't be padded, so we need to pad them as well
-        ctx_attn_mask = [example.pop("ctx_attn_mask") for example in inp_list]
-        ctx_attn_mask = torch.nn.utils.rnn.pad_sequence(
-            ctx_attn_mask,
-            batch_first=True,
-            padding_value=0,
-        )
-
-    chat_ids = None
-    if "chat_ids" in inp_list[0]:
-        chat_ids = [x.pop("chat_ids") for x in inp_list]
-        chat_ids = torch.nn.utils.rnn.pad_sequence(
-            chat_ids,
-            batch_first=True,
-            padding_value=0,
-        )
-        chat_attn_mask = [x.pop("chat_attn_mask") for x in inp_list]
-        chat_attn_mask = torch.nn.utils.rnn.pad_sequence(
-            chat_attn_mask,
-            batch_first=True,
-            padding_value=0,
-        )
-        chat_labels = [x.pop("chat_labels") for x in inp_list]
-        chat_labels = torch.nn.utils.rnn.pad_sequence(
-            chat_labels,
-            batch_first=True,
-            padding_value=-100,
-        )
-        chat_labels = torch.where(chat_attn_mask == 0, -100, chat_labels)
-
-    labels = [x.pop("labels") for x in inp_list]
-    padded_seq = tokenizer.pad(inp_list, **padding_kwargs)
-
-    # hacky explicit padding since the labels are not padded by default
-    labels = tokenizer.pad({"input_ids": labels}, **padding_kwargs)["input_ids"]
-    labels = torch.where(padded_seq["attention_mask"] == 0, -100, labels)
-    out = {**padded_seq, "labels": labels}
-    if ctx_ids is not None:
-        out["ctx_ids"] = ctx_ids
-        out["ctx_attn_mask"] = ctx_attn_mask
-    if chat_ids is not None:
-        out["chat_ids"] = chat_ids
-        out["chat_attn_mask"] = chat_attn_mask
-        out["chat_labels"] = chat_labels
-    return out
-
-
-def generation_collator(inp_list, tokenizer):
-    padding_kwargs = dict(padding=True, padding_side="left", return_tensors="pt")
-    input_ids = [x.pop("input_ids") for x in inp_list]
-    attn_mask = [x.pop("attention_mask") for x in inp_list]
-    labels = [x.pop("labels") for x in inp_list]
-    for i, label in enumerate(labels):
-        # remove the response tokens
-        idx = np.argmax(label != -100)
-        input_ids[i] = input_ids[i][:idx]
-        attn_mask[i] = attn_mask[i][:idx]
-    out = tokenizer.pad(
-        {"input_ids": input_ids, "attention_mask": attn_mask}, **padding_kwargs
-    )
-
-    # we don't include the labels in the output
-    # since we don't want to compute the loss on the labels
-    # during generation
-    if "ctx_ids" in inp_list[0]:
-        # have to be manual since it has [ctx_len, features] shape
-        # pad to the longest ctx_len in the batch
-        # which can have a different length from the input_ids, attn_mask, labels
-        ctx_ids = [example.pop("ctx_ids") for example in inp_list]
-        ctx_ids = torch.nn.utils.rnn.pad_sequence(
-            ctx_ids,
-            batch_first=True,
-            padding_value=0,
-        )
-        # exotic keys won't be padded, so we need to pad them as well
-        ctx_attn_mask = [example.pop("ctx_attn_mask") for example in inp_list]
-        ctx_attn_mask = torch.nn.utils.rnn.pad_sequence(
-            ctx_attn_mask,
-            batch_first=True,
-            padding_value=0,
-        )
-        out["ctx_ids"] = ctx_ids
-        out["ctx_attn_mask"] = ctx_attn_mask
-    return out
-
-
 def evaluate(
     checkpoint_path,
     model_name_or_path,
@@ -793,7 +559,7 @@ def evaluate(
 
     model.eval()
 
-    collator = generation_collator if generative else train_collator
+    collator = generation_collator if generative else eval_collator
 
     trainer_kwargs = {
         "model": model,
@@ -835,67 +601,25 @@ def evaluate(
     return out
 
 
-if __name__ == "__main__":
-    import argparse
-
-    parser = argparse.ArgumentParser(description="Evaluate a checkpoint")
-    parser.add_argument(
-        "--model_name_or_path",
-        type=str,
-        default=None,
-        help="Evaluate a base model from HuggingFace Hub, without loading checkpoint",
-    )
-    parser.add_argument(
-        "--checkpoint_path",
-        type=str,
-        default=None,
-        help="Path to the checkpoint to evaluate",
-    )
-    parser.add_argument(
-        "--split",
-        type=str,
-        choices=["validation", "test"],
-        default="validation",
-        help="Which split to evaluate on",
-    )
-    parser.add_argument(
-        "--datasets",
-        type=str,
-        nargs="+",
-        help=(
-            "Specific datasets to evaluate on."
-            "If not provided, uses default from args.yaml"
-        ),
-    )
-    parser.add_argument(
-        "--eval_batch_size",
-        type=int,
-        default=8,
-        help="Eval batch size for teacher forcing",
-    )
-    parser.add_argument(
-        "--eval_batch_size_gen",
-        type=int,
-        default=32,
-        help="Eval batch size for generation",
-    )
-    parser.add_argument(
-        "--remove_context",
-        action="store_true",
-        help="Remove context when evaluating the base model.",
-    )
-
-    cli_args = parser.parse_args()
+def run_eval(
+    checkpoint_path: str = None,
+    model_name_or_path: str = None,
+    datasets: list[str] = None,
+    split: str = "validation",
+    eval_batch_size: int = 8,
+    remove_context: bool = False,
+    generative: bool = False,
+):
     # setup_logging(output_dir, debug=os.getenv("DEBUG", False))
 
-    assert bool(cli_args.model_name_or_path) ^ bool(cli_args.checkpoint_path), (
+    assert bool(model_name_or_path) ^ bool(checkpoint_path), (
         "Either --model_name_or_path or --checkpoint_path must be provided"
     )
 
     os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
     os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "true"
     os.environ["FLASH_ATTENTION_DETERMINISTIC"] = "1"
-    os.environ["WANDB_MODE"] = "disabled"
+    # os.environ["WANDB_MODE"] = "disabled"
     disable_caching()
     set_seed(42)
     torch.use_deterministic_algorithms(True, warn_only=True)
@@ -905,51 +629,46 @@ if __name__ == "__main__":
     torch.backends.cuda.matmul.allow_tf32 = False
     torch.backends.cudnn.allow_tf32 = False
 
-    if cli_args.checkpoint_path:
-        checkpoint_dir = "/".join(cli_args.checkpoint_path.split("/")[:-1])
-        run_dir = "/".join(cli_args.checkpoint_path.split("/")[:-2])
-        cur_it = int(cli_args.checkpoint_path.split("checkpoint-")[1].split("/")[0])
+    if checkpoint_path:
+        checkpoint_dir = "/".join(checkpoint_path.split("/")[:-1])
+        run_dir = "/".join(checkpoint_path.split("/")[:-2])
+        cur_it = int(checkpoint_path.split("checkpoint-")[1].split("/")[0])
         args = Namespace(**yaml.unsafe_load(open(f"{run_dir}/args.yaml")))
-        print(f"checkpoint_path: {cli_args.checkpoint_path}")
+        print(f"checkpoint_path: {checkpoint_path}")
         print(f"run_dir: {run_dir}")
 
         args.output_dir = f"{run_dir}/eval-results-{cur_it}"
         args.logging_dir = f"{run_dir}/eval-results-{cur_it}"
         args.run_name = run_dir.split("/")[-1]
-        args.remove_context = False  # modulated model doesn't see ctx by default
+        # modulated model doesn't see ctx by default
+        # but remove_context has to be false for correct file naming
+        args.remove_context = False
     else:
         args = Namespace(
-            model_name_or_path=cli_args.model_name_or_path,
-            output_dir=f"eval_results/{cli_args.model_name_or_path}",
-            logging_dir=f"eval_results/{cli_args.model_name_or_path}",
-            run_name=f"eval_results/{cli_args.model_name_or_path}",
+            model_name_or_path=model_name_or_path,
+            output_dir=f"eval_results/{model_name_or_path}",
+            logging_dir=f"eval_results/{model_name_or_path}",
+            run_name=f"eval_results/{model_name_or_path}",
             # max_base_len=2**13,
             # max_ctx_len=-1,  # not used
             val_ds_names=[],
             test_ds_names=[],
-            remove_context=cli_args.remove_context,
+            remove_context=remove_context,
         )
+    setup_logging(args.logging_dir)
 
     # Override dataset names if provided via CLI
-    if cli_args.datasets:
-        if cli_args.split == "validation":
-            args.val_ds_names = cli_args.datasets
+    if datasets:
+        if split == "validation":
+            args.val_ds_names = datasets
         else:
-            args.test_ds_names = cli_args.datasets
+            args.test_ds_names = datasets
 
-    # evaluate(
-    #     cli_args.checkpoint_path,
-    #     cli_args.model_name_or_path,
-    #     cli_args.eval_batch_size,
-    #     args,
-    #     split=cli_args.split,
-    #     generative=False,
-    # )
     evaluate(
-        cli_args.checkpoint_path,
-        cli_args.model_name_or_path,
-        cli_args.eval_batch_size_gen,
+        checkpoint_path,
+        model_name_or_path,
+        eval_batch_size,
         args,
-        split=cli_args.split,
-        generative=True,
+        split,
+        generative=generative,
     )
diff --git a/src/ctx_to_lora/metrics.py b/src/ctx_to_lora/metrics.py
index 5caf8d9..c62807a 100644
--- a/src/ctx_to_lora/metrics.py
+++ b/src/ctx_to_lora/metrics.py
@@ -3,9 +3,23 @@ from collections.abc import Callable
 
 import numpy as np
 import torch
+from rouge_score import rouge_scorer
 from transformers import EvalPrediction
 
 
+def compute_rouge(pred_texts, label_texts):
+    out = defaultdict(list)
+    scorer = rouge_scorer.RougeScorer(["rougeL"], use_stemmer=True)
+    for pred_text, label_text in zip(pred_texts, label_texts):
+        scores = scorer.score(pred_text, label_text)
+        for k, v in scores.items():
+            out[f"{k}.f1"].append(v.fmeasure)
+    for k in out:
+        out[k] = np.mean(out[k])
+    return out
+
+
+@torch.inference_mode()
 def compute_per_token_acc(shift_logits, shift_labels, valid_masks):
     indices = torch.where(valid_masks)
     acc = (shift_logits.argmax(-1) == shift_labels)[indices].float()
@@ -15,6 +29,7 @@ def compute_per_token_acc(shift_logits, shift_labels, valid_masks):
     }
 
 
+@torch.inference_mode()
 def compute_prefix_matching(shift_logits, shift_labels, valid_masks):
     lengths = valid_masks.sum(dim=1)
 
@@ -34,7 +49,7 @@ def compute_prefix_matching(shift_logits, shift_labels, valid_masks):
     }
 
 
-@torch.no_grad()
+@torch.inference_mode()
 def compute_perplexity(shift_logits, shift_labels, valid_masks):
     loss_fct = torch.nn.CrossEntropyLoss(reduction="none")
     loss = loss_fct(shift_logits.transpose(1, 2), shift_labels)
diff --git a/watcher.py b/watcher.py
index 794adcb..a0742f9 100644
--- a/watcher.py
+++ b/watcher.py
@@ -1,19 +1,13 @@
-# handmade file watcher using glob
-# not using watchdog because there are too many saved files
-# but we want to just watch when these files are created
-# */checkpoints/it_*/hypermod.pt (for HyperLoRA)
-import argparse
-import gc
 import itertools
 import os
 import time
 from glob import glob
 
-import torch
+import wandb
 import yaml
 
-import wandb
-from eval import evaluate
+from ctx_to_lora.eval_utils import run_eval
+from ctx_to_lora.utils import clear_gpu
 
 CP_PATTERN = "train_outputs/runs/*/checkpoint*/pytorch_model.bin"
 
@@ -22,13 +16,9 @@ def flatten(l):
     return itertools.chain.from_iterable(l)
 
 
-def clear_gpu():
-    gc.collect()
-    torch.cuda.empty_cache()
-    torch.cuda.reset_max_memory_allocated()
-    torch.cuda.reset_max_memory_cached()
-
-
+# handmade file watcher using glob
+# not using watchdog because there are too many saved files
+# but we want to just watch CP_PATTERN files
 class Watcher:
     def __init__(self, patterns):
         self.patterns = patterns
@@ -42,7 +32,7 @@ class Watcher:
         self.files = self.get_files()
         new_files = self.files - self.last_files
         self.last_files = self.files
-        return new_files
+        return sorted(list(new_files))
 
     def save_state(self):
         with open("watcher_state.yaml", "w") as f:
@@ -56,26 +46,11 @@ class Watcher:
         self.last_files = state["last_files"]
 
 
-# TODO: run eval on "validation" split during training
-# using function from `src/ctx-to-lora/eval.py`
 if __name__ == "__main__":
     os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "true"
     os.environ["FLASH_ATTENTION_DETERMINISTIC"] = "1"
     os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
     os.environ["WANDB_PROJECT"] = "ctx_to_lora"
-    os.environ["WANDB_WATCH"] = ""  # "all"
-    os.environ["WANDB_CONSOLE"] = "off"
-    # checkpoint_path = sys.argv[1]
-    # checkpoint_dir = "/".join(checkpoint_path.split("/")[:-1])
-
-    # run_dir = "/".join(checkpoint_path.split("/")[:-2])
-    # args = Namespace(**yaml.unsafe_load(open(f"{run_dir}/args.yaml", "r")))
-    # print(f"checkpoint_path: {checkpoint_path}")
-    # print(f"run_dir: {run_dir}")
-    # # print(f"args: {args}")
-
-    # evaluate(checkpoint_path, args, split="validation", generative=False)
-    # evaluate(checkpoint_path, args, split="validation", generative=True)
 
     watcher = Watcher([CP_PATTERN])
     watcher.load_state()
@@ -90,25 +65,33 @@ if __name__ == "__main__":
                 # cp is delete before we can read it
                 continue
             run_dir = file.split("/checkpoint")[0]
-            cur_it = int(file.split("checkpoint-")[1].split("/")[0])
-            checkpoint_dir = os.path.dirname(file)
-            args = argparse.Namespace(**yaml.unsafe_load(open(f"{run_dir}/args.yaml")))
-            args.output_dir = f"{run_dir}/eval-results-{cur_it}"
-            args.logging_dir = f"{run_dir}/eval-results-{cur_it}"
-            args.run_name = run_dir.split("/")[-1]
-            print("Evaluating...")
-            print(f"checkpoint_dir: {checkpoint_dir}")
+            run_name = run_dir.split("/")[-1]
+            print(f"Evaluating {file}")
             curstep = int(file.split("checkpoint-")[1].split("/")[0])
             wandb_kwargs = {
                 "project": os.getenv("WANDB_PROJECT"),
-                "group": args.run_name,
-                "name": f"{args.run_name}-eval",
-                "id": f"{args.run_name}-eval",
+                "group": run_name,
+                "name": f"{run_name}-eval",
+                "id": f"{run_name}-eval",
                 "resume": "allow",
             }
             wandb.init(**wandb_kwargs)
-            metrics = evaluate(file, args, split="validation", generative=False)
-            gen_metrics = evaluate(file, args, split="validation", generative=True)
+
+            # TODO: have to change this for bigger models
+            eval_batch_size = 8
+            eval_batch_size_gen = 32
+            metrics = run_eval(
+                checkpoint_path=file,
+                eval_batch_size=eval_batch_size,
+                split="validation",
+                generative=False,
+            )
+            gen_metrics = run_eval(
+                checkpoint_path=file,
+                eval_batch_size=eval_batch_size_gen,
+                split="validation",
+                generative=True,
+            )
             metrics.update(gen_metrics)
             wandb.log(metrics, step=curstep)
             wandb.finish()