Upload run_evaluate_loco.py

scripts/evaluate/run_evaluate_loco.py

@@ -0,0 +1,304 @@
+import json
+import random
+
+import pandas as pd
+import torch
+import os
+
+os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
+from sentence_transformers import SentenceTransformer
+import tqdm
+import numpy as np
+import faiss
+from sklearn.metrics import ndcg_score
+from os.path import join
+from sklearn.preprocessing import normalize
+from transformers import AutoTokenizer, AutoModel
+
+faiss.omp_set_num_threads(16)
+
+
+def find_topk_by_vecs(source_vecs: np.ndarray, target_vecs: np.ndarray, topk: int):
+    if topk > len(target_vecs):
+        topk = len(target_vecs)
+    faiss_index = faiss.IndexFlatIP(target_vecs.shape[1])
+    faiss_index.add(target_vecs)
+
+    res_distance, res_index = faiss_index.search(source_vecs, topk)
+    return res_index, res_distance
+
+
+def get_loco_path_info(q_dir, d_dir):
+    names = []
+    for name in sorted(os.listdir(q_dir)):
+        if name.endswith(".jsonl"):
+            names.append(name)
+    for name in os.listdir(d_dir):
+        if name.endswith(".jsonl"):
+            assert name in names
+    infos = []
+    for name in names:
+        infos.append(["LOCO-V1", name, join(q_dir, name), join(d_dir, name)])
+    infos.sort(key=lambda x: x[1])
+    return infos
+
+
+def get_loco_data(q_path, d_path):
+    passage_list, query2passage_list = [], {}
+
+    original_doc_id2doc = {}
+
+    with open(d_path, "r", encoding="utf8") as fr:
+        for line in fr:
+            item = json.loads(line)
+            if item["passage"].strip():
+                original_doc_id2doc[item["pid"]] = item["passage"].strip()
+                passage_list.append(item["passage"].strip())
+
+    with open(q_path, "r", encoding="utf8") as fr:
+        for line in fr:
+            item = json.loads(line)
+            if item["query"].strip():
+                query2passage_list[item["query"].strip()] = [
+                    original_doc_id2doc[answer_pid]
+                    for answer_pid in item["answer_pids"]
+                    if answer_pid in original_doc_id2doc
+                ]
+    query2passage_list = {k: list(set(v)) for k, v in query2passage_list.items() if list(set(v))}
+    passage_list = list(set(passage_list))
+    passage2id = {passage: idx for idx, passage in enumerate(passage_list)}
+    query2id_list = {k: list(set([passage2id[i] for i in v])) for k, v in query2passage_list.items()}
+    query_list = list(query2id_list.keys())
+    return query_list, passage_list, query2id_list
+
+
+def get_ndcg_score(query_list, passage_list, query2passage_id_list, topk=10, error_data_save_path: str = None):
+    chunk_id2passage_id = {}
+    q_vecs = model.encode(
+        sentences=query_list,
+        batch_size=batch_size,
+        chunk_size=chunk_size,
+        chunk_overlap=chunk_overlap,
+        max_seq_length=max_seq_length,
+        is_q=True,
+    )
+    p_vecs = model.encode(
+        sentences=passage_list,
+        batch_size=batch_size,
+        chunk_size=chunk_size,
+        chunk_overlap=chunk_overlap,
+        max_seq_length=max_seq_length,
+        is_q=False,
+    )
+    # according query2id_list get labels_list
+    query_id_list = [query2passage_id_list[query] for query in query_list]
+    max_doc = max((len(id_list) for id_list in query_id_list))
+
+    labels = np.array([(id_list * max_doc)[:max_doc] for id_list in query_id_list])
+    if isinstance(p_vecs, list):
+        for idx, vec in enumerate(p_vecs):
+            if multi_vec_strategy == "full_text":
+                p_vecs[idx] = normalize(np.mean(vec[1:2, :], axis=0, keepdims=True), axis=1)
+            elif multi_vec_strategy == "full_text+chunks":
+                n_chunk = (vec.shape[0] - 2) // 2
+                if n_chunk > 0:
+                    p_vecs[idx] = np.vstack(
+                        (
+                            normalize(np.mean(vec[:2, :], axis=0, keepdims=True), axis=1),
+                            vec[2:2 + n_chunk, :],
+                        )
+                    )
+                else:
+                    p_vecs[idx] = normalize(np.mean(vec[:2, :], axis=0, keepdims=True), axis=1)
+        p_vecs = np.vstack(p_vecs)
+
+    if isinstance(q_vecs, list):
+        for idx, vec in enumerate(q_vecs):
+            q_vecs[idx] = normalize(np.mean(vec[0:2, :], axis=0, keepdims=True), axis=1)
+        q_vecs = np.vstack(q_vecs)
+    print("q_vecs.shape and dtype", q_vecs.shape, q_vecs.dtype)
+    print("p_vecs.shape and dtype", p_vecs.shape, p_vecs.dtype)
+    # search topk
+    # we calculate ndcg@10
+    topk_index, topk_scores = find_topk_by_vecs(q_vecs, p_vecs, topk * 100)
+    # print("topk_index", topk_index.shape, topk_index)
+    # print("topk_scores", topk_scores.shape, topk_scores)
+    ### we may use multi vectors, so we should modify topk_index and topk_scores
+    if chunk_id2passage_id:
+        new_topk_index, new_topk_scores = [], []
+        # print("chunk_id2passage_id")
+        for chunk_ids, chunk_scores in tqdm.tqdm(zip(topk_index, topk_scores),
+                                                 desc="modify topk_index and topk_scores", disable=True):
+            # processed by row
+            row_ids, row_scores, passage_id_set = [], [], set()
+            for idx, chunk_id in enumerate(chunk_ids):
+                passage_id = chunk_id2passage_id[chunk_id]
+                if passage_id not in passage_id_set:
+                    passage_id_set.add(passage_id)
+                    row_ids.append(passage_id)
+                    row_scores.append(chunk_scores[idx])
+            new_topk_index.append(row_ids[:topk])
+            new_topk_scores.append(row_scores[:topk])
+        topk_index = np.array(new_topk_index)
+        # print("topk_index", topk_index)
+        topk_scores = np.array(new_topk_scores)
+    topk_index, topk_scores = topk_index[:, :topk], topk_scores[:, :topk]
+    is_match = (topk_index == labels[:, :1])
+    for idx in range(1, max_doc):
+        # the or operator means that only one positive doc in pred topk, we think it is recalled
+        is_match = is_match | (topk_index == labels[:, idx:idx + 1])
+
+    # compute recall at topk
+    print("is_match.shape", is_match.shape)
+    # recall_at_k = is_match.sum(axis=1).astype(bool).mean()
+    ndcg = ndcg_score(is_match.astype(dtype=np.float32), topk_scores)
+
+    if error_data_save_path:
+        in_top_k = is_match.sum(axis=1).astype(bool)
+        err_data = []
+        for idx, pred_res in enumerate(in_top_k):
+            if not pred_res:
+                query = query_list[idx]
+                label_doc = passage_list[query2passage_id_list[query][0]]
+                pred_doc = passage_list[topk_index[idx][0]]
+                err_data.append([query, label_doc, pred_doc])
+        pd.DataFrame(err_data, columns=["Query", "Label", "Pred"]).to_excel(error_data_save_path, index=False)
+    return float(ndcg)
+
+
+class ModelWrapper:
+    def __init__(self, model_dir, model_type, max_seq_length):
+        assert model_type in ["dewey", "sentence_transformer"]
+        self.model_type = model_type
+        self.tokenizer = AutoTokenizer.from_pretrained(model_dir)
+        if model_type == "dewey":
+            self.model = AutoModel.from_pretrained(
+                model_dir,
+                attn_implementation="flash_attention_2",
+                trust_remote_code=True,
+            ).cuda().bfloat16().eval()
+            self.model.tokenizer = self.tokenizer
+        else:
+            self.model = SentenceTransformer(
+                model_dir,
+                trust_remote_code=True,
+                device="cpu",
+                model_kwargs={
+                    "torch_dtype": torch.bfloat16,  # fp16
+                    "attn_implementation": "flash_attention_2"
+                },
+            )
+            self.model.max_seq_length = max_seq_length
+            if "NV-Embed-v2" in model_dir:
+                self.model.tokenizer.padding_side = "right"
+            self.pool = self.model.start_multi_process_pool()
+
+    def encode(
+            self,
+            sentences,
+            batch_size,
+            chunk_size,
+            chunk_overlap,
+            max_seq_length,
+            is_q,
+    ):
+        if self.model_type == "dewey":
+            if is_q:
+                prompt = "<|START_INSTRUCTION|>Answer the question<|END_INSTRUCTION|>"
+            else:
+                prompt = "<|START_INSTRUCTION|>Candidate document<|END_INSTRUCTION|>"
+            return self.model.encode(
+                sentences=sentences,
+                batch_size=batch_size,
+                use_cuda=True,
+                show_progress_bar=True,
+                chunk_size=chunk_size,
+                chunk_overlap=chunk_overlap,
+                convert_to_tensor=False,
+                max_seq_length=max_seq_length,
+                normalize_embeddings=True,
+                prompt=prompt,
+                fast_chunk=True,
+            )[0]
+        self.model.max_seq_length = max_seq_length
+        prompt = None
+        if is_q and (
+                "Linq-Embed-Mistral" in model_dir or "e5-mistral-7b-instruct" in model_dir or "SFR-Embedding-Mistral" in model_dir):
+            prompt = PROMPT_E5
+        if is_q and ("NV-Embed-v2" in model_dir):
+            prompt = PROMPT_NV
+        if "chunk_alignment" in model_dir or "dewey" in model_dir:
+            if is_q:
+                prompt = "<|START_INSTRUCTION|>Answer the question<|END_INSTRUCTION|>"
+            else:
+                prompt = "<|START_INSTRUCTION|>Candidate document<|END_INSTRUCTION|>"
+        vecs = self.model.encode_multi_process(
+            add_eos(sentences) if "NV-Embed-v2" in model_dir else sentences,
+            pool=self.pool,
+            show_progress_bar=True,
+            batch_size=batch_size,
+            normalize_embeddings=True,
+            prompt=prompt
+        )
+        return vecs
+
+
+def add_eos(input_examples):
+    input_examples = [input_example + model.tokenizer.eos_token for input_example in input_examples]
+    return input_examples
+
+
+PROMPT_BGE = "Represent this sentence for searching relevant passages:"
+PROMPT_E5 = "Instruct: Given a web search query, retrieve relevant passages that answer the query.\nQuery: "
+PROMPT_NV = "Instruct: Given a question, retrieve passages that answer the question\nQuery: "
+if __name__ == "__main__":
+    chunk_size = -1
+    chunk_overlap = 32
+    batch_size = 2
+    max_seq_length = 8 * 1024
+    multi_vec_strategy = "full_text"  # full_text; full_text+chunks
+    err_data_save_path = None
+    topk = 10
+
+    model_dir = "infgrad/dewey_en_beta"
+    # model_dir = "/home/zd/public_models/Linq-Embed-Mistral/"
+    # model_dir = "/home/zd/public_models/SFR-Embedding-Mistral"
+    # model_dir = "/home/zd/public_models/e5-mistral-7b-instruct"
+    # model_dir = "/home/zd/public_models/bge-m3"
+    # model_dir = "/home/zd/public_models/gte-modernbert-base"
+    # model_dir = "/home/zd/public_models/NV-Embed-v2"
+
+    # sentence_transformer dewey
+    model_type = "sentence_transformer"
+    ## get data info
+    # TODO Please download LOCOV1 data first!
+    data_info = get_loco_path_info(
+        "/home/zd/public_data/LoCoV1-Queries/documents/",
+        "/home/zd/public_data/LoCoV1-Documents/documents/",
+    )
+
+    # load model
+    model = ModelWrapper(model_dir=model_dir, model_type=model_type, max_seq_length=max_seq_length)
+    # model = zd()
+    ndcg_score_list = []
+    for item in data_info:
+        print("\n\n\n\n" + "=" * 20)
+        print(f"evaluate {item[:2]}...")
+        query_list, passage_list, query2passage_id_list = get_loco_data(*item[2:])
+        print("number of all queries", len(query_list))
+        print("number of all passages", len(passage_list))
+        ndcg = get_ndcg_score(query_list, passage_list, query2passage_id_list, topk=topk,
+                              error_data_save_path=err_data_save_path)
+        print(f"{ndcg}")
+        ndcg_score_list.append(ndcg)
+
+    for i in data_info:
+        print(i[0])
+    print("\n\n\n")
+    for i in data_info:
+        print(i[1].replace(".jsonl", ""))
+    print("\n\n\n")
+
+    print(os.path.basename(model_dir))
+    for i in ndcg_score_list:
+        print(i)