gte-small

thenlper

Similitud de oraciones

Modelo General de Embeddings de Texto (GTE). Hacia Embeddings de Textos Generales con Aprendizaje Contrastivo Multiescenario. Los modelos GTE son entrenados por la Academia DAMO de Alibaba. Están principalmente basados en el marco de BERT y actualmente ofrecen tres tamaños diferentes de modelos, incluyendo GTE-large, GTE-base y GTE-small. Los modelos GTE están entrenados en un gran corpus de pares de texto relevantes, cubriendo una amplia gama de dominios y escenarios. Esto permite que los modelos GTE se apliquen a diversas tareas de texto embedido, incluyendo la recuperación de información, similitud textual semántica, clasificación de texto, ranking de texto, etc.

Como usar

import torch.nn.functional as F
from torch import Tensor
from transformers import AutoTokenizer, AutoModel

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]

input_texts = [
    "what is the capital of China?",
    "how to implement quick sort in python?",
    "Beijing",
    "sorting algorithms"
]

tokenizer = AutoTokenizer.from_pretrained("thenlper/gte-small")
model = AutoModel.from_pretrained("thenlper/gte-small")

# Tokenize the input texts
batch_dict = tokenizer(input_texts, max_length=512, padding=True, truncation=True, return_tensors='pt')

outputs = model(**batch_dict)
embeddings = average_pool(outputs.last_hidden_state, batch_dict['attention_mask'])

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

Usando sentence-transformers:
from sentence_transformers import SentenceTransformer
from sentence_transformers.util import cos_sim

sentences = ['That is a happy person', 'That is a very happy person']

model = SentenceTransformer('thenlper/gte-large')
embeddings = model.encode(sentences)
print(cos_sim(embeddings[0], embeddings[1]))

Funcionalidades

Basado en el marco BERT
Tres tamaños de modelos: GTE-large, GTE-base, GTE-small
Entrenado en un gran corpus de pares de texto relevantes
Aplicable a diversas tareas de texto embedido

Casos de uso

Recuperación de información
Similitud textual semántica
Clasificación de texto
Ranking de texto