accessibility-audit-env / inference.py
dixo8055's picture
Clamp rewards strictly inside (0,1) with EPS=0.01
b5f90e7
Raw
History Blame Contribute Delete
8.02 kB
#!/usr/bin/env python3
"""
Inference Script for AccessibilityAudit-Env
============================================
MANDATORY STDOUT FORMAT:
[START] task=<task_name> env=<benchmark> model=<model_name>
[STEP] step=<n> action=<action_str> reward=<0.00> done=<true|false> error=<msg|null>
[END] success=<true|false> steps=<n> score=<score> rewards=<r1,r2,...,rn>
"""
from __future__ import annotations
import asyncio
import os
import textwrap
from typing import List, Optional
from openai import OpenAI
from accessibility_audit_env import AccessibilityAuditAction, AccessibilityAuditEnv
IMAGE_NAME = os.getenv("LOCAL_IMAGE_NAME") or os.getenv("IMAGE_NAME", "accessibility_audit_env:latest")
HF_TOKEN = os.getenv("HF_TOKEN")
API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")
MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")
BENCHMARK = "accessibility_audit"
MAX_STEPS = 5
TEMPERATURE = 0.3
MAX_TOKENS = 4096
TASKS = ["easy", "medium", "hard", "expert"]
SYSTEM_PROMPT = textwrap.dedent(
"""
You are a web accessibility expert. You receive HTML pages with WCAG 2.1
accessibility violations. Your job is to fix ALL violations while preserving
the page's visual appearance and functionality.
You will receive:
1. The current HTML source
2. A list of accessibility violations detected by axe-core (with rule IDs,
impact levels, CSS selectors, and fix guidance)
You must respond with the COMPLETE corrected HTML. Do not explain — just
output the fixed HTML. Make surgical, minimal fixes. Do not rewrite the
entire page.
Common fixes:
- Missing alt text: add descriptive alt attributes to <img>
- Missing labels: add <label for="..."> associated with form inputs
- Color contrast: change colors to meet WCAG AA (4.5:1 ratio)
- Heading hierarchy: do not skip heading levels
- ARIA roles: add role/aria-* on dialogs, tabs, etc.
- Landmark regions: wrap content in <main>, <nav>, <header>, <footer>
- Link text: replace "Click here" with descriptive link text
- Skip navigation: add a skip-to-content link
- <html> must have a lang attribute
- Address fields need autocomplete attributes
Output ONLY the complete HTML. No markdown, no explanation, no code fences.
"""
).strip()
def log_start(task: str, env_name: str, model: str) -> None:
print(f"[START] task={task} env={env_name} model={model}", flush=True)
def log_step(step: int, action: str, reward: float, done: bool, error: Optional[str]) -> None:
error_val = error if error else "null"
done_val = str(done).lower()
action_short = action[:100].replace("\n", " ")
if len(action) > 100:
action_short += "..."
print(
f"[STEP] step={step} action={action_short} reward={reward:.4f} "
f"done={done_val} error={error_val}",
flush=True,
)
def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None:
rewards_str = ",".join(f"{r:.4f}" for r in rewards)
print(
f"[END] success={str(success).lower()} steps={steps} score={score:.4f} rewards={rewards_str}",
flush=True,
)
def _strip_fences(text: str) -> str:
text = text.strip()
if text.startswith("```html"):
text = text[7:]
elif text.startswith("```"):
text = text[3:]
if text.endswith("```"):
text = text[:-3]
return text.strip()
def get_fix_from_llm(
client: OpenAI, html: str, violations_summary: str, step: int
) -> str:
user_prompt = (
f"Step {step}: Fix the following accessibility violations in this HTML.\n\n"
f"VIOLATIONS DETECTED:\n{violations_summary}\n\n"
f"CURRENT HTML:\n{html}\n\n"
"Output the COMPLETE fixed HTML only. No explanation."
)
try:
completion = client.chat.completions.create(
model=MODEL_NAME,
messages=[
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": user_prompt},
],
temperature=TEMPERATURE,
max_tokens=MAX_TOKENS,
stream=False,
)
text = (completion.choices[0].message.content or "").strip()
return _strip_fences(text) or html
except Exception as exc:
print(f"[DEBUG] Model request failed: {exc}", flush=True)
return html
def _obs(result):
"""Support both `result.observation` and result-as-observation shapes."""
return getattr(result, "observation", result)
async def run_task(client: OpenAI, env, task_id: str) -> float:
rewards: List[float] = []
steps_taken = 0
EPS = 0.01
score = EPS
success = False
log_start(task=task_id, env_name=BENCHMARK, model=MODEL_NAME)
try:
result = await env.reset(task_id=task_id)
obs = _obs(result)
for step in range(1, MAX_STEPS + 1):
if getattr(obs, "done", False):
break
fixed_html = get_fix_from_llm(
client,
getattr(obs, "html_source", ""),
getattr(obs, "violation_summary", ""),
step,
)
action_str = f"submit_html(len={len(fixed_html)})"
try:
result = await env.step(AccessibilityAuditAction(fixed_html=fixed_html))
except Exception as step_exc:
steps_taken = step
rewards.append(0.01)
log_step(step, action_str, 0.01, True, str(step_exc))
raise
obs = _obs(result)
reward = float(getattr(result, "reward", None) or getattr(obs, "reward", 0.01) or 0.01)
# Hard guarantee: every reward we report stays strictly inside (0, 1).
EPS_LOCAL = 0.01
if reward <= EPS_LOCAL:
reward = EPS_LOCAL
elif reward >= 1.0 - EPS_LOCAL:
reward = 1.0 - EPS_LOCAL
done = bool(getattr(result, "done", False) or getattr(obs, "done", False))
rewards.append(reward)
steps_taken = step
log_step(step, action_str, reward, done, None)
if done:
break
raw = rewards[-1] if rewards else EPS
score = max(EPS, min(1.0 - EPS, raw))
success = score > 0.1
except Exception as exc:
print(f"[DEBUG] Task {task_id} error: {exc}", flush=True)
success = False
finally:
score = max(EPS, min(1.0 - EPS, score))
log_end(success=success, steps=steps_taken, score=score, rewards=rewards)
return score
async def main() -> None:
client = OpenAI(base_url=API_BASE_URL, api_key=HF_TOKEN)
remote_url = os.getenv("AAE_REMOTE_URL", "").strip().rstrip("/")
provider = None
if remote_url:
base_url = remote_url
print(f"[DEBUG] Using remote env at {base_url}", flush=True)
else:
import subprocess
host_port = 8765
container_name = "aae_inference_run"
subprocess.run(["docker", "rm", "-f", container_name], capture_output=True)
subprocess.run(
["docker", "run", "-d", "--rm", "--name", container_name,
"-p", f"{host_port}:7860", IMAGE_NAME],
check=True, capture_output=True,
)
from openenv.core.containers.runtime.providers import LocalDockerProvider
provider = LocalDockerProvider()
base_url = f"http://localhost:{host_port}"
provider.wait_for_ready(base_url, timeout_s=180.0)
env = AccessibilityAuditEnv(base_url=base_url, provider=provider)
await env.connect()
try:
for task_id in TASKS:
score = await run_task(client, env, task_id)
print(f"[DEBUG] Task {task_id} final score: {score:.3f}", flush=True)
finally:
try:
await env.close()
except Exception as exc:
print(f"[DEBUG] env.close() error: {exc}", flush=True)
if __name__ == "__main__":
asyncio.run(main())