LLM2Vec-Meta-Llama-3-8B-Instruct-mntp-unsup-simcse

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 modelos de solo decodificación. Los pesos de LoRA MNTP están fusionados en el modelo base.
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")
model = PeftModel.from_pretrained(model, "McGill-NLP/LLM2Vec-Meta-Llama-3-8B-Instruct-mntp")
model = model.merge_and_unload()  # Esto puede tardar varios minutos en cpu

# Cargando el modelo SimCSE no supervisado. Esto carga los pesos de LoRA entrenados 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-Meta-Llama-3-8B-Instruct-mntp-unsup-simcse")

# Envoltura 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, recuperar 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. No se requieren instrucciones para los documentos
documents = [
    "Como guía general, el requerimiento promedio de proteína del CDC para mujeres de 19 a 70 años es de 46 gramos por día. Pero, como se puede ver en este gráfico, necesitarás aumentar eso si estás embarazada o entrenando para un maratón. Consulta el gráfico a continuación para ver cuánta proteína deberías estar comiendo cada día.",
    "Definición de cumbre para estudiantes del idioma inglés. : 1. el punto más alto de una montaña: la parte superior 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 cosín 
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)