Skip to content

EasyDel implementation of Clip JAX model#357

Draft
ctr-pmuruganTT wants to merge 2 commits intomainfrom
pmurugan/add_easydel_for_clip
Draft

EasyDel implementation of Clip JAX model#357
ctr-pmuruganTT wants to merge 2 commits intomainfrom
pmurugan/add_easydel_for_clip

Conversation

@ctr-pmuruganTT
Copy link
Copy Markdown
Contributor

Problem description

Modify loader of CLIP Jax model to use EasyDel Instead of Transformers library

What's changed

  • Modified the model source
  • Removed load_processor and used AutoImageProcessor
  • Added get_input_activations_partition_spec and load_parameters_partition_spec

Checklist

  • New/Existing tests provide coverage for changes

"""

from transformers import FlaxCLIPModel
from easydel import CLIPModel
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Using CLIPModel is appropriate when we already have EasyDeL formatted weights. In our case, we are loading weights directly from HuggingFace and converting them into JAX NumPy arrays.
Easydel provides AutoModel family of classes which is exactly for this purposes (model params taking from HF and transfer to jax arrays; after that calling model implementation e.g. ClipModel).
For this use case we should use AutoEasyDeLModelForZeroShotImageClassification (from easydel import AutoEasyDeLModelForZeroShotImageClassification).

framework=Framework.JAX,
)

def _load_processor(self, dtype_override=None):
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why did we deleted this _load_processor? It was better to use AutoImageProcessor instead of CLIPProcessor?

@@ -132,7 +110,7 @@ def load_model(self, dtype_override=None):
from_pt = pretrained_model_name == "openai/clip-vit-large-patch14-336"
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This flag from_pt is unnecessary in the easydel

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants