num_z = 50 for all

detect-pretrain-code-contamination/src/run.py

@@ -86,11 +86,7 @@ def sample_generation(sentence, model, tokenizer, args,data_name):
     input_ids = torch.tensor(tokenizer.encode(prefix)).unsqueeze(0)
     input_ids = input_ids.to(model.device)
 
-    output = None
-    if data_name != "cais/mmlu" and data_name != "gsm8k":
-        output = model.generate(input_ids, max_new_tokens=len(sentence.split())-half_sentence_index, min_new_tokens=1, num_return_sequences=args['num_z'], pad_token_id=tokenizer.eos_token_id, **args['generate_args'])
-    else:
-        output = model.generate(input_ids, max_new_tokens=(len(sentence.split())-half_sentence_index)/2, min_new_tokens=1, num_return_sequences=int(args['num_z']/2), pad_token_id=tokenizer.eos_token_id, **args['generate_args'])
+    output = model.generate(input_ids, max_new_tokens=(len(sentence.split())-half_sentence_index), min_new_tokens=1, num_return_sequences=int(args['num_z']), pad_token_id=tokenizer.eos_token_id, **args['generate_args'])
     # print(output)
     complete_generated_text = tokenizer.batch_decode(output, skip_special_tokens=True)
  
@@ -103,7 +99,7 @@ def RMIA_1(text,target_loss,ref_loss,model1,tokenizer1,ratio_gen,neighbors_dl):
     return result
 
 def get_neighbors(text,ref_loss,model2,tokenizer2,ratio_gen,data_name):
-    cur_args = {'prefix_length': ratio_gen, 'num_z': 100, 'generate_args': {'do_sample': True}}
+    cur_args = {'prefix_length': ratio_gen, 'num_z': 50, 'generate_args': {'do_sample': True}}
     neighbors = sample_generation(text, model2, tokenizer2, cur_args,data_name)
     neighbors_dl = DataLoader(neighbors, batch_size=32, shuffle=False)
     return neighbors_dl