Function read_episodes_from_jsonl

eval/eval_data/process_ac.py:246–311

(jsonl_file_path, train_episodes, output_test_dir)

        return set()

def read_episodes_from_jsonl(jsonl_file_path, train_episodes, output_test_dir):
    episodes = {'in_train': [], 'not_in_train': []}
    
    os.makedirs(output_test_dir, exist_ok=True)
    
    total_lines = sum(1 for _ in open(jsonl_file_path, 'r', encoding='utf-8'))
    
    with open(jsonl_file_path, 'r', encoding='utf-8') as file:
        for line in tqdm(file, total=total_lines, desc="saving data"):
            data = json.loads(line.strip())
            episode_id = data['episode_id'][0]
            
            episode_data = []
            goal = data['goal'][0]
            actions = data['actions']
            screenshot_widths = data['screenshot_widths']
            screenshot_heights = data['screenshot_heights']
            screenshot_path = data['screenshot_path']
            
            for index, action in enumerate(actions):
                action_data = {
                    'action': action,
                    'screenshot_width': screenshot_widths[index],
                    'screenshot_height': screenshot_heights[index],
                    'screenshot_path': screenshot_path[index],
                    'low_instruction': data['step_instructions'][index],
                    'goal': goal,
                    'episode_id': episode_id,
                    'step': index,
                    'episode_length': len(screenshot_path),
                    'ui_trees': data['accessibility_tree'][index]
                }
                test_data = transform_action_data_and_build_test_data(action_data)
                if test_data is not None:
                    episode_data.append(test_data)
            
            # add last finish
            action_data={
                'action': "{\"action_type\":\"finish\"}",
                'screenshot_width': screenshot_widths[-1],
                'screenshot_height': screenshot_heights[-1],
                'screenshot_path': screenshot_path[-1],
                'goal': goal,
                'episode_id': episode_id,
                'step': len(actions),
                'episode_length': len(screenshot_path),
                'ui_trees': data['accessibility_tree'][-1],
                'low_instruction': "finish the task"
            }
            finish_data =transform_action_data_and_build_test_data(action_data)
            episode_data.append(finish_data)
                                    
            if episode_id in train_episodes:
                episodes['in_train'].append(episode_id)
                continue
            else:
                episodes['not_in_train'].append(episode_id)