mixedbread-ai

/

mxbai-embed-large-v1

like 819

Follow

Mixedbread 408

Hello!

This pull request adds an exported onnx model (model_O3.onnx).

Config

OptimizationConfig(
    optimization_level=2,
    enable_transformers_specific_optimizations=True,
    optimize_for_gpu=False,
    fp16=False,
    disable_gelu_fusion=False,
    disable_attention_fusion=False,
    disable_bias_gelu_fusion=False,
    disable_layer_norm_fusion=False,
    disable_rotary_embeddings=False,
    disable_skip_layer_norm_fusion=False,
    disable_bias_skip_layer_norm_fusion=False,
    disable_skip_group_norm_fusion=False,
    disable_bias_splitgelu_fusion=False,
    disable_bias_add_fusion=False,
    disable_group_norm_fusion=True,
    disable_embed_layer_norm_fusion=True,
    enable_gemm_fast_gelu_fusion=False,
    use_mask_index=False,
    disable_packed_kv=True,
    no_attention_mask=False,
    use_raw_attention_mask=False,
    disable_shape_inference=False,
    use_multi_head_attention=False,
    enable_gelu_approximation=True,
    use_group_norm_channels_first=False,
    disable_packed_qkv=False,
    disable_nhwc_conv=False
)

Testing this pull request

You can test this pull request before merging by loading the model from this PR with the revision argument:

from sentence_transformers import SentenceTransformer

# NOTE: Update this to the number of your pull request
pr_number = 2
model = SentenceTransformer(
    "mixedbread-ai/mxbai-embed-large-v1",
    revision=f"refs/pr/{pr_number}",
    backend="onnx",
    model_kwargs={"file_name": "model_O3.onnx"},
)

# Verify that everything works as expected
embeddings = model.encode(["The weather is lovely today.", "It's so sunny outside!", "He drove to the stadium."])
print(embeddings.shape)

similarities = model.similarity(embeddings, embeddings)
print(similarities)

This PR was auto-generated with export_optimized_onnx_model.