vision-coder-openenv / src /server /rewards /text_block_rewards.py
amaljoe88's picture
deploy: sync 712e5bc -> HF
cf6c0e0
Raw
History Blame Contribute Delete
6.35 kB
"""Text block reward: match text blocks between reference and predicted HTML.
Phase 1 of the Design2Code metrics extension.
Uses Playwright to extract text element bounding boxes and content,
then matches blocks using the Hungarian algorithm (scipy.optimize.linear_sum_assignment).
Score = 0.5 * block_match_rate + 0.5 * mean_text_similarity
"""
from __future__ import annotations
import logging
from difflib import SequenceMatcher
from typing import Optional
from openenv.server.rewards import extract_html
logger = logging.getLogger(__name__)
_VIEWPORT_W = 640
_VIEWPORT_H = 480
_IOU_MATCH_THRESHOLD = 0.1 # minimum IoU to count a block as "matched"
def _get_text_blocks(html: str, width: int = _VIEWPORT_W, height: int = _VIEWPORT_H) -> list[dict]:
"""Extract text blocks from HTML by rendering with Playwright.
Returns a list of dicts with keys: 'text', 'x', 'y', 'width', 'height'.
Only leaf-level DOM elements with direct text content are included.
"""
try:
from playwright.sync_api import sync_playwright
with sync_playwright() as p:
browser = p.chromium.launch(args=["--no-sandbox", "--disable-dev-shm-usage"])
page = browser.new_page(viewport={"width": width, "height": height})
page.set_content(html, wait_until="networkidle")
blocks = page.evaluate("""() => {
const results = [];
const walker = document.createTreeWalker(
document.body || document.documentElement,
NodeFilter.SHOW_ELEMENT,
null
);
let node;
while ((node = walker.nextNode())) {
// Collect only direct text content (ignore child elements)
const directText = Array.from(node.childNodes)
.filter(n => n.nodeType === Node.TEXT_NODE)
.map(n => n.textContent.trim())
.join(' ')
.trim();
if (!directText) continue;
const rect = node.getBoundingClientRect();
if (rect.width <= 0 || rect.height <= 0) continue;
if (rect.top < 0 || rect.left < 0) continue;
results.push({
text: directText,
x: rect.left,
y: rect.top,
width: rect.width,
height: rect.height,
});
}
return results;
}""")
browser.close()
return blocks or []
except Exception as exc:
logger.warning("Text block extraction failed: %s", exc)
return []
def _bbox_iou(a: dict, b: dict) -> float:
"""Compute IoU (Intersection over Union) between two bounding boxes."""
ax1, ay1 = a["x"], a["y"]
ax2, ay2 = ax1 + a["width"], ay1 + a["height"]
bx1, by1 = b["x"], b["y"]
bx2, by2 = bx1 + b["width"], by1 + b["height"]
inter_x1 = max(ax1, bx1)
inter_y1 = max(ay1, by1)
inter_x2 = min(ax2, bx2)
inter_y2 = min(ay2, by2)
if inter_x2 <= inter_x1 or inter_y2 <= inter_y1:
return 0.0
inter_area = (inter_x2 - inter_x1) * (inter_y2 - inter_y1)
a_area = a["width"] * a["height"]
b_area = b["width"] * b["height"]
union_area = a_area + b_area - inter_area
return inter_area / union_area if union_area > 0 else 0.0
def _text_similarity(a: str, b: str) -> float:
"""Character-level similarity between two strings using SequenceMatcher."""
if not a and not b:
return 1.0
if not a or not b:
return 0.0
return SequenceMatcher(None, a, b).ratio()
def text_block_reward(
completions: list[list[dict]],
solution: Optional[list[str]] = None,
) -> list[float]:
"""Score text block matching between rendered HTML and reference.
Uses the Hungarian algorithm to match text blocks optimally by IoU,
then scores both the block match rate and text content similarity.
Args:
completions: List of completion message lists.
solution: List of reference HTML strings (one per completion).
Returns:
List of float scores in [0.0, 1.0].
"""
results = []
for i, completion in enumerate(completions):
content = completion[0]["content"]
html = extract_html(content)
ref_html = solution[i] if solution and i < len(solution) else None
if not ref_html:
results.append(0.0)
continue
try:
import numpy as np
from scipy.optimize import linear_sum_assignment
ref_blocks = _get_text_blocks(ref_html)
pred_blocks = _get_text_blocks(html)
if not ref_blocks:
# If reference has no text blocks, penalise non-empty predictions
results.append(1.0 if not pred_blocks else 0.5)
continue
if not pred_blocks:
results.append(0.0)
continue
n_ref = len(ref_blocks)
n_pred = len(pred_blocks)
# Build IoU cost matrix (cost = 1 - IoU so lower = better match)
cost_matrix = np.zeros((n_ref, n_pred), dtype=np.float64)
for r, ref_block in enumerate(ref_blocks):
for p, pred_block in enumerate(pred_blocks):
cost_matrix[r, p] = 1.0 - _bbox_iou(ref_block, pred_block)
row_ind, col_ind = linear_sum_assignment(cost_matrix)
matched = 0
text_scores = []
for r, p in zip(row_ind, col_ind):
iou = 1.0 - cost_matrix[r, p]
if iou > _IOU_MATCH_THRESHOLD:
matched += 1
text_scores.append(
_text_similarity(ref_blocks[r]["text"], pred_blocks[p]["text"])
)
block_match_score = matched / n_ref
text_sim_score = sum(text_scores) / n_ref if text_scores else 0.0
score = 0.5 * block_match_score + 0.5 * text_sim_score
results.append(score)
except Exception as exc:
logger.warning("Text block reward failed: %s", exc)
results.append(0.0)
return results