File tree 1 file changed +9
-8
lines changed 1 file changed +9
-8
lines changed Original file line number Diff line number Diff line change @@ -489,21 +489,22 @@ def sample(
489
489
prefix_lines = self .data [:num_prefix_lines ]
490
490
491
491
samples = []
492
- for _ in range ( num_requests ) :
492
+ while len ( samples ) < num_requests :
493
493
extra_lines = random .choices (self .data ,
494
494
k = num_input_lines - num_prefix_lines )
495
495
prompt = f"{ base_prompt } { '' .join (prefix_lines + extra_lines )} "
496
496
msg = [{"role" : "user" , "content" : prompt }]
497
497
prompt_formatted = tokenizer .apply_chat_template (
498
498
msg , add_generation_prompt = True , tokenize = False )
499
499
prompt_len = len (tokenizer (prompt_formatted ).input_ids )
500
- samples .append (
501
- SampleRequest (
502
- prompt = prompt_formatted
503
- if return_prompt_formatted else prompt ,
504
- prompt_len = prompt_len ,
505
- expected_output_len = output_len ,
506
- ))
500
+ if prompt_len <= input_len :
501
+ samples .append (
502
+ SampleRequest (
503
+ prompt = prompt_formatted
504
+ if return_prompt_formatted else prompt ,
505
+ prompt_len = prompt_len ,
506
+ expected_output_len = output_len ,
507
+ ))
507
508
return samples
508
509
509
510
You can’t perform that action at this time.
0 commit comments