Release GUI-Swift-Mobile-2B: Mobile105 SFT step4658

b7fc81a verified 1 day ago

12.3 kB

"""Versioned Android-only 105-slot adapter; frozen 133-slot models are unchanged.
The user-selected layout is action(1), two coordinate pairs(4), text(100).
Reuse the pinned GUI-Decisions coordinate codebook/targets/readout and the
existing exact text/EOS codec. The one-slot action code is a new, explicit
categorical mapping, not a claim that action names tokenize to one token.
"""

import hashlib
import json
from dataclasses import dataclass, field, replace

from trace2task.mobile133_codec import BUNDLE, Mobile133Codec
from trace2task.mobile133_protocol import MobileExecutionConfig

PROTOCOL = "gui-swift-mobile105-v1"
ACTIONS = (
    "tap",
    "long_press",
    "swipe",
    "scroll_up",
    "scroll_down",
    "scroll_left",
    "scroll_right",
    "type_text",
    "open_app",
    "back",
    "home",
    "enter",
    "menu",
    "wait",
    "done",
    "failed",
)
ACTION_LABELS = tuple("ABCDEFGHIJKLMNOP")
FIELD_WIDTHS = (
    ("action", 1),
    ("coordinate.x", 1),
    ("coordinate.y", 1),
    ("coordinate2.x", 1),
    ("coordinate2.y", 1),
    ("text", 100),
)
LITERALS = (
    '<action>{"action":"',
    '","coordinate":[',
    ",",
    '],"coordinate2":[',
    ",",
    '],"text":"',
    '"}</action>',
)
# Only coordinate/text value positions are copied from the pinned parameter codec.
PARAMETER_POSITIONS = {**{p: p - 1 for p in range(2, 6)}, **{p: p - 28 for p in range(33, 133)}}
SYSTEM = (
    "Execute the current Android instruction using the screenshot and executed-action history. "
    "Predict the 105 masked values in one forward: action, x, y, x2, y2, and 100 text slots. "
    "Action codes: " + ", ".join(f"{label}={a}" for label, a in zip(ACTION_LABELS, ACTIONS)) + ". "
    "tap/long_press use x,y; swipe uses both coordinate pairs; type_text/open_app use text. "
    "Other actions have no parameters. Coordinates and historical coordinates use the same "
    "256-bin codebook; the decoder maps them to 0-1000. Scroll direction means browsing direction. "
    "Copy text exactly, at most 99 native tokens followed by EOS. Timing is an executor setting. "
    "Predict only applicable fields in the supplied <action> table; do not output explanations."
)


def from_source_action(action):
    """Lossless semantic projection of an admitted Android action, also for history."""
    a = action.get("action")
    if a in ("click", "long_press", "swipe", "type", "open", "wait"):
        expected = {"action"}
        if a in ("click", "long_press", "swipe"):
            expected.add("coordinate")
        if a == "swipe":
            expected.add("coordinate2")
        if a in ("type", "open"):
            expected.add("text")
        if set(action) != expected:
            raise ValueError("mobile105_unrepresentable_parameters:" + str(a))
        return dict(
            action, action={"click": "tap", "type": "type_text", "open": "open_app"}.get(a, a)
        )
    if a == "scroll" and set(action) == {"action", "direction"}:
        name = "scroll_" + str(action["direction"])
        if name in ACTIONS:
            return {"action": name}
    if a == "system_button" and set(action) == {"action", "button"}:
        names = {"Back": "back", "Home": "home", "Enter": "enter", "Menu": "menu"}
        if action["button"] in names:
            return {"action": names[action["button"]]}
    if a == "key" and set(action) == {"action", "text"}:
        names = {"BACK": "back", "HOME": "home", "ENTER": "enter", "MENU": "menu"}
        if action["text"] in names:
            return {"action": names[action["text"]]}
    if (
        a == "terminate"
        and set(action) == {"action", "status"}
        and action["status"] in ("success", "failure")
    ):
        return {"action": "done" if action["status"] == "success" else "failed"}
    raise ValueError("unsupported_mobile105_source_action:" + str(action))


def to_source_action(action):
    a = action.get("action")
    if a not in ACTIONS:
        raise ValueError("unknown_mobile105_action")
    if a in ("tap", "long_press", "swipe", "type_text", "open_app", "wait"):
        result = dict(
            action, action={"tap": "click", "type_text": "type", "open_app": "open"}.get(a, a)
        )
    elif set(action) != {"action"}:
        raise ValueError("mobile105_unexpected_parameters")
    elif a.startswith("scroll_"):
        result = {"action": "scroll", "direction": a.removeprefix("scroll_")}
    elif a in ("back", "home", "enter", "menu"):
        result = {"action": "system_button", "button": a.capitalize()}
    else:
        result = {"action": "terminate", "status": "success" if a == "done" else "failure"}
    # Reject time or any silently ignored extra fields, including for history.
    if from_source_action(result) != action:
        raise ValueError("mobile105_noncanonical_action")
    return result


@dataclass(frozen=True)
class Mobile105ExecutionConfig:
    protocol: str = PROTOCOL
    primitives: MobileExecutionConfig = field(default_factory=MobileExecutionConfig)

    def __post_init__(self):
        if self.protocol != PROTOCOL or not isinstance(self.primitives, MobileExecutionConfig):
            raise ValueError("mobile105_execution_contract_mismatch")


class Mobile105Codec:
    variant = PROTOCOL
    output_slots = 105

    def __init__(self, tokenizer, *, bundle=BUNDLE):
        self.source_codec = Mobile133Codec(tokenizer, bundle=bundle)
        self.upstream = self.source_codec.upstream
        for name in (
            "tokenizer",
            "base_codec",
            "bin_ids",
            "codebook_sha256",
            "mask_id",
            "pad_id",
            "eos_id",
            "special_ids",
            "vocab_size",
        ):
            setattr(self, name, getattr(self.source_codec, name))
        self.action_ids = {
            a: self.base_codec._single(label)
            for label, a in zip(ACTION_LABELS, ACTIONS, strict=True)
        }
        if len(set(self.action_ids.values())) != len(ACTIONS):
            raise ValueError("mobile105_action_code_collision")
        built = self.build_parallel_input()
        self.input_ids, self.slot_positions = built.input_ids, built.slot_positions
        self.field_slices = {}
        start = 0
        for name, width in FIELD_WIDTHS:
            self.field_slices[name] = slice(start, start + width)
            start += width
        assert start == self.output_slots

    def build_parallel_input(self, prefix_ids=()):
        from trace2task.gui_owl_short_parallel_input import ParallelInput

        ids, positions, fields = list(prefix_ids), [], []
        for literal, (name, width) in zip(LITERALS[:-1], FIELD_WIDTHS, strict=True):
            ids.extend(self.tokenizer.encode(literal, add_special_tokens=False).ids)
            pos = tuple(range(len(ids), len(ids) + width))
            positions.extend(pos)
            fields.append((name, pos))
            ids.extend([self.mask_id] * width)
        ids.extend(self.tokenizer.encode(LITERALS[-1], add_special_tokens=False).ids)
        return ParallelInput(tuple(ids), tuple(positions), tuple(fields), len(prefix_ids))

    def contract(self):
        return {
            "protocol": PROTOCOL,
            "output_slots": 105,
            "layout": [[name, width] for name, width in FIELD_WIDTHS],
            "action_codes": [
                {"action": a, "label": label, "token_id": self.action_ids[a]}
                for label, a in zip(ACTION_LABELS, ACTIONS, strict=True)
            ],
            "coordinate_codebook_sha256": self.codebook_sha256,
            "scaffold_sha256": hashlib.sha256(json.dumps(self.input_ids).encode()).hexdigest(),
            "max_text_tokens": 99,
            "text_slots_including_eos": 100,
            "timing": "executor defaults, no duration labels",
            "scroll_direction": "browsing direction; no predicted endpoints",
            "inactive_slots": "ignored in supervision and execution",
            "readout": "native causal logits at slot_position-1; single forward",
        }

    def validate_contract(self, value):
        if value != self.contract():
            raise ValueError("mobile105_protocol_contract_mismatch")

    def adapt_target(self, action, *, coordinate_supervision=None):
        source = to_source_action(action)
        t = self.source_codec.adapt_target(
            source, "android", coordinate_supervision=coordinate_supervision
        )
        canonical = from_source_action(t.canonical_action)
        ids, hard = [self.pad_id] * 105, [-100] * 105
        ids[0] = hard[0] = self.action_ids[canonical["action"]]
        soft, fields = {}, [("action", 0, 1)]
        for name, start, stop in t.fields:
            if name not in self.field_slices or name == "action":
                continue  # direction/button/status are encoded by the action class.
            for old in range(start, stop):
                new = PARAMETER_POSITIONS[old]
                ids[new], hard[new] = t.target_ids[old], t.hard_labels[old]
                if old in t.soft_labels:
                    soft[new] = t.soft_labels[old]
            fields.append((name, PARAMETER_POSITIONS[start], PARAMETER_POSITIONS[stop - 1] + 1))
        return replace(
            t,
            canonical_action=canonical,
            target_ids=tuple(ids),
            hard_labels=tuple(hard),
            soft_labels=soft,
            fields=tuple(fields),
        )

    def adapt_training_target(self, row):
        self.source_codec.check_row(row)
        if row["platform"] != "android":
            raise ValueError("mobile105_requires_android")
        return self.adapt_target(
            from_source_action(row["target"]),
            coordinate_supervision=self.upstream.coordinate_metadata_from_row(row),
        )

    def messages(self, row):
        self.source_codec.check_row(row)
        if row["platform"] != "android":
            raise ValueError("mobile105_requires_android")
        history = [
            dict(step, actions=[from_source_action(a) for a in step["actions"]])
            for step in row["history"]
        ]
        # Reuse the pinned image order, complete-history formatter and escaping.
        result = self.upstream.messages(dict(row, history=history))
        result[0]["content"] = SYSTEM
        result[1]["content"][0]["text"] = result[1]["content"][0]["text"].replace(
            "\nTask: ", "\nCurrent instruction: ", 1
        )
        return result

    def bind_history_tokens(self, prefix_ids, row):
        # The action-name projection preserves coordinate order and every history item.
        return self.source_codec.bind_history_tokens(prefix_ids, row)

    def decode(self, logits):
        import torch

        if logits.ndim != 2 or logits.shape[0] != 105 or logits.shape[1] < self.vocab_size:
            raise ValueError("expected_105_by_vocab_logits")
        if not torch.isfinite(logits).all():
            raise ValueError("nonfinite_logits")
        a = ACTIONS[int(logits[0, list(self.action_ids.values())].argmax())]
        if a not in ("tap", "long_press", "swipe", "type_text", "open_app"):
            return {"action": a}
        # Parameter-only compatibility adapter: no second model call, no new readout rule.
        full = logits.new_full((133, logits.shape[1]), -60)
        old = list(PARAMETER_POSITIONS)
        full[old] = logits[list(PARAMETER_POSITIONS.values())]
        old_name = {"tap": "click", "type_text": "type", "open_app": "open"}.get(a, a)
        for p, token in enumerate(self.source_codec.action_encodings[old_name]):
            full[p, token] = 60
        if a == "long_press":
            full[self.source_codec.field_slices["time"], self.pad_id] = 60
        decoded = self.source_codec.decode(full, platform="android")
        return from_source_action(decoded)

    def execution_plan(self, action, size, *, config, installed_apps=None):
        if not isinstance(config, Mobile105ExecutionConfig) or config.protocol != PROTOCOL:
            raise ValueError("mobile105_execution_contract_mismatch")
        canonical = self.adapt_target(action).canonical_action
        return self.source_codec.execution_plan(
            to_source_action(canonical),
            size,
            config=config.primitives,
            installed_apps=installed_apps,
        )

	"""Versioned Android-only 105-slot adapter; frozen 133-slot models are unchanged.

	The user-selected layout is action(1), two coordinate pairs(4), text(100).
	Reuse the pinned GUI-Decisions coordinate codebook/targets/readout and the
	existing exact text/EOS codec. The one-slot action code is a new, explicit
	categorical mapping, not a claim that action names tokenize to one token.
	"""

	import hashlib
	import json
	from dataclasses import dataclass, field, replace

	from trace2task.mobile133_codec import BUNDLE, Mobile133Codec
	from trace2task.mobile133_protocol import MobileExecutionConfig

	PROTOCOL = "gui-swift-mobile105-v1"
	ACTIONS = (
	"tap",
	"long_press",
	"swipe",
	"scroll_up",
	"scroll_down",
	"scroll_left",
	"scroll_right",
	"type_text",
	"open_app",
	"back",
	"home",
	"enter",
	"menu",
	"wait",
	"done",
	"failed",
	)
	ACTION_LABELS = tuple("ABCDEFGHIJKLMNOP")
	FIELD_WIDTHS = (
	("action", 1),
	("coordinate.x", 1),
	("coordinate.y", 1),
	("coordinate2.x", 1),
	("coordinate2.y", 1),
	("text", 100),
	)
	LITERALS = (
	'<action>{"action":"',
	'","coordinate":[',
	",",
	'],"coordinate2":[',
	",",
	'],"text":"',
	'"}</action>',
	)
	# Only coordinate/text value positions are copied from the pinned parameter codec.
	PARAMETER_POSITIONS = {{p: p - 1 for p in range(2, 6)}, {p: p - 28 for p in range(33, 133)}}
	SYSTEM = (
	"Execute the current Android instruction using the screenshot and executed-action history. "
	"Predict the 105 masked values in one forward: action, x, y, x2, y2, and 100 text slots. "
	"Action codes: " + ", ".join(f"{label}={a}" for label, a in zip(ACTION_LABELS, ACTIONS)) + ". "
	"tap/long_press use x,y; swipe uses both coordinate pairs; type_text/open_app use text. "
	"Other actions have no parameters. Coordinates and historical coordinates use the same "
	"256-bin codebook; the decoder maps them to 0-1000. Scroll direction means browsing direction. "
	"Copy text exactly, at most 99 native tokens followed by EOS. Timing is an executor setting. "
	"Predict only applicable fields in the supplied <action> table; do not output explanations."
	)


	def from_source_action(action):
	"""Lossless semantic projection of an admitted Android action, also for history."""
	a = action.get("action")
	if a in ("click", "long_press", "swipe", "type", "open", "wait"):
	expected = {"action"}
	if a in ("click", "long_press", "swipe"):
	expected.add("coordinate")
	if a == "swipe":
	expected.add("coordinate2")
	if a in ("type", "open"):
	expected.add("text")
	if set(action) != expected:
	raise ValueError("mobile105_unrepresentable_parameters:" + str(a))
	return dict(
	action, action={"click": "tap", "type": "type_text", "open": "open_app"}.get(a, a)
	)
	if a == "scroll" and set(action) == {"action", "direction"}:
	name = "scroll_" + str(action["direction"])
	if name in ACTIONS:
	return {"action": name}
	if a == "system_button" and set(action) == {"action", "button"}:
	names = {"Back": "back", "Home": "home", "Enter": "enter", "Menu": "menu"}
	if action["button"] in names:
	return {"action": names[action["button"]]}
	if a == "key" and set(action) == {"action", "text"}:
	names = {"BACK": "back", "HOME": "home", "ENTER": "enter", "MENU": "menu"}
	if action["text"] in names:
	return {"action": names[action["text"]]}
	if (
	a == "terminate"
	and set(action) == {"action", "status"}
	and action["status"] in ("success", "failure")
	):
	return {"action": "done" if action["status"] == "success" else "failed"}
	raise ValueError("unsupported_mobile105_source_action:" + str(action))


	def to_source_action(action):
	a = action.get("action")
	if a not in ACTIONS:
	raise ValueError("unknown_mobile105_action")
	if a in ("tap", "long_press", "swipe", "type_text", "open_app", "wait"):
	result = dict(
	action, action={"tap": "click", "type_text": "type", "open_app": "open"}.get(a, a)
	)
	elif set(action) != {"action"}:
	raise ValueError("mobile105_unexpected_parameters")
	elif a.startswith("scroll_"):
	result = {"action": "scroll", "direction": a.removeprefix("scroll_")}
	elif a in ("back", "home", "enter", "menu"):
	result = {"action": "system_button", "button": a.capitalize()}
	else:
	result = {"action": "terminate", "status": "success" if a == "done" else "failure"}
	# Reject time or any silently ignored extra fields, including for history.
	if from_source_action(result) != action:
	raise ValueError("mobile105_noncanonical_action")
	return result


	@dataclass(frozen=True)
	class Mobile105ExecutionConfig:
	protocol: str = PROTOCOL
	primitives: MobileExecutionConfig = field(default_factory=MobileExecutionConfig)

	def __post_init__(self):
	if self.protocol != PROTOCOL or not isinstance(self.primitives, MobileExecutionConfig):
	raise ValueError("mobile105_execution_contract_mismatch")


	class Mobile105Codec:
	variant = PROTOCOL
	output_slots = 105

	def __init__(self, tokenizer, *, bundle=BUNDLE):
	self.source_codec = Mobile133Codec(tokenizer, bundle=bundle)
	self.upstream = self.source_codec.upstream
	for name in (
	"tokenizer",
	"base_codec",
	"bin_ids",
	"codebook_sha256",
	"mask_id",
	"pad_id",
	"eos_id",
	"special_ids",
	"vocab_size",
	):
	setattr(self, name, getattr(self.source_codec, name))
	self.action_ids = {
	a: self.base_codec._single(label)
	for label, a in zip(ACTION_LABELS, ACTIONS, strict=True)
	}
	if len(set(self.action_ids.values())) != len(ACTIONS):
	raise ValueError("mobile105_action_code_collision")
	built = self.build_parallel_input()
	self.input_ids, self.slot_positions = built.input_ids, built.slot_positions
	self.field_slices = {}
	start = 0
	for name, width in FIELD_WIDTHS:
	self.field_slices[name] = slice(start, start + width)
	start += width
	assert start == self.output_slots

	def build_parallel_input(self, prefix_ids=()):
	from trace2task.gui_owl_short_parallel_input import ParallelInput

	ids, positions, fields = list(prefix_ids), [], []
	for literal, (name, width) in zip(LITERALS[:-1], FIELD_WIDTHS, strict=True):
	ids.extend(self.tokenizer.encode(literal, add_special_tokens=False).ids)
	pos = tuple(range(len(ids), len(ids) + width))
	positions.extend(pos)
	fields.append((name, pos))
	ids.extend([self.mask_id] * width)
	ids.extend(self.tokenizer.encode(LITERALS[-1], add_special_tokens=False).ids)
	return ParallelInput(tuple(ids), tuple(positions), tuple(fields), len(prefix_ids))

	def contract(self):
	return {
	"protocol": PROTOCOL,
	"output_slots": 105,
	"layout": [[name, width] for name, width in FIELD_WIDTHS],
	"action_codes": [
	{"action": a, "label": label, "token_id": self.action_ids[a]}
	for label, a in zip(ACTION_LABELS, ACTIONS, strict=True)
	],
	"coordinate_codebook_sha256": self.codebook_sha256,
	"scaffold_sha256": hashlib.sha256(json.dumps(self.input_ids).encode()).hexdigest(),
	"max_text_tokens": 99,
	"text_slots_including_eos": 100,
	"timing": "executor defaults, no duration labels",
	"scroll_direction": "browsing direction; no predicted endpoints",
	"inactive_slots": "ignored in supervision and execution",
	"readout": "native causal logits at slot_position-1; single forward",
	}

	def validate_contract(self, value):
	if value != self.contract():
	raise ValueError("mobile105_protocol_contract_mismatch")

	def adapt_target(self, action, *, coordinate_supervision=None):
	source = to_source_action(action)
	t = self.source_codec.adapt_target(
	source, "android", coordinate_supervision=coordinate_supervision
	)
	canonical = from_source_action(t.canonical_action)
	ids, hard = [self.pad_id] * 105, [-100] * 105
	ids[0] = hard[0] = self.action_ids[canonical["action"]]
	soft, fields = {}, [("action", 0, 1)]
	for name, start, stop in t.fields:
	if name not in self.field_slices or name == "action":
	continue # direction/button/status are encoded by the action class.
	for old in range(start, stop):
	new = PARAMETER_POSITIONS[old]
	ids[new], hard[new] = t.target_ids[old], t.hard_labels[old]
	if old in t.soft_labels:
	soft[new] = t.soft_labels[old]
	fields.append((name, PARAMETER_POSITIONS[start], PARAMETER_POSITIONS[stop - 1] + 1))
	return replace(
	t,
	canonical_action=canonical,
	target_ids=tuple(ids),
	hard_labels=tuple(hard),
	soft_labels=soft,
	fields=tuple(fields),
	)

	def adapt_training_target(self, row):
	self.source_codec.check_row(row)
	if row["platform"] != "android":
	raise ValueError("mobile105_requires_android")
	return self.adapt_target(
	from_source_action(row["target"]),
	coordinate_supervision=self.upstream.coordinate_metadata_from_row(row),
	)

	def messages(self, row):
	self.source_codec.check_row(row)
	if row["platform"] != "android":
	raise ValueError("mobile105_requires_android")
	history = [
	dict(step, actions=[from_source_action(a) for a in step["actions"]])
	for step in row["history"]
	]
	# Reuse the pinned image order, complete-history formatter and escaping.
	result = self.upstream.messages(dict(row, history=history))
	result[0]["content"] = SYSTEM
	result[1]["content"][0]["text"] = result[1]["content"][0]["text"].replace(
	"\nTask: ", "\nCurrent instruction: ", 1
	)
	return result

	def bind_history_tokens(self, prefix_ids, row):
	# The action-name projection preserves coordinate order and every history item.
	return self.source_codec.bind_history_tokens(prefix_ids, row)

	def decode(self, logits):
	import torch

	if logits.ndim != 2 or logits.shape[0] != 105 or logits.shape[1] < self.vocab_size:
	raise ValueError("expected_105_by_vocab_logits")
	if not torch.isfinite(logits).all():
	raise ValueError("nonfinite_logits")
	a = ACTIONS[int(logits[0, list(self.action_ids.values())].argmax())]
	if a not in ("tap", "long_press", "swipe", "type_text", "open_app"):
	return {"action": a}
	# Parameter-only compatibility adapter: no second model call, no new readout rule.
	full = logits.new_full((133, logits.shape[1]), -60)
	old = list(PARAMETER_POSITIONS)
	full[old] = logits[list(PARAMETER_POSITIONS.values())]
	old_name = {"tap": "click", "type_text": "type", "open_app": "open"}.get(a, a)
	for p, token in enumerate(self.source_codec.action_encodings[old_name]):
	full[p, token] = 60
	if a == "long_press":
	full[self.source_codec.field_slices["time"], self.pad_id] = 60
	decoded = self.source_codec.decode(full, platform="android")
	return from_source_action(decoded)

	def execution_plan(self, action, size, *, config, installed_apps=None):
	if not isinstance(config, Mobile105ExecutionConfig) or config.protocol != PROTOCOL:
	raise ValueError("mobile105_execution_contract_mismatch")
	canonical = self.adapt_target(action).canonical_action
	return self.source_codec.execution_plan(
	to_source_action(canonical),
	size,
	config=config.primitives,
	installed_apps=installed_apps,
	)