File size: 1,777 Bytes
ccf46b4
3b9f0aa
a8fed56
 
 
 
 
ccf46b4
 
a8fed56
 
 
b290091
 
a8fed56
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ca83acc
 
 
 
 
 
 
a8fed56
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
---
pipeline_tag: feature-extraction
tags:
  - sentence-transformers
  - feature-extraction
  - sentence-similarity
language: en
license: apache-2.0
---

# ONNX Conversion of [sentence-transformers/all-MiniLM-L6-v2](https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2)

- ONNX model for CPU with O3 optimisation
- This is a [sentence-transformers](https://www.SBERT.net) model: It maps sentences & paragraphs to a 384 dimensional dense vector space and can be used for tasks like clustering or semantic search.

## Usage

```python
import torch
import torch.nn.functional as F
from optimum.onnxruntime import ORTModelForFeatureExtraction
from transformers import AutoTokenizer

sentences = [
    "The llama (/ˈlɑːmə/) (Lama glama) is a domesticated South American camelid.",
    "The alpaca (Lama pacos) is a species of South American camelid mammal.",
    "The vicuña (Lama vicugna) (/vɪˈkuːnjə/) is one of the two wild South American camelids.",
]

model_name = "EmbeddedLLM/all-MiniLM-L6-v2-onnx-o3-cpu"
device = "cpu"
provider = "CPUExecutionProvider"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = ORTModelForFeatureExtraction.from_pretrained(
    model_name, use_io_binding=True, provider=provider, device_map=device
)
inputs = tokenizer(
    sentences,
    padding=True,
    truncation=True,
    return_tensors="pt",
    max_length=model.config.max_position_embeddings,
)
inputs = inputs.to(device)
token_embeddings = model(**inputs).last_hidden_state
# Pool
att_mask = inputs["attention_mask"].unsqueeze(-1).expand(token_embeddings.size()).float()
embeddings = torch.sum(token_embeddings * att_mask, 1) / torch.clamp(att_mask.sum(1), min=1e-9)
embeddings = F.normalize(embeddings, p=2, dim=1)
print(embeddings.cpu().numpy().shape)
```