McGill-NLP/LLM2Vec-Sheared-LLaMA-mntp

from llm2vec import LLM2Vec

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

# Cargar modelo base Mistral, junto con el código personalizado que habilita conexiones bidireccionales en LLMs solo decodificadores.
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",
)

# Cargar modelo MNTP (Predicción de la siguiente máscara de token).
model = PeftModel.from_pretrained(
model,
"McGill-NLP/LLM2Vec-Sheared-LLaMA-mntp",
)

# Envoltorio para 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 web, recupera los 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. Las instrucciones no son necesarias para los documentos
documents = [
"Como pauta general, el requerimiento 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á aumentar eso si está esperando o entrenando para un maratón. Consulte el gráfico a continuación para ver cuánta proteína debe comer 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)

# Calcule la similitud del 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.8180, 0.5825],
[0.1069, 0.1931]])