mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-30 13:09:40 +00:00
fix(config): accept pretrained_model dir for --config_path on resume (#4023)
Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
This commit is contained in:
@@ -66,3 +66,27 @@ def test_from_pretrained_raises_when_no_root_config_and_no_checkpoints(monkeypat
|
||||
|
||||
with pytest.raises(FileNotFoundError, match="train_config.json not found"):
|
||||
TrainPipelineConfig.from_pretrained("user/empty-repo")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("pass_dir", [False, True])
|
||||
def test_resolve_resume_checkpoint_accepts_file_or_pretrained_model_dir(tmp_path, monkeypatch, pass_dir):
|
||||
"""`--config_path` may point at the checkpoint's train_config.json or at its
|
||||
pretrained_model/ directory; both must resolve `policy.pretrained_path` to the
|
||||
pretrained_model/ directory (regression test for the directory case, which
|
||||
previously resolved one level too high and failed on model.safetensors)."""
|
||||
pretrained_dir = tmp_path / "checkpoints" / "000002" / "pretrained_model"
|
||||
pretrained_dir.mkdir(parents=True)
|
||||
(pretrained_dir / "train_config.json").touch()
|
||||
target = pretrained_dir if pass_dir else pretrained_dir / "train_config.json"
|
||||
|
||||
from lerobot.policies.act.configuration_act import ACTConfig
|
||||
|
||||
cfg = tc.draccus.parse(TrainPipelineConfig, args=["--dataset.repo_id", "u/d"])
|
||||
cfg.policy = ACTConfig()
|
||||
cfg.resume = True
|
||||
monkeypatch.setattr(tc.parser, "parse_arg", lambda name: str(target) if name == "config_path" else None)
|
||||
|
||||
cfg._resolve_resume_checkpoint()
|
||||
|
||||
assert cfg.policy.pretrained_path == pretrained_dir
|
||||
assert cfg.checkpoint_path == pretrained_dir.parent
|
||||
|
||||
Reference in New Issue
Block a user