NeuML/pubmedbert-base-embeddings
NeuML
Similitud de oraciones
Este es un modelo PubMedBERT-base afinado utilizando sentence-transformers. Mapea oraciones y párrafos a un espacio vectorial denso de 768 dimensiones y puede usarse para tareas como agrupamiento o búsqueda semántica. El conjunto de datos de entrenamiento se generó utilizando una muestra aleatoria de pares de títulos-resumen de PubMed junto con pares de títulos similares. PubMedBERT Embeddings produce embeddings de mayor calidad que los modelos generalizados para literatura médica. Un ajuste fino adicional para un subdominio médico resultará en un rendimiento aún mejor.
Como usar
Uso con txtai
import txtai
embeddings = txtai.Embeddings(path="neuml/pubmedbert-base-embeddings", content=True)
embeddings.index(documents())
# Ejecutar una consulta
embeddings.search("consulta para ejecutar")
Uso con Sentence-Transformers
from sentence_transformers import SentenceTransformer
sentences = ["Esta es una oración de ejemplo", "Cada oración se convierte"]
model = SentenceTransformer("neuml/pubmedbert-base-embeddings")
embeddings = model.encode(sentences)
print(embeddings)
Uso con Transformers de Hugging Face
from transformers import AutoTokenizer, AutoModel
import torch
# Pooling medio: tener en cuenta la máscara de atención para un promedio correcto
def meanpooling(output, mask):
embeddings = output[0] # El primer elemento de model_output contiene todos los embeddings de los tokens
mask = mask.unsqueeze(-1).expand(embeddings.size()).float()
return torch.sum(embeddings * mask, 1) / torch.clamp(mask.sum(1), min=1e-9)
# Oraciones para las que queremos embeddings de oraciones
sentences = ['Esta es una oración de ejemplo', 'Cada oración se convierte']
# Cargar modelo del HuggingFace Hub
tokenizer = AutoTokenizer.from_pretrained("neuml/pubmedbert-base-embeddings")
model = AutoModel.from_pretrained("neuml/pubmedbert-base-embeddings")
# Tokenizar oraciones
inputs = tokenizer(sentences, padding=True, truncation=True, return_tensors='pt')
# Calcular embeddings de los tokens
with torch.no_grad():
output = model(**inputs)
# Realizar pooling. En este caso, pooling medio.
embeddings = meanpooling(output, inputs['attention_mask'])
print("Embeddings de oraciones:")
print(embeddings)
Funcionalidades
- Genera embeddings de alta calidad para literatura médica
- 768 dimensiones en el espacio vectorial
- Entrenado con datos de títulos-resumen de PubMed
- Se puede usar para búsqueda semántica y agrupamiento
- Compatible con múltiples bibliotecas, como sentence-transformers y transformers
Casos de uso
- Búsqueda semántica en literatura médica
- Agrupamiento de documentos médicos
- Fuente de conocimiento para generación aumentada por recuperación (RAG)