pxyyy/NuminaMath-CoT-smp20k-removed-top500-by-logix-for-MATH-Correct-2k

import numpy as np
import torch
from tqdm import tqdm
from datasets import load_dataset, DatasetDict, Dataset
import datasets

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)

def calculate_overlap(score, n):
    num_queries = score.shape[0]
    total_overlap = 0
    top_docs=[]
    for q in tqdm(range(num_queries)):
        top_docs.append(get_top_n_docs(score[q], n))
    common_docs = set.intersection(*top_docs)
    total_overlap += len(common_docs)
    
    return total_overlap 


def top_k_idx(score, k):
    """Return top-k document indices for a query, ignoring negative scores."""
    doc_contributions = np.sum(score, axis=0)
    # print(doc_contributions)
    sorted_indices = np.argsort(-doc_contributions)  # Descending order
    top_k_indices = sorted_indices[:k]    # Take top k
    return top_k_indices

if __name__ == "__main__":
    
    remove_ks=[500, 1000, 1500, 2000]

    score=torch.load('/shared/storage-01/xp12/mp-llm/logix/examples/language_modeling/save-MATH-Correct-2k/scores.pt').cpu().numpy()
    raw_dataset = load_dataset('pxyyy/NuminaMath-CoT-smp20k', split='train')

    # dedplicate the dataset on a new column (problem + solution)
    raw_dataset = raw_dataset.map(lambda x: {'problem': x['problem'], 'solution': x['solution'], 'text': x['problem'] + ' ' + x['solution']})
    
    seen_texts = set()
    keep_indices = []

    for idx, example in enumerate(raw_dataset):
        text = example["text"]
        if text not in seen_texts:
            seen_texts.add(text)
            keep_indices.append(idx)

    raw_dataset = raw_dataset.select(keep_indices)
    print(raw_dataset)
    assert len(raw_dataset) == score.shape[1]

    for k in remove_ks:
        top_k_indices = top_k_idx(score, k)
        # print(top_k_indices)

        removed_topk = set(range(len(raw_dataset))) - set(top_k_indices)

        removed_top_k_docs = raw_dataset.select(removed_topk)

        print(removed_top_k_docs)
        removed_top_k_docs.push_to_hub(f'pxyyy/NuminaMath-CoT-smp20k-removed-top{k}-by-logix-for-MATH-Correct-2k')