File size: 3,173 Bytes
e9cd0c7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 | import argparse
from typing import Any
import mteb
import json
import torch
import numpy as np
from mteb.models.instructions import task_to_instruction
from mteb.models.text_formatting_utils import corpus_to_texts
from llm2vec import LLM2Vec
def llm2vec_instruction(instruction):
if len(instruction) > 0 and instruction[-1] != ":":
instruction = instruction.strip(".") + ":"
return instruction
class LLM2VecWrapper:
def __init__(self, model=None, task_to_instructions=None):
self.task_to_instructions = task_to_instructions
self.model = model
def encode(
self,
sentences: list[str],
*,
prompt_name: str = None,
**kwargs: Any, # noqa
) -> np.ndarray:
if prompt_name is not None:
instruction = (
self.task_to_instructions[prompt_name]
if self.task_to_instructions
and prompt_name in self.task_to_instructions
else llm2vec_instruction(task_to_instruction(prompt_name))
)
else:
instruction = ""
sentences = [[instruction, sentence] for sentence in sentences]
return self.model.encode(sentences, **kwargs)
def encode_corpus(
self,
corpus: list[dict[str, str]] | dict[str, list[str]] | list[str],
prompt_name: str = None,
**kwargs: Any,
) -> np.ndarray:
sentences = corpus_to_texts(corpus, sep=" ")
sentences = [["", sentence] for sentence in sentences]
if "request_qid" in kwargs:
kwargs.pop("request_qid")
return self.model.encode(sentences, **kwargs)
def encode_queries(self, queries: list[str], **kwargs: Any) -> np.ndarray:
return self.encode(queries, **kwargs)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
"--base_model_name_or_path",
type=str,
default="McGill-NLP/LLM2Vec-Meta-Llama-3-8B-Instruct-mntp",
)
parser.add_argument(
"--peft_model_name_or_path",
type=str,
default="McGill-NLP/LLM2Vec-Meta-Llama-3-8B-Instruct-mntp-supervised",
)
parser.add_argument("--task_name", type=str, default="STS16")
parser.add_argument(
"--task_to_instructions_fp",
type=str,
default="test_configs/mteb/task_to_instructions.json",
)
parser.add_argument("--output_dir", type=str, default="results")
args = parser.parse_args()
task_to_instructions = None
if args.task_to_instructions_fp is not None:
with open(args.task_to_instructions_fp, "r") as f:
task_to_instructions = json.load(f)
l2v_model = LLM2Vec.from_pretrained(
args.base_model_name_or_path,
peft_model_name_or_path=args.peft_model_name_or_path,
device_map="cuda" if torch.cuda.is_available() else "cpu",
torch_dtype=torch.bfloat16,
)
model = LLM2VecWrapper(model=l2v_model, task_to_instructions=task_to_instructions)
tasks = mteb.get_tasks(tasks=[args.task_name])
evaluation = mteb.MTEB(tasks=tasks)
results = evaluation.run(model, output_folder=args.output_dir)
|