add omegaconf to _JSON_SERIALIZABLE_TYPES

[not-lain]

I'm working on integrating an AI model with PyTorchModelHubMixin and it seems that I cannot serialize the input because it does not belong to the _JSON_SERIALIZABLE_TYPES .
some of my parameters are of type omegaconf i tried converting them to a python standard dict, but again i am still limited by _JSON_SERIALIZABLE_TYPES.

also pinging @Wauplin for feedback

[Wauplin]

@not-lain could you share which library you are trying to integrate and what you've tried that didn't work? I'm not familiar with omegaconf objects but indeed, only json-serializable values can be saved in the config.json (by definition). It should be the case for dicts if the values themselves are json-serializable. With a more concrete example I would be able to look into it.

[not-lain]

@Wauplin yeh sure :
https://colab.research.google.com/drive/1rStTsfAqoT7wXijt4a2wPbd6ZUYG3jTc?usp=sharing
I also tried converting omegaconf to dict using OmegaConf.to_container(cfg) but it didn't work either

[Wauplin]

I feel that there are 2 different parts to the problem:

• being able to serialize OmegaConf objects => as you said, something like OmegaConf.to_container should work and put everything as a dict in config.json
• being able to deserialize to omegaconf => that part PytorchModelHubMixin cannot really "guess" it. What I would do is to add some logic to the __init__ method which does if isinstance(config, dict): config = omegaconf.create(config) => this way the model can be initialized either from a dictionary or a omegaconf directly.

That been said, I do like how they handle configs with proper yaml files. Here we are not talking about a few hyperparameters but entire yaml files which are tens of lines long. What PytorchHubMixin will do in that case is to put everything in a single config.json which was not the first intention of the authors I think when separating configs in https://github.com/NeuralCarver/Michelangelo/tree/main/configs. This integration will require more work I think as it would be good to save yaml files independently in the model repo along the weights (so inherit from PytorchHubMixin but then slightly overwrite _from_pretrained and _save_pretrained to correctly use/create the config files).

[not-lain]

@Wauplin

I think we can do something like this in the deserialization :
1 - inspect init method and extract variable type (example: def __init__(self, param : OmegaConf ) )
2 - convert the parameter into an omegaconf
3- instantiate the model

in the serialization :

• check if the parameter passed to the model is either of type dict or either of type omegaconf, in the case it is omegaconf we apply temp_variable = OmegaConf.to_container(param)

[not-lain]

I can create a pr about this implementing the same logic, should I do it ?

[Wauplin]

inspect init method and extract variable type (example: def init(self, param : OmegaConf ) )

This is what we already do for dataclasses so it's possible to do it. However, I'm not sure we want to do it's not really standard library (and not so popular compared to other libs). I don't want us to add more and more special cases and would prefer that special cases are handled by individual repos instead.

OR we could add a way to define a custom json serializer/deserializer when inheriting from PytorchModelHubMixin. Definitely more complex/less straightforward but at least we wouldn't make exceptions in huggingface_hub for each and every possible ways to structure data.

[not-lain]

fundamentally omegaconf and dict are different on a miniature state which is that you can do both omg.key and omg[key] with omega conf, so it's really not far off

would prefer that special cases are handled by individual repos instead.

i can do it all i need to do it transform the input to dict before passing it to the init and then retransform it back to omegaconf after the init method

define a custom json serializer/deserializer when inheriting from PytorchModelHubMixin

nahh, too much work me, i'll just work aroud it by the solution i mentioned in this comment (convert to dict before passing to init and then reconvert it back) shortcuts are my thing 😎👉👉

[not-lain]

I made it work with the solution above.
feel free to keep it open or close this issue.
to any future readers, here's what i did.
1-

params = OmegaConf.to_container(omg_params)
model(**params)

2- reconvert them back in the init method using

def __init__(self,params) : 
  omg_params =  OmegaConf.create(params)

[Wauplin]

Great! This is what should be done yes! Maybe in __init__ you could support the case if a omega config is provided so that it doesn't break existing usage? Something like this:

def __init__(self,params) : 
  omg_params =  OmegaConf.create(params) if isinstance(params, dict) else params

Apart from that, I think it's the right approach. I'll close this issue as there's nothing to fix/add in huggingface_hub directly.