♻️ Style Transfer: add unique checkpoint folder for each run

lib/style-transfer/style_transfer/rb_gen/steps/dpo.py

@@ -1,5 +1,3 @@
-import glob
-
 import hydra
 import numpy as np
 import pandas as pd
@@ -61,7 +59,7 @@ def dpo_train(
     dataset = dataset.select_columns(["prompts", "chosen", "rejected"])
     dataset = dataset.rename_column("prompts", "prompt")
 
-    cfg.dpo.training_args.output_dir = f"models/dpo/{step}"
+    cfg.dpo.training_args.output_dir = f"models/{wandb.run.id}/dpo/{step}"
     args = hydra.utils.instantiate(cfg.dpo.training_args)
     args.padding_value = tokenizer.eos_token_id
     model = AutoPeftModelForCausalLM.from_pretrained(pretrained_model_name_or_path=model_path)
@@ -74,10 +72,7 @@ def dpo_train(
         train_dataset=dataset,
         callbacks=[CustomWandbCallback],
     )
-    if glob.glob(f"{args.output_dir}/*"):
-        dpo_trainer.train(resume_from_checkpoint=True)
-    else:
-        dpo_trainer.train()
+    dpo_trainer.train()
 
     dpo_path = args.output_dir
     dpo_trainer.save_model(dpo_path)

lib/style-transfer/style_transfer/rb_gen/steps/generate.py

@@ -51,14 +51,14 @@ def generate(
         torch_dtype=torch.bfloat16,
         trust_remote_code=True,
     ).merge_and_unload()
-    model.save_pretrained("models/merged/")
-    tokenizer.save_pretrained("models/merged/")
+    model.save_pretrained(f"models/{wandb.run.id}/merged/")
+    tokenizer.save_pretrained(f"models/{wandb.run.id}/merged/")
     del model
     del tokenizer
     gc.collect()
     torch.cuda.empty_cache()
     logging.info("🫧 Building VLLM Pipeline ...")
-    llm = hydra.utils.instantiate(cfg.gen.llm)
+    llm = hydra.utils.instantiate(cfg.gen.llm, model=f"models/{wandb.run.id}/merged/")
 
     logging.info("🎉 And it's done!")
 
@@ -79,7 +79,7 @@ def generate(
     del llm
     gc.collect()
     torch.cuda.empty_cache()
-    shutil.rmtree("models/merged/")
+    shutil.rmtree(f"models/{wandb.run.id}/merged/")
     return gen_pred_dataset