mirror of
https://github.com/qdrant/fastembed.git
synced 2026-07-25 12:11:05 -05:00
57 lines
2.1 KiB
Python
57 lines
2.1 KiB
Python
from pathlib import Path
|
|
from typing import List
|
|
|
|
import click
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from optimum.onnxruntime import ORTModelForFeatureExtraction
|
|
from optimum.pipelines import pipeline
|
|
from torch import Tensor
|
|
from transformers import AutoModel, AutoTokenizer
|
|
|
|
|
|
def average_pool(last_hidden_states: Tensor, attention_mask: Tensor) -> Tensor:
|
|
last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
|
|
return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
|
|
|
|
|
|
def hf_embed(model_id: str, texts: List[str], tokenizer):
|
|
# Tokenize the input texts
|
|
model = AutoModel.from_pretrained(model_id)
|
|
model.eval()
|
|
encoded_input = tokenizer(texts, padding=True, truncation=True, return_tensors="pt")
|
|
|
|
model_output = model(**encoded_input)
|
|
sentence_embeddings = model_output[0][:, 0]
|
|
sentence_embeddings = torch.nn.functional.normalize(sentence_embeddings, p=2, dim=1)
|
|
return sentence_embeddings
|
|
|
|
|
|
@click.command()
|
|
@click.option("--model_id", help="model id from huggingface.co/models")
|
|
@click.option("--model_dir", help="The person to greet.")
|
|
def setup(model_id, model_dir):
|
|
text = "This is a test sentence"
|
|
tokenizer = AutoTokenizer.from_pretrained(model_id)
|
|
output_dir = Path(model_dir)
|
|
model = ORTModelForFeatureExtraction.from_pretrained(output_dir)
|
|
onnx_quant_embed = pipeline(
|
|
"feature-extraction", model=model, accelerator="ort", tokenizer=tokenizer, return_tensors=True
|
|
)
|
|
quant_embeddings = onnx_quant_embed([text])
|
|
quant_embeddings = F.normalize(quant_embeddings[0][:,0], p=2, dim=1)
|
|
quant_embeddings = quant_embeddings.detach().numpy()
|
|
print(quant_embeddings.shape)
|
|
|
|
torch_embeddings = hf_embed(model_id, texts=[text], tokenizer=tokenizer)
|
|
torch_embeddings = F.normalize(torch_embeddings, p=2, dim=1)
|
|
torch_embeddings = torch_embeddings.detach().numpy()
|
|
print(torch_embeddings.shape)
|
|
assert quant_embeddings.shape == torch_embeddings.shape
|
|
print(np.allclose(quant_embeddings, torch_embeddings, atol=1e-5))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
setup()
|