McGill-NLP/LLM2Vec-Sheared-LLaMA-mntp-unsup-simcse

pip install llm2vec

from llm2vec import LLM2Vec

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

# Cargando modelo base Mistral, junto con código personalizado que habilita conexiones bidireccionales en LLMs sólo de decodificación. Los pesos de MNTP LoRA se fusionan en el modelo base.
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 la cpu

# Cargando modelo SimCSE no supervisado. Esto carga los pesos de MNTP entrenados en el modelo SimCSE (LoRA).
model = PeftModel.from_pretrained(
model, "McGill-NLP/LLM2Vec-Sheared-LLaMA-mntp-unsup-simcse")

# Envoltorio 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, recupere pasajes relevantes que respondan a la consulta:")
queries = [
[instruction, "cuánta proteína debe consumir una mujer"],
[instruction, "definición de cumbre"],
]
q_reps = l2v.encode(queries)

# Codificación de documentos. Instrucciones no son necesarias para documentos
documents = [
"Como regla 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 se puede ver en esta tabla, necesitarás aumentar eso si estás esperando un bebé o entrenando para un maratón. Consulta la tabla a continuación para ver cuánta proteína debes consumir 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 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.5964, 0.1270],[0.0698, 0.2394]]) ```