diff --git a/finetune_moss.py b/finetune_moss.py index 2f65533..156b855 100644 --- a/finetune_moss.py +++ b/finetune_moss.py @@ -287,7 +287,7 @@ def train(args): parser.add_argument('--eval_bsz_per_gpu', default=4, type=int) parser.add_argument('--weight_decay', default=0.1, type=float) parser.add_argument('--learning_rate', default=9e-6, type=float) - parser.add_argument('--warmup_rates', default=0.05, type=int) + parser.add_argument('--warmup_rates', default=0.05, type=float) parser.add_argument('--n_epochs', default=2, type=int) # Other Args diff --git a/tests/test_finetune_cli.py b/tests/test_finetune_cli.py new file mode 100644 index 0000000..d32658a --- /dev/null +++ b/tests/test_finetune_cli.py @@ -0,0 +1,29 @@ +import subprocess +import sys +import unittest +from pathlib import Path + + +REPOSITORY_ROOT = Path(__file__).resolve().parents[1] + + +class FinetuneCliTest(unittest.TestCase): + def test_accepts_fractional_warmup_rate(self): + result = subprocess.run( + [ + sys.executable, + str(REPOSITORY_ROOT / "finetune_moss.py"), + "--warmup_rates", + "0.1", + "--help", + ], + cwd=REPOSITORY_ROOT, + capture_output=True, + text=True, + ) + + self.assertEqual(result.returncode, 0, result.stderr) + + +if __name__ == "__main__": + unittest.main()