File size: 4,583 Bytes
a61addb | 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 | """
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}")
|