mirror of
https://github.com/SakanaAI/doc-to-lora.git
synced 2026-07-23 17:01:04 +02:00
891c0bd256
* faster slice * add facts + ctx_qa * new configs * new scripts * intx_sft.py to train.py * add kaggle for downloading facts * max_new_tokens cli for eval * generate negative_nq * scripts + configs * default vals * max_val_samples_per_ds=500 * more efficient layer-to-layer ctx encoder * use_per_ctx_average_loss * faster processing * small exp distill * scripts * more robust watcher * per-module l1_norm avg * per-ctx average loss * clear_gpu
119 lines
3.8 KiB
Python
119 lines
3.8 KiB
Python
import itertools
|
|
import os
|
|
import time
|
|
from glob import glob
|
|
|
|
import wandb
|
|
import yaml
|
|
|
|
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"
|
|
|
|
|
|
def flatten(l):
|
|
return itertools.chain.from_iterable(l)
|
|
|
|
|
|
# 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
|
|
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
|
|
return sorted(list(new_files))
|
|
|
|
def update(self, file):
|
|
if file in self.last_files:
|
|
return
|
|
self.last_files.add(file)
|
|
print(f"Added {file} to evaluated 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") 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["FLASH_ATTENTION_DETERMINISTIC"] = "1"
|
|
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
|
|
os.environ["WANDB_PROJECT"] = "ctx_to_lora"
|
|
|
|
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(20)
|
|
if not os.path.exists(file):
|
|
# cp is delete before we can read it
|
|
continue
|
|
run_dir = file.split("/checkpoint")[0]
|
|
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": run_name,
|
|
"name": f"{run_name}-eval",
|
|
"id": f"{run_name}-eval",
|
|
"resume": "allow",
|
|
}
|
|
wandb.init(**wandb_kwargs)
|
|
|
|
# TODO: have to change this for bigger models
|
|
eval_batch_size = 8
|
|
eval_batch_size_gen = 16
|
|
try:
|
|
# metrics = run_eval(
|
|
# checkpoint_path=file,
|
|
# eval_batch_size=eval_batch_size,
|
|
# split="validation",
|
|
# generative=False,
|
|
# )
|
|
metrics = {}
|
|
except FileNotFoundError as e:
|
|
print(f"Error evaluating {file}: {e}. The checkpoint might be deleted.")
|
|
continue
|
|
try:
|
|
gen_metrics = run_eval(
|
|
checkpoint_path=file,
|
|
eval_batch_size=eval_batch_size_gen,
|
|
split="validation",
|
|
generative=True,
|
|
)
|
|
except FileNotFoundError as e:
|
|
print(f"Error evaluating {file}: {e}. The checkpoint might be deleted.")
|
|
gen_metrics = {}
|
|
file = ""
|
|
metrics.update(gen_metrics)
|
|
for k in metrics:
|
|
wandb.log(metrics[k], step=curstep)
|
|
wandb.finish()
|
|
print(f"Logged metrics: {metrics}")
|
|
print("=" * 80)
|
|
clear_gpu()
|
|
watcher.update(file)
|
|
|
|
watcher.save_state()
|