Add sweep command

w2s/sweep.py

@@ -0,0 +1,17 @@
+from argparse import ArgumentParser
+
+from w2s.ds_registry import _REGISTRY
+from w2s.train import TrainConfig, train
+
+
+def main():
+ parser = ArgumentParser()
+ parser.add_argument("rank", type=int)
+ args = parser.parse_args()
+
+ DATASETS = sorted(_REGISTRY.keys())
+ train(TrainConfig("Qwen/Qwen1.5-0.5B", DATASETS[args.rank]))
+
+
+if __name__ == "__main__":
+ main()

w2s/train.py

@@ -66,8 +66,7 @@ def move_best_ckpt(trainer: Trainer):
  print(f"Best model (loss {perf:.3f}) saved at: {dest}")
 
 
-def main():
- cfg = parse(TrainConfig)
+def train(cfg: TrainConfig):
  lora_cfg = LoraConfig(target_modules=LORA_MODULES)
 
  STRONG_NAME = "meta-llama/Meta-Llama-3-8B"
@@ -241,4 +240,4 @@ def strong_processor(examples):
 
 
 if __name__ == "__main__":
- main()
+ train(parse(TrainConfig))