McGill-NLP/LLM2Vec-Llama-2-7b-chat-hf-mntp-supervised

from llm2vec import LLM2Vec

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

# Cargando el modelo base Mistral, junto con el código personalizado que habilita conexiones bidireccionales en LLMs de solo decodificador. Los pesos de MNTP LoRA se fusionan en el modelo base.
tokenizer = AutoTokenizer.from_pretrained(
"McGill-NLP/LLM2Vec-Llama-2-7b-chat-hf-mntp")
config = AutoConfig.from_pretrained(
"McGill-NLP/LLM2Vec-Llama-2-7b-chat-hf-mntp", trust_remote_code=True
)
model = AutoModel.from_pretrained(
"McGill-NLP/LLM2Vec-Llama-2-7b-chat-hf-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-Llama-2-7b-chat-hf-mntp",
)
model = model.merge_and_unload() # Esto puede tomar varios minutos en CPU

# Cargando el modelo supervisado. Esto carga los pesos de LoRA entrenados sobre el modelo MNTP. Por lo tanto, los pesos finales son -- Modelo base + MNTP (LoRA) + supervisado (LoRA).
model = PeftModel.from_pretrained(
model, "McGill-NLP/LLM2Vec-Llama-2-7b-chat-hf-mntp-supervised")

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

# Codificación de consultas utilizando instrucciones
instruction = (
"Dada una consulta de búsqueda web, recupere los pasajes relevantes que respondan a la consulta:")
queries = [
[instruction, "cuánta proteína 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 guía general, la CDC recomienda un promedio de 46 gramos de proteína al día para mujeres de 19 a 70 años. Pero, como puede ver en este gráfico, necesitará aumentar eso si está esperando un hijo o entrenando para un maratón. Consulte el gráfico a continuación para ver cuánta proteína debe comer cada día.",
"Definición de cumbre para los estudiantes 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 series de reuniones entre los líderes de dos o más gobiernos.",
]
d_reps = l2v.encode(documents)

# Calcular la similitud del 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)

Esto imprimirá algo similar a:

tensor([[0.5417, 0.0780],
[0.0627, 0.5726]])