--- license: mit tags: - sentence-embeddings - endpoints-template - optimum library_name: generic --- # Optimized and Quantized [sentence-transformers/all-MiniLM-L6-v2](https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2) with a custom pipeline.py This repository implements a `custom` task for `sentence-embeddings` for 🤗 Inference Endpoints for accelerated inference using [🤗 Optimum](https://huggingface.co/docs/optimum/index). The code for the customized pipeline is in the [pipeline.py](https://huggingface.co/philschmid/all-MiniLM-L6-v2-optimum-embeddings/blob/main/pipeline.py). In the [how to create your own optimized and quantized model](#how-to-create-your-own-optimized-and-quantized-model) you will learn how the model was converted & optimized, it is based on the [Accelerate Sentence Transformers with Hugging Face Optimum](https://www.philschmid.de/optimize-sentence-transformers) blog post. It also includes how to create your custom pipeline and test it. There is also a [notebook](https://huggingface.co/philschmid/all-MiniLM-L6-v2-optimum-embeddings/blob/main/convert.ipynb) included. To use deploy this model a an Inference Endpoint you have to select `Custom` as task to use the `pipeline.py` file. -> _double check if it is selected_ ### expected Request payload ```json { "inputs": "The sky is a blue today and not gray", } ``` below is an example on how to run a request using Python and `requests`. ## Run Request ```python import json from typing import List import requests as r import base64 ENDPOINT_URL = "" HF_TOKEN = "" def predict(document_string:str=None): payload = {"inputs": document_string} response = r.post( ENDPOINT_URL, headers={"Authorization": f"Bearer {HF_TOKEN}"}, json=payload ) return response.json() prediction = predict( path_to_image="The sky is a blue today and not gray" ) ``` expected output ```python {'embeddings': [[-0.021580450236797333, 0.021715054288506508, 0.00979710929095745, -0.0005379787762649357, 0.04682469740509987, -0.013600599952042103, ... } ``` ## How to create your own optimized and quantized model Steps: [1. Convert model to ONNX](#1-convert-model-to-onnx) [2. Optimize & quantize model with Optimum](#2-optimize--quantize-model-with-optimum) [3. Create Custom Handler for Inference Endpoints](#3-create-custom-handler-for-inference-endpoints) Helpful links: * [Accelerate Sentence Transformers with Hugging Face Optimum](https://www.philschmid.de/optimize-sentence-transformers) * [Create Custom Handler Endpoints](https://link-to-docs) ## Setup & Installation ```python %%writefile requirements.txt optimum[onnxruntime]==1.3.0 mkl-include mkl ``` install requirements ```python !pip install -r requirements.txt ``` ## 1. Convert model to ONNX ```python from optimum.onnxruntime import ORTModelForFeatureExtraction from transformers import AutoTokenizer from pathlib import Path model_id="sentence-transformers/all-MiniLM-L6-v2" onnx_path = Path(".") # load vanilla transformers and convert to onnx model = ORTModelForFeatureExtraction.from_pretrained(model_id, from_transformers=True) tokenizer = AutoTokenizer.from_pretrained(model_id) # save onnx checkpoint and tokenizer model.save_pretrained(onnx_path) tokenizer.save_pretrained(onnx_path) ``` ## 2. Optimize & quantize model with Optimum ```python from optimum.onnxruntime import ORTOptimizer, ORTQuantizer from optimum.onnxruntime.configuration import OptimizationConfig, AutoQuantizationConfig # create ORTOptimizer and define optimization configuration optimizer = ORTOptimizer.from_pretrained(model_id, feature=model.pipeline_task) optimization_config = OptimizationConfig(optimization_level=99) # enable all optimizations # apply the optimization configuration to the model optimizer.export( onnx_model_path=onnx_path / "model.onnx", onnx_optimized_model_output_path=onnx_path / "model-optimized.onnx", optimization_config=optimization_config, ) # create ORTQuantizer and define quantization configuration dynamic_quantizer = ORTQuantizer.from_pretrained(model_id, feature=model.pipeline_task) dqconfig = AutoQuantizationConfig.avx512_vnni(is_static=False, per_channel=False) # apply the quantization configuration to the model model_quantized_path = dynamic_quantizer.export( onnx_model_path=onnx_path / "model-optimized.onnx", onnx_quantized_model_output_path=onnx_path / "model-quantized.onnx", quantization_config=dqconfig, ) ``` ## 3. Create Custom Handler for Inference Endpoints ```python %%writefile pipeline.py from typing import Dict, List, Any from optimum.onnxruntime import ORTModelForFeatureExtraction from transformers import AutoTokenizer import torch.nn.functional as F import torch # copied from the model card def mean_pooling(model_output, attention_mask): token_embeddings = model_output[0] #First element of model_output contains all token embeddings input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float() return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9) class PreTrainedPipeline(): def __init__(self, path=""): # load the optimized model self.model = ORTModelForFeatureExtraction.from_pretrained(path, file_name="model-quantized.onnx") self.tokenizer = AutoTokenizer.from_pretrained(path) def __call__(self, data: Any) -> List[List[Dict[str, float]]]: """ Args: data (:obj:): includes the input data and the parameters for the inference. Return: A :obj:`list`:. The list contains the embeddings of the inference inputs """ inputs = data.get("inputs", data) # tokenize the input encoded_inputs = self.tokenizer(inputs, padding=True, truncation=True, return_tensors='pt') # run the model outputs = self.model(**encoded_inputs) # Perform pooling sentence_embeddings = mean_pooling(outputs, encoded_inputs['attention_mask']) # Normalize embeddings sentence_embeddings = F.normalize(sentence_embeddings, p=2, dim=1) # postprocess the prediction return {"embeddings": sentence_embeddings.tolist()} ``` test custom pipeline ```python from pipeline import PreTrainedPipeline # init handler my_handler = PreTrainedPipeline(path=".") # prepare sample payload request = {"inputs": "I am quite excited how this will turn out"} # test the handler %timeit my_handler(request) ``` results ``` 1.55 ms ± 2.04 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each) ```