LLM2Vec-Meta-Llama-3-8B-Instruct-mntp-supervised

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

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()

model = PeftModel.from_pretrained(
model, "McGill-NLP/LLM2Vec-Meta-Llama-3-8B-Instruct-mntp-supervised")

l2v = LLM2Vec(model, tokenizer, pooling_mode="mean", max_length=512)

instruction = (
"Dada una consulta de búsqueda 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)

documents = [
"Como pauta general, el requisito promedio de proteína del CDC para mujeres de 19 a 70 años es de 46 gramos por día. Pero, como puede ver en este gráfico, necesitará 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 consumir cada día.",
"Definición de cumbre para aprendices 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 una serie de reuniones entre los líderes de dos o más gobiernos.",
]
d_reps = l2v.encode(documents)

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)