Trouble using `10MPaper` model for inference

[psaegert]

Hi there,

I'm trying to use the model pretrained on the 10M dataset for inference via the fit_func.ipynb notebook.

Since the notebook loads the 100M-dataset model by default, I changed

## Load equation configuration and architecture configuration
import omegaconf
- with open('100M/eq_setting.json', 'r') as json_file:
+ with open('10MPaper/equation_config.json', 'r') as json_file:
  eq_setting = json.load(json_file)

- cfg = omegaconf.OmegaConf.load("100M/config.yaml")
+ cfg = omegaconf.OmegaConf.load("10MPaper/config.yaml")

but now this cell

params_fit = FitParams(word2id=eq_setting["word2id"], 
                            id2word={int(k): v for k,v in eq_setting["id2word"].items()}, 
                            una_ops=eq_setting["una_ops"], 
                            bin_ops=eq_setting["bin_ops"], 
                            total_variables=list(eq_setting["total_variables"]),  
                            total_coefficients=list(eq_setting["total_coefficients"]),
                            rewrite_functions=list(eq_setting["rewrite_functions"]),
                            bfgs=bfgs,
                            beam_size=cfg.inference.beam_size #This parameter is a tradeoff between accuracy and fitting time
                            )

results in a KeyError:

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
/tmp/ipykernel_10858/898078201.py in <module>
----> 1 params_fit = FitParams(word2id=eq_setting["word2id"], 
      2                             id2word={int(k): v for k,v in eq_setting["id2word"].items()},
      3                             una_ops=eq_setting["una_ops"],
      4                             bin_ops=eq_setting["bin_ops"],
      5                             total_variables=list(eq_setting["total_variables"]),

KeyError: 'word2id'

Can anyone replicate this error? How can I use the 10M-dataset model for inference?