| 310 | return res, lines |
| 311 | |
| 312 | def wrap_conv(self, data: list[dict]): # add |
| 313 | lines = [] |
| 314 | res = [] |
| 315 | for line in data: |
| 316 | # print(line) |
| 317 | collated, partial_qa = self.conv_collater(line) |
| 318 | # collated: ['Q', 'QAQ', 'QAQAQ', ...] |
| 319 | # partial_qa: [ |
| 320 | # [{'q': 'q'}], |
| 321 | # [{'q': 'q', 'a': 'a'}, {'q'}], |
| 322 | # [{'q': 'q', 'a': 'a'}, {'q': 'q', 'a': 'a'}, {'q': 'q'}] |
| 323 | # ] |
| 324 | res.extend(collated) # 1d list |
| 325 | lines.extend(partial_qa) |
| 326 | return res, lines |
| 327 | |
| 328 | def unwrap(self, outputs, num_return_sequences): |
| 329 | batch_return = [] |