deepvk/USER-base

Utilizar este modelo se vuelve sencillo cuando tienes sentence-transformers instalado:

pip install -U sentence-transformers

Luego puedes usar el modelo así:

from sentence_transformers import SentenceTransformer

input_texts = [
"query: Cuando se lanzó al agua el primer torpedero 'Spokoinyy'?",
"query: ¿Hay petróleo en Udmurtia?",
"pasaje: Spokoinyy (destructor)\nListado en la lista de la Armada Soviética el 19 de agosto de 1952.",
"pasaje: Las obras de búsqueda de petróleo en Udmurtia comenzaron inmediatamente después de la Segunda Guerra Mundial en 1945 y continúan hasta hoy. La extracción de petróleo comenzó en 1967."
]

model = SentenceTransformer("deepvk/USER-base")
embeddings = model.encode(input_texts, normalize_embeddings=True)

Sin embargo, también puedes usar el modelo directamente con Transformers:

import torch.nn.functional as F
from torch import Tensor, inference_mode
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 = [
    "query: Cuando se lanzó al agua el primer torpedero 'Spokoinyy'?",
    "query: ¿Hay petróleo en Udmurtia?",
    "pasaje: Spokoinyy (destructor)\nListado en la lista de la Armada Soviética el 19 de agosto de 1952.",
    "pasaje: Las obras de búsqueda de petróleo en Udmurtia comenzaron inmediatamente después de la Segunda Guerra Mundial en 1945 y continúan hasta hoy. La extracción de petróleo comenzó en 1967."
]

tokenizer = AutoTokenizer.from_pretrained("deepvk/USER-base")
model = AutoModel.from_pretrained("deepvk/USER-base")

batch_dict = tokenizer(
    input_texts, padding=True, truncation=True, return_tensors="pt")
with inference_mode():
    outputs = model(**batch_dict)
embeddings = average_pool(
  outputs.last_hidden_state, batch_dict["attention_mask"]
)
embeddings = F.normalize(embeddings, p=2, dim=1)

scores = (embeddings[:2] @ embeddings[2:].T) * 100
print(scores.round(decimals=2))

⚠️ Atención ⚠️ Cada texto de entrada debe comenzar con el prefijo "query: " o "pasaje: ". Para tareas distintas de la recuperación, simplemente utiliza el prefijo "query: ".