"""Clear the<start>and<end>generated by the model in inference"""

import json

model_name = ''
task = 'block_completion'

result_path = f'../data/result_data/{task}/{model_name}/VersiCode_block_completion.json'  # Modify the file according to the task format


with open(result_path, 'r', encoding='utf-8') as fr:
    lodict = json.load(fr)
data_dict = lodict
data_list = data_dict

for data in data_list:
    temp_list = []
    model_output_list = eval(data['model_output'])
    for output in model_output_list:
        if '<start>' in output and '<end>' in output:
            start_index = output.find('<start>') + len('<start>')
            end_index = output.find('<end>')
            content = (
                output[start_index:end_index]
                .replace('```python', '')
                .replace('```', '')
            )
        else:
            content = 'no_answer'

        temp_list.append(content)

    data['model_output_clear'] = str(temp_list)

with open(result_path, 'w', encoding='utf-8') as fw:
    json.dump(data_dict, fw, indent=4, ensure_ascii=False)