add watcher

hyperlora/eval.py

@@ -153,7 +153,7 @@ def decode_test_result(test_dataset, test_result, tokenizer):
 def eval_generation(eval_trainer, tokenizer, datasets, split, gen_kwargs):
     if not isinstance(datasets, dict):
         datasets = {"": datasets}
-
+    out = {}
     for ds_name, ds in datasets.items():
         split_name = f"{split}_{ds_name}" if ds_name else split
         eval_result = eval_trainer.predict(
@@ -173,22 +173,26 @@ def eval_generation(eval_trainer, tokenizer, datasets, split, gen_kwargs):
             split=split_name,
             output_dir=eval_trainer.args.output_dir,
         )
+        out[split_name] = eval_result.metrics
         eval_trainer.log_metrics(split_name, eval_result.metrics)
         eval_trainer.save_metrics(split_name, eval_result.metrics)
         clear_gpu()
+    return out
 
 
 @torch.no_grad()
 def eval_teacher_forcing(eval_trainer, datasets, split):
     if not isinstance(datasets, dict):
         datasets = {"": datasets}
-
+    out = {}
     for ds_name, ds in datasets.items():
         split_name = f"{split}_{ds_name}" if ds_name else split
         metrics = eval_trainer.evaluate(ds, metric_key_prefix=split_name)
+        out[split_name] = metrics
         eval_trainer.log_metrics(split_name, metrics)
         eval_trainer.save_metrics(split_name, metrics)
         clear_gpu()
+    return out
 
 
 def train_collator(inp_list, tokenizer):
@@ -376,6 +380,7 @@ def evaluate(checkpoint_path, args, split, generative):
         "args": eval_trainer_args,
         "data_collator": partial(collator, tokenizer=tokenizer),
     }
+    out = {}
     if not generative:
 
         trainer_kwargs["compute_metrics"] = partial(
@@ -390,12 +395,15 @@ def evaluate(checkpoint_path, args, split, generative):
         # spents a few hours on this but couldn't find the reason
         for ds_name, ds in datasets.items():
             eval_trainer = Trainer(**trainer_kwargs)
-            eval_teacher_forcing(eval_trainer, {ds_name: ds}, split)
+            metrics = eval_teacher_forcing(eval_trainer, {ds_name: ds}, split)
+            out.update(metrics)
     else:
         eval_trainer = Seq2SeqTrainer(**trainer_kwargs)
-        eval_generation(eval_trainer, tokenizer, datasets, split, gen_kwargs)
+        metrics = eval_generation(eval_trainer, tokenizer, datasets, split, gen_kwargs)
+        out.update(metrics)
 
     clear_gpu()
+    return out
 
 
 if __name__ == "__main__":

hyperlora/watcher.py

@@ -0,0 +1,109 @@
+# 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 itertools
+import shutil
+import time
+import os
+import argparse
+from glob import glob
+
+import numpy as np
+import pandas as pd
+import wandb
+import yaml
+
+from eval import evaluate
+
+
+CP_PATTERN = "train_outputs/runs/*/checkpoint*/pytorch_model.bin"
+
+
+def flatten(l):
+    return itertools.chain.from_iterable(l)
+
+
+class Watcher:
+    def __init__(self, patterns):
+        self.patterns = patterns
+        self.files = self.get_files()
+        self.last_files = self.files
+
+    def get_files(self):
+        return set(flatten(glob(pattern) for pattern in self.patterns))
+
+    def watch(self):
+        self.files = self.get_files()
+        new_files = self.files - self.last_files
+        self.last_files = self.files
+        return new_files
+
+    def save_state(self):
+        with open("watcher_state.yaml", "w") as f:
+            yaml.dump({"last_files": self.last_files}, f)
+
+    def load_state(self):
+        if not os.path.exists("watcher_state.yaml"):
+            return
+        with open("watcher_state.yaml", "r") as f:
+            state = yaml.safe_load(f)
+        self.last_files = state["last_files"]
+
+
+if __name__ == "__main__":
+    os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "true"
+    os.environ["TOKENIZERS_PARALLELISM"] = "true"
+    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()
+    print("Watching for new files...")
+    while True:
+        time.sleep(10)
+        new_files = watcher.watch()
+        for file in new_files:
+            # workaround to prevent loading incomplete checkpoints
+            time.sleep(10)
+            if not os.path.exists(file):
+                # cp is delete before we can read it
+                continue
+            run_dir = file.split("/checkpoint")[0]
+            checkpoint_dir = os.path.dirname(file)
+            args = argparse.Namespace(
+                **yaml.safe_load(open(f"{run_dir}/args.yaml", "r"))
+            )
+            args.output_dir = checkpoint_dir
+            args.logging_dir = checkpoint_dir
+            args.run_name = run_dir.split("/")[-1]
+            print("Evaluating...")
+            print(f"checkpoint_dir: {checkpoint_dir}")
+            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",
+                "resume": "allow",
+            }
+            wandb.init(**wandb_kwargs)
+            metrics = evaluate(file, args, split="validation", generative=False)
+            gen_metrics = evaluate(file, args, split="validation", generative=True)
+            metrics.update(gen_metrics)
+            wandb.log(metrics, step=curstep)
+            wandb.finish()
+
+        watcher.save_state()