google

/

paligemma-3b-ft-ai2d-224-jax

like 4

Follow

Google 8.87k

Hi @bsmani , To load the .npz model and perform inference in the paligemma-3b-ft-ai2d-224-jax model, i have tried the steps you can follow in this google colab link.

1. Load the model:

• Use the jax.tree_util.tree_map function to load the weights from the .npz file.
• Modify the state dictionary to match the shape of the input data.

2. Perform inference:

• Call the predict_step function with the processed input and the loaded model.
• Extract the logits and apply softmax to obtain probability distributions.

Kindly try these steps and let me know if you are facing any issue. Thank you.

@lkv : Could you please share the Colab notebook? I'm trying to inference a fine-tuned PaliGemma 2 model (.npz file). There are very scare resources on how to load the fine-tuned model for inferencing.