Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Fix emb model export and load with trfrs (#756)
# What does this PR do? Fixes #744 With the PR, we should be once again able to export embedding model via transformers library or sentence transformer library depending on the class called: * With Transformers ```python import torch from optimum.neuron import NeuronModelForFeatureExtraction from transformers import AutoConfig, AutoTokenizer compiler_args = {"auto_cast": "matmul", "auto_cast_type": "fp16"} input_shapes = {"batch_size": 4, "sequence_length": 512} model = NeuronModelForFeatureExtraction.from_pretrained( model_id="TaylorAI/bge-micro-v2", # BERT SMALL export=True, disable_neuron_cache=True, **compiler_args, **input_shapes, ) ``` * With Sentence Transformers ```python import torch from optimum.neuron import NeuronModelForSentenceTransformers from transformers import AutoConfig, AutoTokenizer compiler_args = {"auto_cast": "matmul", "auto_cast_type": "fp16"} input_shapes = {"batch_size": 4, "sequence_length": 512} model = NeuronModelForSentenceTransformers.from_pretrained( model_id="TaylorAI/bge-micro-v2", # BERT SMALL export=True, disable_neuron_cache=True, **compiler_args, **input_shapes, ) ```
- Loading branch information