Views
No views yet
task: search result query format. I have no idea why this model benchmarks better than the base model on most retrieval tasks, but I'll take it.

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/embeddinggemma-300m-ONNX-uint8")
10 self.session = rt.InferenceSession("C:/LLM/embeddinggemma-300m-ONNX-uint8/onnx/model.onnx", providers=["CPUExecutionProvider"])
11 self.scale = 0.22116543352603912 / 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 if prompt_type == PromptType.query:
28 sentences = [f"task: search result | query: {s}" for s in sentences]
29 inputs = self.tokenizer(sentences, padding=True, truncation=True, return_tensors="np")
30 q = self.session.run(["sentence_embedding"], dict(inputs))
31 return self.dequantize(q, self.scale)
32
33
34model = CustomModel()
35benchmark = mteb.get_benchmark("NanoBEIR")
36evaluation = mteb.MTEB(tasks=benchmark)
37results = evaluation.run(model, corpus_chunk_size=128)
38for r in results:
39 print(r)
401from fastembed import TextEmbedding
2from fastembed.common.model_description import PoolingType, ModelSource
3
4TextEmbedding.add_custom_model(
5 model="embeddinggemma-300m-ONNX-uint8",
6 pooling=PoolingType.DISABLED,
7 normalization=False,
8 sources=ModelSource(hf="electroglyph/embeddinggemma-300m-ONNX-uint8"),
9 dim=768,
10 model_file="onnx/model.onnx",
11)
12
13model = TextEmbedding(model_name="embeddinggemma-300m-ONNX-uint8")
14embeddings = list(model.embed("test"))
15print(embeddings)