File size: 2,605 Bytes
e587c3d |
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 53 54 55 56 57 58 59 |
from sentence_transformers import models
import torch
from transformers import AutoTokenizer
from optimum.onnxruntime import ORTModelForFeatureExtraction
import numpy as np
# Load tokenizer and ONNX model
model_path = "./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_path).to(device)
class ONNXSentenceTransformer:
def __init__(self, model, tokenizer):
self.model = model
self.tokenizer = tokenizer
self.word_embedding_dimension = 768
self.pooling = models.Pooling(word_embedding_dimension=self.word_embedding_dimension, pooling_mode_mean_tokens=True)
def encode(self, sentences, batch_size=32):
if isinstance(sentences, str):
sentences = [sentences]
embeddings = []
for i in range(0, len(sentences), batch_size):
batch = sentences[i:i+batch_size]
inputs = self.tokenizer(batch, return_tensors="pt", padding=True, truncation=True)
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
with torch.no_grad():
outputs = self.model(**inputs)
last_hidden = outputs.last_hidden_state
attention_mask = inputs['attention_mask'].to(last_hidden.device)
features = {'token_embeddings': last_hidden, 'attention_mask': attention_mask}
pooled = self.pooling(features)['sentence_embedding']
embeddings.append(pooled)
return torch.cat(embeddings, dim=0).cpu().detach().numpy()
# Usage example
onnx_st = ONNXSentenceTransformer(onnx_model, tokenizer)
words = ["apple", "banana", "car"]
embeddings = onnx_st.encode(words)
print(embeddings)
for idx, embedding in enumerate(embeddings):
print(f"Embedding {idx+1}: {embedding.shape}")
# Cosine similarity demonstration
def cosine_similarity(a, b):
a = a.flatten()
b = b.flatten()
return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))
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}")
|