McGill-NLP/LLM2Vec-Sheared-LLaMA-mntp-supervised

pip install llm2vec

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

tokenizer = AutoTokenizer.from_pretrained("McGill-NLP/LLM2Vec-Sheared-LLaMA-mntp")
config = AutoConfig.from_pretrained("McGill-NLP/LLM2Vec-Sheared-LLaMA-mntp", trust_remote_code=True)
model = AutoModel.from_pretrained("McGill-NLP/LLM2Vec-Sheared-LLaMA-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-Sheared-LLaMA-mntp")
model = model.merge_and_unload() # Esto puede tomar varios minutos en CPU

# Cargando el modelo supervisado
model = PeftModel.from_pretrained(model, "McGill-NLP/LLM2Vec-Sheared-LLaMA-mntp-supervised")
l2v = LLM2Vec(model, tokenizer, pooling_mode="mean", max_length=512)

# Codificación de consultas usando instrucciones
instruction = ("Dada una consulta de búsqueda web, recupere 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)

documents = [
    "Como guía general, el requisito promedio de proteína de los CDC para mujeres de 19 a 70 años es de 46 gramos por día. Pero, como puede ver en este cuadro, deberá aumentar eso si está esperando o entrenando para un maratón. Consulte la tabla a continuación para ver cuánta proteína debe comer cada día.",
    "Definición de cumbre para estudiantes de 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)

# Calcular 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.6500, 0.1291], [0.0916, 0.4733]])