dunzhang/stella_en_400M_v5

from sentence_transformers import SentenceTransformer

# Este modelo soporta dos prompts: "s2p_query" y "s2s_query" para tareas de oración a pasaje y de oración a oración, respectivamente.
# Están definidos en 'config_sentence_transformers.json'
query_prompt_name = "s2p_query"
queries = [
"¿Cuáles son algunas formas de reducir el estrés?",
"¿Cuáles son los beneficios de beber té verde?",
]
# Los documentos no necesitan prompts
docs = [
"Hay muchas formas efectivas de reducir el estrés. Algunas técnicas comunes incluyen respiración profunda, meditación y actividad física. Participar en pasatiempos, pasar tiempo en la naturaleza y conectarse con seres queridos también puede ayudar a aliviar el estrés. Además, establecer límites, practicar el autocuidado y aprender a decir no puede evitar que el estrés se acumule.",
"El té verde se ha consumido durante siglos y es conocido por sus potenciales beneficios para la salud. Contiene antioxidantes que pueden ayudar a proteger al cuerpo contra el daño causado por los radicales libres. El consumo regular de té verde se ha asociado con una mejor salud del corazón, una función cognitiva mejorada y un riesgo reducido de ciertos tipos de cáncer. Los polifenoles del té verde también pueden tener propiedades antiinflamatorias y de pérdida de peso.",
]

# ¡La dimensión predeterminada es 1024, si necesitas otras dimensiones, clona el modelo y modifica 'modules.json' para reemplazar '2_Dense_1024' por otra dimensión, por ejemplo, '2_Dense_256' o '2_Dense_8192'!
# en gpu
model = SentenceTransformer('dunzhang/stella_en_400M_v5', trust_remote_code=True).cuda()
# también puedes usar este modelo sin las características de 'use_memory_efficient_attention' y 'unpad_inputs'. Puede funcionar en CPU.
# model = SentenceTransformer(
#     'dunzhang/stella_en_400M_v5',
#     trust_remote_code=True,
#     device='cpu',
#     config_kwargs={'use_memory_efficient_attention': False, 'unpad_inputs': False}
# )
query_embeddings = model.encode(queries, prompt_name=query_prompt_name)
doc_embeddings = model.encode(docs)
print(query_embeddings.shape, doc_embeddings.shape)
# (2, 1024) (2, 1024)

similarities = model.similarity(query_embeddings, doc_embeddings)
print(similarities)
# tensor([[0.8398, 0.2990],
#         [0.3282, 0.8095]])

import os
import torch
from transformers import AutoModel, AutoTokenizer
from sklearn.preprocessing import normalize

query_prompt = "Instrucción: Dada una consulta de búsqueda web, recuperar pasajes relevantes que respondan a la consulta.\nConsulta: "
queries = [
"¿Cuáles son algunas formas de reducir el estrés?",
"¿Cuáles son los beneficios de beber té verde?",
]
queries = [query_prompt + query for query in queries]
# Los documentos no necesitan prompts
docs = [
"Hay muchas formas efectivas de reducir el estrés. Algunas técnicas comunes incluyen respiración profunda, meditación y actividad física. Participar en pasatiempos, pasar tiempo en la naturaleza y conectarse con seres queridos también puede ayudar a aliviar el estrés. Además, establecer límites, practicar el autocuidado y aprender a decir no puede evitar que el estrés se acumule.",
"El té verde se ha consumido durante siglos y es conocido por sus potenciales beneficios para la salud. Contiene antioxidantes que pueden ayudar a proteger al cuerpo contra el daño causado por los radicales libres. El consumo regular de té verde se ha asociado con una mejor salud del corazón, una función cognitiva mejorada y un riesgo reducido de ciertos tipos de cáncer. Los polifenoles del té verde también pueden tener propiedades antiinflamatorias y de pérdida de peso.",
]

# La ruta de tu modelo después de clonarlo
model_dir = "{Tu RUTA_DEL_MODELO}"

vector_dim = 1024
vector_linear_directory = f"2_Dense_{vector_dim}"
model = AutoModel.from_pretrained(model_dir, trust_remote_code=True).cuda().eval()
# también puedes usar este modelo sin las características de 'use_memory_efficient_attention' y 'unpad_inputs'. Puede funcionar en CPU.
# model = AutoModel.from_pretrained(model_dir, trust_remote_code=True,use_memory_efficient_attention=False,unpad_inputs=False).cuda().eval()
tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
vector_linear = torch.nn.Linear(in_features=model.config.hidden_size, out_features=vector_dim)
vector_linear_dict = {
k.replace('linear.', ''): v for k, v in
torch.load(os.path.join(model_dir, f"{vector_linear_directory}/pytorch_model.bin")).items()
}
vector_linear.load_state_dict(vector_linear_dict)
vector_linear.cuda()

# Insertar las consultas
with torch.no_grad():
input_data = tokenizer(queries, padding="longest", truncation=True, max_length=512, return_tensors="pt")
input_data = {k: v.cuda() for k, v in input_data.items()}
attention_mask = input_data['attention_mask']
last_hidden_state = model(**input_data)[0]
last_hidden = last_hidden_state.masked_fill(~attention_mask[..., None].bool(), 0.0)
query_vectors = last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
query_vectors = normalize(vector_linear(query_vectors).cpu().numpy())

# Insertar los documentos
with torch.no_grad():
input_data = tokenizer(docs, padding="longest", truncation=True, max_length=512, return_tensors="pt")
input_data = {k: v.cuda() for k, v in input_data.items()}
attention_mask = input_data['attention_mask']
last_hidden_state = model(**input_data)[0]
last_hidden = last_hidden_state.masked_fill(~attention_mask[..., None].bool(), 0.0)
docs_vectors = last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
docs_vectors = normalize(vector_linear(docs_vectors).cpu().numpy())

print(query_vectors.shape, docs_vectors.shape)
# (2, 1024) (2, 1024)

similarities = query_vectors @ docs_vectors.T
print(similarities)
# [[0.8397531  0.29900077]
#  [0.32818374 0.80954516]]