Views
No views yet


1import mteb
2from mteb.encoder_interface import PromptType
3import numpy as np
4import onnxruntime as rt
5from transformers import AutoTokenizer
6
7class CustomModel:
8 def __init__(self) -> None:
9 self.tokenizer = AutoTokenizer.from_pretrained("C:/LLM/snowflake-arctic-embed-m-v2.0")
10 self.session = rt.InferenceSession("snowflake-arctic-embed-m-v2.0-uint8.onnx", providers=["CPUExecutionProvider"])
11 self.scale = 0.15288913249969482 / 127.0
12
13 def dequantize(self, quantized: list | np.ndarray, scale: float) -> np.ndarray:
14 quantized = np.array(quantized)
15 dequant = (quantized.astype(np.float32) - 128) * scale
16 if dequant.ndim == 3 and dequant.shape[0] == 1:
17 return np.squeeze(dequant, axis=0)
18 return dequant
19
20 def encode(
21 self,
22 sentences: list[str],
23 task_name: str,
24 prompt_type: PromptType | None = None,
25 **kwargs,
26 ) -> np.ndarray:
27 inputs = self.tokenizer(sentences, padding=True, truncation=True, return_tensors="np", max_length=8192)
28 q = self.session.run(["sentence_embedding"], dict(inputs))
29 return self.dequantize(q, self.scale)
30
31
32model = CustomModel()
33benchmark = mteb.get_benchmark("NanoBEIR")
34evaluation = mteb.MTEB(tasks=benchmark)
35results = evaluation.run(model, corpus_chunk_size=4)
36for r in results:
37 print(r)
381from fastembed import TextEmbedding
2from fastembed.common.model_description import PoolingType, ModelSource
3
4TextEmbedding.add_custom_model(
5 model="snowflake-arctic-embed-m-v2.0-ONNX-uint8",
6 pooling=PoolingType.DISABLED,
7 normalization=False,
8 sources=ModelSource(hf="electroglyph/snowflake-arctic-embed-m-v2.0-ONNX-uint8"),
9 dim=768,
10 model_file="onnx/model.onnx",
11)
12
13model = TextEmbedding(model_name="snowflake-arctic-embed-m-v2.0-ONNX-uint8")
14embeddings = list(model.embed("test"))
15print(embeddings)