Function test

scripts/eval_pipeline_attack.py:282–476

(args: argparse.Namespace, config_file_list: list[str], use_caption=True)

@beartype
def test(args: argparse.Namespace, config_file_list: list[str], use_caption=True) -> None:
    scores = {}
    max_steps = args.max_steps

    early_stop_thresholds = {
        "parsing_failure": args.parsing_failure_th,
        "repeating_action": args.repeating_action_failure_th,
    }

    if (
        args.observation_type
        in [
            "accessibility_tree_with_captioner",
            "image_som",
        ]
        and use_caption
    ):
        device = torch.device("cuda") if torch.cuda.is_available() else "cpu"
        dtype = torch.float16 if torch.cuda.is_available() else torch.float32
        # caption_image_fn = image_utils.get_captioning_fn(device, dtype, args.captioning_model)
        caption_image_fn = get_captioning_model(args.captioning_model)
    else:
        caption_image_fn = None

    # Load a (possibly different) captioning model for running VQA evals.
    if caption_image_fn and args.eval_captioning_model == args.captioning_model:
        eval_caption_image_fn = caption_image_fn
    else:
        eval_caption_image_fn = image_utils.get_captioning_fn(
            args.eval_captioning_model_device,
            torch.float16
            if (torch.cuda.is_available() and args.eval_captioning_model_device == "cuda")
            else torch.float32,
            args.eval_captioning_model,
        )

    agent = construct_agent(
        args,
        captioning_fn=caption_image_fn if args.observation_type == "accessibility_tree_with_captioner" else None,
    )  # NOTE: captioning_fn here is used for captioning input images.

    env = ScriptBrowserEnv(
        headless=not args.render,
        slow_mo=args.slow_mo,
        observation_type=args.observation_type,
        current_viewport_only=args.current_viewport_only,
        viewport_size={
            "width": args.viewport_width,
            "height": args.viewport_height,
        },
        save_trace_enabled=args.save_trace_enabled,
        sleep_after_execution=args.sleep_after_execution,
        # NOTE: captioning_fn here is used for LLM + captioning baselines.
        # This can be different from the captioning model used for evals.
        captioning_fn=caption_image_fn,
    )

    for config_file in config_file_list: