ONNX
English
vons
research
candidate-selection
File size: 8,355 Bytes
a76a881
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
"""Label-free context compression for offline Mind2Web selection.

The adapter is deliberately independent of the evaluator and gold labels. It
keeps only the last two interaction-history entries, removes presentation-only
DOM attributes, limits each candidate with the exact runtime tokenizer, and
then reduces candidate/state text until the serialized inference prompt is at
most 512 tokens.
"""

from __future__ import annotations

import json
import re
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from typing import Any, Protocol

CONTEXT_ADAPTER_VERSION = "vons.mind2web-context/v1"
MAX_HISTORY_ACTIONS = 2
MAX_CANDIDATE_TOKENS = 15
MAX_PROMPT_TOKENS = 512

_SVG_BLOCK = re.compile(r"<svg\b[^>]*>.*?</svg\s*>", re.IGNORECASE | re.DOTALL)
_STYLE_BLOCK = re.compile(r"<style\b[^>]*>.*?</style\s*>", re.IGNORECASE | re.DOTALL)
_ATTRIBUTE = re.compile(
    r"(?<!\S)([A-Za-z_:][-A-Za-z0-9_:.]*)=(\"[^\"]*\"|'[^']*'|\S+)"
)
_WHITESPACE = re.compile(r"\s+")
_KEPT_ATTRIBUTES = frozenset({"type", "aria-label", "id"})


class _Encoding(Protocol):
    ids: Sequence[int]


class ExactTokenizer(Protocol):
    def encode(self, text: str) -> _Encoding: ...

    def decode(self, ids: Sequence[int], skip_special_tokens: bool = True) -> str: ...


@dataclass(frozen=True)
class CompressionStats:
    version: str
    original_prompt_tokens: int
    prompt_tokens: int
    candidate_token_budget: int
    max_candidate_tokens: int
    history_actions_before: int
    history_actions_after: int
    candidates: int


def strip_dom_presentation(text: str) -> str:
    """Keep visible text, tag names, and the small semantic attribute allowlist."""
    value = _SVG_BLOCK.sub(" ", text)
    value = _STYLE_BLOCK.sub(" ", value)

    def replace_attribute(match: re.Match[str]) -> str:
        name = match.group(1).lower()
        return match.group(0) if name in _KEPT_ATTRIBUTES else ""

    value = _ATTRIBUTE.sub(replace_attribute, value)
    return _WHITESPACE.sub(" ", value).strip()


def _token_ids(tokenizer: ExactTokenizer, text: str) -> list[int]:
    return list(tokenizer.encode(text).ids)


def _decode_prefix(tokenizer: ExactTokenizer, ids: Sequence[int], limit: int) -> str:
    if limit <= 0 or not ids:
        return ""
    prefix = list(ids[:limit])
    value = tokenizer.decode(prefix, skip_special_tokens=True).strip()
    while value and len(_token_ids(tokenizer, value)) > limit:
        prefix.pop()
        value = tokenizer.decode(prefix, skip_special_tokens=True).strip()
    return value


def _state_mapping(value: object) -> dict[str, object]:
    if isinstance(value, str):
        parsed = json.loads(value)
    else:
        parsed = value
    if not isinstance(parsed, Mapping):
        raise TypeError("Mind2Web state must be a JSON object")
    goal = parsed.get("goal", "")
    history = parsed.get("previous_actions", [])
    if not isinstance(goal, str):
        raise TypeError("Mind2Web goal must be a string")
    if not isinstance(history, list) or any(not isinstance(item, str) for item in history):
        raise TypeError("Mind2Web previous_actions must be a string array")
    return {"goal": goal, "previous_actions": history[-MAX_HISTORY_ACTIONS:]}


def _canonical(value: object) -> str:
    return json.dumps(
        value,
        ensure_ascii=False,
        sort_keys=True,
        separators=(",", ":"),
        allow_nan=False,
    )


def _prompt_text(state: Mapping[str, object], question: str, candidates: Sequence[str]) -> str:
    prefix = f"{_canonical(state)}\nQuestion: {question}"
    return prefix + "\nCandidates:\n" + "\n".join(candidates)


def _truncate_goal_to_fit(
    tokenizer: ExactTokenizer,
    state: dict[str, object],
    question: str,
    candidates: Sequence[str],
    maximum: int,
) -> dict[str, object]:
    candidate_state = {"goal": str(state["goal"]), "previous_actions": []}
    if len(_token_ids(tokenizer, _prompt_text(candidate_state, question, candidates))) <= maximum:
        return candidate_state
    goal_ids = _token_ids(tokenizer, str(state["goal"]))
    left, right = 0, len(goal_ids)
    best = {"goal": "", "previous_actions": []}
    while left <= right:
        middle = (left + right) // 2
        trial = {
            "goal": _decode_prefix(tokenizer, goal_ids, middle),
            "previous_actions": [],
        }
        if len(_token_ids(tokenizer, _prompt_text(trial, question, candidates))) <= maximum:
            best = trial
            left = middle + 1
        else:
            right = middle - 1
    return best


def compress_request(
    request: Mapping[str, Any],
    tokenizer: ExactTokenizer,
    *,
    max_prompt_tokens: int = MAX_PROMPT_TOKENS,
    max_candidate_tokens: int = MAX_CANDIDATE_TOKENS,
) -> tuple[dict[str, object], CompressionStats]:
    """Return an inference request that is guaranteed to fit the hard budget."""
    if max_prompt_tokens <= 0 or max_candidate_tokens < 0:
        raise ValueError("token budgets must be positive")
    question = request.get("question")
    candidates = request.get("candidates")
    if not isinstance(question, str) or not isinstance(candidates, list):
        raise TypeError("Mind2Web request requires question and candidates")
    if any(
        not isinstance(item, Mapping)
        or not isinstance(item.get("id"), str)
        or not isinstance(item.get("text"), str)
        for item in candidates
    ):
        raise TypeError("Mind2Web candidates require string id and text")

    raw_state = _state_mapping(request.get("state"))
    raw_history = request.get("state")
    if isinstance(raw_history, str):
        raw_history = json.loads(raw_history)
    history_before = (
        len(raw_history.get("previous_actions", []))
        if isinstance(raw_history, Mapping)
        and isinstance(raw_history.get("previous_actions", []), list)
        else 0
    )
    raw_texts = [str(item["text"]) for item in candidates]
    original_prompt_tokens = len(
        _token_ids(tokenizer, _prompt_text(raw_state, question, raw_texts))
    )
    cleaned = [strip_dom_presentation(text) for text in raw_texts]
    candidate_ids = [_token_ids(tokenizer, text) for text in cleaned]

    def texts_for_budget(budget: int) -> list[str]:
        return [_decode_prefix(tokenizer, ids, budget) for ids in candidate_ids]

    left, right = 0, max_candidate_tokens
    selected_budget = 0
    selected_texts = texts_for_budget(0)
    while left <= right:
        middle = (left + right) // 2
        trial_texts = texts_for_budget(middle)
        token_count = len(
            _token_ids(tokenizer, _prompt_text(raw_state, question, trial_texts))
        )
        if token_count <= max_prompt_tokens:
            selected_budget = middle
            selected_texts = trial_texts
            left = middle + 1
        else:
            right = middle - 1

    state = raw_state
    if len(_token_ids(tokenizer, _prompt_text(state, question, selected_texts))) > max_prompt_tokens:
        selected_budget = 0
        selected_texts = texts_for_budget(0)
        state = _truncate_goal_to_fit(
            tokenizer,
            raw_state,
            question,
            selected_texts,
            max_prompt_tokens,
        )

    prompt_tokens = len(_token_ids(tokenizer, _prompt_text(state, question, selected_texts)))
    if prompt_tokens > max_prompt_tokens:
        raise ValueError("fixed Mind2Web prompt scaffolding exceeds the hard token budget")
    compressed_candidates = [
        {"id": str(item["id"]), "text": text}
        for item, text in zip(candidates, selected_texts, strict=True)
    ]
    max_observed_candidate_tokens = max(
        (len(_token_ids(tokenizer, text)) for text in selected_texts),
        default=0,
    )
    compressed = {
        "state": _canonical(state),
        "question": question,
        "candidates": compressed_candidates,
    }
    stats = CompressionStats(
        version=CONTEXT_ADAPTER_VERSION,
        original_prompt_tokens=original_prompt_tokens,
        prompt_tokens=prompt_tokens,
        candidate_token_budget=selected_budget,
        max_candidate_tokens=max_observed_candidate_tokens,
        history_actions_before=history_before,
        history_actions_after=len(state["previous_actions"]),
        candidates=len(compressed_candidates),
    )
    return compressed, stats