Trying to prevent OOM on mmlu and gsm8k by halving seq len

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

@@ -75,7 +75,7 @@ def calculatePerplexity(sentence, model, tokenizer, gpu):
         all_prob.append(probability)
     return torch.exp(loss).item(), all_prob, loss.item()
 
-def sample_generation(sentence, model, tokenizer, args):
+def sample_generation(sentence, model, tokenizer, args,data_name):
     half_sentence_index = math.ceil(len(sentence.split())*args['prefix_length'])
 
     if half_sentence_index > 0:
@@ -86,7 +86,11 @@ def sample_generation(sentence, model, tokenizer, args):
     input_ids = torch.tensor(tokenizer.encode(prefix)).unsqueeze(0)
     input_ids = input_ids.to(model.device)
 
-    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'])
+    output = None
+    if data_name != "cais/mmlu" or 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=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)
  
@@ -98,9 +102,9 @@ def RMIA_1(text,target_loss,ref_loss,model1,tokenizer1,ratio_gen,neighbors_dl):
     result = torch.count_nonzero(target_losses_z < target_loss).item() / len(target_losses_z)
     return result
 
-def get_neighbors(text,ref_loss,model2,tokenizer2,ratio_gen):
+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}}
-    neighbors = sample_generation(text, model2, tokenizer2, cur_args)
+    neighbors = sample_generation(text, model2, tokenizer2, cur_args,data_name)
     neighbors_dl = DataLoader(neighbors, batch_size=32, shuffle=False)
     return neighbors_dl
 
@@ -134,7 +138,7 @@ def evaluate_data(test_data, col_name, target_model, ref_model, ratio_gen, data_
         counter = 0
         for ex in tqdm(test_data):
             text = ex[col_name]
-            new_ex = get_neighbors(text,inference2_pass[counter][2],model2,tokenizer2,ratio_gen)
+            new_ex = get_neighbors(text,inference2_pass[counter][2],model2,tokenizer2,ratio_gen,data_name)
             counter = counter + 1
             neighbors_dls.append(new_ex)
         unload_model(model2,tokenizer2)