(args)
| 172 | return results |
| 173 | |
| 174 | def generate_response(args): |
| 175 | path = args.input_path |
| 176 | path_save = args.output_path |
| 177 | block_size = args.block_size |
| 178 | model_name = args.model_name |
| 179 | tensor_parallel_size = args.tensor_parallel_size |
| 180 | enable_thinking = args.enable_thinking |
| 181 | print('path:', path) |
| 182 | dataset = construct_dataset(path, args.train_size) |
| 183 | model = LLM(model_name, gpu_memory_utilization=0.95, tensor_parallel_size=tensor_parallel_size, |
| 184 | enable_expert_parallel=False, max_model_len=32768) |
| 185 | sampling_params = SamplingParams(max_tokens=32768, temperature=0.7, top_p=0.8, top_k=20, min_p=0) |
| 186 | print('path_save:', path_save) |
| 187 | processed_ids = load_processed_ids(path_save) |
| 188 | messages, samples = [], [] |
| 189 | count = 0 |
| 190 | for sample in tqdm(dataset): |
| 191 | if sample['id_ddm'] in processed_ids: |
| 192 | continue |
| 193 | text = sample['text'] |
| 194 | prompt = build_prompt(text) |
| 195 | messages.append([ |
| 196 | {"role": "user", "content": prompt} |
| 197 | ]) |
| 198 | |
| 199 | samples.append(sample) |
| 200 | if len(samples) >= block_size: |
| 201 | responses = model.chat(messages, sampling_params=sampling_params, |
| 202 | chat_template_kwargs={"enable_thinking": enable_thinking}) |
| 203 | outputs = [r.outputs[0].text for r in responses] |
| 204 | for sample, output in zip(samples, outputs): |
| 205 | think_text = sample['text'] |
| 206 | |
| 207 | pred_result = parse_and_locate_spans(output, think_text) |
| 208 | if pred_result is None: |
| 209 | continue |
| 210 | count+= 1 |
| 211 | sample_output = { |
| 212 | 'id_ddm': sample['id_ddm'], |
| 213 | 'think_text': think_text, |
| 214 | 'pred_result': pred_result |
| 215 | } |
| 216 | append_jsonl(path_save, sample_output) |
| 217 | messages, samples = [], [] |
| 218 | |
| 219 | if len(messages) > 0: |
| 220 | responses = model.chat(messages, sampling_params=sampling_params, |
| 221 | chat_template_kwargs={"enable_thinking": enable_thinking}) |
| 222 | outputs = [r.outputs[0].text for r in responses] |
| 223 | for sample, output in zip(samples, outputs): |
| 224 | think_text = sample['text'] |
| 225 | pred_result = parse_and_locate_spans(output, think_text) |
| 226 | if pred_result is None: |
| 227 | continue |
| 228 | sample_output = { |
| 229 | 'id_ddm': sample['id_ddm'], |
| 230 | 'think_text': think_text, |
| 231 | 'pred_result': pred_result |
no test coverage detected