| 73 | |
| 74 | @TEXT_POSTPROCESSORS.register_module() |
| 75 | def answer_cleansing( |
| 76 | method: str, |
| 77 | prediction: str, |
| 78 | options: list, |
| 79 | label: str, |
| 80 | ) -> str: |
| 81 | options_str = r'\b(' + '|'.join(options) + r')\b' |
| 82 | prediction = re.findall(options_str, prediction) |
| 83 | |
| 84 | if len(prediction) == 0: |
| 85 | prediction = [] |
| 86 | else: |
| 87 | # If there is a "label" and its length is 1, |
| 88 | # process prediction accordingly |
| 89 | if len(label) == 1: |
| 90 | if method == 'few-shot': |
| 91 | answer_flag = True if len(prediction) > 1 else False |
| 92 | # choose the first or last element based on the answer_flag |
| 93 | if answer_flag: |
| 94 | prediction = [prediction[0]] |
| 95 | else: |
| 96 | prediction = [prediction[-1]] |
| 97 | elif method == 'zero-shot': |
| 98 | # choose the first element in list |
| 99 | prediction = [prediction[0]] |
| 100 | else: |
| 101 | raise ValueError('Method is not properly defined ...') |
| 102 | |
| 103 | # Remove trailing period if it exists |
| 104 | if prediction[0] and prediction[0].endswith('.'): |
| 105 | prediction[0] = prediction[0][:-1] |
| 106 | |
| 107 | return prediction[0] |