McGill-NLP/LLM2Vec-Mistral-7B-Instruct-v2-mntp-unsup-simcse

from llm2vec import LLM2Vec

import torch
from transformers import AutoTokenizer, AutoModel, AutoConfig
from peft import PeftModel

# Cargar el modelo Mistral base, junto con código personalizado que habilita conexiones bidireccionales en LLMs que solo son decodificadores. Los pesos de MNTP LoRA están fusionados en el modelo base.
tokenizer = AutoTokenizer.from_pretrained(
"McGill-NLP/LLM2Vec-Mistral-7B-Instruct-v2-mntp")
config = AutoConfig.from_pretrained(
"McGill-NLP/LLM2Vec-Mistral-7B-Instruct-v2-mntp", trust_remote_code=True
)
model = AutoModel.from_pretrained(
"McGill-NLP/LLM2Vec-Mistral-7B-Instruct-v2-mntp",
trust_remote_code=True,
config=config,
torch_dtype=torch.bfloat16,
device_map="cuda" if torch.cuda.is_available() else "cpu",
)
model = PeftModel.from_pretrained(
model,
"McGill-NLP/LLM2Vec-Mistral-7B-Instruct-v2-mntp",
)
model = model.merge_and_unload() # Esto puede tardar varios minutos en CPU

# Cargando el modelo SimCSE no supervisado. Esto carga los pesos entrenados de LoRA sobre el modelo MNTP. Por lo tanto, los pesos finales son -- Modelo base + MNTP (LoRA) + SimCSE (LoRA).
model = PeftModel.from_pretrained(
model, "McGill-NLP/LLM2Vec-Mistral-7B-Instruct-v2-mntp-unsup-simcse")

# Envoltorio para operaciones de codificación y agrupamiento
l2v = LLM2Vec(model, tokenizer, pooling_mode="mean", max_length=512)

# Codificación de consultas usando instrucciones
instruction = (
"Dada una consulta de búsqueda web, recupera pasajes relevantes que respondan a la consulta:")
queries = [
[instruction, "cuánto proteínas debe comer una mujer"],
[instruction, "definir cumbre"],
]
q_reps = l2v.encode(queries)

# Codificación de documentos. Las instrucciones no son necesarias para los documentos
documents = [
"Como una guía general, el requerimiento promedio de proteínas del CDC para mujeres de 19 a 70 años es de 46 gramos por día. Pero, como puedes ver en esta tabla, necesitarás aumentar si estás esperando o entrenando para un maratón. Consulta la tabla a continuación para ver cuántas proteínas deberías comer cada día.",
"Definición de cumbre para aprendices de inglés: 1 el punto más alto de una montaña: la cima de una montaña. 2 el nivel más alto. 3 una reunión o serie de reuniones entre los líderes de dos o más gobiernos.",
]
d_reps = l2v.encode(documents)

# Calcular similitud de coseno
q_reps_norm = torch.nn.functional.normalize(q_reps, p=2, dim=1)
d_reps_norm = torch.nn.functional.normalize(d_reps, p=2, dim=1)
cos_sim = torch.mm(q_reps_norm, d_reps_norm.transpose(0, 1))

print(cos_sim)

[0.2298, 0.5792]])```