(response, num_rounds)
| 71 | |
| 72 | |
| 73 | def process_response(response, num_rounds): |
| 74 | output = str(response[0]["generated_text"]) |
| 75 | output = output.replace("<|endoftext|></s>", "") |
| 76 | all_positions = [m.start() for m in re.finditer("Human: ", output)] |
| 77 | place_of_second_q = -1 |
| 78 | if len(all_positions) > num_rounds: |
| 79 | place_of_second_q = all_positions[num_rounds] |
| 80 | if place_of_second_q != -1: |
| 81 | output = output[0:place_of_second_q] |
| 82 | return output |
| 83 | |
| 84 | |
| 85 | def main(args): |