|
|
from sentence_transformers import models
|
|
|
import torch
|
|
|
from transformers import AutoTokenizer
|
|
|
from optimum.onnxruntime import ORTModelForFeatureExtraction
|
|
|
import numpy as np
|
|
|
import os
|
|
|
import onnxruntime
|
|
|
|
|
|
|
|
|
model_dir = "embeddinggemma-300m"
|
|
|
tokenizer = AutoTokenizer.from_pretrained("google/embeddinggemma-300m-qat-q4_0-unquantized")
|
|
|
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
|
onnx_model = ORTModelForFeatureExtraction.from_pretrained(
|
|
|
model_dir,
|
|
|
file_name="model.onnx"
|
|
|
).to(device)
|
|
|
|
|
|
class ONNXTransformer:
|
|
|
def __init__(self, onnx_model, tokenizer, max_seq_length=2048):
|
|
|
self.onnx_model = onnx_model
|
|
|
self.tokenizer = tokenizer
|
|
|
self.max_seq_length = max_seq_length
|
|
|
def encode(self, sentences):
|
|
|
inputs = self.tokenizer(sentences, return_tensors="pt", padding=True, truncation=True, max_length=self.max_seq_length)
|
|
|
input_ids = inputs['input_ids']
|
|
|
sequence_length = input_ids.shape[1]
|
|
|
position_ids = torch.arange(sequence_length)[None, :].expand(input_ids.shape[0], sequence_length)
|
|
|
inputs['position_ids'] = position_ids.to(input_ids.device)
|
|
|
with torch.no_grad():
|
|
|
outputs = self.onnx_model(**inputs)
|
|
|
return outputs.last_hidden_state
|
|
|
|
|
|
modules = []
|
|
|
onnx_transformer = ONNXTransformer(onnx_model, tokenizer, max_seq_length=2048)
|
|
|
modules.append(onnx_transformer)
|
|
|
for idx, name in [(1, "Pooling"), (2, "Dense"), (3, "Dense"), (4, "Normalize")]:
|
|
|
module_path = os.path.join(model_dir, f"{idx}_{name}")
|
|
|
if name == "Pooling":
|
|
|
modules.append(models.Pooling(module_path))
|
|
|
elif name == "Dense":
|
|
|
|
|
|
dense_onnx_path = os.path.join(model_dir, "onnx", f"dense{idx-1}.onnx")
|
|
|
modules.append(onnxruntime.InferenceSession(dense_onnx_path, providers=["CPUExecutionProvider"]))
|
|
|
elif name == "Normalize":
|
|
|
modules.append(models.Normalize())
|
|
|
|
|
|
class ONNXSentenceTransformer:
|
|
|
def __init__(self, modules):
|
|
|
self.modules = modules
|
|
|
def encode(self, sentences):
|
|
|
features = self.modules[0].encode(sentences)
|
|
|
for module in self.modules[1:]:
|
|
|
if isinstance(module, models.Pooling):
|
|
|
features = module({'token_embeddings': features, 'attention_mask': torch.ones(features.shape[:2], device=features.device)})['sentence_embedding']
|
|
|
elif isinstance(module, onnxruntime.InferenceSession):
|
|
|
|
|
|
if isinstance(features, torch.Tensor):
|
|
|
features = features.cpu().detach().numpy()
|
|
|
outputs = []
|
|
|
for vec in features:
|
|
|
ort_inputs = {module.get_inputs()[0].name: vec.reshape(1, -1)}
|
|
|
out = module.run(None, ort_inputs)[0]
|
|
|
outputs.append(out.squeeze(0))
|
|
|
features = np.stack(outputs, axis=0)
|
|
|
elif isinstance(module, models.Normalize):
|
|
|
|
|
|
if not isinstance(features, torch.Tensor):
|
|
|
features = torch.from_numpy(features)
|
|
|
features = module({'sentence_embedding': features})['sentence_embedding']
|
|
|
if isinstance(features, torch.Tensor):
|
|
|
return features.cpu().detach().numpy()
|
|
|
return features
|
|
|
|
|
|
onnx_st = ONNXSentenceTransformer(modules)
|
|
|
|
|
|
def cosine_similarity(a, b):
|
|
|
a = a.flatten()
|
|
|
b = b.flatten()
|
|
|
return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
words = ["apple", "banana", "car"]
|
|
|
embeddings = onnx_st.encode(words)
|
|
|
print(embeddings)
|
|
|
for idx, embedding in enumerate(embeddings):
|
|
|
print(f"Embedding {idx+1}: {embedding.shape}")
|
|
|
|
|
|
print("\nCosine similarities:")
|
|
|
print(f"apple vs banana: {cosine_similarity(embeddings[0], embeddings[1]):.4f}")
|
|
|
print(f"apple vs car: {cosine_similarity(embeddings[0], embeddings[2]):.4f}")
|
|
|
print(f"banana vs car: {cosine_similarity(embeddings[1], embeddings[2]):.4f}")
|
|
|
|
