WebExperT/prepare_test.py at master · Luohh5/WebExperT · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137

import json
import pdb
import jsonlines
import os
from tqdm import tqdm
import sys
sys.path.insert(0, '../')
# from SeeAct.src.data_utils.format_prompt_utils import get_choices

def build_example(problem, id):
    answer = 'Answer: ' + problem['choices'][problem['answer']] + '\n'
    question = problem['question']
    rationale = problem['lecture'] + problem['solution']
    choices = problem['choices'].copy()
    choices.remove(problem['choices'][problem['answer']])
    distractors = ''
    for i, choice in enumerate(choices):
        distractors += f'({i + 1}) {choice}\n'
    if problem['image'] != None:
        image = 'Picture: <img>' + \
                os.path.join(data_root, problem['split'], str(id), problem[
                    'image']) + f'</img>\n'
        context = f'Context: ' + problem['hint'] + '\n' if problem['hint'] != '' else ''
        # qg_prompt = 'Please generate a question from this picture, context and the corresponding answer.' if \
        # problem['hint'] != '' else 'Please generate a question from this picture and the corresponding answer.'
        # rg_prompt = 'According to this picture, context and question with its answer, how to make a reasoning to answer the question?' if \
        # problem[
        #     'hint'] != '' else 'According to this picture and question with its answer, how to make a reasoning to answer the question?'
        # dg_prompt = 'According to this picture, context and question with its answer obtained through reasoning, please generate at least 1 plausible yet incorrect answers which should be similar and grammatically consistent with the correct answer and seperate them with numbers like (1) (2) (3). \n' if \
        # problem[
        #     'hint'] != '' else 'According to this picture and question with its answer obtained through reasoning, please generate at least 1 plausible yet incorrect answers which should be similar and grammatically consistent with the correct answer and seperate them with numbers like (1) (2) (3).\n'
        qg = '\nExample:\n' + image + context + answer + f'Question: {question}'
        rg = '\nExample:\n' + image + context + f'Question: {question}\n' + answer + f'Reasoning: {rationale}'
        dg = '\nExample:\n' + image + context + f'Question: {question}\n' + answer + f'Distractors: {distractors}'
    else:
        context = 'Context: ' + problem['hint'] + '\n' if problem['hint'] != '' else ''
        # qg_prompt = 'Please generate a question from this context and the corresponding answer.' if \
        #     problem[
        #         'hint'] != '' else 'Please generate a question from the corresponding answer.'
        # rg_prompt = 'According to this context and question with its answer, how to make a reasoning to answer the question?' if \
        #     problem[
        #         'hint'] != '' else 'According to this question with its answer, how to make a reasoning to answer the question?'
        # dg_prompt = 'According to this context and question with its answer obtained through reasoning, please generate at least 1 plausible yet incorrect answers which should be similar and grammatically consistent with the correct answer and seperate them with numbers like (1) (2) (3). \n' if \
        #     problem[
        #         'hint'] != '' else 'According to this question with its answer obtained through reasoning, please generate at least 1 plausible yet incorrect answers which should be similar and grammatically consistent with the correct answer and seperate them with numbers like (1) (2) (3).\n'
        qg = '\nExample:\n' + context + answer + f'Question: {question}'
        rg = '\nExample:\n' + context + f'Question: {question}\n' + answer + f'Reasoning: {rationale}'
        dg = '\nExample:\n' + context + f'Question: {question}\n' + answer + f'Distractors: {distractors}'
    return qg, rg, dg


save_list = []
split = 'test_website'
save_root = f'playground/data/flan-t5-finetune/mind2web_{split}_noplan_origin.json'
data_root = f'playground/data/{split}_kp'
query_path = 'playground/grounding_results/30_selected_tasks/image_annotation_3images_website'
plan_path = f'playground/mind2web_train_pl/{split}.jsonl'
predict_plans = open(plan_path,'r').readlines()
plan_id = 0
previous_k = 5

for root, dirs, files in os.walk(data_root):
    for id, file in enumerate(files):
        # if file == 'train_5.json' or file == 'train_10.json':
        #     continue
        examples = json.load(open(os.path.join(data_root, file)))

        for example in tqdm(examples):
            task = f"{example['confirmed_task']}"
            keypoints = example['keypoints']
            for a_id, action in enumerate(example['actions']):
                predict_plan = json.loads(predict_plans[plan_id])['response']
                plan_id += 1
                task_action = f"{example['annotation_id']}_{action['action_uid']}"
                # print(os.path.join(query_path, task_action))
                # pdb.set_trace()
                if not os.path.exists(os.path.join(query_path, task_action, 'queries.jsonl')):
                    continue
                quries = json.load(open(os.path.join(query_path, task_action, 'queries.jsonl')))

                # element_list = list(action["score"].keys())
                if len(action['pos_candidates']) > 0:
                    correct_element_id = action['pos_candidates'][0]['backend_node_id']
                else:
                    correct_element_id = -1
                # dict = {"id": f"{example['annotation_id']}-{action['action_uid']}"}
                image = f"{query_path}/{task_action}/images/0_labeled.jpg"
                # if example['plan'].find()
                # elements = []
                # positive_elements = action['pos_candidates']
                # negative_elements = action['neg_candidates']
                # elements.extend(positive_elements)
                # elements.extend(negative_elements)
                #
                # all_elements_id = [item['backend_node_id'] for item in elements]
                # all_elements = get_choices(action, all_elements_id, -1,
                #                            keep_html_brackets=True)
                # element_dict = {}
                # for element in all_elements:
                #     element_dict.update({element[0]:element[1]})
                if action['operation']['value'] == "":
                    action['operation']['value'] = "None"
                assistant_value = f"\nACTION: {action['operation']['op']}\nVALUE: {action['operation']['value']}"

                # prompt_2 = f"{predict_plan}\nPrevious Actions:\n"
                prompt_2 = f"Previous Actions:\n"
                if len(quries["previous_actions"]) > 0:
                    for pre_act_id, action in enumerate(quries["previous_actions"][-previous_k:]):
                        prompt_2 += f"{pre_act_id+1}. {action}\n"
                else:
                    prompt_2 += "None\n"
                prompt_2 += f"""Combined with the image and previous action, what should be the next action to complete the task: {task} Please select from the following choices:\n"""

                # pdb.set_trace()
                # print(action["score"])
                # pdb.set_trace()
                correct_element_latter = "None"
                # element_choice = action["score"][:26]
                for idx, choice in enumerate(quries["choices"]):
                    # convert to ascii A, B, C, D, ...
                    prompt_2 += f"{chr(65 + idx)}. {choice[1]}\n"
                    if choice[0] == correct_element_id:
                        correct_element_latter = chr(65 + idx)
                assistant_value = f"ELEMENT: {correct_element_latter}" + assistant_value
                prompt_2 += """None. None of the other options match the correct element."""
                user_value = "<image>\n" + prompt_2

                # if assistant_value.find('Fishing') != -1:
                #     continue
                conversations = {"id": f"{example['annotation_id']}_{a_id}", "image": image, "conversations": [{"from": "human", "value": user_value}, {"from": "gpt", "value": assistant_value}]}
              # print(dict)
                # pdb.set_trace()
                save_list.append(conversations)
print(len(save_list))
with open(save_root, 'w') as fp:
    json.dump(save_list, fp)