# This script is identical as the one above, except it groups all changes into one pull request for convenience
# requires sentence_transformers>=3.2.0
from sentence_transformers import (
    SentenceTransformer,
    export_optimized_onnx_model,
    export_dynamic_quantized_onnx_model,
)
from huggingface_hub import upload_folder

# The model to export to ONNX (+ optimize, quantize), OpenVINO
model_id = "BAAI/llm-embedder"
# Where to save the exported models locally
output_dir = model_id.replace("/", "-")
# Where to create the PR
output_model_id = "BAAI/llm-embedder"

onnx_model = SentenceTransformer(model_id, backend="onnx", model_kwargs={"export": True})
onnx_model.save_pretrained(output_dir)

for optimization_config in ["O1", "O2", "O3", "O4"]:
    export_optimized_onnx_model(
        onnx_model,
        optimization_config=optimization_config,
        model_name_or_path=output_dir,
    )

for quantization_config in ["arm64", "avx2", "avx512", "avx512_vnni"]:
    export_dynamic_quantized_onnx_model(
        onnx_model,
        quantization_config=quantization_config,
        model_name_or_path=output_dir,
    )

openvino_model = SentenceTransformer(model_id, backend="openvino")
openvino_model.save_pretrained(output_dir)

upload_folder(
    folder_path=output_dir,
    repo_id=output_model_id,
    repo_type="model",
    create_pr=True,
    commit_message="Export model to ONNX and OpenVINO",
)