| """ |
| Build K-step GUI Transition data from AndroidControl dataset for GUI-Shift training. |
| Run this first to construct the training dataset. |
| """ |
| import json |
| import os |
| import random |
| from collections import defaultdict |
| from datasets import load_dataset |
|
|
| def parse_action(action_str): |
| """Parse AndroidControl action string to GUI-Shift action format.""" |
| if not action_str.startswith("!FUNCTIONCALL"): |
| return None |
| json_str = action_str.replace("!FUNCTIONCALL", "") |
| try: |
| action = json.loads(json_str) |
| name = action["name"] |
| params = action.get("parameters", {}) |
| if name == "click": |
| return {"action_type": "click", "x": params["x"], "y": params["y"]} |
| elif name == "long_press": |
| return {"action_type": "long_press", "x": params["x"], "y": params["y"]} |
| elif name == "scroll": |
| return {"action_type": "scroll", "direction": params["direction"]} |
| elif name == "open_app": |
| return {"action_type": "open_app", "app_name": params["app_name"]} |
| elif name == "input_text": |
| return {"action_type": "input_text", "text": params["text"]} |
| elif name == "navigate_back": |
| return {"action_type": "navigate_back"} |
| elif name == "navigate_home": |
| return {"action_type": "navigate_home"} |
| elif name == "wait": |
| return {"action_type": "wait"} |
| elif name == "status": |
| return None |
| else: |
| return None |
| except Exception: |
| return None |
|
|
| def build_k_step_transitions(dataset_split, k=1, max_samples=2000, output_dir="./data"): |
| os.makedirs(output_dir, exist_ok=True) |
| episodes = defaultdict(list) |
| for i, item in enumerate(dataset_split): |
| json_data = item["json"] |
| episode_id = json_data["episode_id"] |
| episodes[episode_id].append({ |
| "index": i, |
| "step_id": json_data["step_id"], |
| "messages": json_data["messages"], |
| "png": item["png"], |
| }) |
| for ep_id in episodes: |
| episodes[ep_id].sort(key=lambda x: x["step_id"]) |
|
|
| transitions = [] |
| for ep_id, steps in episodes.items(): |
| for i in range(len(steps) - k): |
| state_t = steps[i] |
| state_tk = steps[i + k] |
| messages = state_t["messages"] |
| action_str = None |
| for msg in messages: |
| if msg["role"] == "assistant": |
| action_str = msg["content"] |
| break |
| if not action_str: |
| continue |
| action = parse_action(action_str) |
| if action is None: |
| continue |
|
|
| img_t = state_t["png"] |
| img_tk = state_tk["png"] |
| img_t_path = os.path.join(output_dir, f"ep{ep_id}_step{state_t['step_id']}.png") |
| img_tk_path = os.path.join(output_dir, f"ep{ep_id}_step{state_tk['step_id']}.png") |
| img_t.save(img_t_path) |
| img_tk.save(img_tk_path) |
|
|
| problem = ( |
| "You are given two GUI screenshots. The first screenshot shows the current state, " |
| "and the second screenshot shows the state after executing some actions. " |
| "Predict the first action that transitions from the first screenshot to the second screenshot." |
| ) |
| transitions.append({ |
| "image": [img_t_path, img_tk_path], |
| "conversations": [ |
| {"from": "human", "value": f"<image>\n<image>\n{problem}"}, |
| {"from": "gpt", "value": json.dumps(action)} |
| ], |
| "problem": problem, |
| "solution": json.dumps(action), |
| }) |
|
|
| random.shuffle(transitions) |
| transitions = transitions[:max_samples] |
| output_path = os.path.join(output_dir, f"ui_transition_training_{max_samples}_k_{k}_no_reasoning.jsonl") |
| with open(output_path, "w") as f: |
| for item in transitions: |
| f.write(json.dumps(item) + "\n") |
| print(f"Saved {len(transitions)} K-step transitions to {output_path}") |
| return output_path |
|
|
| if __name__ == "__main__": |
| print("Loading AndroidControl dataset...") |
| ds = load_dataset("ckg/AndroidControlParsedWithImages-20k", split="train", streaming=True) |
| items = [] |
| for i, item in enumerate(ds): |
| items.append(item) |
| if i >= 10000: |
| break |
| print(f"Processing {len(items)} items...") |
| output = build_k_step_transitions(items, k=1, max_samples=2000, output_dir="./data") |
| print(f"Done: {output}") |
|
|
