[Inference] Add SentenceTransformers support to pipeline for feature-extration

[philschmid]

What does this PR do?

This PR adds a new, slightly modified FeatureExtractionPipeline from Transformers that allows us to use it with sentence-transformers models. When using the pipeline object from optimum,, the library checks if the requested model for feature-extraction is a sentence-transformers model and if so, it would return the sentence_embeddings instead of the first hidden state.

Thats is done by adding a new is_sentence_transformer_model that checks if the requested model is a transformers or sentence-transformers model. If it is a sentence-transformers model, it uses NeuronModelForSentenceTransformers and the FeatureExtractionPipeline returns _model_outputs.sentence_embedding[0] instead of model_outputs[0]

Example:

from optimum.neuron import pipeline

input_shapes = {"batch_size": 1, "sequence_length": 64} 
p = pipeline("feature-extraction","sentence-transformers/all-MiniLM-L6-v2",export=True, input_shapes=input_shapes)
> # Using Sentence Transformers compatible Feature extraction pipeline

p("test")
> [0.06765521317720413,
> 0.06349243223667145,
> 0.04871273413300514,
> 0.0793028473854065,

Validated with torch.allclose

Implications:

• sentence-transformers models will now always return the sentence_embeddings when initialized with the FeatureExtractionPipeline pipeline.

Alternatives options:

• Instead of modifying the feature-extraction pipeline, we could also introduce a new task sentence-embeddings to optimum, but that might hinder more general adoption since it is unique to optimum-neuron.

[philschmid]

@tomaarsen can you also do a review?

[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.

[dacorvo]

Several pipeline tests are failing. It is unclear to me if they should be triggered or not.

[tomaarsen]

My primary concern at this time is that the Sentence Transformer tokenizer uses this max_seq_length as the "correct" maximum length as opposed to the value defined in the tokenizer_config.json.

Here, we are relying on the tokenizer defined in the Pipeline, which won't use the max_seq_length. As a result, I think this ST pipeline component will perform differently (worse, to be precise) for longer input texts. A solution is to use model_inputs = self.model.tokenize(inputs) instead.
Do note that the ST tokenize method does not allow for extra tokenize kwargs such as truncation, return_tensors, or padding. These are unfortunately hardcoded at the moment.

[tomaarsen]

Yet another solution is to rely exclusively on self.model.encode(...) in def _forward, but I recognize that this might clash with some requirements of the Pipeline.

[philschmid]

@tomaarsen for inferentia models need to be traced to a sequence length before running inference since we have static shapes. You always need to specify a sequence_length and batch_size before you can compile a model which is then used.
This abstracted away by the user in the NeuronModelForSentenceTransformers class.

[philschmid]

There is no "sentence-transformers" used at all at the end. since the model is traced and it happens with transformers. We should be good here.

[philschmid]

See here:

optimum-neuron/optimum/neuron/modeling.py

Line 201 in 18460aa

class NeuronModelForSentenceTransformers(NeuronBaseModel):

[tomaarsen]

Ahh, I see! Thanks for the heads up. I figured it was more like Intel Gaudi, which does just work with regular Sentence Transformers (as long as the padding is "max_length" to also get static shapes & the device is "hpu").

Then my concern still stands: I think the max_seq_length might not be taken into account correctly.

[JingyaHuang]

The sentence_bert_config.json is not taken into consideration during the export, and maybe we should. @tomaarsen where can we usually find the max_seq_length? Is there a specific name / path that it's stored, if so could we add it to config.json?

[JingyaHuang]

fyi, max_seq_length is not taken into account at all by the export of Neuron model. Shall we prevent users from setting static seq len higher than this value?

[tomaarsen]

For Sentence Transformer models, max_seq_length should have priority. In 99% of cases, this is stored in the sentence_bert_config.json file in the root of the model directory/repository. You might indeed want to store it in a config.json when exporting for Neuron, or override model_max_length in tokenizer_config.json as that should work in a more "expected" fashion.

[dacorvo]

Disclaimer: I don't know sentence transformers, so this review is on the general outlook of the code and test, which looks good to me. @JingyaHuang and @tomaarsen reviews will be more relevant.

[JingyaHuang]

Thanks @philschmid, left some small nits.

[JingyaHuang]

The sentence_bert_config.json is not taken into consideration during the export, and maybe we should. @tomaarsen where can we usually find the max_seq_length? Is there a specific name / path that it's stored, if so could we add it to config.json?

[JingyaHuang]

LGTM, thanks @philschmid for adding it!

And thanks @tomaarsen for raising the concern of max_seq_length, it's on the backlog, will improve it in a coming PR.