McGill-NLP/LLM2Vec-Llama-2-7b-chat-hf-mntp

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 modelos de lenguaje de solo decodificación.
tokenizer = AutoTokenizer.from_pretrained(
"McGill-NLP/LLM2Vec-Llama-2-7b-chat-hf-mntp")
config = AutoConfig.from_pretrained(
"McGill-NLP/LLM2Vec-Llama-2-7b-chat-hf-mntp", trust_remote_code=True
)
model = AutoModel.from_pretrained(
"McGill-NLP/LLM2Vec-Llama-2-7b-chat-hf-mntp",
trust_remote_code=True,
config=config,
torch_dtype=torch.bfloat16,
device_map="cuda" if torch.cuda.is_available() else "cpu",
)

# Cargando el modelo de Predicción de la Siguiente Token Enmascarada (MNTP).
model = PeftModel.from_pretrained(
model,
"McGill-NLP/LLM2Vec-Llama-2-7b-chat-hf-mntp",
)

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

# Codificación de consultas usando instrucciones
instruction = (
"Given a web search query, retrieve relevant passages that answer the query:")
queries = [
[instruction, "how much protein should a female eat"],
[instruction, "summit define"],
]
q_reps = l2v.encode(queries)

# Codificación de documentos. No se requieren instrucciones para documentos
documents = [
"As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day. But, as you can see from this chart, you'll need to increase that if you're expecting or training for a marathon. Check out the chart below to see how much protein you should be eating each day.",
"Definition of summit for English Language Learners. : 1  the highest point of a mountain : the top of a mountain. : 2  the highest level. : 3  a meeting or series of meetings between the leaders of two or more governments.",
]
d_reps = l2v.encode(documents)

# Calculo de 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)