`OCPCalculator` not compatible with ocp 2.0 format configs

[EricZQu]

I'm trying to use OCPCalculator to run some relaxations, but it stopped working after the config update.

MWE:

from fairchem.core import OCPCalculator

calc = OCPCalculator(
    config_yml='configs/s2ef/2M/equiformer_v2/equiformer_v2_N@12_L@6_M@2.yml',
    checkpoint_path='checkpoints/equiformerv2/eq2_83M_2M.pt'
)

Error:

KeyError                                  Traceback (most recent call last)
Cell In[3], line 3
      1 from fairchem.core import OCPCalculator
----> 3 calc = OCPCalculator(
      4     config_yml='configs/s2ef/2M/equiformer_v2/equiformer_v2_N@12_L@6_M@2.yml',
      5     checkpoint_path='checkpoints/equiformerv2/eq2_83M_2M.pt'
      6 )

File ~/fairchem/src/fairchem/core/common/relaxation/ase_utils.py:164, in OCPCalculator.__init__(self, config_yml, checkpoint_path, model_name, local_cache, trainer, cutoff, max_neighbors, cpu, seed)
    160     config["model"] = config["model_attributes"]
    162 # for checkpoints with relaxation datasets defined, remove to avoid
    163 # unnecesarily trying to load that dataset
--> 164 if "relax_dataset" in config["task"]:
    165     del config["task"]["relax_dataset"]
    167 # Calculate the edge indices on the fly

KeyError: 'task'

If you guys are open to contributions, I could try to fix it 😃

[mshuaibii]

Thanks for flagging this!

Always open to contributions :)

FWIW - when constructing the calculator you don't need to specify the config, the checkpoint is sufficient. But we certainly should get this resolved.

[EricZQu]

Thanks for working on this! I think it should be resolved because it won't support the checkpoints trained with new configs as well.

Nit: this checkpoint loading method is also outdated (if I'm understanding this correctly). Now the trainer expects a None if checkpoint is not passed. So the line 231-232 should be removed.

https://github.com/FAIR-Chem/fairchem/blob/1b4b62b4b71dfe755c94c60839f5d810e5d14a91/src/fairchem/core/common/relaxation/ase_utils.py#L221-L234

https://github.com/FAIR-Chem/fairchem/blob/1b4b62b4b71dfe755c94c60839f5d810e5d14a91/src/fairchem/core/trainers/base_trainer.py#L449-L459

[mshuaibii]

Yup that should be reflected in the PR above. Feel free to check out the branch and test it on your end. It should be working fine now.

[EricZQu]

Oh thanks! I missed that. And it is working for me now. Feel free to close this issue when merged in.