import numpy as np
import torch
from tqdm import tqdm
from datasets import load_dataset, DatasetDict, Dataset
import datasets
import pickle
def get_top_n_docs(scores, n):
"""Return top-n document indices for a query, ignoring negative scores."""
valid_docs = np.where(scores >= 0)[0] # Filter out negative scores
sorted_indices = np.argsort(-scores[valid_docs]) # Descending order
top_n_indices = valid_docs[sorted_indices][:n] # Take top n
return set(top_n_indices)
