README.md

---
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)
```