File size: 20,090 Bytes
78ac017
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
#!/usr/bin/env python3
"""Build a source-disjoint, pixel-grounded VST prefix-32 semantic pack.

The upstream VST metadata provides cumulative count states.  A recent-only
window cannot in general support those labels, so this compiler exposes a
uniform causal prefix ending strictly before the query time.  It emits only
plain semantic answers (open integer and balanced MCQ); no agent action or
private state is ever an actor target.
"""

from __future__ import annotations

import argparse
import collections
import concurrent.futures
import hashlib
import io
import json
import math
import os
import shutil
import tempfile
from collections.abc import Mapping, Sequence
from pathlib import Path
from typing import Any

from PIL import Image

from data_pipeline.temporal_scope.semantic_mcq import (
    DIRECT_PROMPT_PREFIX,
    DIRECT_PROMPT_SUFFIX,
    rebalance_answer_slot,
)
from viss_core.data_v0.canonical import (
    attach_content_sha256,
    canonical_json_line,
    file_sha256,
    stable_id,
    write_jsonl,
)
from viss_core.data_v0.schema import validate_row


FORMAT = "streameva.vst_prefix32_semantic.v1"
PREFLIGHT_FORMAT = "streameva.vst_count_metadata_preflight.v1"
OPTION_LETTERS = "ABCD"
TARGET_COUNTS = tuple(range(11))
PREFIX_FRAME_MAX_PIXELS = 128 * 28 * 28
PREFIX_FRAME_MAX_SIDE = 504


class VSTPrefixSemanticError(ValueError):
    """Raised when a prefix-32 semantic candidate violates its contract."""


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--selection", type=Path, required=True)
    parser.add_argument("--preflight", type=Path, required=True)
    parser.add_argument("--media-root", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--frames", type=int, default=32)
    parser.add_argument("--train-per-count", type=int, default=92)
    parser.add_argument("--holdout-per-count", type=int, default=5)
    parser.add_argument("--workers", type=int, default=16)
    parser.add_argument("--seed", type=int, default=173)
    args = parser.parse_args()
    print(json.dumps(materialize(args), indent=2, sort_keys=True))


def select_available_states(
    rows: Sequence[Mapping[str, Any]],
    *,
    media_root: Path,
    train_per_count: int,
    holdout_per_count: int,
    seed: int,
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
    """Select balanced, source-disjoint train and holdout states."""

    if train_per_count <= 0 or holdout_per_count <= 0:
        raise VSTPrefixSemanticError("train/holdout quotas must be positive")
    by_count: dict[int, list[dict[str, Any]]] = {
        value: [] for value in TARGET_COUNTS
    }
    seen_videos: set[str] = set()
    seen_ids: set[str] = set()
    for raw in rows:
        row = dict(raw)
        answer = int(row.get("answer", -1))
        relative = Path(str(row.get("source_video_relative") or ""))
        candidate_id = str(row.get("candidate_id") or "")
        if (
            answer not in by_count
            or not candidate_id
            or relative.is_absolute()
            or ".." in relative.parts
            or relative.suffix.casefold() != ".mp4"
        ):
            raise VSTPrefixSemanticError("selection row is malformed")
        if candidate_id in seen_ids or relative.as_posix() in seen_videos:
            raise VSTPrefixSemanticError("selection IDs/videos must be unique")
        seen_ids.add(candidate_id)
        seen_videos.add(relative.as_posix())
        source = (media_root / relative).resolve()
        if not source.is_file() or source.is_symlink():
            continue
        row["source_video"] = str(source)
        by_count[answer].append(row)

    train: list[dict[str, Any]] = []
    holdout: list[dict[str, Any]] = []
    required = train_per_count + holdout_per_count
    for answer in TARGET_COUNTS:
        values = by_count[answer]
        values.sort(
            key=lambda row: hashlib.sha256(
                f"{seed}|{answer}|{row['candidate_id']}".encode()
            ).digest()
        )
        if len(values) < required:
            raise VSTPrefixSemanticError(
                f"count {answer} has {len(values)} available videos; need {required}"
            )
        for row in values[:holdout_per_count]:
            value = dict(row)
            value["split"] = "policy_holdout"
            holdout.append(value)
        for row in values[holdout_per_count:required]:
            value = dict(row)
            value["split"] = "train"
            train.append(value)
    return (
        sorted(train, key=lambda row: str(row["candidate_id"])),
        sorted(holdout, key=lambda row: str(row["candidate_id"])),
    )


def uniform_causal_indices(
    *, source_total: int, source_fps: float, query_time: float, frames: int
) -> tuple[list[int], list[float]]:
    """Return exactly ``frames`` uniform indices whose timestamps are causal."""

    if (
        source_total <= 0
        or not math.isfinite(source_fps)
        or source_fps <= 0
        or not math.isfinite(query_time)
        or query_time <= 0
        or frames <= 1
    ):
        raise VSTPrefixSemanticError("invalid video sampling metadata")
    # ceil(t * fps) - 1 is strictly earlier than t, including exact boundaries.
    last = min(source_total - 1, max(0, math.ceil(query_time * source_fps) - 1))
    if last == 0:
        indices = [0] * frames
    else:
        indices = [
            int(round(position * last / (frames - 1))) for position in range(frames)
        ]
    timestamps = [index / source_fps for index in indices]
    if any(timestamp >= query_time + 1e-9 for timestamp in timestamps):
        raise VSTPrefixSemanticError("prefix sampler exposed a future frame")
    return indices, timestamps


def project_semantic_rows(
    row: Mapping[str, Any],
    *,
    relative_paths: Sequence[str],
    frame_timestamps: Sequence[float],
    source_sha256: str,
    selection_sha256: str,
) -> list[dict[str, Any]]:
    """Project one grounded count state into open and balanced-MCQ rows."""

    if len(relative_paths) != 32 or len(frame_timestamps) != 32:
        raise VSTPrefixSemanticError("prefix projection requires exactly 32 frames")
    answer = int(row["answer"])
    if answer not in TARGET_COUNTS:
        raise VSTPrefixSemanticError("count answer is outside 0--10")
    candidate_id = str(row["candidate_id"])
    split = str(row["split"])
    if split not in {"train", "policy_holdout"}:
        raise VSTPrefixSemanticError("invalid semantic split")
    images = [
        {
            "type": "image",
            "image": path,
            "frame_id": stable_id(
                "vstprefixframe", candidate_id, index, timestamp, digest_length=20
            ),
        }
        for index, (path, timestamp) in enumerate(
            zip(relative_paths, frame_timestamps, strict=True)
        )
    ]
    parent = {
        "kind": "vst_count_metadata_state",
        "id": stable_id("vstprefixsource", candidate_id, digest_length=20),
        "sha256": selection_sha256,
    }
    group_id = stable_id("vstprefixgroup", candidate_id, digest_length=20)
    question = " ".join(str(row["question"]).split())
    answer_text = str(answer)
    open_row = attach_content_sha256(
        {
            "schema": "streameva.sft_row.v1",
            "uid": stable_id("sftvstprefixopen", candidate_id, digest_length=20),
            "component": "preservation",
            "group_id": group_id,
            "split": split,
            "decision_loss_profile": "full",
            "messages": [
                {
                    "role": "user",
                    "content": [
                        *images,
                        {
                            "type": "text",
                            "text": (
                                f"{question}\nAnswer with one integer from 0 to 10 "
                                "and no other text."
                            ),
                        },
                    ],
                    "loss": False,
                },
                {
                    "role": "assistant",
                    "content": [{"type": "text", "text": answer_text}],
                    "loss": True,
                },
            ],
            "sample_loss_weight": 1.0,
            "parents": [parent],
        }
    )

    distractors = [value for value in TARGET_COUNTS if value != answer]
    distractors.sort(
        key=lambda value: (
            abs(value - answer),
            hashlib.sha256(f"{candidate_id}|{value}".encode()).digest(),
        )
    )
    options, mcq_answer = rebalance_answer_slot(
        [answer_text, *[str(value) for value in distractors[:3]]],
        answer="A",
        answer_slot=int(hashlib.sha256(candidate_id.encode()).hexdigest(), 16) % 4,
        seed=f"{candidate_id}|prefix32-mcq-v1",
    )
    option_text = "\n".join(
        f"{letter}. {value}"
        for letter, value in zip(OPTION_LETTERS, options, strict=True)
    )
    mcq_prompt = (
        f"{DIRECT_PROMPT_PREFIX}\n\nQuestion: {question}\n\n"
        f"Options:\n{option_text}\n\n{DIRECT_PROMPT_SUFFIX}"
    )
    mcq_row = attach_content_sha256(
        {
            "schema": "streameva.sft_row.v1",
            "uid": stable_id("sftvstprefixmcq", candidate_id, digest_length=20),
            "component": "preservation",
            "group_id": group_id,
            "split": split,
            "decision_loss_profile": "full",
            "messages": [
                {
                    "role": "user",
                    "content": [*images, {"type": "text", "text": mcq_prompt}],
                    "loss": False,
                },
                {
                    "role": "assistant",
                    "content": [{"type": "text", "text": mcq_answer}],
                    "loss": True,
                },
            ],
            "sample_loss_weight": 1.0,
            "parents": [parent],
        }
    )
    for output in (open_row, mcq_row):
        validate_row(output)
    # Bind source pixels in the manifest-facing auxiliary state without leaking
    # an absolute path into the strict trainer row.
    del source_sha256
    return [open_row, mcq_row]


def materialize(args: argparse.Namespace) -> dict[str, Any]:
    selection = args.selection.resolve()
    preflight_path = args.preflight.resolve()
    media_root = args.media_root.resolve()
    output = args.output.resolve()
    if output.exists():
        raise FileExistsError(output)
    if args.frames != 32:
        raise VSTPrefixSemanticError("this release is frozen to exactly 32 frames")
    if not 1 <= args.workers <= 32:
        raise VSTPrefixSemanticError("workers must lie in [1, 32]")
    preflight = json.loads(preflight_path.read_text(encoding="utf-8"))
    declaration = preflight.get("selection_artifact", {})
    if (
        preflight.get("format") != PREFLIGHT_FORMAT
        or preflight.get("status") != "pass_metadata_firewall_and_selection"
        or preflight.get("contracts", {}).get("basename_hits_admitted") != 0
        or preflight.get("contracts", {}).get("question_hits_admitted") != 0
        or declaration.get("sha256") != file_sha256(selection)
    ):
        raise VSTPrefixSemanticError("metadata preflight is not valid or hash-bound")
    rows = [
        json.loads(line)
        for line in selection.read_text(encoding="utf-8").splitlines()
        if line.strip()
    ]
    train, holdout = select_available_states(
        rows,
        media_root=media_root,
        train_per_count=args.train_per_count,
        holdout_per_count=args.holdout_per_count,
        seed=args.seed,
    )
    selected = [*train, *holdout]
    output.parent.mkdir(parents=True, exist_ok=True)
    scratch = Path(tempfile.mkdtemp(prefix=f".{output.name}.partial.", dir=output.parent))
    try:
        media_out = scratch / "media"

        def decode(row: Mapping[str, Any]) -> tuple[str, dict[str, Any]]:
            return str(row["candidate_id"]), _decode_prefix(row, media_out, args.frames)

        decoded: dict[str, dict[str, Any]] = {}
        errors: list[dict[str, str]] = []
        with concurrent.futures.ThreadPoolExecutor(max_workers=args.workers) as pool:
            futures = {pool.submit(decode, row): row for row in selected}
            for future in concurrent.futures.as_completed(futures):
                row = futures[future]
                try:
                    candidate_id, state = future.result()
                    decoded[candidate_id] = state
                except Exception as exc:  # pragma: no cover - remote media failures
                    errors.append(
                        {
                            "candidate_id": str(row["candidate_id"]),
                            "reason": f"{type(exc).__name__}: {exc}",
                        }
                    )
        write_jsonl(errors, scratch / "decode_errors.jsonl")
        if errors or len(decoded) != len(selected):
            raise VSTPrefixSemanticError(f"prefix decode failed for {len(errors)} rows")

        train_rows: list[dict[str, Any]] = []
        holdout_rows: list[dict[str, Any]] = []
        source_inventory = []
        selection_sha = file_sha256(selection)
        for row in selected:
            state = decoded[str(row["candidate_id"])]
            source_sha = file_sha256(str(row["source_video"]))
            source_inventory.append(
                {
                    "candidate_id": row["candidate_id"],
                    "source_video_relative": row["source_video_relative"],
                    "source_sha256": source_sha,
                    "split": row["split"],
                    "query_time": row["query_time"],
                    "answer": row["answer"],
                    "frame_indices": state["frame_indices"],
                    "frame_timestamps": state["frame_timestamps"],
                }
            )
            projected = project_semantic_rows(
                row,
                relative_paths=state["relative_paths"],
                frame_timestamps=state["frame_timestamps"],
                source_sha256=source_sha,
                selection_sha256=selection_sha,
            )
            (train_rows if row["split"] == "train" else holdout_rows).extend(projected)
        write_jsonl(train_rows, scratch / "trainer/sft.jsonl", sort_key="uid")
        write_jsonl(
            holdout_rows, scratch / "holdout/policy_holdout.jsonl", sort_key="uid"
        )
        write_jsonl(
            source_inventory,
            scratch / "source_inventory.jsonl",
            sort_key="candidate_id",
        )
        media_files = sorted(path for path in media_out.rglob("*.jpg") if path.is_file())
        manifest = {
            "format": FORMAT,
            "status": "materialized_internal_research_candidate",
            "training_allowed": True,
            "public_release_allowed": False,
            "protocol": {
                "causal_prefix_only": True,
                "frames": 32,
                "per_frame_max_pixels": PREFIX_FRAME_MAX_PIXELS,
                "sampling": "exact_uniform_indices_over_[0,query_time)",
                "visual_input_mode": "independent_image_frames",
            },
            "counts": {
                "train_source_states": len(train),
                "holdout_source_states": len(holdout),
                "train_rows": len(train_rows),
                "holdout_rows": len(holdout_rows),
                "media_files": len(media_files),
                "train_by_answer": dict(
                    sorted(collections.Counter(str(row["answer"]) for row in train).items())
                ),
                "holdout_by_answer": dict(
                    sorted(
                        collections.Counter(str(row["answer"]) for row in holdout).items()
                    )
                ),
            },
            "inputs": {
                "selection_sha256": selection_sha,
                "preflight_sha256": file_sha256(preflight_path),
                "builder_sha256": file_sha256(Path(__file__)),
            },
            "artifacts": {
                "trainer/sft.jsonl": file_sha256(scratch / "trainer/sft.jsonl"),
                "holdout/policy_holdout.jsonl": file_sha256(
                    scratch / "holdout/policy_holdout.jsonl"
                ),
                "source_inventory.jsonl": file_sha256(
                    scratch / "source_inventory.jsonl"
                ),
                "decode_errors.jsonl": file_sha256(scratch / "decode_errors.jsonl"),
            },
            "contracts": {
                "all_labels_metadata_derived": True,
                "all_pixels_strictly_before_query": True,
                "benchmark_firewall_inherited": True,
                "holdout_source_disjoint": True,
                "no_agent_actions": True,
                "no_private_state_actor_targets": True,
                "source_videos_globally_unique": True,
                "teacher_used_for_labels": False,
            },
            "license": {
                "public_media_redistribution_claimed": False,
                "release_mode": "recipe_plus_source_checksums",
            },
        }
        (scratch / "manifest.json").write_bytes(canonical_json_line(manifest))
        os.replace(scratch, output)
        return {**manifest, "output": str(output)}
    except Exception:
        shutil.rmtree(scratch, ignore_errors=True)
        raise


def _decode_prefix(
    row: Mapping[str, Any], media_root: Path, frames: int
) -> dict[str, Any]:
    from decord import VideoReader

    source = Path(str(row["source_video"]))
    reader = VideoReader(str(source))
    try:
        fps = float(reader.get_avg_fps())
        total = len(reader)
        indices, timestamps = uniform_causal_indices(
            source_total=total,
            source_fps=fps,
            query_time=float(row["query_time"]),
            frames=frames,
        )
        arrays = reader.get_batch(indices).asnumpy()
    finally:
        del reader
    relative_paths = []
    for array in arrays:
        image = Image.fromarray(array)
        height, width = image.height, image.width
        resized_height, resized_width = _rounded_size(height, width)
        if (resized_height, resized_width) != (height, width):
            image = image.resize((resized_width, resized_height), Image.Resampling.LANCZOS)
        buffer = io.BytesIO()
        image.save(buffer, format="JPEG", quality=95, subsampling=0, optimize=False)
        payload = buffer.getvalue()
        digest = hashlib.sha256(payload).hexdigest()
        relative = Path("media") / digest[:2] / f"{digest}.jpg"
        destination = media_root.parent / relative
        destination.parent.mkdir(parents=True, exist_ok=True)
        try:
            with destination.open("xb") as handle:
                handle.write(payload)
        except FileExistsError:
            pass
        if file_sha256(destination) != digest:
            raise VSTPrefixSemanticError("decoded frame hash mismatch")
        relative_paths.append(relative.as_posix())
    return {
        "relative_paths": relative_paths,
        "frame_indices": indices,
        "frame_timestamps": timestamps,
    }


def _rounded_size(height: int, width: int) -> tuple[int, int]:
    factor = 28
    max_pixels = PREFIX_FRAME_MAX_PIXELS
    max_side = PREFIX_FRAME_MAX_SIDE
    scale = min(1.0, max_side / max(height, width))
    if height * width * scale * scale > max_pixels:
        scale = min(scale, math.sqrt(max_pixels / (height * width)))
    new_height = max(factor, int(round(height * scale / factor)) * factor)
    new_width = max(factor, int(round(width * scale / factor)) * factor)
    while new_height * new_width > max_pixels:
        if new_height >= new_width and new_height > factor:
            new_height -= factor
        elif new_width > factor:
            new_width -= factor
        else:
            break
    return new_height, new_width


if __name__ == "__main__":
    main()