Views
No views yet
1# Requires transformers>=4.36.0
2import onnxruntime as ort
3import numpy as np
4from transformers import AutoTokenizer
5input_texts = [
6 "what is the capital of China?",
7 "how to implement quick sort in python?",
8 "北京",
9 "快排算法介绍"
10]
11# Load the tokenizer (using the original model for tokenizer)
12tokenizer = AutoTokenizer.from_pretrained('Alibaba-NLP/gte-multilingual-base')
13# Load the ONNX model
14session = ort.InferenceSession("model.onnx")
15# Tokenize the input texts
16batch_dict = tokenizer(input_texts, max_length=8192, padding=True, truncation=True, return_tensors='np')
17# Run inference
18outputs = session.run(None, {
19 "input_ids": batch_dict["input_ids"],
20 "attention_mask": batch_dict["attention_mask"]
21})
22# Get embeddings from the second output (last hidden states)
23# Extract the [CLS] token embedding (first token) for each sequence
24last_hidden_states = outputs[1] # Shape: (batch_size, seq_len, hidden_size)
25dimension = 768 # The output dimension of the output embedding, should be in [128, 768]
26embeddings = last_hidden_states[:, 0, :dimension] # Shape: (batch_size, dimension)
27# Debug: Check embeddings
28print(f"Embeddings shape: {embeddings.shape}")
29print(f"First few values of first embedding: {embeddings[0][:5]}")
30print(f"First few values of second embedding: {embeddings[1][:5]}")
31# Normalize embeddings
32embeddings = embeddings / np.linalg.norm(embeddings, axis=1, keepdims=True)
33# Calculate similarity scores
34scores = (embeddings[:1] @ embeddings[1:].T) * 100
35print(scores.tolist())
36# [[0.3016996383666992, 0.7503870129585266, 0.3203084468841553]]