McGill-NLP/LLM2Vec-Meta-Llama-3-8B-Instruct-mntp

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 únicamente de decodificación.
tokenizer = AutoTokenizer.from_pretrained(
"McGill-NLP/LLM2Vec-Meta-Llama-3-8B-Instruct-mntp")
config = AutoConfig.from_pretrained(
"McGill-NLP/LLM2Vec-Meta-Llama-3-8B-Instruct-mntp", trust_remote_code=True
)
model = AutoModel.from_pretrained(
"McGill-NLP/LLM2Vec-Meta-Llama-3-8B-Instruct-mntp",
trust_remote_code=True,
config=config,
torch_dtype=torch.bfloat16,
device_map="cuda" if torch.cuda.is_available() else "cpu",
)

# Cargando el modelo MNTP (Predicción enmascarada del siguiente token).
model = PeftModel.from_pretrained(
model,
"McGill-NLP/LLM2Vec-Meta-Llama-3-8B-Instruct-mntp",
)

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

# Codificando consultas usando instrucciones
instruction = (
"Dada una consulta de búsqueda web, recupera pasajes relevantes que respondan a la consulta:")
queries = [
[instruction, "cuánta proteína debería consumir una mujer"],
[instruction, "definir cumbre"],
]
q_reps = l2v.encode(queries)

# Codificando documentos. No se requieren instrucciones para los documentos
documents = [
"Como pauta general, el requerimiento promedio de proteínas según el CDC para mujeres de 19 a 70 años es de 46 gramos por día. Sin embargo, como puedes ver en esta tabla, necesitarás aumentar eso si estás embarazada o entrenando para un maratón. Consulta la tabla a continuación para ver cuánta proteína deberías estar consumiendo cada día.",
"Definición de cumbre para aprendices del idioma 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)

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

tensor([[0.7740, 0.5580], [0.4845, 0.4993]])