wehub-resource-sync
229 lines
8.3 KiB
Python
229 lines
8.3 KiB
Python
import faiss
|
|
import torch
|
|
import logging
|
|
import numpy as np
|
|
import pytrec_eval
|
|
from tqdm import tqdm
|
|
from collections import defaultdict
|
|
from typing import Dict, List, Tuple, Optional
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
# Modified from https://github.com/beir-cellar/beir/blob/f062f038c4bfd19a8ca942a9910b1e0d218759d4/beir/retrieval/custom_metrics.py#L4
|
|
def evaluate_mrr(
|
|
qrels: Dict[str, Dict[str, int]],
|
|
results: Dict[str, Dict[str, float]],
|
|
k_values: List[int],
|
|
) -> Tuple[Dict[str, float]]:
|
|
"""Compute mean reciprocal rank (MRR).
|
|
|
|
Args:
|
|
qrels (Dict[str, Dict[str, int]]): Ground truth relevance.
|
|
results (Dict[str, Dict[str, float]]): Search results to evaluate.
|
|
k_values (List[int]): Cutoffs.
|
|
|
|
Returns:
|
|
Tuple[Dict[str, float]]: MRR results at provided k values.
|
|
"""
|
|
mrr = defaultdict(list)
|
|
|
|
k_max, top_hits = max(k_values), {}
|
|
|
|
for query_id, doc_scores in results.items():
|
|
top_hits[query_id] = sorted(
|
|
doc_scores.items(), key=lambda item: item[1], reverse=True
|
|
)[0:k_max]
|
|
|
|
for query_id in top_hits:
|
|
query_relevant_docs = {
|
|
doc_id for doc_id in qrels[query_id] if qrels[query_id][doc_id] > 0
|
|
}
|
|
for k in k_values:
|
|
rr = 0
|
|
for rank, hit in enumerate(top_hits[query_id][0:k], 1):
|
|
if hit[0] in query_relevant_docs:
|
|
rr = 1.0 / rank
|
|
break
|
|
mrr[f"MRR@{k}"].append(rr)
|
|
|
|
for k in k_values:
|
|
mrr[f"MRR@{k}"] = round(sum(mrr[f"MRR@{k}"]) / len(qrels), 5)
|
|
return mrr
|
|
|
|
|
|
# Modified from https://github.com/beir-cellar/beir/blob/f062f038c4bfd19a8ca942a9910b1e0d218759d4/beir/retrieval/custom_metrics.py#L33
|
|
def evaluate_recall_cap(
|
|
qrels: Dict[str, Dict[str, int]],
|
|
results: Dict[str, Dict[str, float]],
|
|
k_values: List[int]
|
|
) -> Tuple[Dict[str, float]]:
|
|
"""Compute capped recall.
|
|
|
|
Args:
|
|
qrels (Dict[str, Dict[str, int]]): Ground truth relevance.
|
|
results (Dict[str, Dict[str, float]]): Search results to evaluate.
|
|
k_values (List[int]): Cutoffs.
|
|
|
|
Returns:
|
|
Tuple[Dict[str, float]]: Capped recall results at provided k values.
|
|
"""
|
|
capped_recall = {}
|
|
|
|
for k in k_values:
|
|
capped_recall[f"R_cap@{k}"] = 0.0
|
|
|
|
k_max = max(k_values)
|
|
logging.info("\n")
|
|
|
|
for query_id, doc_scores in results.items():
|
|
top_hits = sorted(doc_scores.items(), key=lambda item: item[1], reverse=True)[0:k_max]
|
|
query_relevant_docs = [doc_id for doc_id in qrels[query_id] if qrels[query_id][doc_id] > 0]
|
|
for k in k_values:
|
|
retrieved_docs = [row[0] for row in top_hits[0:k] if qrels[query_id].get(row[0], 0) > 0]
|
|
denominator = min(len(query_relevant_docs), k)
|
|
capped_recall[f"R_cap@{k}"] += (len(retrieved_docs) / denominator)
|
|
|
|
for k in k_values:
|
|
capped_recall[f"R_cap@{k}"] = round(capped_recall[f"R_cap@{k}"]/len(qrels), 5)
|
|
logging.info("R_cap@{}: {:.4f}".format(k, capped_recall[f"R_cap@{k}"]))
|
|
|
|
return capped_recall
|
|
|
|
|
|
# Modified from https://github.com/embeddings-benchmark/mteb/blob/18f730696451a5aaa026494cecf288fd5cde9fd0/mteb/evaluation/evaluators/RetrievalEvaluator.py#L501
|
|
def evaluate_metrics(
|
|
qrels: Dict[str, Dict[str, int]],
|
|
results: Dict[str, Dict[str, float]],
|
|
k_values: List[int],
|
|
) -> Tuple[
|
|
Dict[str, float],
|
|
Dict[str, float],
|
|
Dict[str, float],
|
|
Dict[str, float],
|
|
]:
|
|
"""Evaluate the main metrics.
|
|
|
|
Args:
|
|
qrels (Dict[str, Dict[str, int]]): Ground truth relevance.
|
|
results (Dict[str, Dict[str, float]]): Search results to evaluate.
|
|
k_values (List[int]): Cutoffs.
|
|
|
|
Returns:
|
|
Tuple[ Dict[str, float], Dict[str, float], Dict[str, float], Dict[str, float], ]: Results of different metrics at
|
|
different provided k values.
|
|
"""
|
|
all_ndcgs, all_aps, all_recalls, all_precisions = defaultdict(list), defaultdict(list), defaultdict(list), defaultdict(list)
|
|
|
|
map_string = "map_cut." + ",".join([str(k) for k in k_values])
|
|
ndcg_string = "ndcg_cut." + ",".join([str(k) for k in k_values])
|
|
recall_string = "recall." + ",".join([str(k) for k in k_values])
|
|
precision_string = "P." + ",".join([str(k) for k in k_values])
|
|
evaluator = pytrec_eval.RelevanceEvaluator(
|
|
qrels, {map_string, ndcg_string, recall_string, precision_string}
|
|
)
|
|
scores = evaluator.evaluate(results)
|
|
|
|
for query_id in scores.keys():
|
|
for k in k_values:
|
|
all_ndcgs[f"NDCG@{k}"].append(scores[query_id]["ndcg_cut_" + str(k)])
|
|
all_aps[f"MAP@{k}"].append(scores[query_id]["map_cut_" + str(k)])
|
|
all_recalls[f"Recall@{k}"].append(scores[query_id]["recall_" + str(k)])
|
|
all_precisions[f"P@{k}"].append(scores[query_id]["P_" + str(k)])
|
|
|
|
ndcg, _map, recall, precision = (
|
|
all_ndcgs.copy(),
|
|
all_aps.copy(),
|
|
all_recalls.copy(),
|
|
all_precisions.copy(),
|
|
)
|
|
|
|
for k in k_values:
|
|
ndcg[f"NDCG@{k}"] = round(sum(ndcg[f"NDCG@{k}"]) / len(scores), 5)
|
|
_map[f"MAP@{k}"] = round(sum(_map[f"MAP@{k}"]) / len(scores), 5)
|
|
recall[f"Recall@{k}"] = round(sum(recall[f"Recall@{k}"]) / len(scores), 5)
|
|
precision[f"P@{k}"] = round(sum(precision[f"P@{k}"]) / len(scores), 5)
|
|
|
|
return ndcg, _map, recall, precision
|
|
|
|
|
|
def index(
|
|
index_factory: str = "Flat",
|
|
corpus_embeddings: Optional[np.ndarray] = None,
|
|
load_path: Optional[str] = None,
|
|
device: Optional[str] = None
|
|
):
|
|
"""Create and add embeddings into a Faiss index.
|
|
|
|
Args:
|
|
index_factory (str, optional): Type of Faiss index to create. Defaults to "Flat".
|
|
corpus_embeddings (Optional[np.ndarray], optional): The embedding vectors of the corpus. Defaults to None.
|
|
load_path (Optional[str], optional): Path to load embeddings from. Defaults to None.
|
|
device (Optional[str], optional): Device to hold Faiss index. Defaults to None.
|
|
|
|
Returns:
|
|
faiss.Index: The Faiss index that contains all the corpus embeddings.
|
|
"""
|
|
if corpus_embeddings is None:
|
|
corpus_embeddings = np.load(load_path)
|
|
|
|
logger.info(f"Shape of embeddings: {corpus_embeddings.shape}")
|
|
# create faiss index
|
|
logger.info(f'Indexing {corpus_embeddings.shape[0]} documents...')
|
|
faiss_index = faiss.index_factory(corpus_embeddings.shape[-1], index_factory, faiss.METRIC_INNER_PRODUCT)
|
|
|
|
if device is None and torch.cuda.is_available():
|
|
try:
|
|
co = faiss.GpuMultipleClonerOptions()
|
|
co.shard = True
|
|
co.useFloat16 = True
|
|
faiss_index = faiss.index_cpu_to_all_gpus(faiss_index, co)
|
|
except:
|
|
print('faiss do not support GPU, please uninstall faiss-cpu, faiss-gpu and install faiss-gpu again.')
|
|
|
|
logger.info('Adding embeddings ...')
|
|
corpus_embeddings = corpus_embeddings.astype(np.float32)
|
|
faiss_index.train(corpus_embeddings)
|
|
faiss_index.add(corpus_embeddings)
|
|
logger.info('Embeddings add over...')
|
|
return faiss_index
|
|
|
|
|
|
def search(
|
|
faiss_index: faiss.Index,
|
|
k: int = 100,
|
|
query_embeddings: Optional[np.ndarray] = None,
|
|
load_path: Optional[str] = None
|
|
):
|
|
"""
|
|
1. Encode queries into dense embeddings;
|
|
2. Search through faiss index
|
|
|
|
Args:
|
|
faiss_index (faiss.Index): The Faiss index that contains all the corpus embeddings.
|
|
k (int, optional): Top k numbers of closest neighbours. Defaults to :data:`100`.
|
|
query_embeddings (Optional[np.ndarray], optional): The embedding vectors of queries. Defaults to :data:`None`.
|
|
load_path (Optional[str], optional): Path to load embeddings from. Defaults to :data:`None`.
|
|
|
|
Returns:
|
|
Tuple[np.ndarray, np.ndarray]: The scores of search results and their corresponding indices.
|
|
"""
|
|
if query_embeddings is None:
|
|
query_embeddings = np.load(load_path)
|
|
|
|
query_size = len(query_embeddings)
|
|
|
|
all_scores = []
|
|
all_indices = []
|
|
|
|
for i in tqdm(range(0, query_size, 32), desc="Searching"):
|
|
j = min(i + 32, query_size)
|
|
query_embedding = query_embeddings[i: j]
|
|
score, indice = faiss_index.search(query_embedding.astype(np.float32), k=k)
|
|
all_scores.append(score)
|
|
all_indices.append(indice)
|
|
|
|
all_scores = np.concatenate(all_scores, axis=0)
|
|
all_indices = np.concatenate(all_indices, axis=0)
|
|
return all_scores, all_indices
|