sentence-transformers-transformer not supported yet with neuron backend

[austinmw]

This blog post by HuggingFace says that compiling the BGE model works:
https://www.philschmid.de/inferentia2-embeddings

Yet when I try it, I get the below error.

%pip install git+https://github.com/huggingface/optimum-neuron.git
%pip install "sagemaker>=2.197.0"  --upgrade

from optimum.neuron import NeuronModelForFeatureExtraction
from transformers import AutoTokenizer

# model id you want to compile
vanilla_model_id = "BAAI/bge-base-en-v1.5"

# configs for compiling model
input_shapes = {
  "sequence_length": 384, # max length of the document (max 512)
  "batch_size": 1 # batch size for the model
  }

emb = NeuronModelForFeatureExtraction.from_pretrained(vanilla_model_id, export=True, **input_shapes)
tokenizer = AutoTokenizer.from_pretrained(vanilla_model_id)

# Save locally or upload to the HuggingFace Hub
save_directory = "embedding_model"
emb.save_pretrained(save_directory)
tokenizer.save_pretrained(save_directory)

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
Cell In[6], [line 27](vscode-notebook-cell:?execution_count=6&line=27)
     [21](vscode-notebook-cell:?execution_count=6&line=21) # configs for compiling model
     [22](vscode-notebook-cell:?execution_count=6&line=22) input_shapes = {
     [23](vscode-notebook-cell:?execution_count=6&line=23)   "sequence_length": 384, # max length of the document (max 512)
     [24](vscode-notebook-cell:?execution_count=6&line=24)   "batch_size": 1 # batch size for the model
     [25](vscode-notebook-cell:?execution_count=6&line=25)   }
---> [27](vscode-notebook-cell:?execution_count=6&line=27) emb = NeuronModelForFeatureExtraction.from_pretrained(vanilla_model_id, export=True, **input_shapes)
     [28](vscode-notebook-cell:?execution_count=6&line=28) tokenizer = AutoTokenizer.from_pretrained(vanilla_model_id)
     [30](vscode-notebook-cell:?execution_count=6&line=30) # Save locally or upload to the HuggingFace Hub

File [~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/modeling_base.py:372](https://vscode-remote+ssh-002dremote-002bsagemaker-002dnotebook-002disengard-002dspgi.vscode-resource.vscode-cdn.net/home/ec2-user/SageMaker/finetune_bge_bedrock/~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/modeling_base.py:372), in OptimizedModel.from_pretrained(cls, model_id, export, force_download, use_auth_token, cache_dir, subfolder, config, local_files_only, trust_remote_code, revision, **kwargs)
    [369](https://vscode-remote+ssh-002dremote-002bsagemaker-002dnotebook-002disengard-002dspgi.vscode-resource.vscode-cdn.net/home/ec2-user/SageMaker/finetune_bge_bedrock/~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/modeling_base.py:369)     trust_remote_code = False
    [371](https://vscode-remote+ssh-002dremote-002bsagemaker-002dnotebook-002disengard-002dspgi.vscode-resource.vscode-cdn.net/home/ec2-user/SageMaker/finetune_bge_bedrock/~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/modeling_base.py:371) from_pretrained_method = cls._from_transformers if export else cls._from_pretrained
--> [372](https://vscode-remote+ssh-002dremote-002bsagemaker-002dnotebook-002disengard-002dspgi.vscode-resource.vscode-cdn.net/home/ec2-user/SageMaker/finetune_bge_bedrock/~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/modeling_base.py:372) return from_pretrained_method(
    [373](https://vscode-remote+ssh-002dremote-002bsagemaker-002dnotebook-002disengard-002dspgi.vscode-resource.vscode-cdn.net/home/ec2-user/SageMaker/finetune_bge_bedrock/~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/modeling_base.py:373)     model_id=model_id,
    [374](https://vscode-remote+ssh-002dremote-002bsagemaker-002dnotebook-002disengard-002dspgi.vscode-resource.vscode-cdn.net/home/ec2-user/SageMaker/finetune_bge_bedrock/~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/modeling_base.py:374)     config=config,
    [375](https://vscode-remote+ssh-002dremote-002bsagemaker-002dnotebook-002disengard-002dspgi.vscode-resource.vscode-cdn.net/home/ec2-user/SageMaker/finetune_bge_bedrock/~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/modeling_base.py:375)     revision=revision,
    [376](https://vscode-remote+ssh-002dremote-002bsagemaker-002dnotebook-002disengard-002dspgi.vscode-resource.vscode-cdn.net/home/ec2-user/SageMaker/finetune_bge_bedrock/~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/modeling_base.py:376)     cache_dir=cache_dir,
    [377](https://vscode-remote+ssh-002dremote-002bsagemaker-002dnotebook-002disengard-002dspgi.vscode-resource.vscode-cdn.net/home/ec2-user/SageMaker/finetune_bge_bedrock/~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/modeling_base.py:377)     force_download=force_download,
    [378](https://vscode-remote+ssh-002dremote-002bsagemaker-002dnotebook-002disengard-002dspgi.vscode-resource.vscode-cdn.net/home/ec2-user/SageMaker/finetune_bge_bedrock/~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/modeling_base.py:378)     use_auth_token=use_auth_token,
    [379](https://vscode-remote+ssh-002dremote-002bsagemaker-002dnotebook-002disengard-002dspgi.vscode-resource.vscode-cdn.net/home/ec2-user/SageMaker/finetune_bge_bedrock/~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/modeling_base.py:379)     subfolder=subfolder,
    [380](https://vscode-remote+ssh-002dremote-002bsagemaker-002dnotebook-002disengard-002dspgi.vscode-resource.vscode-cdn.net/home/ec2-user/SageMaker/finetune_bge_bedrock/~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/modeling_base.py:380)     local_files_only=local_files_only,
    [381](https://vscode-remote+ssh-002dremote-002bsagemaker-002dnotebook-002disengard-002dspgi.vscode-resource.vscode-cdn.net/home/ec2-user/SageMaker/finetune_bge_bedrock/~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/modeling_base.py:381)     trust_remote_code=trust_remote_code,
    [382](https://vscode-remote+ssh-002dremote-002bsagemaker-002dnotebook-002disengard-002dspgi.vscode-resource.vscode-cdn.net/home/ec2-user/SageMaker/finetune_bge_bedrock/~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/modeling_base.py:382)     **kwargs,
    [383](https://vscode-remote+ssh-002dremote-002bsagemaker-002dnotebook-002disengard-002dspgi.vscode-resource.vscode-cdn.net/home/ec2-user/SageMaker/finetune_bge_bedrock/~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/modeling_base.py:383) )

File [~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/neuron/modeling_base.py:252](https://vscode-remote+ssh-002dremote-002bsagemaker-002dnotebook-002disengard-002dspgi.vscode-resource.vscode-cdn.net/home/ec2-user/SageMaker/finetune_bge_bedrock/~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/neuron/modeling_base.py:252), in NeuronBaseModel._from_transformers(cls, model_id, config, use_auth_token, revision, force_download, cache_dir, compiler_workdir, optlevel, subfolder, local_files_only, trust_remote_code, task, auto_cast, auto_cast_type, disable_fast_relayout, disable_fallback, dynamic_batch_size, **kwargs_shapes)
    [238](https://vscode-remote+ssh-002dremote-002bsagemaker-002dnotebook-002disengard-002dspgi.vscode-resource.vscode-cdn.net/home/ec2-user/SageMaker/finetune_bge_bedrock/~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/neuron/modeling_base.py:238) model = TasksManager.get_model_from_task(
    [239](https://vscode-remote+ssh-002dremote-002bsagemaker-002dnotebook-002disengard-002dspgi.vscode-resource.vscode-cdn.net/home/ec2-user/SageMaker/finetune_bge_bedrock/~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/neuron/modeling_base.py:239)     task=task,
    [240](https://vscode-remote+ssh-002dremote-002bsagemaker-002dnotebook-002disengard-002dspgi.vscode-resource.vscode-cdn.net/home/ec2-user/SageMaker/finetune_bge_bedrock/~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/neuron/modeling_base.py:240)     model_name_or_path=model_id,
   (...)
    [248](https://vscode-remote+ssh-002dremote-002bsagemaker-002dnotebook-002disengard-002dspgi.vscode-resource.vscode-cdn.net/home/ec2-user/SageMaker/finetune_bge_bedrock/~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/neuron/modeling_base.py:248)     trust_remote_code=trust_remote_code,
    [249](https://vscode-remote+ssh-002dremote-002bsagemaker-002dnotebook-002disengard-002dspgi.vscode-resource.vscode-cdn.net/home/ec2-user/SageMaker/finetune_bge_bedrock/~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/neuron/modeling_base.py:249) )
    [251](https://vscode-remote+ssh-002dremote-002bsagemaker-002dnotebook-002disengard-002dspgi.vscode-resource.vscode-cdn.net/home/ec2-user/SageMaker/finetune_bge_bedrock/~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/neuron/modeling_base.py:251) task = TasksManager.map_from_synonym(task)
--> [252](https://vscode-remote+ssh-002dremote-002bsagemaker-002dnotebook-002disengard-002dspgi.vscode-resource.vscode-cdn.net/home/ec2-user/SageMaker/finetune_bge_bedrock/~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/neuron/modeling_base.py:252) neuron_config_constructor = TasksManager.get_exporter_config_constructor(
    [253](https://vscode-remote+ssh-002dremote-002bsagemaker-002dnotebook-002disengard-002dspgi.vscode-resource.vscode-cdn.net/home/ec2-user/SageMaker/finetune_bge_bedrock/~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/neuron/modeling_base.py:253)     model=model, exporter="neuron", task=task
    [254](https://vscode-remote+ssh-002dremote-002bsagemaker-002dnotebook-002disengard-002dspgi.vscode-resource.vscode-cdn.net/home/ec2-user/SageMaker/finetune_bge_bedrock/~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/neuron/modeling_base.py:254) )
    [256](https://vscode-remote+ssh-002dremote-002bsagemaker-002dnotebook-002disengard-002dspgi.vscode-resource.vscode-cdn.net/home/ec2-user/SageMaker/finetune_bge_bedrock/~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/neuron/modeling_base.py:256) input_shapes = {}
    [257](https://vscode-remote+ssh-002dremote-002bsagemaker-002dnotebook-002disengard-002dspgi.vscode-resource.vscode-cdn.net/home/ec2-user/SageMaker/finetune_bge_bedrock/~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/neuron/modeling_base.py:257) for name in neuron_config_constructor.func.get_mandatory_axes_for_task(task):

File [~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/exporters/tasks.py:1881](https://vscode-remote+ssh-002dremote-002bsagemaker-002dnotebook-002disengard-002dspgi.vscode-resource.vscode-cdn.net/home/ec2-user/SageMaker/finetune_bge_bedrock/~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/exporters/tasks.py:1881), in TasksManager.get_exporter_config_constructor(exporter, model, task, model_type, model_name, exporter_config_kwargs, library_name)
   [1878](https://vscode-remote+ssh-002dremote-002bsagemaker-002dnotebook-002disengard-002dspgi.vscode-resource.vscode-cdn.net/home/ec2-user/SageMaker/finetune_bge_bedrock/~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/exporters/tasks.py:1878)     model_type = model_type.replace("_", "-")
   [1879](https://vscode-remote+ssh-002dremote-002bsagemaker-002dnotebook-002disengard-002dspgi.vscode-resource.vscode-cdn.net/home/ec2-user/SageMaker/finetune_bge_bedrock/~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/exporters/tasks.py:1879)     model_name = getattr(model, "name", model_name)
-> [1881](https://vscode-remote+ssh-002dremote-002bsagemaker-002dnotebook-002disengard-002dspgi.vscode-resource.vscode-cdn.net/home/ec2-user/SageMaker/finetune_bge_bedrock/~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/exporters/tasks.py:1881) model_tasks = TasksManager.get_supported_tasks_for_model_type(
   [1882](https://vscode-remote+ssh-002dremote-002bsagemaker-002dnotebook-002disengard-002dspgi.vscode-resource.vscode-cdn.net/home/ec2-user/SageMaker/finetune_bge_bedrock/~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/exporters/tasks.py:1882)     model_type, exporter, model_name=model_name, library_name=library_name
   [1883](https://vscode-remote+ssh-002dremote-002bsagemaker-002dnotebook-002disengard-002dspgi.vscode-resource.vscode-cdn.net/home/ec2-user/SageMaker/finetune_bge_bedrock/~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/exporters/tasks.py:1883) )
   [1885](https://vscode-remote+ssh-002dremote-002bsagemaker-002dnotebook-002disengard-002dspgi.vscode-resource.vscode-cdn.net/home/ec2-user/SageMaker/finetune_bge_bedrock/~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/exporters/tasks.py:1885) if task not in model_tasks:
   [1886](https://vscode-remote+ssh-002dremote-002bsagemaker-002dnotebook-002disengard-002dspgi.vscode-resource.vscode-cdn.net/home/ec2-user/SageMaker/finetune_bge_bedrock/~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/exporters/tasks.py:1886)     synonyms = TasksManager.synonyms_for_task(task)

File [~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/exporters/tasks.py:1146](https://vscode-remote+ssh-002dremote-002bsagemaker-002dnotebook-002disengard-002dspgi.vscode-resource.vscode-cdn.net/home/ec2-user/SageMaker/finetune_bge_bedrock/~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/exporters/tasks.py:1146), in TasksManager.get_supported_tasks_for_model_type(model_type, exporter, model_name, library_name)
   [1140](https://vscode-remote+ssh-002dremote-002bsagemaker-002dnotebook-002disengard-002dspgi.vscode-resource.vscode-cdn.net/home/ec2-user/SageMaker/finetune_bge_bedrock/~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/exporters/tasks.py:1140)         raise KeyError(
   [1141](https://vscode-remote+ssh-002dremote-002bsagemaker-002dnotebook-002disengard-002dspgi.vscode-resource.vscode-cdn.net/home/ec2-user/SageMaker/finetune_bge_bedrock/~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/exporters/tasks.py:1141)             f"{model_type_and_model_name} is not supported yet for {library_name}. "
   [1142](https://vscode-remote+ssh-002dremote-002bsagemaker-002dnotebook-002disengard-002dspgi.vscode-resource.vscode-cdn.net/home/ec2-user/SageMaker/finetune_bge_bedrock/~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/exporters/tasks.py:1142)             f"Only {list(TasksManager._SUPPORTED_MODEL_TYPE.keys())} are supported. "
   [1143](https://vscode-remote+ssh-002dremote-002bsagemaker-002dnotebook-002disengard-002dspgi.vscode-resource.vscode-cdn.net/home/ec2-user/SageMaker/finetune_bge_bedrock/~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/exporters/tasks.py:1143)             f"If you want to support {model_type} please propose a PR or open up an issue."
   [1144](https://vscode-remote+ssh-002dremote-002bsagemaker-002dnotebook-002disengard-002dspgi.vscode-resource.vscode-cdn.net/home/ec2-user/SageMaker/finetune_bge_bedrock/~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/exporters/tasks.py:1144)         )
   [1145](https://vscode-remote+ssh-002dremote-002bsagemaker-002dnotebook-002disengard-002dspgi.vscode-resource.vscode-cdn.net/home/ec2-user/SageMaker/finetune_bge_bedrock/~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/exporters/tasks.py:1145) if exporter not in TasksManager._SUPPORTED_MODEL_TYPE[model_type]:
-> [1146](https://vscode-remote+ssh-002dremote-002bsagemaker-002dnotebook-002disengard-002dspgi.vscode-resource.vscode-cdn.net/home/ec2-user/SageMaker/finetune_bge_bedrock/~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/exporters/tasks.py:1146)     raise KeyError(
   [1147](https://vscode-remote+ssh-002dremote-002bsagemaker-002dnotebook-002disengard-002dspgi.vscode-resource.vscode-cdn.net/home/ec2-user/SageMaker/finetune_bge_bedrock/~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/exporters/tasks.py:1147)         f"{model_type_and_model_name} is not supported yet with the {exporter} backend. "
   [1148](https://vscode-remote+ssh-002dremote-002bsagemaker-002dnotebook-002disengard-002dspgi.vscode-resource.vscode-cdn.net/home/ec2-user/SageMaker/finetune_bge_bedrock/~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/exporters/tasks.py:1148)         f"Only {list(TasksManager._SUPPORTED_MODEL_TYPE[model_type].keys())} are supported. "
   [1149](https://vscode-remote+ssh-002dremote-002bsagemaker-002dnotebook-002disengard-002dspgi.vscode-resource.vscode-cdn.net/home/ec2-user/SageMaker/finetune_bge_bedrock/~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/exporters/tasks.py:1149)         f"If you want to support {exporter} please propose a PR or open up an issue."
   [1150](https://vscode-remote+ssh-002dremote-002bsagemaker-002dnotebook-002disengard-002dspgi.vscode-resource.vscode-cdn.net/home/ec2-user/SageMaker/finetune_bge_bedrock/~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/exporters/tasks.py:1150)     )
   [1152](https://vscode-remote+ssh-002dremote-002bsagemaker-002dnotebook-002disengard-002dspgi.vscode-resource.vscode-cdn.net/home/ec2-user/SageMaker/finetune_bge_bedrock/~/anaconda3/envs/rerank/lib/python3.11/site-packages/optimum/exporters/tasks.py:1152) return TasksManager._SUPPORTED_MODEL_TYPE[model_type][exporter]

KeyError: 'sentence-transformers-transformer is not supported yet with the neuron backend. Only [] are supported. If you want to support neuron please propose a PR or open up an issue.'

[philschmid]

@austinmw could you please share some details on the versions you use, including optium-neuron, transformers, torch-neuronx?

cc @JingyaHuang

[JingyaHuang]

Hi @austinmw,

Thanks for reporting the issue, the error is introduced by a PR in Optimum main repo: huggingface/optimum#1589, in which the model_type of all sentence-transformers models are patched to so-call sentence-transformers-transformer which is not an officially supported task in optimum-neuron...

TL,DR

• A quick workaround would be downgrade optimum to 1.15.0 Meanwhile, I will discuss with the team how to better tackle the support for sentence-transformers models.

[JingyaHuang]

Small update on this, the PR for better supporting sentence transformers is merged in optimum-neuron. We built a new modeling API NeuronModelForSenetenceTransformers to replace the previous NeuronModelForFeatureExtraction, which it will output token_embeddings and sentence_embedding.

An example snippet here:

from transformers import AutoTokenizer
from optimum.neuron import NeuronModelForSenetenceTransformers

tokenizer = AutoTokenizer.from_pretrained("optimum/bge-base-en-v1.5-neuronx")
model = NeuronModelForSenetenceTransformers.from_pretrained("optimum/bge-base-en-v1.5-neuronx")

inputs = tokenizer("In the smouldering promise of the fall of Troy, a mythical world of gods and mortals rises from the ashes.", return_tensors="pt")

outputs = model(**inputs)
token_embeddings = outputs.token_embeddings
sentence_embedding = = outputs.sentence_embedding

Please feel free if you want to give it a try by installing optimum-neuron from source. THX

[mmcclean-aws]

If you find any further issues please raise in the https://github.com/huggingface/optimum-neuron repo

[austinmw]

@JingyaHuang Awesome, thanks!

Btw, very minor, but I noticed a spelling error in the naming of the NeuronModelForSenetenceTransformers class