| 196 | |
| 197 | |
| 198 | def parse_input_token_extra_ids(prompt_table_path, kv_cache_enable_block_reuse, |
| 199 | input_token_extra_ids, |
| 200 | input_token_extra_ids_file, max_input_length): |
| 201 | batch_extra_ids = None |
| 202 | if prompt_table_path and kv_cache_enable_block_reuse: |
| 203 | assert input_token_extra_ids or input_token_extra_ids_file, \ |
| 204 | "Input token extra ids must be provided when p-tuning and KV Cache reuse are both enabled" |
| 205 | batch_extra_ids = [] |
| 206 | if input_token_extra_ids_file: |
| 207 | if input_token_extra_ids_file.endswith('.csv'): |
| 208 | with open(input_token_extra_ids_file, 'r') as csv_file: |
| 209 | csv_reader = csv.reader(csv_file, delimiter=',') |
| 210 | for line in csv_reader: |
| 211 | extra_ids = [int(num) for num in line] |
| 212 | batch_extra_ids.append(extra_ids[-max_input_length:]) |
| 213 | elif input_token_extra_ids_file.endswith('.npy'): |
| 214 | inputs = np.load(input_token_extra_ids_file) |
| 215 | for extra_ids in inputs: |
| 216 | batch_extra_ids.append(extra_ids[-max_input_length:]) |
| 217 | else: |
| 218 | print('Input file format not supported.') |
| 219 | raise SystemExit |
| 220 | else: |
| 221 | batch_extra_ids.append(input_token_extra_ids) |
| 222 | return batch_extra_ids |
| 223 | |
| 224 | |
| 225 | def print_output(tokenizer, |