[Inference]Support sentence transformers clip

[JingyaHuang]

What does this PR do?

Follow up for #408, support the inference of sentence transformers clip.

[Compilation]

• With CLI

optimum-cli export neuron -m sentence-transformers/clip-ViT-B-32 --sequence_length 64 --text_batch_size 3 --image_batch_size 1 --num_channels 3 --height 224 --width 224 --task feature-extraction --library-name sentence_transformers --subfolder 0_CLIPModel clip_emb/

• Code

from optimum.neuron import NeuronModelForSentenceTransformers

# [Compile]
model_id = "sentence-transformers/clip-ViT-B-32"

# configs for compiling model
input_shapes = {
    "num_channels": 3,
    "height": 224,
    "width": 224,
    "text_batch_size": 3,
    "image_batch_size": 1,
    "sequence_length": 64,
}

emb_model = NeuronModelForSentenceTransformers.from_pretrained(
    model_id, subfolder="0_CLIPModel", export=True, library_name="sentence_transformers", dynamic_batch_size=False, **input_shapes
)

# Save locally or upload to the HuggingFace Hub
save_directory = "clip_emb/"
emb_model.save_pretrained(save_directory)

[Inference]

from PIL import Image
from sentence_transformers import util
from transformers import CLIPProcessor

from optimum.neuron import NeuronModelForSentenceTransformers

save_directory = "clip_emb"
emb_model = NeuronModelForSentenceTransformers.from_pretrained(save_directory)

processor = CLIPProcessor.from_pretrained(save_directory)
inputs = processor(
    text=["Two dogs in the snow", 'A cat on a table', 'A picture of London at night'], images=Image.open("two_dogs_in_snow.jpg"), return_tensors="pt", padding=True
)  

outputs = emb_model(**inputs)


# Compute cosine similarities
cos_scores = util.cos_sim(outputs.image_embeds, outputs.text_embeds)
print(cos_scores)

# tensor([[0.3072, 0.1016, 0.1095]])

Caveat

Since compiled models with dynamic batch size enabled only accept tensors with the same batch size, we cannot set dynamic_batch_size=True if the input texts and images have different batch sizes. And as NeuronModelForSentenceTransformers pad the inputs to the batch size used during the compilation, you could use a relatively larger batch_size during the compilation for flexibility with the trade-off of compute).

eg. if you want to encode 3 or 4 or 5 texts and 1 image, you could set text_batch_size = 5 = max(3, 4, 5) and image_batch_size = 1 during the compilation.

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you make sure to update the documentation with your changes?
• Did you write any new necessary tests?

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[michaelbenayoun]

Maybe put a check of the transformers version, if it's above the coming release, then fail. This way we can actually catch that and remove it.

[JingyaHuang]

Sounds good!

[JingyaHuang]

Just checked the PR I submited to transformers, it's not merged yet and I'm not 100% percent sure that it should be accepted by trfrs maintainers, thus completely don't know what trfr version to put here. Let's keep it this way, I will keep my eyes on how it goes...

[JingyaHuang]

huggingface/transformers#29169

[michaelbenayoun]

LGTM