51 lines
2.4 KiB
Python
51 lines
2.4 KiB
Python
# A dependency-light way to run the onnx model
|
|
|
|
|
|
import numpy as np
|
|
from typing import List
|
|
import os
|
|
|
|
__location__ = os.path.realpath(os.path.join(os.getcwd(), os.path.dirname(__file__)))
|
|
MODEL_ID = "sentence-transformers/all-MiniLM-L6-v2"
|
|
|
|
def normalize(v):
|
|
norm = np.linalg.norm(v, axis=1)
|
|
norm[norm == 0] = 1e-12
|
|
return v / norm[:, np.newaxis]
|
|
|
|
# Sampel implementation of the default sentence-transformers model using ONNX
|
|
class DefaultEmbeddingModel():
|
|
|
|
def __init__(self):
|
|
from tokenizers import Tokenizer
|
|
import onnxruntime as ort
|
|
# max_seq_length = 256, for some reason sentence-transformers uses 256 even though the HF config has a max length of 128
|
|
# https://github.com/UKPLab/sentence-transformers/blob/3e1929fddef16df94f8bc6e3b10598a98f46e62d/docs/_static/html/models_en_sentence_embeddings.html#LL480
|
|
self.tokenizer = Tokenizer.from_file(os.path.join(__location__, "models/onnx/tokenizer.json"))
|
|
self.tokenizer.enable_truncation(max_length=256)
|
|
self.tokenizer.enable_padding(pad_id=0, pad_token="[PAD]", length=256)
|
|
self.model = ort.InferenceSession(os.path.join(__location__,"models/onnx/model.onnx"))
|
|
|
|
|
|
def __call__(self, documents: List[str], batch_size: int = 32):
|
|
all_embeddings = []
|
|
for i in range(0, len(documents), batch_size):
|
|
batch = documents[i:i + batch_size]
|
|
encoded = [self.tokenizer.encode(d) for d in batch]
|
|
input_ids = np.array([e.ids for e in encoded])
|
|
attention_mask = np.array([e.attention_mask for e in encoded])
|
|
onnx_input = {
|
|
"input_ids": np.array(input_ids, dtype=np.int64),
|
|
"attention_mask": np.array(attention_mask, dtype=np.int64),
|
|
"token_type_ids": np.array([np.zeros(len(e), dtype=np.int64) for e in input_ids], dtype=np.int64),
|
|
}
|
|
model_output = self.model.run(None, onnx_input)
|
|
last_hidden_state = model_output[0]
|
|
# Perform mean pooling with attention weighting
|
|
input_mask_expanded = np.broadcast_to(np.expand_dims(attention_mask, -1), last_hidden_state.shape)
|
|
embeddings = np.sum(last_hidden_state * input_mask_expanded, 1) / np.clip(input_mask_expanded.sum(1), a_min=1e-9, a_max=None)
|
|
embeddings = normalize(embeddings).astype(np.float32)
|
|
all_embeddings.append(embeddings)
|
|
return np.concatenate(all_embeddings)
|
|
|