Buckets:
| """Contracts and aggregation for offline Visual-Jigsaw rollout quality scores.""" | |
| from __future__ import annotations | |
| import hashlib | |
| import json | |
| import math | |
| import os | |
| from collections import OrderedDict, defaultdict | |
| from pathlib import Path | |
| import numpy as np | |
| import torch | |
| from PIL import Image | |
| from .bsq_codes import unpack_visual_codes | |
| SAMPLE_SCHEMA = "zebra_visual_quality_sample_v1" | |
| SUMMARY_SCHEMA = "zebra_visual_quality_summary_v1" | |
| DEFAULT_ASPECT_TOLERANCE = 0.10 | |
| DEFAULT_BLANK_STD_THRESHOLD = 0.01 | |
| DEFAULT_DUPLICATE_DINO_THRESHOLD = 0.995 | |
| DEFAULT_DUPLICATE_LPIPS_THRESHOLD = 0.02 | |
| def read_json(path: Path) -> dict: | |
| return json.loads(path.read_text(encoding="utf-8")) | |
| def atomic_write_json(path: Path, payload: dict) -> None: | |
| path.parent.mkdir(parents=True, exist_ok=True) | |
| temporary = path.with_name(f".{path.name}.{os.getpid()}.tmp") | |
| try: | |
| temporary.write_text( | |
| json.dumps(payload, ensure_ascii=False, sort_keys=True) + "\n", | |
| encoding="utf-8", | |
| ) | |
| os.replace(temporary, path) | |
| finally: | |
| temporary.unlink(missing_ok=True) | |
| def file_sha256(path: Path, *, chunk_size: int = 1024 * 1024) -> str: | |
| digest = hashlib.sha256() | |
| with path.open("rb") as handle: | |
| for chunk in iter(lambda: handle.read(chunk_size), b""): | |
| digest.update(chunk) | |
| return digest.hexdigest() | |
| def completed_interleaved_test_runs(eval_root: Path) -> list[Path]: | |
| """Return complete full-test UniAR rollout directories in stable order.""" | |
| runs = [] | |
| for run in sorted(path for path in eval_root.iterdir() if path.is_dir()): | |
| arguments = run / "arguments.json" | |
| summary = run / "summary.json" | |
| predictions = run / "predictions.jsonl" | |
| if not (arguments.exists() and summary.exists() and predictions.exists()): | |
| continue | |
| args = read_json(arguments) | |
| result = read_json(summary) | |
| if ( | |
| args.get("mode") == "interleaved" | |
| and args.get("selection") == "test" | |
| and int(result.get("samples", 0)) == 1095 | |
| ): | |
| runs.append(run) | |
| return runs | |
| def prediction_rows(eval_run: Path) -> list[dict]: | |
| rows = [ | |
| json.loads(line) | |
| for line in (eval_run / "predictions.jsonl").read_text( | |
| encoding="utf-8" | |
| ).splitlines() | |
| if line | |
| ] | |
| sample_ids = [row["sample_id"] for row in rows] | |
| if len(sample_ids) != len(set(sample_ids)): | |
| raise ValueError(f"duplicate prediction sample IDs in {eval_run}") | |
| indices = [int(row["evaluation_index"]) for row in rows] | |
| if indices != list(range(len(rows))): | |
| raise ValueError(f"prediction evaluation indices are not contiguous: {eval_run}") | |
| return rows | |
| def sample_artifact_dirs(eval_run: Path) -> dict[str, Path]: | |
| result = {} | |
| for rank in sorted((eval_run / "shards").glob("rank-*")): | |
| predictions = rank / "predictions.jsonl" | |
| if not predictions.exists(): | |
| continue | |
| with predictions.open(encoding="utf-8") as handle: | |
| for line in handle: | |
| row = json.loads(line) | |
| sample_id = row["sample_id"] | |
| directory = rank / "samples" / sample_id | |
| if sample_id in result: | |
| raise ValueError(f"duplicate sample artifact directory: {sample_id}") | |
| if not (directory / "generated_codes.pt").exists(): | |
| raise FileNotFoundError(directory / "generated_codes.pt") | |
| result[sample_id] = directory | |
| expected = {row["sample_id"] for row in prediction_rows(eval_run)} | |
| if set(result) != expected: | |
| raise ValueError( | |
| f"sample artifact IDs differ from merged predictions: {eval_run}" | |
| ) | |
| return result | |
| class PackedItemCache: | |
| """Manifest-indexed mmap reader that preserves raw trajectory provenance.""" | |
| def __init__(self, root: Path, *, max_open_shards: int = 2): | |
| self.root = root | |
| if max_open_shards <= 0: | |
| raise ValueError("max_open_shards must be positive") | |
| rows = [ | |
| json.loads(line) | |
| for line in (root / "manifest.jsonl").read_text( | |
| encoding="utf-8" | |
| ).splitlines() | |
| if line | |
| ] | |
| self.by_id = {row["sample_id"]: row for row in rows} | |
| if len(self.by_id) != len(rows): | |
| raise ValueError("cache manifest contains duplicate sample IDs") | |
| self.max_open_shards = max_open_shards | |
| self.shards: OrderedDict[str, list[dict]] = OrderedDict() | |
| def item(self, sample_id: str) -> dict: | |
| row = self.by_id[sample_id] | |
| relative = row["cache_shard"] | |
| if relative not in self.shards: | |
| payload = torch.load( | |
| self.root / relative, | |
| map_location="cpu", | |
| weights_only=False, | |
| mmap=True, | |
| ) | |
| if payload.get("schema") != "zebra_uniar_cache_v1": | |
| raise ValueError(f"unsupported cache shard schema: {relative}") | |
| self.shards[relative] = payload["items"] | |
| while len(self.shards) > self.max_open_shards: | |
| self.shards.popitem(last=False) | |
| self.shards.move_to_end(relative) | |
| item = self.shards[relative][row["cache_index"]] | |
| if item["sample_id"] != sample_id: | |
| raise ValueError(f"cache item mismatch for {sample_id}") | |
| return item | |
| def generated_state_events(item: dict) -> list[dict]: | |
| return [ | |
| event | |
| for event in item["trajectory"]["assistant"] | |
| if event["role"] == "generated_state" | |
| ] | |
| def target_code_images(item: dict) -> list[torch.Tensor]: | |
| codes = unpack_visual_codes( | |
| item["visual_codes_packed"], item["visual_codes_shape"] | |
| ) | |
| generated = codes[item["generated_flat_mask"].to(torch.bool)] | |
| counts = [int(value) for value in item["generated_image_token_counts"].tolist()] | |
| if sum(counts) != generated.shape[0]: | |
| raise ValueError(f"generated target-code count mismatch: {item['sample_id']}") | |
| return list(torch.split(generated, counts)) | |
| def geometry_metrics( | |
| reference_size: tuple[int, int], | |
| generated_size: tuple[int, int], | |
| *, | |
| aspect_tolerance: float = DEFAULT_ASPECT_TOLERANCE, | |
| ) -> dict: | |
| """Return explicit size/aspect diagnostics before any RGB alignment.""" | |
| ref_width, ref_height = reference_size | |
| gen_width, gen_height = generated_size | |
| if min(ref_width, ref_height, gen_width, gen_height) <= 0: | |
| raise ValueError("image dimensions must be positive") | |
| ref_aspect = ref_width / ref_height | |
| gen_aspect = gen_width / gen_height | |
| log_error = abs(math.log(gen_aspect / ref_aspect)) | |
| relative_error = math.exp(log_error) - 1.0 | |
| return { | |
| "reference_width": ref_width, | |
| "reference_height": ref_height, | |
| "generated_width": gen_width, | |
| "generated_height": gen_height, | |
| "reference_aspect_ratio": ref_aspect, | |
| "generated_aspect_ratio": gen_aspect, | |
| "aspect_ratio_relative_error": relative_error, | |
| "area_ratio": (gen_width * gen_height) / (ref_width * ref_height), | |
| "exact_resolution": reference_size == generated_size, | |
| "geometry_comparable": relative_error <= aspect_tolerance, | |
| } | |
| def align_generated_to_reference( | |
| generated: Image.Image, reference: Image.Image | |
| ) -> tuple[torch.Tensor, torch.Tensor]: | |
| """Resize prediction to reference coordinates with fixed bicubic filtering.""" | |
| reference = reference.convert("RGB") | |
| generated = generated.convert("RGB") | |
| if generated.size != reference.size: | |
| generated = generated.resize(reference.size, Image.Resampling.BICUBIC) | |
| return image_tensor(generated), image_tensor(reference) | |
| def image_tensor(image: Image.Image) -> torch.Tensor: | |
| values = np.asarray(image.convert("RGB"), dtype=np.float32) / 255.0 | |
| return torch.from_numpy(values).permute(2, 0, 1).unsqueeze(0) | |
| def blank_image_metrics( | |
| image: Image.Image, | |
| *, | |
| std_threshold: float = DEFAULT_BLANK_STD_THRESHOLD, | |
| ) -> dict: | |
| values = np.asarray(image.convert("RGB"), dtype=np.float32) / 255.0 | |
| luma = 0.2126 * values[..., 0] + 0.7152 * values[..., 1] + 0.0722 * values[..., 2] | |
| std = float(luma.std()) | |
| dynamic_range = float(luma.max() - luma.min()) | |
| return { | |
| "luma_mean": float(luma.mean()), | |
| "luma_std": std, | |
| "luma_dynamic_range": dynamic_range, | |
| "blank": std < std_threshold, | |
| } | |
| def native_code_metrics( | |
| generated: torch.Tensor, target: torch.Tensor | |
| ) -> dict | None: | |
| """Compare exact-grid grouped BSQ states without decoder involvement.""" | |
| generated = generated.detach().cpu().to(torch.bool) | |
| target = target.detach().cpu().to(torch.bool) | |
| if generated.ndim == target.ndim + 1 and generated.shape[0] == 1: | |
| generated = generated[0] | |
| if generated.shape != target.shape: | |
| return None | |
| equal = generated == target | |
| positions = equal.reshape(equal.shape[0], -1) | |
| equal_bits = int(equal.sum()) | |
| total_bits = equal.numel() | |
| return { | |
| "bits": total_bits, | |
| "equal_bits": equal_bits, | |
| "hamming_bits": total_bits - equal_bits, | |
| "hamming_fraction": (total_bits - equal_bits) / total_bits, | |
| "bit_accuracy": equal_bits / total_bits, | |
| "token_agreement": float(positions.all(-1).float().mean()), | |
| } | |
| def finite_number(value) -> float | None: | |
| value = float(value) | |
| return value if math.isfinite(value) else None | |
| def distribution(values: list[float]) -> dict: | |
| array = np.asarray(values, dtype=np.float64) | |
| if not len(array): | |
| return {"count": 0, "mean": None, "q05": None, "q50": None, "q95": None} | |
| return { | |
| "count": len(array), | |
| "mean": float(array.mean()), | |
| "q05": float(np.quantile(array, 0.05)), | |
| "q50": float(np.quantile(array, 0.50)), | |
| "q95": float(np.quantile(array, 0.95)), | |
| } | |
| _QUALITY_METRICS = { | |
| "luma_mean", | |
| "luma_std", | |
| "luma_dynamic_range", | |
| "aspect_ratio_relative_error", | |
| "area_ratio", | |
| "hamming_fraction", | |
| "bit_accuracy", | |
| "token_agreement", | |
| "mse", | |
| "psnr_db", | |
| "ssim", | |
| "ms_ssim", | |
| "lpips", | |
| "dists", | |
| "dinov2_cosine", | |
| "clip_cosine", | |
| } | |
| def _flatten_metrics(value: dict, prefix: str = ""): | |
| for key, item in value.items(): | |
| name = f"{prefix}.{key}" if prefix else key | |
| if isinstance(item, dict): | |
| yield from _flatten_metrics(item, name) | |
| elif ( | |
| key in _QUALITY_METRICS | |
| and isinstance(item, (int, float)) | |
| and not isinstance(item, bool) | |
| and math.isfinite(float(item)) | |
| ): | |
| yield name, float(item) | |
| def _metric_distributions(samples: list[dict], section: str) -> dict: | |
| by_metric: dict[str, list[tuple[str, float]]] = defaultdict(list) | |
| for sample in samples: | |
| sample_id = sample["sample_id"] | |
| for pair in sample.get(section, []): | |
| for key, value in _flatten_metrics(pair): | |
| by_metric[key].append((sample_id, value)) | |
| output = {} | |
| for metric, rows in sorted(by_metric.items()): | |
| per_sample: dict[str, list[float]] = defaultdict(list) | |
| for sample_id, value in rows: | |
| per_sample[sample_id].append(value) | |
| trajectory_values = [float(np.mean(values)) for values in per_sample.values()] | |
| output[metric] = { | |
| "trajectory_balanced": distribution(trajectory_values), | |
| "image_weighted": distribution([value for _sample_id, value in rows]), | |
| } | |
| return output | |
| def summarize_samples(samples: list[dict], *, eval_run: str) -> dict: | |
| expected_images = sum(int(row["expected_images"]) for row in samples) | |
| generated_images = sum(int(row["generated_images"]) for row in samples) | |
| blank_images = sum( | |
| int(image["blank"]) for row in samples for image in row["generated_quality"] | |
| ) | |
| adjacent = [pair for row in samples for pair in row["adjacent_generated"]] | |
| near_duplicates = sum(int(pair["near_duplicate"]) for pair in adjacent) | |
| pairs = [pair for row in samples for pair in row["paired_images"]] | |
| return { | |
| "schema": SUMMARY_SCHEMA, | |
| "eval_run": eval_run, | |
| "samples": len(samples), | |
| "count_matched_samples": sum(int(row["count_match"]) for row in samples), | |
| "expected_images": expected_images, | |
| "generated_images": generated_images, | |
| "paired_images": len(pairs), | |
| "exact_resolution_pairs": sum( | |
| int(pair["geometry"]["exact_resolution"]) for pair in pairs | |
| ), | |
| "geometry_comparable_pairs": sum( | |
| int(pair["geometry"]["geometry_comparable"]) for pair in pairs | |
| ), | |
| "native_comparable_pairs": sum(pair.get("native") is not None for pair in pairs), | |
| "blank_images": blank_images, | |
| "blank_fraction": blank_images / generated_images if generated_images else None, | |
| "adjacent_generated_pairs": len(adjacent), | |
| "near_duplicate_pairs": near_duplicates, | |
| "near_duplicate_fraction": near_duplicates / len(adjacent) if adjacent else None, | |
| "generated_quality": _metric_distributions(samples, "generated_quality"), | |
| "paired_metrics": _metric_distributions(samples, "paired_images"), | |
| "adjacent_metrics": _metric_distributions(samples, "adjacent_generated"), | |
| } | |
| def campaign_row(summary: dict) -> dict: | |
| """Flatten headline trajectory-balanced metrics for cross-run comparison.""" | |
| row = { | |
| "eval_run": summary["eval_run"], | |
| "samples": summary["samples"], | |
| "count_matched_samples": summary["count_matched_samples"], | |
| "generated_images": summary["generated_images"], | |
| "paired_images": summary["paired_images"], | |
| "exact_resolution_pairs": summary["exact_resolution_pairs"], | |
| "geometry_comparable_pairs": summary["geometry_comparable_pairs"], | |
| "blank_fraction": summary["blank_fraction"], | |
| "near_duplicate_fraction": summary["near_duplicate_fraction"], | |
| } | |
| for name, values in summary["paired_metrics"].items(): | |
| row[name] = values["trajectory_balanced"]["mean"] | |
| return row | |
| __all__ = [ | |
| "DEFAULT_ASPECT_TOLERANCE", | |
| "DEFAULT_BLANK_STD_THRESHOLD", | |
| "DEFAULT_DUPLICATE_DINO_THRESHOLD", | |
| "DEFAULT_DUPLICATE_LPIPS_THRESHOLD", | |
| "PackedItemCache", | |
| "SAMPLE_SCHEMA", | |
| "SUMMARY_SCHEMA", | |
| "align_generated_to_reference", | |
| "atomic_write_json", | |
| "blank_image_metrics", | |
| "campaign_row", | |
| "completed_interleaved_test_runs", | |
| "distribution", | |
| "file_sha256", | |
| "finite_number", | |
| "generated_state_events", | |
| "geometry_metrics", | |
| "image_tensor", | |
| "native_code_metrics", | |
| "prediction_rows", | |
| "read_json", | |
| "sample_artifact_dirs", | |
| "summarize_samples", | |
| "target_code_images", | |
| ] | |
Xet Storage Details
- Size:
- 14.9 kB
- Xet hash:
- dfc89be3f2f50142da221e4a6af38517b04a0b23a036f744a90fe333f2f8d89a
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.