Instalación

pip install llm2vec

Uso

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 que solo funcionan como decodificadores. Los pesos MNTP LoRA son 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 tomar varios minutos en la CPU

# Cargando el modelo supervisado. Esto carga los pesos 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-Mistral-7B-Instruct-v2-mntp-supervised")

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

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

# Codificando documentos. Las instrucciones no son necesarias para los documentos
documents = [
  "Como guía general, el requisito promedio de proteínas del CDC para mujeres de 19 a 70 años es de 46 gramos por día. Pero, como puede ver en este cuadro, deberá aumentarlo si está esperando o entrenando para un maratón. Consulte el cuadro a continuación para ver cuánta proteína debe consumir cada día.",
  "Definición de cumbre para 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 serie de reuniones entre los líderes de dos o más gobiernos.",
]
d_reps = l2v.encode(documents)

# Calculando la similitud coseno
aq_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)