evaluate.py

from collections import OrderedDict
from datetime import timedelta
from functools import partial
import os
import json
import numpy as np
from tqdm import tqdm
import torch
import torch.nn.functional as F
import torch.distributed as dist
from transformers import set_seed
from argparse import ArgumentParser
import pandas as pd
import src.custom_utils as utils
from PIL import Image
from natsort import natsorted
import math
from pathlib import Path
import faiss
from src.models import Ret2Model

logger = utils.get_logger(__name__)

PLACEHOLDER_IMG = Image.new('RGB', (336, 336), color='black')
PLACEHOLDER_TXT = ''


def parse_command_line():
    parser = ArgumentParser()
    parser.add_argument('--dataset_path', type=str)
    parser.add_argument('--dataset_passages_path', type=str)
    parser.add_argument('--image_root', type=str)
    parser.add_argument('--checkpoint_path', type=str)
    parser.add_argument('--action', type=str, default='index',
                        choices=['index', 'search', 'create_index'])
    parser.add_argument('--index_root', type=str, default='evaluation')
    parser.add_argument('--index_model_name', type=str)
    parser.add_argument('--index_dataset_name', type=str)
    parser.add_argument('--batch_size', type=int, default=4)
    parser.add_argument('--top_k', type=int, default=100)
    parser.add_argument('--dataloader_num_workers', type=int, default=0)
    parser.add_argument('--skip_metrics', action='store_true')
    parser.add_argument('--no_instruction', action='store_true')
    parser.add_argument('--fp16', action='store_true')

    args = parser.parse_args()

    return args


class SingleTokenIndexDataset(torch.utils.data.Dataset):
    def __init__(self, args: ArgumentParser):
        self.args = args

        jsonl_path = args.dataset_passages_path if args.action == 'index' else args.dataset_path
        data = pd.read_json(jsonl_path, lines=True)
        logger.info(f"[Rank {args.rank}] Loaded {len(data)} lines from {jsonl_path}")
        local_len = math.ceil(len(data) / dist.get_world_size())
        data = data.iloc[args.rank * local_len:(args.rank + 1) * local_len]
        data.reset_index(drop=True, inplace=True)
        logger.info(f"[Rank {args.rank}] Local samples: {len(data)}")

        if args.image_root:
            img_root = Path(args.image_root)
            if args.action == 'index':
                data.passage_image_path = data.passage_image_path.apply(lambda x: x if pd.isna(x) else str(img_root.joinpath(x))) 
            elif args.action == 'search':
                data.image_path = data.image_path.apply(lambda x: x if pd.isna(x) else str(img_root.joinpath(x))) 
        self.data = data

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        sample = self.data.iloc[idx]
        ret = {}

        if self.args.action == 'index':
            ret['passage_id'] = str(sample.passage_id)
            text = PLACEHOLDER_TXT if pd.isna(sample.passage_text) else sample.passage_text
            image_path = None if pd.isna(sample.passage_image_path) else sample.passage_image_path
        
        elif self.args.action == 'search':
            passage_id = sample.passage_id
            if not isinstance(passage_id, list):
                passage_id = [passage_id]
            ret['passage_id'] = [(None if pid is None else str(pid)) for pid in passage_id]
            if pd.isna(sample.question):
                if self.args.no_instruction:
                    text = ''
                    text_mask = 0
                else:
                    text = sample.instruction
                    text_mask = 1 if text else 0
            elif self.args.no_instruction:
                text = sample.question
                text_mask = 1 if text else 0
            else:
                text = f"{sample.instruction} {sample.question}"
                text_mask = 1
            image_path = None if pd.isna(sample.image_path) else sample.image_path
            ret['query_id'] = sample.data_id
            ret['text_mask'] = text_mask
            if (isinstance(sample.answer, list) and not pd.isna(sample.answer).all()) or not pd.isna(sample.answer):
                answer = sample.answer
                if not isinstance(answer, list):
                    answer = [answer]
                answer = [x for x in answer if not pd.isna(x)]
                assert len(answer) > 0
                ret['answer'] = answer

        ret['text'] = text.lower() if self.args.lowercase_text else text
        if 'text_mask' not in ret:
            ret['text_mask'] = 1 if text else 0
        ret['image'] = Image.open(image_path).convert('RGB') if image_path else PLACEHOLDER_IMG
        ret['vision_mask'] = 1 if image_path else 0

        return ret


def collate_fn(samples, tokenizer, image_processor, dtype):
    batch = {}
    for k in samples[0]:
        values = [x[k] for x in samples]
        if k == 'text':
            padding = 'max_length' if utils.is_siglip_tokenizer(tokenizer) else True
            text_inputs = tokenizer(values, return_tensors='pt', padding=padding, truncation=True)
            batch['input_ids'] = text_inputs['input_ids']
            batch['attention_mask'] = text_inputs.get('attention_mask')
            batch['text'] = values 
        elif k == 'image':
            batch['pixel_values'] = image_processor(images=values, return_tensors='pt').pixel_values
        elif k == 'text_mask' or k == 'vision_mask':
            batch[k] = torch.tensor(values, dtype=dtype)
        else:
            batch[k] = values
    return batch


def to_device(batch: dict, device: torch.device):
    return {k: (v.to(device) if isinstance(v, torch.Tensor) else v) for k, v in batch.items()}


if __name__ == '__main__':
    set_seed(42)
    args = parse_command_line()

    output_path = Path(args.index_root).joinpath(args.index_model_name, args.index_dataset_name, 'embeddings')

    if args.action == 'create_index':
        embedding_files = [x for x in os.scandir(output_path) 
                        if x.is_file() and x.name.endswith('.npy') and 'embedding' in x.name]
        embedding_files = natsorted(embedding_files, key=lambda x: x.name)
        embeddings = []
        for emb_file in embedding_files:
            emb_data = np.load(emb_file.path)
            embeddings.append(emb_data)
        embeddings = np.vstack(embeddings).astype(np.float32)
        
        dimension = embeddings.shape[1]
        num_vectors = embeddings.shape[0]
        logger.info(f"Creating flat IP index with {num_vectors} vectors of dimension {dimension}")
        
        index = faiss.IndexFlatIP(dimension)
        index.add(embeddings)
        index_path = str(output_path.joinpath('knn.index'))
        faiss.write_index(index, index_path)
        
        index_infos = {
            "index_type": "IndexFlatIP",
            "dimension": dimension,
            "num_vectors": num_vectors,
            "metric": "inner_product"
        }
        
        index_infos_path = str(output_path.joinpath('index_infos.json'))
        with open(index_infos_path, 'w') as f:
            json.dump(index_infos, f, indent=2)
        
        logger.info(f"Created index: {index_path}")
        logger.info(f"Created index info: {index_infos_path}")
        
        # Create final JSON with all text data
        json_name = output_path.joinpath('knn.json')
        all_text = []
        info_files = [x for x in os.scandir(output_path) if x.is_file(
        ) and x.name.endswith('.json') and x.name != 'index_infos.json' and x.name != 'knn.json']
        info_files = natsorted(info_files, key=lambda x: x.name)
        for item in tqdm(info_files, desc='Creating final json with IDs ...'):
            text = json.load(open(item.path, 'r'))
            all_text.extend(text)
        json.dump(all_text, open(json_name, 'w'))
        logger.info(f"Created json knn {json_name}")

    else:

        dist.init_process_group(
            backend='nccl', timeout=timedelta(seconds=1800))

        model = Ret2Model.from_pretrained(args.checkpoint_path)
        model.eval()

        rank = dist.get_rank()
        args.rank = rank
        torch.cuda.set_device(rank)
        dtype = torch.float16 if args.fp16 else torch.float32
        model.to(dtype=dtype, device=rank)

        args.lowercase_text = utils.is_siglip2_tokenizer(model.tokenizer)
        dataset = SingleTokenIndexDataset(args)

        dataloader = torch.utils.data.DataLoader(
            dataset,
            batch_size=args.batch_size,
            shuffle=False,
            num_workers=args.dataloader_num_workers,
            pin_memory=False,
            drop_last=False,
            collate_fn=partial(
                collate_fn,
                tokenizer=model.tokenizer,
                image_processor=model.image_processor,
                dtype=dtype
            )
        )

        if args.action == 'index':
            passage_data = []
            all_ret_feats = []

            for i, batch in enumerate(tqdm(dataloader, desc=f"[Rank {rank}] Indexing")):
                pids = batch.pop('passage_id')
                passage_text = batch.pop('text')
                batch = to_device(batch, rank)

                with torch.inference_mode():
                    ret_feats = model.get_ret_features(**batch)
                    ret_feats = F.normalize(ret_feats, p=2, dim=-1).cpu()

                passage_data.extend(zip(pids, passage_text))
                all_ret_feats.append(ret_feats)

            all_ret_feats = torch.cat(all_ret_feats, dim=0).numpy()
            logger.info(
                f"[Rank {rank}] Local embeddings: {len(all_ret_feats)} (shape {all_ret_feats.shape})")

            if rank == 0:
                output_path.mkdir(parents=True, exist_ok=True)
            dist.barrier()
            embedding_path = output_path.joinpath(f"embeddings_{str(rank).zfill(5)}.npy")
            np.save(embedding_path, all_ret_feats)
            passage_data_path = output_path.joinpath(f"embeddings_{str(rank).zfill(5)}.json")
            with open(passage_data_path, 'w') as f:
                json.dump(passage_data, f)
            logger.info(f"[Rank {rank}] Done saving embeddings")

            # check for non-unique passage_ids, required for M2KR
            dist.barrier()
            if rank == 0:
                unique_pids = set()
                for embd_info_path in output_path.glob('embeddings_*.json'):
                    local_data = json.load(open(embd_info_path, 'r'))
                    pidxs_to_delete = []
                    for i, (pid, _) in enumerate(local_data):
                        if pid in unique_pids:
                            pidxs_to_delete.append(i)
                        else:
                            unique_pids.add(pid)
                    if len(pidxs_to_delete):
                        logger.info(
                            f"Deleting {len(pidxs_to_delete)} non-unique passages from {embd_info_path}")
                        embd_path = embd_info_path.with_suffix('.npy')
                        local_embds = np.load(embd_path)
                        assert len(local_embds) == len(local_data)
                        local_data = [x for i, x in enumerate(local_data) if i not in set(pidxs_to_delete)]
                        assert len(local_data) == len(dict(local_data))
                        local_embds = np.delete(local_embds, np.array(
                            pidxs_to_delete, dtype=np.int64), axis=0)
                        assert len(local_embds) == len(local_data)
                        with open(embd_info_path, 'w') as f:
                            json.dump(local_data, f)
                        np.save(embd_path, local_embds)

        elif args.action == 'search':                    
            Ks = [x for x in [1, 2, 3, 5, 10, 50, 100] if x <= args.top_k]
            recall_metrics = np.zeros(len(Ks))
            precall_metrics = np.zeros(len(Ks))
            compute_pseudo_recall = not args.skip_metrics and any(~ dataset.data.answer.isna())
            N = 0

            logger.info("Loading retrieval index")
            retrieval_index = faiss.read_index(
                os.path.join(output_path, 'knn.index'))
            passage_data = json.load(
                open(os.path.join(output_path, 'knn.json')))
            passage_data = OrderedDict(passage_data)
            pid2pidx = {x: i for i, x in enumerate(passage_data)}
            pidx2pid = {i: x for i, x in enumerate(passage_data)}
            logger.info(f"Done loading retrieval index with {len(passage_data)} samples")

            all_predictions = []

            for batch in tqdm(dataloader, desc='Searching'):
                batch_tgt_pids = batch.pop('passage_id', None)
                batch_qids = batch.pop('query_id')
                batch_answers = batch.pop('answer', None)
                batch.pop('text')
                batch = to_device(batch, rank)

                with torch.inference_mode():
                    ret_feats = model.get_ret_features(**batch)
                    ret_feats = F.normalize(ret_feats, p=2, dim=-1).cpu()

                batch_ret_sim, batch_ret_idxs = retrieval_index.search(ret_feats.numpy(), args.top_k)
                N += len(batch_ret_idxs)

                for i, ret_idxs in enumerate(batch_ret_idxs):
                    ret_pids = [
                        pidx2pid[x] for x in ret_idxs if x > -1]
                    if not args.skip_metrics:
                        tgt_pidxs = [pid2pidx[x]
                                            for x in batch_tgt_pids[i] if x in pid2pidx]
                        # in M2KR, some passage_ids are not unique
                        for j, k in enumerate(Ks):
                            recall_metrics[j] += int(
                                any([1 for x in tgt_pidxs if x in ret_idxs[:k]]))
                        if compute_pseudo_recall:
                            hit_list = []
                            for pid in ret_pids:
                                ret_text = passage_data[pid].lower()
                                found = False
                                for answer in batch_answers[i]:
                                    safe_answer = answer
                                    if isinstance(safe_answer, dict):
                                        safe_answer = str(safe_answer['wikidata'])
                                    if safe_answer.strip().lower() in ret_text:
                                        found = True
                                        break
                                if found:
                                    hit_list.append(1)
                                    break
                                else:
                                    hit_list.append(0)
                            for j, k in enumerate(Ks):
                                pr = float(np.max(np.array(hit_list[:k])))
                                precall_metrics[j] += pr

                    list_ret_idxs = [
                        x for x in ret_idxs.tolist() if x > -1][:10]
                    all_predictions.extend((
                        batch_qids[i],
                        idx,
                        j,
                        batch_ret_sim[i, j]
                    ) for j, idx in enumerate(list_ret_idxs))

            if not args.skip_metrics:
                recall_metrics = recall_metrics / N

                recall_dict = {f"R@{k}": v.item()
                            for k, v in zip(Ks, recall_metrics)}

                if compute_pseudo_recall:
                    precall_metrics = precall_metrics / N
                    recall_dict.update((f"PR@{k}", v.item())
                                    for k, v in zip(Ks, precall_metrics))

                # Get the maximum length for each column (header or value) for proper alignment
                column_widths = {key: max(len(key), len(f"{value:.8f}"))
                                for key, value in recall_dict.items()}

                # Print the headers (the keys as columns)
                header_row = "| " + \
                    " | ".join(
                        [f"{key:<{column_widths[key]}}" for key in recall_dict.keys()]) + " |"
                print(header_row)

                # Print the separator line for the table
                separator_row = "| " + \
                    " | ".join(["-" * column_widths[key]
                            for key in recall_dict.keys()]) + " |"
                print(separator_row)

                # Print the values (the data in a single row)
                value_row = "| " + " | ".join([f"{round(value, 5):<{column_widths[key]}.5f}".replace(
                    '.', ',') for key, value in recall_dict.items()]) + " |"
                print(value_row)

                metrics_filename = os.path.join(
                    args.index_root, args.index_model_name, args.index_dataset_name, "metrics.json")
                with open(metrics_filename, 'w') as f:
                    json.dump(recall_dict, f, indent=4)
                logger.info(f"JSON metrics saved to {metrics_filename}")
                metrics_filename = os.path.join(
                    args.index_root, args.index_model_name, args.index_dataset_name, "metrics.txt")
                with open(metrics_filename, 'w') as f:
                    f.writelines('\n'.join([header_row, separator_row, value_row]))
                    f.write('\n')
                logger.info(f"Markdown metrics saved to {metrics_filename}")
            
            ranking_filename = os.path.join(
                args.index_root, args.index_model_name, args.index_dataset_name, "ranking.tsv")
            pd.DataFrame(all_predictions).to_csv(
                ranking_filename, sep='\t', header=['qid', 'pidx', 'rank', 'score'], index=False)
            logger.info(f"TSV ranking saved to {ranking_filename}")