More aggressive OOM prevention, num_z reduced to 50 for gsm8k and mmlu

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

@@ -90,7 +90,7 @@ def sample_generation(sentence, model, tokenizer, args,data_name):
     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'])
+        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'])
     # print(output)
     complete_generated_text = tokenizer.batch_decode(output, skip_special_tokens=True)