laya-browser / code /finetune /build_items.py

v19s: WebChain real-site trajectories, format v5, webgym x7 + DAgger, harness fixes; replaces v17s

454b3e6 verified 29 days ago

9.41 kB

"""cases.jsonl + pages.jsonl -> tokenized training items (train split) and eval cases (held-out pages).
    python finetune/build_items.py out/pages.jsonl out/cases.jsonl out/
"""
import json, random, os, sys
from array import array
import torch
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from common_ft import build_request, gold_for, goal_done_question
from transformers import AutoTokenizer
from laya.common import QTYPES, build_sequence, render_options

MAX_LEN, HEAD_MAX_LEN = int(os.environ.get('LAYA_MAXLEN', '1024')), int(os.environ.get('LAYA_HEAD', '512'))
EVAL_EVERY = 5   # pages with index % 5 == 0 are held out
MAX_TARGETS = int(os.environ.get('MAX_TARGETS', '40'))
FINAL_P = float(os.environ.get('FINAL_P', '0.2'))   # share of target items built like the server's final round
# NOUL=1: also emit the yes/no question "is the whole goal visibly done on this page?" (common_ft.GOAL_DONE).
# yes = states whose gold is DONE; no = every mid-trajectory state (incl. webgym's results-shown-but-filter-missing
# states).  NOUL_NEG_P keeps the negatives at roughly 60:40 (no-heavy, to counter the DONE-too-early bias).
NOUL = os.environ.get('NOUL') == '1'
NOUL_NEG_P = float(os.environ.get('NOUL_NEG_P', '0.3'))


def main():
    pages_f, cases_f, out = sys.argv[1:4]
    pages = [json.loads(l) for l in open(pages_f)]
    cases = [json.loads(l) for l in open(cases_f)]
    for extra in sys.argv[4:]:
        cases += [json.loads(l) for l in open(extra)]
    snap = os.environ.get("LAYA_BASE")
    tok = AutoTokenizer.from_pretrained(os.path.join(snap, "tokenizer"))
    items, n_skip, ev = [], 0, []
    import hashlib
    STOPS = [" Stop when it is open.", " Stop once that page is visible.", " Then stop.", " Finish when it has loaded.", ""]
    # evaluation hygiene: never train on a suite task's start page (suite A and B start URLs), whatever the source
    import re as _re
    def _norm(u): return (u or "").split("#")[0].rstrip("/").replace("://www.", "://")
    suite_urls = set()
    for f in ("apps/browser_suite.py", "apps/browser_suite_b.py"):
        if os.path.exists(f):
            suite_urls |= {_norm(u) for u in _re.findall(r'\("[a-z0-9-]+", "(https?://[^"]+)"', open(f).read())}
    n_suite = 0
    if os.environ.get("LAYA_FMT") == "v6":
        # v6 history carries each action's result = the page the NEXT step starts on (same trajectory, one more step)
        traj = lambda c: (c.get("source"), c.get("task_id") or c.get("goal"), c.get("website"))
        at = {}
        for c in cases:
            pg = c.get("page_obj") or (pages[c["page"]] if c.get("page", -1) >= 0 else {})
            at.setdefault((traj(c), len(c.get("history") or [])), (pg.get("url") or c.get("url"), pg.get("title") or c.get("title")))
        filled = 0
        for c in cases:
            for i, h in enumerate(c.get("history") or []):
                if not h.get("url") and (traj(c), i + 1) in at:
                    h["url"], h["title"] = at[(traj(c), i + 1)]; filled += 1
        print("v6: history steps given their result page:", filled, flush=True)
    keep = float(os.environ.get("CASE_KEEP", "1"))   # CASE_KEEP=p: build from a random p of the cases (smaller, faster)
    if keep < 1:
        cases = [c for c in cases if int(hashlib.md5((c["goal"] + str(len(c.get("history") or []))).encode()).hexdigest()[:8], 16) / 0xffffffff < keep]
        print("CASE_KEEP", keep, "->", len(cases), "cases", flush=True)
    for c in cases:
        if _norm(c.get("url") or (c.get("page_obj") or {}).get("url")) in suite_urls or (c.get("page", -1) >= 0 and _norm(pages[c["page"]]["url"]) in suite_urls):
            n_suite += 1; continue
        if c.get("source") == "dagger":
            continue
        if c["kind"] == "done" and "page_obj" not in c:
            continue  # template DONE cases leak phrasing; real ones come from make_done_cases.py
        page = c.get("page_obj") or pages[c["page"]]
        h = int(hashlib.md5(c["goal"].encode()).hexdigest(), 16)
        is_sub = str(c.get("skill", "")).endswith("@sub") or c.get("skill") == "subgoal_done"
        goal = c["goal"] + ("" if is_sub or c.get("source") == "webgym" else STOPS[h % len(STOPS)])
        c = {**c, "goal": goal}
        state, questions, targets, controls = build_request(page, goal, c.get("history", []))
        if c.get("noul_only"):
            # a state where a failed run wrongly stopped: only the completion question, answered "not done"
            if NOUL and not ((hashlib.md5(c["website"].encode()).digest()[0] % EVAL_EVERY == 0)):
                nq = goal_done_question(goal)
                qq = {"t": "noul", "ins": json.dumps(nq["instructions"]), "crit": None}
                seq, markers = build_sequence(tok, state, qq, MAX_LEN, HEAD_MAX_LEN)
                if len(markers) == 2:
                    items.extend([{"ids": array("i", seq), "markers": markers, "qtype": QTYPES["noul"], "target": [1.0, 0.0], "label": 0,
                                   "qid": "goal_done", "gold_op": "NOT_DONE", "src": c.get("source")}] * 2)
            elif NOUL:
                ev.append({**c, "gold_index": None})
            continue
        gop, gidx = gold_for(c, targets, controls)
        if gop is None:
            n_skip += 1; continue
        held = (hashlib.md5(c["website"].encode()).digest()[0] % EVAL_EVERY == 0) if c.get("source") in ("mind2web", "nnetnav", "webchain", "gobrowse") else (c["page"] % EVAL_EVERY == 0)
        if held:
            ev.append({**c, "gold_index": gidx}); continue
        golds = {"operation": gop}
        if gidx is not None:
            golds[gop.lower() + "_target"] = gidx
        for qid, gold in golds.items():
            q = questions[qid]; keys = list(q["criteria"])
            rs = random.Random(h + len(keys) + len(qid))
            if qid != "operation" and len(keys) > 8 and rs.random() < FINAL_P:
                width = rs.randint(2, 6)          # the server's final round: only the chunk winners (2-6 options)
            else:
                width = MAX_TARGETS
            if qid != "operation" and len(keys) > width:
                # inference splits wide choices into chunks of <= MAX_TARGETS (systemone_server MAXOPT); train on the same
                # widths, so option labels are never truncated to a few tokens: gold + a random sample, page order kept
                keep = set(rs.sample([k for k in keys if k != gold], width - 1)) | {gold}
                keys = [k for k in keys if k in keep]
                q = {**q, "criteria": {k: q["criteria"][k] for k in keys}}
            target = [1.0 if k == gold else 0.0 for k in keys]
            qq = {"t": "choice", "ins": json.dumps(q["instructions"]), "crit": q["criteria"]}
            seq, markers = build_sequence(tok, state, qq, MAX_LEN, HEAD_MAX_LEN)
            if len(markers) != len(render_options(qq)):
                n_skip += 1; continue
            item = {"ids": array("i", seq), "markers": markers, "qtype": QTYPES["choice"], "target": target, "label": keys.index(gold), "qid": qid, "gold_op": gop, "src": c.get("source")}
            # class balance: CLICK dominates the operation question, so repeat the rare operations
            reps = {"DONE": 4, "TYPE_TEXT": 3, "SELECT": 3, "PRESS_ENTER": 4, "SCROLL_DOWN": 2}.get(gop, 1) if qid == "operation" else 1
            if c.get("source") == "dagger":
                reps *= 5
            if c.get("source") == "webgym" or str(c.get("skill", "")).endswith("@sub"):
                reps = 1      # webgym already pairs every DONE with its "keep going" twin; no op-level oversampling
            if c.get("source") == "rollout":
                reps *= 3   # scripted scroll / search-submit / select trajectories: the skills the model lacked   # on-policy teacher corrections from real tasks: few but exactly where the policy fails
            items.extend([item] * reps)
        if NOUL:
            yes = gop == "DONE"
            rn = random.Random(h + 7)
            if yes or rn.random() < NOUL_NEG_P:
                nq = goal_done_question(goal)
                qq = {"t": "noul", "ins": json.dumps(nq["instructions"]), "crit": None}
                seq, markers = build_sequence(tok, state, qq, MAX_LEN, HEAD_MAX_LEN)
                if len(markers) == 2:
                    items.append({"ids": array("i", seq), "markers": markers, "qtype": QTYPES["noul"], "target": [0.0, 1.0] if yes else [1.0, 0.0],
                                  "label": int(yes), "qid": "goal_done", "gold_op": gop, "src": c.get("source")})
    torch.save(items, os.path.join(out, "train_items.pt"))
    with open(os.path.join(out, "eval_cases.jsonl"), "w") as f:
        for c in ev: f.write(json.dumps(c, ensure_ascii=False) + "\n")
    print(f"dropped {n_suite} cases on suite start pages; dagger cases are never used")
    lens = [len(i["ids"]) for i in items]
    import collections
    print("goal_done (noul) items:", collections.Counter(i["label"] for i in items if i["qid"] == "goal_done"))
    print("operation label counts:", dict(collections.Counter(i["gold_op"] for i in items if i["qid"] == "operation")))
    print(f"train items {len(items)} (op {sum(i['qid']=='operation' for i in items)}, target {sum(i['qid']!='operation' for i in items)}), "
          f"eval cases {len(ev)}, skipped {n_skip}, seq len mean {sum(lens)/len(lens):.0f} max {max(lens)}")

if __name__ == "__main__":
    main()

	"""cases.jsonl + pages.jsonl -> tokenized training items (train split) and eval cases (held-out pages).

	python finetune/build_items.py out/pages.jsonl out/cases.jsonl out/
	"""
	import json, random, os, sys
	from array import array
	import torch
	sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
	from common_ft import build_request, gold_for, goal_done_question
	from transformers import AutoTokenizer
	from laya.common import QTYPES, build_sequence, render_options

	MAX_LEN, HEAD_MAX_LEN = int(os.environ.get('LAYA_MAXLEN', '1024')), int(os.environ.get('LAYA_HEAD', '512'))
	EVAL_EVERY = 5 # pages with index % 5 == 0 are held out
	MAX_TARGETS = int(os.environ.get('MAX_TARGETS', '40'))
	FINAL_P = float(os.environ.get('FINAL_P', '0.2')) # share of target items built like the server's final round
	# NOUL=1: also emit the yes/no question "is the whole goal visibly done on this page?" (common_ft.GOAL_DONE).
	# yes = states whose gold is DONE; no = every mid-trajectory state (incl. webgym's results-shown-but-filter-missing
	# states). NOUL_NEG_P keeps the negatives at roughly 60:40 (no-heavy, to counter the DONE-too-early bias).
	NOUL = os.environ.get('NOUL') == '1'
	NOUL_NEG_P = float(os.environ.get('NOUL_NEG_P', '0.3'))


	def main():
	pages_f, cases_f, out = sys.argv[1:4]
	pages = [json.loads(l) for l in open(pages_f)]
	cases = [json.loads(l) for l in open(cases_f)]
	for extra in sys.argv[4:]:
	cases += [json.loads(l) for l in open(extra)]
	snap = os.environ.get("LAYA_BASE")
	tok = AutoTokenizer.from_pretrained(os.path.join(snap, "tokenizer"))
	items, n_skip, ev = [], 0, []
	import hashlib
	STOPS = [" Stop when it is open.", " Stop once that page is visible.", " Then stop.", " Finish when it has loaded.", ""]
	# evaluation hygiene: never train on a suite task's start page (suite A and B start URLs), whatever the source
	import re as _re
	def _norm(u): return (u or "").split("#")[0].rstrip("/").replace("://www.", "://")
	suite_urls = set()
	for f in ("apps/browser_suite.py", "apps/browser_suite_b.py"):
	if os.path.exists(f):
	suite_urls \|= {_norm(u) for u in _re.findall(r'\("[a-z0-9-]+", "(https?://[^"]+)"', open(f).read())}
	n_suite = 0
	if os.environ.get("LAYA_FMT") == "v6":
	# v6 history carries each action's result = the page the NEXT step starts on (same trajectory, one more step)
	traj = lambda c: (c.get("source"), c.get("task_id") or c.get("goal"), c.get("website"))
	at = {}
	for c in cases:
	pg = c.get("page_obj") or (pages[c["page"]] if c.get("page", -1) >= 0 else {})
	at.setdefault((traj(c), len(c.get("history") or [])), (pg.get("url") or c.get("url"), pg.get("title") or c.get("title")))
	filled = 0
	for c in cases:
	for i, h in enumerate(c.get("history") or []):
	if not h.get("url") and (traj(c), i + 1) in at:
	h["url"], h["title"] = at[(traj(c), i + 1)]; filled += 1
	print("v6: history steps given their result page:", filled, flush=True)
	keep = float(os.environ.get("CASE_KEEP", "1")) # CASE_KEEP=p: build from a random p of the cases (smaller, faster)
	if keep < 1:
	cases = [c for c in cases if int(hashlib.md5((c["goal"] + str(len(c.get("history") or []))).encode()).hexdigest()[:8], 16) / 0xffffffff < keep]
	print("CASE_KEEP", keep, "->", len(cases), "cases", flush=True)
	for c in cases:
	if _norm(c.get("url") or (c.get("page_obj") or {}).get("url")) in suite_urls or (c.get("page", -1) >= 0 and _norm(pages[c["page"]]["url"]) in suite_urls):
	n_suite += 1; continue
	if c.get("source") == "dagger":
	continue
	if c["kind"] == "done" and "page_obj" not in c:
	continue # template DONE cases leak phrasing; real ones come from make_done_cases.py
	page = c.get("page_obj") or pages[c["page"]]
	h = int(hashlib.md5(c["goal"].encode()).hexdigest(), 16)
	is_sub = str(c.get("skill", "")).endswith("@sub") or c.get("skill") == "subgoal_done"
	goal = c["goal"] + ("" if is_sub or c.get("source") == "webgym" else STOPS[h % len(STOPS)])
	c = {**c, "goal": goal}
	state, questions, targets, controls = build_request(page, goal, c.get("history", []))
	if c.get("noul_only"):
	# a state where a failed run wrongly stopped: only the completion question, answered "not done"
	if NOUL and not ((hashlib.md5(c["website"].encode()).digest()[0] % EVAL_EVERY == 0)):
	nq = goal_done_question(goal)
	qq = {"t": "noul", "ins": json.dumps(nq["instructions"]), "crit": None}
	seq, markers = build_sequence(tok, state, qq, MAX_LEN, HEAD_MAX_LEN)
	if len(markers) == 2:
	items.extend([{"ids": array("i", seq), "markers": markers, "qtype": QTYPES["noul"], "target": [1.0, 0.0], "label": 0,
	"qid": "goal_done", "gold_op": "NOT_DONE", "src": c.get("source")}] * 2)
	elif NOUL:
	ev.append({**c, "gold_index": None})
	continue
	gop, gidx = gold_for(c, targets, controls)
	if gop is None:
	n_skip += 1; continue
	held = (hashlib.md5(c["website"].encode()).digest()[0] % EVAL_EVERY == 0) if c.get("source") in ("mind2web", "nnetnav", "webchain", "gobrowse") else (c["page"] % EVAL_EVERY == 0)
	if held:
	ev.append({**c, "gold_index": gidx}); continue
	golds = {"operation": gop}
	if gidx is not None:
	golds[gop.lower() + "_target"] = gidx
	for qid, gold in golds.items():
	q = questions[qid]; keys = list(q["criteria"])
	rs = random.Random(h + len(keys) + len(qid))
	if qid != "operation" and len(keys) > 8 and rs.random() < FINAL_P:
	width = rs.randint(2, 6) # the server's final round: only the chunk winners (2-6 options)
	else:
	width = MAX_TARGETS
	if qid != "operation" and len(keys) > width:
	# inference splits wide choices into chunks of <= MAX_TARGETS (systemone_server MAXOPT); train on the same
	# widths, so option labels are never truncated to a few tokens: gold + a random sample, page order kept
	keep = set(rs.sample([k for k in keys if k != gold], width - 1)) \| {gold}
	keys = [k for k in keys if k in keep]
	q = {**q, "criteria": {k: q["criteria"][k] for k in keys}}
	target = [1.0 if k == gold else 0.0 for k in keys]
	qq = {"t": "choice", "ins": json.dumps(q["instructions"]), "crit": q["criteria"]}
	seq, markers = build_sequence(tok, state, qq, MAX_LEN, HEAD_MAX_LEN)
	if len(markers) != len(render_options(qq)):
	n_skip += 1; continue
	item = {"ids": array("i", seq), "markers": markers, "qtype": QTYPES["choice"], "target": target, "label": keys.index(gold), "qid": qid, "gold_op": gop, "src": c.get("source")}
	# class balance: CLICK dominates the operation question, so repeat the rare operations
	reps = {"DONE": 4, "TYPE_TEXT": 3, "SELECT": 3, "PRESS_ENTER": 4, "SCROLL_DOWN": 2}.get(gop, 1) if qid == "operation" else 1
	if c.get("source") == "dagger":
	reps *= 5
	if c.get("source") == "webgym" or str(c.get("skill", "")).endswith("@sub"):
	reps = 1 # webgym already pairs every DONE with its "keep going" twin; no op-level oversampling
	if c.get("source") == "rollout":
	reps *= 3 # scripted scroll / search-submit / select trajectories: the skills the model lacked # on-policy teacher corrections from real tasks: few but exactly where the policy fails
	items.extend([item] * reps)
	if NOUL:
	yes = gop == "DONE"
	rn = random.Random(h + 7)
	if yes or rn.random() < NOUL_NEG_P:
	nq = goal_done_question(goal)
	qq = {"t": "noul", "ins": json.dumps(nq["instructions"]), "crit": None}
	seq, markers = build_sequence(tok, state, qq, MAX_LEN, HEAD_MAX_LEN)
	if len(markers) == 2:
	items.append({"ids": array("i", seq), "markers": markers, "qtype": QTYPES["noul"], "target": [0.0, 1.0] if yes else [1.0, 0.0],
	"label": int(yes), "qid": "goal_done", "gold_op": gop, "src": c.get("source")})
	torch.save(items, os.path.join(out, "train_items.pt"))
	with open(os.path.join(out, "eval_cases.jsonl"), "w") as f:
	for c in ev: f.write(json.dumps(c, ensure_ascii=False) + "\n")
	print(f"dropped {n_suite} cases on suite start pages; dagger cases are never used")
	lens = [len(i["ids"]) for i in items]
	import collections
	print("goal_done (noul) items:", collections.Counter(i["label"] for i in items if i["qid"] == "goal_done"))
	print("operation label counts:", dict(collections.Counter(i["gold_op"] for i in items if i["qid"] == "operation")))
	print(f"train items {len(items)} (op {sum(i['qid']=='operation' for i in items)}, target {sum(i['qid']!='operation' for i in items)}), "
	f"eval cases {len(ev)}, skipped {n_skip}, seq len mean {sum(lens)/len(lens):.0f} max {max(lens)}")

	if __name__ == "__main__":
	main()