gte-base

thenlper

Similitud de oraciones

Un modelo de General Text Embeddings (GTE). Hacia General Text Embeddings con Aprendizaje Contrastivo Multietapa. Los modelos GTE son entrenados por la Academia DAMO de Alibaba. Se basan principalmente en el marco 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 corpus a gran escala de pares de textos relevantes, cubriendo una amplia gama de dominios y escenarios. Esto permite que los modelos GTE se apliquen a varias tareas de downstream de embeddings de texto, incluyendo recuperación de información, similitud textual semántica, reordenamiento 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-base")
model = AutoModel.from_pretrained("thenlper/gte-base")

# 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())

Uso con 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-base')
embeddings = model.encode(sentences)
print(cos_sim(embeddings[0], embeddings[1]))

Funcionalidades

Basado en la estructura BERT
Disponible en tres tamaños: GTE-large, GTE-base, GTE-small
Entrenado en un corpus a gran escala de pares de textos relevantes
Applicable a tareas de downstream como la recuperación de información, similitud textual semántica y reordenamiento de textos
Ofrecido en varios tamaños y dimensiones

Casos de uso

Recuperación de información
Similitud textual semántica
Reordenamiento de textos