gte-base-en-v1.5

# Requiere transformers>=4.36.0
import torch.nn.functional as F
from transformers import AutoModel, AutoTokenizer

input_texts = [
"¿Cuál es la capital de China?",
"¿Cómo implementar quick sort en python?",
"Beijing",
"algoritmos de clasificación"
]

model_path = 'Alibaba-NLP/gte-base-en-v1.5'
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModel.from_pretrained(model_path, trust_remote_code=True)

# Tokenizar los textos de entrada
batch_dict = tokenizer(input_texts, max_length=8192, padding=True, truncation=True, return_tensors='pt')

outputs = model(**batch_dict)
embeddings = outputs.last_hidden_state[:, 0]

# (Opcional) normalizar las incrustaciones
embeddings = F.normalize(embeddings, p=2, dim=1)
scores = (embeddings[:1] @ embeddings[1:].T) * 100
print(scores.tolist())

Es recomendable instalar xformers y habilitar unpadding para la aceleración, consulte enable-unpadding-and-xformers.

# Requiere sentence_transformers>=2.7.0
from sentence_transformers import SentenceTransformer
from sentence_transformers.util import cos_sim

sentences = ['Esa es una persona feliz', 'Esa es una persona muy feliz']

model = SentenceTransformer('Alibaba-NLP/gte-base-en-v1.5', trust_remote_code=True)
embeddings = model.encode(sentences)
print(cos_sim(embeddings[0], embeddings[1]))

// npm i @xenova/transformers
import { pipeline, dot } from '@xenova/transformers';

// Crear tubería de extracción de características
const extractor = await pipeline('feature-extraction', 'Alibaba-NLP/gte-base-en-v1.5', {
quantized: false, // Comente esta línea para usar la versión cuantizada
});

// Generar incrustaciones de oraciones
const sentences = [
"¿Cuál es la capital de China?",
"¿Cómo implementar quick sort en python?",
"Beijing",
"algoritmos de clasificación"
]
const output = await extractor(sentences, { normalize: true, pooling: 'cls' });

// Calcular puntuaciones de similitud
const [source_embeddings, ...document_embeddings ] = output.tolist();
const similarities = document_embeddings.map(x => 100 * dot(source_embeddings, x));
console.log(similarities); // [34.504930869007296, 64.03973265120138, 19.520042686034362]