| 176 | |
| 177 | |
| 178 | def load(example: dict, use_caption: bool) -> tuple: |
| 179 | intent = example["intent"] |
| 180 | images = [] |
| 181 | if example["query_image"]: |
| 182 | images.append(example["query_image"].convert("RGB")) |
| 183 | trajectory = [ |
| 184 | { |
| 185 | "observation": { |
| 186 | "text": example["obs_text"] if use_caption else example["no_cap_obs_text"], |
| 187 | "image": np.array(example["obs_screenshot"]), |
| 188 | }, |
| 189 | "info": {"page": DetachedPage(example["obs_url"], "")}, |
| 190 | } |
| 191 | ] |
| 192 | meta_data = {"action_history": [example["previous_action"]]} |
| 193 | |
| 194 | return intent, images, trajectory, meta_data |
| 195 | |
| 196 | |
| 197 | @beartype |