Model card specific features of Keras mixin don't work with JAX model converted to TF

[merveenoyan]

Describe the bug

I need to investigate why JAX models converted to TF and then pushed using push_to_hub_keras don't have automatically generated sections in model card. The bare minimum should have model plot and training hyperparameters. Metrics aren't written if the history isn't kept. (which is normal for this case)
Also pinging @sayakpaul

Reproduction

Code is here: https://github.com/sayakpaul/maxim-tf/blob/main/convert_to_tf.py#L204

Logs

No response

System Info

Need to ask but is insignificant for this case.

[sayakpaul]

@merveenoyan thanks for reporting it.

I did some investigations:

• I created a simple model using the functional API.
• I then used push_to_hub_keras() to host it on Hub.

Here's the Colab: https://colab.research.google.com/gist/sayakpaul/1bea35f00acaa9926d2beb744058de1c/scratchpad.ipynb.

We see the model plot in the corresponding model repository but not the other stuff you mentioned. I wonder why the model plot doesn't appear for the MAXIM model. I can initialize the MAXIM model and actually plot the model.

You can replicate it with the following code:

from huggingface_hub.keras_mixin import from_pretrained_keras
import tensorflow as tf

model = from_pretrained_keras("sayakpaul/S-3_denoising_sidd")
tf.keras.utils.plot_model(model)

But it takes an incredibly long amount of time to plot it. Probably that's the reason the model plot is not there in the model repository.

[merveenoyan]

@sayakpaul I will debug it and let you know 🙂 I wonder why hyperparams aren't there though, let's see.