collect_result.py

import math
import zlib
import json
import math
import torch
import argparse
from tqdm import tqdm
from pathlib import Path
import concurrent.futures
import torch, torch.nn.functional as F
from transformers import AutoTokenizer, AutoModelForCausalLM
from deepscaler.rewards.math_reward import deepscaler_reward_fn
from deepscaler.rewards.math_utils.utils import get_multiple_choice_answer
from utils import *
from datasets import load_from_disk

parser = argparse.ArgumentParser()
parser.add_argument("--results_path", default="./results_member", type=str)
parser.add_argument("--model_name", default= "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B", help="model name or model path")
parser.add_argument("--ref_model_name", default="bespokelabs/Bespoke-Stratos-7B", help="reference model name or path")
parser.add_argument("--dataset_name", default="olympiadbench", choices=["amc23", "aime24", "aime25", "olympiadbench", "minerva_math", "gpqa_diamond",])
parser.add_argument("--perturb_dataset_path", default="./datasets/perturbed_datasets", type=str)
parser.add_argument("--skip_detection", action='store_true', help="if we apply the detection approach")
parser.add_argument("--mia", default="loss", type=str)
parser.add_argument("--min_k", default=0.2, type=float)
parser.add_argument("--global_size", default=1, type=int)
parser.add_argument("--sharding", default=1, type=int)
parser.add_argument("--num_responses", default=1, type=int)
args = parser.parse_args()

loss_dir = Path("{}/{}/{}/{}".format(
            args.results_path, args.model_name.split("/")[-1], args.dataset_name, "loss"
        ))
jsonl_paths = sorted(loss_dir.glob("*.jsonl"))
records = []
for path in jsonl_paths:
    with path.open("r", encoding="utf-8") as f:
        for line in f:
            if line.strip():
                records.append(json.loads(line))
                
if args.sharding == 1:
    total_responses = sum(len(record["response"]) for record in records)
    print("\n============")
    print(f"Loaded {total_responses:,} total responses from {len(jsonl_paths)} shard(s).")

correct_token_len = 0; less_16k = 0; total_num = 0; token_len = 0; acc = 0
for record in records:
    answer = record["answer"]
    for token_ids, response in zip(record["token_ids"], record["response"]):
        if "gpqa_diamond" in args.dataset_name:
            acc += (get_multiple_choice_answer(response) == answer)
            if get_multiple_choice_answer(response) == answer:
                correct_token_len += len(token_ids)
        else:
            acc += deepscaler_reward_fn(solution_str=response, ground_truth=answer)
            if deepscaler_reward_fn(solution_str=response, ground_truth=answer):
                correct_token_len += len(token_ids)
                    
        if len(token_ids)<=16384:
            less_16k += 1
            
        token_len += len(token_ids)
        total_num += 1

if args.sharding == 1:
    print("Dataset: {} | Model path: {}".format(args.dataset_name,  args.model_name))
    print(" Average Acc: {:.2f}% | Ratio of less than 16k: {:.2f}% | Correct avg token length: {} | Wrong avg token length: {} | Average response token length: {} | ".format(
                acc/total_num*100, less_16k/total_num*100, correct_token_len/(acc+1e-3), (token_len-correct_token_len)/(total_num-acc+1e-3), token_len/total_num,
            ))
    print("============")

if args.skip_detection or args.mia=="distillation":
    exit()

records = group_records_by_prompt(records)
if args.sharding == 1:
    print("Len of records after the group: ", len(records))
records = prepare_sharding_records_for_mia(records, args)

if args.mia == "loss":
    tokenizer = AutoTokenizer.from_pretrained(args.model_name)
    model  = AutoModelForCausalLM.from_pretrained(
                                        args.model_name, dtype=torch.bfloat16,
                                        attn_implementation="flash_attention_2",
                                        device_map="auto"
                                    ).eval()

    loss_value = []; min_k_value = []; max_k_value = []; zlib_value = []; labels = []
    for record in tqdm(records):
        loss_r = 0; zlib_r= 0; min_k_r = 0;  max_k_r = 0
        for i, (response, response_ids) in enumerate(zip(record["response"], record['token_ids'])):
            if i == args.num_responses:
                break
            
            prompt_response = [record["prompt"]+response]
            ids = tokenizer(prompt_response, return_tensors="pt", padding=True,
                            truncation=True, max_length=32768).input_ids.cuda()
     
            with torch.inference_mode():
                logits = model(ids).logits 
                all_logprob = F.log_softmax(logits, dim=-1)
            
            target_ids  = ids[:, 1:]
            tok_lp = all_logprob.gather(2, target_ids.unsqueeze(-1)).squeeze(-1) 
            response_logprob = tok_lp[0, len(record['prompt_token_ids'])-1:len(record['prompt_token_ids'])-1+len(response_ids)]

            # Loss
            loss_r += response_logprob.mean(dim=0)
        
            # zlip
            compression_r = zlib.compress(response.encode('utf-8')) 
            zlib_r += response_logprob.mean(dim=0)/len(compression_r)
    
            # min-k on ques
            min_k_r += torch.mean(torch.topk(response_logprob, math.ceil(args.min_k*len(response_logprob)), largest=False).values)
            max_k_r += torch.mean(torch.topk(response_logprob, math.ceil(args.min_k*len(response_logprob)), largest=True).values)
            del response_logprob, tok_lp,  all_logprob, logits,

        loss_r /= max(i, args.num_responses); zlib_r /= max(i, args.num_responses)
        min_k_r /= max(i, args.num_responses); max_k_r /= max(i, args.num_responses)

        labels.append(record["membership"])
        loss_value.append(loss_r); zlib_value.append(zlib_r)
        min_k_value.append(min_k_r); max_k_value.append(max_k_r)
        
    print("DATASET: ", args.dataset_name, " | MODEL: ", args.model_name)
    save_sharding_records_result(labels, loss_value, "LOSS", args)
    save_sharding_records_result(labels, zlib_value, "ZLIB", args)
    save_sharding_records_result(labels, min_k_value, "MIN-K", args)
    save_sharding_records_result(labels, max_k_value, "MAX-K", args)

elif args.mia == "min-k++":
    tokenizer = AutoTokenizer.from_pretrained(args.model_name)
    model  = AutoModelForCausalLM.from_pretrained(
                                        args.model_name, dtype=torch.bfloat16,
                                        attn_implementation="flash_attention_2",
                                        device_map="auto"
                                    ).eval() 
    
    min_k_plus2_value = []; labels = []
    for record in tqdm(records):
        min_k_plus2_r = 0
        for i, (response, response_ids) in enumerate(zip(record["response"], record['token_ids'])):
            if i == args.num_responses:
                break
            
            prompt_response = [record["prompt"]+response]
            ids = tokenizer(prompt_response, return_tensors="pt", padding=True,
                            truncation=True, max_length=32768).input_ids.cuda()
            
            with torch.inference_mode():
                logits = model(ids).logits

                all_probs = F.softmax(logits, dim=-1)
                all_logprob = F.log_softmax(logits, dim=-1)
                all_logprob_mean = (all_probs * all_logprob).sum(-1)
            
                # reuse probs to hold p log²p to save memory
                all_probs.mul_(all_logprob)            
                all_probs.mul_(all_logprob)            
                all_logprob_std = (all_probs.sum(-1) - all_logprob_mean.pow(2)).sqrt()
            
            target_ids  = ids[:, 1:]
            tok_lp = all_logprob.gather(2, target_ids.unsqueeze(-1)).squeeze(-1) 
            prompt_logprob = tok_lp[0, :len(record['prompt_token_ids'])-1]
            response_logprob = tok_lp[0, len(record['prompt_token_ids'])-1:len(record['prompt_token_ids'])-1+len(response_ids)]
        
            # min-k_plus plus on ques
            norm_resp_logprob = (response_logprob - all_logprob_mean[0, prompt_logprob.numel():prompt_logprob.numel()+len(response_logprob)]) \
                                / all_logprob_std[0, prompt_logprob.numel():prompt_logprob.numel()+len(response_logprob)]
            min_k_plus2_r += torch.mean(torch.topk(norm_resp_logprob, math.ceil(args.min_k*len(norm_resp_logprob)), largest=False).values)
          
            del response_logprob, tok_lp, all_logprob, norm_resp_logprob, prompt_logprob, all_logprob_mean, all_logprob_std, logits, all_probs
          
        min_k_plus2_r /= max(i, args.num_responses)
        labels.append(record["membership"])
        min_k_plus2_value.append(min_k_plus2_r)

    print("DATASET: ", args.dataset_name, " | MODEL: ", args.model_name)
    save_sharding_records_result(labels, min_k_plus2_value, "MIN-K++", args)

elif args.mia == "ref":
    assert torch.cuda.device_count() >= 2
    
    tokenizer = AutoTokenizer.from_pretrained(args.model_name)
    model  = AutoModelForCausalLM.from_pretrained(
                                        args.model_name, dtype=torch.bfloat16,
                                        attn_implementation="flash_attention_2",
                                        device_map="cuda:0"
                                    ).eval()
    if args.sharding == 1:
        print("Reference model name: ", args.ref_model_name)
    ref_model  = AutoModelForCausalLM.from_pretrained(
                                        args.ref_model_name, dtype=torch.bfloat16,
                                        attn_implementation="flash_attention_2",
                                        device_map="cuda:1"
                                    ).eval()
    ref_tokenizer = AutoTokenizer.from_pretrained(args.ref_model_name)
    
    ref_value = []; lira_value = []; labels = []
    for record in tqdm(records):
        ref_r = 0;  lira_r = 0
        for i, (response, response_ids) in enumerate(zip(record["response"], record['token_ids'])):
            if i == args.num_responses:
                break
            prompt_response = [record["prompt"]+response]
            ids = tokenizer(prompt_response, return_tensors="pt", padding=True,
                            truncation=True, max_length=32768).input_ids
            
            ref_prompt = prepare_reference_model_prompt(args.ref_model_name, args.model_name, record["prompt"], response)
            ref_prompt_response = [ref_prompt + response]

            ref_ids = ref_tokenizer(ref_prompt_response, return_tensors="pt", padding=True,
                            truncation=True, max_length=32768).input_ids
            ref_prompt_token_ids = ref_tokenizer(ref_prompt)["input_ids"]
           
            def run_embedding(model, ids):
                with torch.inference_mode():
                    ids = ids.to(model.device)
                    logits = model(ids).logits
                    all_logprob = F.log_softmax(logits, dim=-1)
                return all_logprob
            
            with concurrent.futures.ThreadPoolExecutor(max_workers=4) as pool:
                model_return = pool.submit(run_embedding, model, ids)
                ref_model_return = pool.submit(run_embedding, ref_model, ref_ids)

                all_logprob = model_return.result()
                ref_all_logprob = ref_model_return.result()
            
            ids = ids.to("cuda:1"); ref_ids = ref_ids.to("cuda:1"); all_logprob = all_logprob.to("cuda:1")
            
            target_ids  = ids[:, 1:]
            ref_target_ids  = ref_ids[:, 1:]
            tok_lp = all_logprob.gather(2, target_ids.unsqueeze(-1)).squeeze(-1) 
            ref_tok_lp = ref_all_logprob.gather(2, ref_target_ids.unsqueeze(-1)).squeeze(-1) 
          
            response_logprob = tok_lp[0, len(record['prompt_token_ids'])-1:]
            ref_response_logprob = ref_tok_lp[0, len(ref_prompt_token_ids)-1:]
          
            ref_r += torch.mean(response_logprob) - torch.mean(ref_response_logprob)
            lira_r += torch.mean(ref_response_logprob)/torch.mean(response_logprob)
            
            del response_logprob, ref_response_logprob, tok_lp, ref_tok_lp, all_logprob, ref_all_logprob, model_return, ref_model_return
            
        ref_r /= max(i, args.num_responses); lira_r /= max(i, args.num_responses)
        labels.append(record["membership"])
        ref_value.append(ref_r); lira_value.append(lira_r)
    
    print("DATASET: ", args.dataset_name, " | MODEL: ", args.model_name)
    save_sharding_records_result(labels, ref_value, "REF", args)
    save_sharding_records_result(labels, lira_value, "LiRA", args)

elif args.mia == "perturb":
    tokenizer = AutoTokenizer.from_pretrained(args.model_name)
    model  = AutoModelForCausalLM.from_pretrained(
                                        args.model_name, dtype=torch.bfloat16,
                                        attn_implementation="flash_attention_2",
                                        device_map="auto"
                                    ).eval()
    dataset_name = args.dataset_name[:-len("_shuffle")] if "shuffle" in args.dataset_name else args.dataset_name
    perturb_ds = load_from_disk(f"{args.perturb_dataset_path}/{dataset_name}_synonyms_results_variations_dataset")
    loss_value = []; labels = []
    for record in tqdm(records):
        neighbor_r = 0
        for i, (response, response_ids) in enumerate(zip(record["response"], record['token_ids'])):
            if i == args.num_responses:
                break
            # Inference for the original prompts
            prompt_response = [record["prompt"]+response]
            ids = tokenizer(prompt_response, return_tensors="pt", padding=True,
                            truncation=True, max_length=32768).input_ids.cuda()
            
            with torch.inference_mode():
                logits = model(ids).logits 
                all_logprob = F.log_softmax(logits, dim=-1)
        
            target_ids  = ids[:, 1:]  
            tok_lp = all_logprob.gather(2, target_ids.unsqueeze(-1)).squeeze(-1) 
            response_logprob = tok_lp[0, len(record['prompt_token_ids'])-1:len(record['prompt_token_ids'])-1+len(response_ids)]
            loss_r = response_logprob.mean(dim=0)
       
            del response_logprob, tok_lp,  all_logprob, logits, ids, target_ids

            perturb_loss_r = 0
            # Inference for the perturbed prompts
            prompt_no_template = extract_user_instruction(record["prompt"], args.model_name)
            perturb_ds_index = perturb_ds["question_with_template"].index(prompt_no_template)
            neighbors = perturb_ds['variations_with_template'][perturb_ds_index]        
            neighbors = neighbors[:5]
            
            for neighbor in neighbors:

                perturb_prompt_response = prepare_perturbed_prompt_response(neighbor, args.model_name, response)
                neighbors_ids = tokenizer(perturb_prompt_response, return_tensors="pt", padding=True, truncation=True, max_length=32768).input_ids.cuda()
                perturb_prompt_input_ids = prepare_perturbed_prompt_token_ids(tokenizer, neighbor, args.model_name)
                
                with torch.inference_mode():
                    logits = model(neighbors_ids).logits 
                    all_logprob = F.log_softmax(logits, dim=-1)
                
                target_ids  = neighbors_ids[:, 1:]  
                tok_lp = all_logprob.gather(2, target_ids.unsqueeze(-1)).squeeze(-1) 
                
                response_logprob = tok_lp[0, len(perturb_prompt_input_ids)-1:len(perturb_prompt_input_ids)-1+len(response_ids)]
                perturb_loss_r += response_logprob.mean(dim=0)
                del response_logprob, tok_lp,  all_logprob, logits, neighbors_ids, target_ids
        
            torch.cuda.empty_cache()
            perturb_loss_r /= len(neighbors)
            neighbor_r += (loss_r - perturb_loss_r)

        labels.append(record["membership"])
        loss_value.append(neighbor_r/max(i, args.num_responses)) 
            
    print("DATASET: ", args.dataset_name, " | MODEL: ", args.model_name)
    save_sharding_records_result(labels, loss_value, "PERTURB", args)

else:
    assert NotImplementedError