sources/qualification-snapshot/scripts/probe_torch_stream_reference.py

#!/usr/bin/env python3
"""Bounded original/reconstructed BF16 MPS quality reference, never a speed test.

Inputs: a SHA-checked mono_f32le/16 kHz manifest with explicit case roles,
references, and sample counts. Relative PCM paths are relative to the manifest.
Each selected case is an independent rolling session; total selected audio is
at most 360 seconds. No native/MLX runtime is imported. One frontend window is
computed per decoder clock, unlike upstream's optional batched conv precompute.

Example (development only):
 .venv/bin/python scripts/probe_torch_stream_reference.py --manifest MANIFEST \
   --case ID --weights original --startup realtime18 --delay-ms 240 --output NEW.json

Before heldout use, freeze the identical profile with --freeze-candidate NEW.json
(without --case/--output), binding --release-candidate-record and its required
SHA256. Then pass both pinned records and --allow-heldout on the actual run.
--self-test exercises tiny CPU-only geometry, cache, and greedy-call contracts.
"""
from __future__ import annotations

import argparse
import codecs
from dataclasses import dataclass
import gc
import hashlib
import importlib.metadata
import json
import os
from pathlib import Path
import platform
import signal
import struct
import subprocess
import sys
import time
import traceback

ROOT = Path(__file__).resolve().parents[1]
sys.path[:0] = [str(ROOT), str(ROOT / 'scripts')]
PINNED = {
    'scripts/probe_original_bf16_rolling.py': 'c08fdc1327a8d9a7b2ff09e582c48a05367f2b73a4ad18ccd203c2c70d3df2f5',
    'scripts/probe_bf16_delay240_short.py': 'c7ef90271c0b6fda32abfdaee8715844233224982fb3361cf36d7c49c55cfa31',
    'scripts/probe_bf16_rolling_control.py': '675e044aaf3e842da3dd8fac06175ef4fb6272da1d90b13c0d8a06c49caa7d70',
    'baseline_probe.py': 'd9fcfb9d4c99868b6f8d3a365411135238959e8833ed85a14f23e95fc05a2f34',
    'mixed_bundle.py': '2fb3802af6bf9191d2ad1af179baed34898a6710b268e0946544c99ee6557dfa',
    'mixed_quant.py': '33dfe17035e064ec7a61ad48290042fe1ad660e1897667d56372d7a7cac0babe',
    'audio8_ternary.py': 'c2077768445ef4fb85a32f0b81fb6e2fcce6bd70ae324d1d3443d7e8090baf5f',
    'recovery_bundle.py': '0706905ea4ef32e9ee4e6eac31b98f2a2f7826268e1eef4747e8bea07c91f47e',
}
PERIOD, LEFT, LOOK_BACK, LOOK_AHEAD = 1280, 23040, 840, 40
MAX_SAMPLES = 360 * 16000
VERSIONS = ('torch', 'transformers', 'numpy', 'safetensors', 'jiwer')
TABLE = ROOT / 'artifacts/native-tokenizer-v1.a8tok'
TABLE_SHA = '7fdc1e903297ea1fa438b77503db993c50a67f8b325b4e55754f08b8579c5410'
ALLOCATOR_ENV = {'PYTORCH_MPS_LOW_WATERMARK_RATIO': '0.65', 'PYTORCH_MPS_HIGH_WATERMARK_RATIO': '0.95'}


def require(ok, message):
    if not ok:
        raise ValueError(message)


def digest(path):
    h = hashlib.sha256()
    with Path(path).open('rb') as stream:
        for block in iter(lambda: stream.read(1024 * 1024), b''):
            h.update(block)
    return h.hexdigest()


def canonical_sha(value):
    return hashlib.sha256(json.dumps(value, sort_keys=True, separators=(',', ':'),
                                     ensure_ascii=False, allow_nan=False).encode()).hexdigest()


def release_identity(path, expected_sha):
    if path is None and expected_sha is None:
        return None
    require(path is not None and expected_sha is not None, 'release record requires path and explicit SHA256')
    path = Path(path).resolve(strict=True)
    require(digest(path) == expected_sha, 'external release candidate SHA mismatch')
    record = json.loads(path.read_text())
    require(isinstance(record, dict) and record, 'external release candidate must be a nonempty JSON object')
    return dict(path=str(path), sha256=expected_sha, canonical_identity_sha256=canonical_sha(record))


def source_helpers():
    for name, expected in PINNED.items():
        require(digest(ROOT / name) == expected, f'frozen helper changed: {name}')
    import probe_original_bf16_rolling as original
    import probe_bf16_delay240_short as short
    for name, expected in original.PINNED.items():
        require(digest(ROOT / name) == expected, f'pinned upstream/backend changed: {name}')
    return original, short


@dataclass(frozen=True)
class Schedule:
    startup: str
    delay_ms: int

    def __post_init__(self):
        require(self.startup in ('hf', 'realtime18'), 'startup must be hf or realtime18')
        require(type(self.delay_ms) is int and self.delay_ms in (80, 160, 240, 480), 'audited delays: 80/160/240/480 ms')

    @property
    def delay(self):
        return self.delay_ms // 80

    @property
    def prefill(self):
        return 19 + self.delay if self.startup == 'hf' else 18

    def bounds(self, index):
        require(type(index) is int and index >= 0, 'invalid emission index')
        start = 0 if index == 0 else (self.prefill + index - 1) * PERIOD - LOOK_BACK
        return start, (self.prefill + index) * PERIOD + LOOK_AHEAD

    def emissions(self, samples):
        require(type(samples) is int and 0 < samples <= MAX_SAMPLES, 'audio must contain 1..5760000 samples')
        if self.startup == 'hf':
            return (samples + PERIOD - 1) // PERIOD + 11
        return (samples + (self.delay + 11) * PERIOD - LOOK_AHEAD) // PERIOD + 1

    def clock(self, index, samples):
        require(index < self.emissions(samples), 'emission beyond EOF policy')
        a, b = self.bounds(index)
        nominal = max(0, b - LEFT)
        return dict(window_start_sample=a, window_end_sample=b,
                    nominal_source_dependency_samples=nominal,
                    source_dependency_samples=min(samples, nominal),
                    source_dependency_seconds=min(samples, nominal) / 16000,
                    source_eof_padding_used=nominal > samples)

    def window(self, source, index):
        import numpy as np
        a, b = self.bounds(index)
        out = np.zeros(b - a, dtype=np.float32)
        first, end = max(0, a - LEFT), min(len(source), b - LEFT)
        if end > first:
            out[LEFT + first - a:LEFT + end - a] = source[first:end]
        return out

    def metadata(self, samples=None):
        result = dict(startup=self.startup, transcription_delay_ms=self.delay_ms,
                      delay_tokens=self.delay, frame_len=4, token_period_samples=PERIOD,
                      token_period_ms=80, left_pad_samples=LEFT, startup_input_tokens=self.prefill,
                      initial_waveform_samples=self.bounds(0)[1], first_batch_windows=1,
                      batch_windows=1, rolling_context=375, trim_interval=38, stable_prefix=16,
                      first_trim_index=376 - self.prefill,
                      eof_policy=('hf_diagnostic_ceil_audio_plus_11_emissions' if self.startup == 'hf'
                                  else 'official_complete_centered_window_delay_plus_11_silence_tokens'),
                      full_official_realtime_backend_parity=False)
        if samples is not None:
            count = self.emissions(samples)
            trims = list(range(376 - self.prefill, count, 38))
            absolute = self.prefill + count - 1
            result.update(expected_emissions=count, trim_indices=trims,
                          final_absolute_text_position=absolute,
                          final_local_text_position=absolute - 38 * len(trims))
        return result


def select_cases(manifest_path, ids):
    manifest_path = Path(manifest_path).resolve(strict=True)
    manifest = json.loads(manifest_path.read_text())
    require(manifest.get('format') in ('mono_f32le', 'mono_16000_f32le') and manifest.get('sample_rate') == 16000,
            'manifest must identify mono F32LE at 16000 Hz')
    rows = manifest.get('cases', [])
    require(isinstance(rows, list) and rows and all(isinstance(x, dict) for x in rows), 'missing cases')
    catalog = {x.get('id'): x for x in rows}
    require(len(catalog) == len(rows) and all(isinstance(k, str) and k for k in catalog), 'duplicate/bad case IDs')
    require(ids and len(set(ids)) == len(ids) and all(k in catalog for k in ids), 'explicit unique known --case IDs required')
    selected = []
    for key in ids:
        row = dict(catalog[key])
        require(row.get('role') in ('development', 'diagnostic', 'heldout'), f'explicit data role required: {key}')
        require(row.get('language') in ('en', 'zh') and isinstance(row.get('reference'), str), 'language/reference missing')
        require(type(row.get('samples')) is int and 0 < row['samples'] <= MAX_SAMPLES, 'case exceeds 360-second scope')
        if 'seconds' in row:
            require(abs(row['seconds'] - row['samples'] / 16000) < 1e-8, 'manifest duration differs from sample count')
        path = Path(row['path'])
        path = (path if path.is_absolute() else manifest_path.parent / path).resolve(strict=True)
        require(path.stat().st_size == row['samples'] * 4 and digest(path) == row.get('sha256'), f'PCM identity changed: {key}')
        row['path'] = str(path)
        selected.append(row)
    require(sum(x['samples'] for x in selected) <= MAX_SAMPLES, 'selected total audio exceeds 360 seconds')
    return selected


def weight_identity(kind, bundle):
    require(kind in ('original', 'reconstructed'), 'unknown weight route')
    if kind == 'original':
        require(bundle is None, 'original weights must not receive a packed bundle')
        return dict(kind=kind, checkpoint_catalog_sha256=digest(ROOT / 'reports/baseline_checkpoint_files.json'))
    require(bundle is not None, 'reconstructed weights require --packed-bundle')
    bundle = Path(bundle).resolve(strict=True)
    manifest = json.loads((bundle / 'manifest.json').read_text())
    require(manifest.get('profile') in ('q4', 'e4_d3_full8_h4', 'e4_da4_dm3_full8_h4'), 'unsupported reconstructed profile')
    require(manifest.get('weights_file') == 'weights.bin' and manifest.get('group_size') == 64,
            'unexpected packed payload/group layout')
    require(manifest.get('source_revision') == 'b4413de154ed6bdef0a4011028b1ebd12aca8152', 'packed source revision differs')
    require((bundle / 'weights.bin').stat().st_size == manifest['weights_bytes'], 'packed payload size changed')
    require(manifest.get('tie_word_embeddings') is True and manifest.get('aliases') == {
        'language_model.lm_head.weight': 'language_model.model.embed_tokens.weight'}, 'packed embedding/head tie changed')
    return dict(kind=kind, bundle=str(bundle), manifest_sha256=digest(bundle / 'manifest.json'),
                profile=manifest['profile'], weights_sha256=manifest['weights_sha256'],
                weights_bytes=manifest['weights_bytes'], group_size=64,
                validation='Full payload/codec validation by the pinned loader before any weight copy')


def candidate_identity(args):
    original, _ = source_helpers()
    require(digest(TABLE) == TABLE_SHA, 'tokenizer byte table changed')
    return dict(manifest_sha256=digest(args.manifest), weights=weight_identity(args.weights, args.packed_bundle),
                checkpoint_catalog_sha256=digest(ROOT / 'reports/baseline_checkpoint_files.json'),
                release_candidate=release_identity(args.release_candidate_record, args.release_candidate_record_sha256),
                schedule=Schedule(args.startup, args.delay_ms).metadata(),
                source_sha256=PINNED | original.PINNED | {
                    str(Path(__file__).relative_to(ROOT)): digest(__file__)},
                hf_revision=original.HF_REVISION, upstream_revision=original.UPSTREAM_REVISION,
                token_table_sha256=TABLE_SHA, scorer='baseline_probe.normalize + jiwer word EN/character ZH',
                versions={name: importlib.metadata.version(name) for name in VERSIONS},
                mps_allocator_environment=ALLOCATOR_ENV,
                frontend='one_window_direct_upstream_no_precomputed_convolution',
                dtype='BF16', device='MPS', attention_backend='sdpa', eos_suppressed=True,
                fitting=False, activation_quantization=False, prefix=False)


def heldout_gate(cases, allow, candidate_path, candidate_sha, identity):
    if not any(row['role'] == 'heldout' for row in cases):
        return None
    require(allow and candidate_path and candidate_sha, 'heldout requires --allow-heldout and SHA-pinned --candidate-record')
    external = identity.get('release_candidate')
    require(external is not None, 'heldout requires a SHA-pinned external release candidate')
    require(release_identity(external['path'], external['sha256']) == external, 'external release identity changed')
    path = Path(candidate_path).resolve(strict=True)
    require(digest(path) == candidate_sha, 'candidate record SHA mismatch')
    record = json.loads(path.read_text())
    require(record.get('format') == 'audio8-torch-stream-candidate-v1' and record.get('status') == 'frozen',
            'candidate record is not a frozen Torch stream profile')
    require(record.get('identity') == identity and record.get('identity_sha256') == canonical_sha(identity),
            'heldout run differs from frozen weights/runtime/decode/scoring')
    return dict(path=str(path), sha256=candidate_sha, status='explicitly_authorized_frozen_profile')


class IncrementalText:
    """Consume the pinned byte-piece oracle table without importing MLX."""
    def __init__(self, table):
        self.table = table
        self.decoder = codecs.getincrementaldecoder('utf-8')('replace')
        self.text = ''
        self.finished = False

    def push(self, token):
        # The table schema is checked by the pinned short helper's parser.
        offsets, flags, pieces, special, model_count = (
            self.table[k] for k in ('offsets', 'flags', 'pieces', 'special', 'model_count'))
        require(0 <= token < model_count, 'token outside vocabulary')
        if self.finished:
            return ''
        if token == special[3]:
            return self.finish()
        if token >= len(flags) or flags[token] or token in special:
            return ''
        text = self.decoder.decode(pieces[offsets[token]:offsets[token + 1]], final=False)
        self.text += text
        return text

    def finish(self):
        if self.finished:
            return ''
        self.finished = True
        text = self.decoder.decode(b'', final=True)
        self.text += text
        return text


def decode_case(torch, model, streaming, feature, source, case, plan, special, language_id,
                text_decoder, emit, memory_check, helpers_module):
    """Only direct pinned model calls, explicit clocks, and pinned cache helpers."""
    device, dtype = next(model.parameters()).device, next(model.parameters()).dtype
    helpers = helpers_module.pure_helpers(torch)
    inv_text = model.language_model.model.rotary_emb.inv_freq
    inv_audio = model.audio_tower.rotary_emb.inv_freq
    audio_cache = text_cache = None
    local = dropped = 0
    ids, trim_indices, first_readable, first_text_delta = [], [], None, None
    t_cond = model.build_t_cond(plan.delay, batch_size=1, device=device, dtype=dtype, frame_len=4)
    maximum_cache_bytes = 0
    for index in range(plan.emissions(len(source))):
        width = plan.prefill if index == 0 else 1
        require(local + dropped == (0 if index == 0 else plan.prefill + index - 1), 'absolute token clock changed')
        if local + width > 375:
            require(local == 375 and width == 1, 'unexpected pre-query trim boundary')
            helpers_module.trim_pair(text_cache, audio_cache, helpers, inv_text, inv_audio)
            local -= 38
            dropped += 38
            trim_indices.append(index)
            emit(dict(kind='torch_stream_trim', case_id=case['id'], index=index,
                      dropped_tokens=38, total_dropped_tokens=dropped, local_text_position=local,
                      encoder_local_position=local * 4, encoder_rebase=-152, stable_prefix=16))
        require((0 if text_cache is None else text_cache.get_seq_length()) == local, 'decoder pre-query cache clock changed')
        require((0 if audio_cache is None else audio_cache.get_seq_length()) == local * 4, 'encoder pre-query cache clock changed')
        prompt = ([special['bos_token_id'], language_id] + [special['streaming_pad_token_id']] * (width - 2)
                  if index == 0 else [ids[-1]])
        audio_window = torch.as_tensor(plan.window(source, index), device=device, dtype=torch.float32).unsqueeze(0)
        audio_positions = torch.arange(local * 4, (local + width) * 4, device=device, dtype=torch.long).unsqueeze(0)
        audio_embeds, audio_cache = streaming._streaming_audio_embeds_batch(
            model=model, feature_extractor=feature, audio_windows=audio_window,
            audio_past_key_values=audio_cache, sampling_rate=16000, num_80ms_tokens=width,
            frame_len=4, position_ids=audio_positions, precomputed_conv_embeds=None, dtype=dtype, device=device)
        input_ids = torch.tensor([prompt], device=device, dtype=torch.long)
        inputs_embeds = model.build_text_inputs_embeds(input_ids=input_ids, audio_embeds=audio_embeds).to(dtype=dtype)
        positions = torch.arange(local, local + width, device=device, dtype=torch.long)
        output = model.forward_language_model_with_delay(
            inputs_embeds=inputs_embeds, attention_mask=None, position_ids=positions.unsqueeze(0),
            cache_position=positions, past_key_values=text_cache, use_cache=True, logits_to_keep=1,
            num_delay_tokens=plan.delay, frame_len=4, t_cond=t_cond)
        text_cache = output.past_key_values
        local += width
        require(text_cache.get_seq_length() == local <= 375, 'decoder post-query cache bound changed')
        require(audio_cache.get_seq_length() == local * 4, 'encoder post-query cumulative clock changed')
        require(all(x.keys.shape[-2] == local and x.keys.shape == x.values.shape for x in text_cache.layers),
                'decoder per-layer cache bound changed')
        require(all(x.keys.shape[-2] == min(749, local * 4) and x.keys.shape == x.values.shape
                    and x.get_seq_length() == local * 4 for x in audio_cache.layers), 'encoder per-layer cache bound changed')
        logits = output.logits[0, -1, :].clone()
        logits[special['eos_token_id']] = float('-inf')
        token = int(logits.argmax(dim=-1).item())
        require(token != special['eos_token_id'], 'EOS suppression failed')
        ids.append(token)
        delta = text_decoder.push(token)
        clock = plan.clock(index, len(source))
        if delta and first_text_delta is None:
            first_text_delta = dict(index=index, text=delta, **clock)
        if delta.strip() and first_readable is None:
            first_readable = dict(index=index, text=delta, criterion='non_whitespace_utf8_delta', **clock)
        cache_bytes = sum(x.keys.numel() * x.keys.element_size() + x.values.numel() * x.values.element_size()
                          for cache in (audio_cache, text_cache) for x in cache.layers)
        maximum_cache_bytes = max(maximum_cache_bytes, cache_bytes)
        event = dict(kind='torch_stream_token', case_id=case['id'], index=index, token_id=token,
                     visible_text_delta=delta, local_text_position=local, absolute_text_position=local + dropped,
                     encoder_local_position=local * 4, encoder_absolute_position=(local + dropped) * 4,
                     decoder_cached_tokens=local, encoder_history_frames=int(audio_cache.layers[0].keys.shape[-2]),
                     cache_tensor_payload_bytes=cache_bytes, trim_count=len(trim_indices), **clock)
        if index == 0 or index in trim_indices or delta:
            values, indices = logits.float().topk(2)
            event.update(top_two_ids=indices.cpu().tolist(), top_two_logits=values.cpu().tolist())
        emit(event)
        del output, logits, inputs_embeds, audio_embeds, audio_window, input_ids, positions, audio_positions
        memory_check(index)
    tail = text_decoder.finish()
    if tail:
        emit(dict(kind='torch_stream_text_flush', case_id=case['id'], text=tail, source_dependency_samples=len(source)))
    expected = plan.metadata(len(source))
    require(trim_indices == expected['trim_indices'] and local == expected['final_local_text_position'], 'final trim schedule differs')
    return dict(generated_token_ids=ids, generated_ids_sha256=hashlib.sha256(struct.pack(f'<{len(ids)}I', *ids)).hexdigest(),
                text=text_decoder.text, first_readable=first_readable, first_text_delta=first_text_delta,
                final_local_text_position=local,
                final_absolute_text_position=local + dropped, trim_indices=trim_indices,
                max_cache_tensor_payload_bytes=maximum_cache_bytes, all_expected_windows_emitted=True,
                full_input_samples_observed=len(source), schedule=expected)


def score(reference, hypothesis, language):
    import jiwer
    from baseline_probe import normalize
    ref, hyp = normalize(reference, language), normalize(hypothesis, language)
    metric = jiwer.process_characters(ref, hyp) if language == 'zh' else jiwer.process_words(ref, hyp)
    return dict(reference=reference, normalized_reference=ref, normalized_hypothesis=hyp,
                metric_name='CER' if language == 'zh' else 'WER',
                error_rate=metric.cer if language == 'zh' else metric.wer,
                substitutions=metric.substitutions, deletions=metric.deletions,
                insertions=metric.insertions, hits=metric.hits,
                reference_units=metric.hits + metric.substitutions + metric.deletions)


def worker(output):
    from baseline_probe import MemorySampler, atomic_json
    report = json.loads(output.read_text())
    sampler, started = MemorySampler(), time.monotonic()
    stage, current_case = 'preflight', None
    events = Path(report['paths']['events']).open('x')
    def emit(event):
        event.update(stage=stage, elapsed_seconds_diagnostic_only=time.monotonic() - started)
        events.write(json.dumps(event, ensure_ascii=False) + '\n')
        events.flush()
    try:
        require(sys.platform == 'darwin' and platform.machine() == 'arm64', 'Apple Silicon MPS required')
        require(Path(sys.prefix).resolve() == (ROOT / '.venv').resolve(), 'use the source-pinned .venv Python')
        original, short = source_helpers()
        identity = report['candidate_identity']
        for path, expected in identity['source_sha256'].items():
            require(digest(ROOT / path) == expected, f'source changed after supervisor preflight: {path}')
        require({k: os.environ.get(k) for k in ALLOCATOR_ENV} == ALLOCATOR_ENV, 'allocator profile differs')
        require(subprocess.check_output(['git', '-C', str(ROOT / 'vendor/Audio8-ASR-Infinite'), 'rev-parse', 'HEAD'], text=True).strip()
                == original.UPSTREAM_REVISION, 'upstream revision changed')
        require(digest(report['manifest']) == identity['manifest_sha256'], 'input manifest changed')
        cases = select_cases(report['manifest'], [x['id'] for x in report['selected_cases']])
        require(cases == report['selected_cases'], 'selected input identity changed')
        if report['heldout_authorization']:
            auth = report['heldout_authorization']
            heldout_gate(cases, True, auth['path'], auth['sha256'], identity)
        checkpoint, report['checkpoint_provenance'] = short.verify_checkpoint()
        require(report['checkpoint_provenance']['catalog_sha256'] == identity['checkpoint_catalog_sha256'],
                'checkpoint catalog changed after profile freeze')
        read_table, oracle = short.token_helpers()
        require(digest(TABLE) == TABLE_SHA, 'token table changed')
        table = read_table(TABLE)
        import numpy as np
        import torch
        from transformers import AutoTokenizer, AutoFeatureExtractor
        sys.path.insert(0, str(ROOT / 'vendor/Audio8-ASR-Infinite'))
        from audio8_asr_infinite import streaming_inference as streaming
        from audio8_asr_infinite.modeling.modeling_audio8_asr_infinite import (
            Audio8ASRInfiniteForConditionalGeneration, resolve_qwen_language_token_id,
            resolve_qwen_streaming_special_token_ids)
        require(torch.backends.mps.is_available(), 'MPS unavailable')
        require({name: importlib.metadata.version(name) for name in VERSIONS} == identity['versions'], 'package versions changed')
        torch.set_num_threads(4)
        torch.manual_seed(0)
        sampler.torch = torch
        def memory_check(index=None):
            original.check_mps_memory(torch.mps, ceiling_bytes=18_000_000_000,
                                      release_threshold_bytes=12_000_000_000, emit=emit, after_token_index=index)
        stage = 'loading_cpu_bf16'
        report['status'] = stage
        atomic_json(output, report)
        tokenizer = AutoTokenizer.from_pretrained(str(checkpoint), trust_remote_code=True, local_files_only=True)
        feature = AutoFeatureExtractor.from_pretrained(str(checkpoint), trust_remote_code=True, local_files_only=True)
        model = Audio8ASRInfiniteForConditionalGeneration.from_pretrained(
            str(checkpoint), torch_dtype=torch.bfloat16, trust_remote_code=True, local_files_only=True)
        require(sum(p.numel() for p in model.parameters()) == 4086290208, 'model inventory changed')
        require(all(p.device.type == 'cpu' and p.dtype == torch.bfloat16 for p in model.parameters()), 'CPU BF16 destination required')
        weights = identity['weights']
        if weights['kind'] == 'reconstructed':
            require(weight_identity('reconstructed', weights['bundle']) == weights, 'packed manifest changed')
            stage = 'reconstructing_cpu_bf16'
            report['status'] = stage
            atomic_json(output, report)
            if weights['profile'] == 'e4_da4_dm3_full8_h4':
                from recovery_bundle import load_into_model
            else:
                from mixed_bundle import load_into_model
            loaded = load_into_model(model, Path(weights['bundle']), chunk_rows=64)
            require(loaded['weights_sha256'] == weights['weights_sha256'] and loaded['weights_bytes'] == weights['weights_bytes']
                    and loaded['verification']['passed'] and loaded['refit_performed'] is False, 'packed load validation mismatch')
            report['dense_reconstruction'] = loaded
        require(model.language_model.lm_head.weight.data_ptr() == model.language_model.model.embed_tokens.weight.data_ptr(), 'tied head changed')
        model.to(device='mps').eval()
        model.config.use_cache = True
        model.gradient_checkpointing_disable()
        report['attention_implementations'] = dict(encoder=model.audio_tower.config._attn_implementation,
                                                  decoder=model.language_model.config._attn_implementation)
        require(report['attention_implementations'] == dict(encoder='sdpa', decoder='sdpa'), 'SDPA profile changed')
        memory_check()
        special = resolve_qwen_streaming_special_token_ids(tokenizer)
        report['special_ids'] = special
        plan = Schedule(report['startup'], report['delay_ms'])
        report['results'] = []
        for case in cases:
            current_case, stage = case['id'], 'decoding'
            report.update(status='decoding', current_case_id=current_case)
            atomic_json(output, report)
            source = np.memmap(case['path'], dtype='<f4', mode='r')
            for first in range(0, len(source), 65536):
                require(np.isfinite(source[first:first + 65536]).all(), 'nonfinite PCM')
            language = resolve_qwen_language_token_id(tokenizer, case['language'])
            emit(dict(kind='torch_stream_case_ready', case_id=case['id'], role=case['role'],
                      input_sha256=case['sha256'], input_samples=case['samples'], language=case['language'],
                      language_token_id=language, schedule=plan.metadata(len(source)), weights=weights,
                      prefix_used=False, activation_quantization=False, speed_claim=False))
            with torch.inference_mode():
                result = decode_case(torch, model, streaming, feature, source, case, plan, special, language,
                                     IncrementalText(table), emit, memory_check, original)
            require(digest(case['path']) == case['sha256'], 'PCM changed during decode')
            require(oracle(result['generated_token_ids'], table)['text'] == result['text'], 'independent byte oracle text mismatch')
            visible, _, _ = streaming.decode_visible_text(tokenizer=tokenizer, generated_token_ids=result['generated_token_ids'], special_ids=special)
            require(visible == result['text'], 'upstream text decoder mismatch')
            result.update(id=case['id'], role=case['role'], language=case['language'], input_sha256=case['sha256'],
                          input_samples=case['samples'], language_token_id=language, **score(case['reference'], result['text'], case['language']))
            report['results'].append(result)
            emit(dict(kind='torch_stream_case_complete', **result))
            atomic_json(output, report)
            del source
            gc.collect()
            torch.mps.synchronize()
            torch.mps.empty_cache()
            memory_check()
        if report.get('expected_short_report'):
            expected_path = Path(report['expected_short_report'])
            require(digest(expected_path) == report['expected_short_report_sha256'], 'short reference changed')
            expected = json.loads(expected_path.read_text())
            require(weights['kind'] == 'original' and plan == Schedule('hf', 240) and len(cases) == 1
                    and cases[0]['sha256'] == short.INPUT_SHA and expected['status'] == 'complete', 'short comparison profile mismatch')
            old = expected['passes']['original_bf16']
            result = report['results'][0]
            mismatches = [dict(index=i, direct=a, saved_batched_conv=b)
                          for i, (a, b) in enumerate(zip(result['generated_token_ids'], old['token_ids'])) if a != b]
            report['short_original_parity'] = dict(token_count_equal=len(result['generated_token_ids']) == len(old['token_ids']),
                all_ids_equal=not mismatches and len(result['generated_token_ids']) == len(old['token_ids']),
                mismatch_rows=mismatches, text_equal=result['text'] == old['result']['final_text'],
                difference='Direct single-window frontend versus saved batched convolution precompute')
            require(report['short_original_parity']['all_ids_equal'] and report['short_original_parity']['text_equal'], 'short original parity gate failed')
        for path, expected in identity['source_sha256'].items():
            require(digest(ROOT / path) == expected, f'implementation changed during reference: {path}')
        require(digest(report['manifest']) == identity['manifest_sha256'], 'manifest changed during reference')
        if identity.get('release_candidate'):
            external = identity['release_candidate']
            require(release_identity(external['path'], external['sha256']) == external, 'release candidate changed during reference')
        report.update(status='complete', assets_unchanged=True, current_case_id=None)
        emit(dict(kind='torch_stream_complete', cases=len(cases), quality_approved=False, speed_claim=False))
    except BaseException as error:
        report.update(status='failed', error=f'{type(error).__name__}: {error}', traceback=traceback.format_exc(), current_case_id=current_case)
        emit(dict(kind='torch_stream_failed', error=report['error'], case_id=current_case))
        raise
    finally:
        sampler.stop.set()
        sampler.thread.join(timeout=1)
        report['memory'] = sampler.report()
        report['finished_at_unix'] = time.time()
        events.close()
        report['events_sha256'] = digest(report['paths']['events'])
        atomic_json(output, report)


def self_test():
    """Small real CPU cache tests; no checkpoint, accelerator, or downloads."""
    import unittest
    import numpy as np
    import torch
    from types import SimpleNamespace
    from transformers.cache_utils import DynamicLayer, DynamicSlidingWindowLayer
    torch.set_num_threads(1)
    original, short = source_helpers()

    class Cache:
        def __init__(self, layers):
            self.layers = layers
        def get_seq_length(self):
            return self.layers[0].get_seq_length()

    class Model:
        def __init__(self):
            self.parameter = torch.nn.Parameter(torch.zeros(1, dtype=torch.bfloat16))
            inv = 1 / (1e6 ** (torch.arange(0, 8, 2).float() / 8))
            self.language_model = SimpleNamespace(model=SimpleNamespace(rotary_emb=SimpleNamespace(inv_freq=inv)))
            self.audio_tower = SimpleNamespace(rotary_emb=SimpleNamespace(inv_freq=inv))
            self.prompts, self.audio_positions, self.positions = [], [], []
        def parameters(self):
            yield self.parameter
        def build_t_cond(self, *args, **kwargs):
            return torch.zeros(1, 1, 8, dtype=torch.bfloat16)
        def build_text_inputs_embeds(self, input_ids, audio_embeds):
            self.prompts.append(input_ids.tolist()[0])
            return audio_embeds
        def audio(self, **kw):
            width = kw['num_80ms_tokens'] * 4
            cache = kw['audio_past_key_values'] or Cache([DynamicSlidingWindowLayer(sliding_window=750)])
            pos = kw['position_ids'].flatten().tolist()
            assert pos == list(range(cache.get_seq_length(), cache.get_seq_length() + width))
            self.audio_positions.append(pos)
            values = torch.zeros(1, 1, width, 8, dtype=torch.bfloat16)
            cache.layers[0].update(values, values)
            return torch.zeros(1, width // 4, 8, dtype=torch.bfloat16), cache
        def forward_language_model_with_delay(self, **kw):
            cache = kw['past_key_values'] or Cache([DynamicLayer()])
            width = kw['inputs_embeds'].shape[1]
            pos = kw['position_ids'].flatten().tolist()
            assert pos == kw['cache_position'].tolist() == list(range(cache.get_seq_length(), cache.get_seq_length() + width))
            self.positions.append(pos)
            values = torch.zeros(1, 1, width, 8, dtype=torch.bfloat16)
            cache.layers[0].update(values, values)
            logits = torch.zeros(1, 1, 8, dtype=torch.bfloat16)
            logits[0, 0, 3], logits[0, 0, 5] = 100, 1
            return SimpleNamespace(past_key_values=cache, logits=logits)

    class Text:
        text = ''
        def push(self, token):
            assert token == 5
            self.text += 'x'
            return 'x'
        def finish(self):
            return ''

    class Tests(unittest.TestCase):
        def test_geometry_and_eof(self):
            for startup in ('hf', 'realtime18'):
                for delay in (80, 160, 240, 480):
                    plan = Schedule(startup, delay)
                    for samples in (1, 39, 40, 41, 1279, 1280, 1281, 93680):
                        source = np.arange(samples, dtype=np.float32)
                        right = ((samples + PERIOD - 1) // PERIOD * PERIOD - samples + (plan.delay + 12) * PERIOD
                                 if startup == 'hf' else (plan.delay + 11) * PERIOD)
                        padded = np.concatenate((np.zeros(LEFT, np.float32), source, np.zeros(right, np.float32)))
                        count = 0
                        while plan.bounds(count)[1] <= len(padded):
                            a, b = plan.bounds(count)
                            np.testing.assert_array_equal(plan.window(source, count), padded[a:b])
                            count += 1
                        self.assertEqual(count, plan.emissions(samples))
                        self.assertEqual(plan.clock(count - 1, samples)['source_dependency_samples'], samples)
            self.assertEqual(Schedule('hf', 240).bounds(0), (0, 28200))
            self.assertEqual(Schedule('hf', 240).metadata()['first_trim_index'], 354)
            self.assertEqual(Schedule('realtime18', 80).bounds(0), (0, 23080))

        def test_direct_call_feedback_repeated_trims_and_short_eof(self):
            special = dict(bos_token_id=0, pad_token_id=1, streaming_word_token_id=2,
                           eos_token_id=3, streaming_pad_token_id=7)
            for startup in ('hf', 'realtime18'):
                for delay in (80, 160, 240, 480):
                    for samples in (1, 40, 1281, 40 * 16000):
                        plan, model, events = Schedule(startup, delay), Model(), []
                        with torch.inference_mode():
                            result = decode_case(torch, model, SimpleNamespace(_streaming_audio_embeds_batch=model.audio),
                                None, np.zeros(samples, np.float32), {'id': 'tiny'}, plan, special, 4, Text(),
                                events.append, lambda index: None, original)
                        self.assertEqual(result['generated_token_ids'], [5] * plan.emissions(samples))
                        self.assertEqual(model.prompts[0], [0, 4] + [7] * (plan.prefill - 2))
                        self.assertEqual(model.prompts[1:], [[5]] * (plan.emissions(samples) - 1))
                        for index in result['trim_indices']:
                            self.assertEqual(model.positions[index], [337])
                            self.assertEqual(model.audio_positions[index], list(range(1348, 1352)))
                        self.assertEqual(result['final_absolute_text_position'], plan.prefill + len(model.prompts) - 1)

        def test_invalid_geometry_and_heldout_closed(self):
            for delay in (0, 81, 320, True):
                with self.assertRaises(ValueError):
                    Schedule('hf', delay)
            for samples in (0, -1, MAX_SAMPLES + 1, True):
                with self.assertRaises(ValueError):
                    Schedule('hf', 240).emissions(samples)
            with self.assertRaisesRegex(ValueError, 'heldout requires'):
                heldout_gate([{'role': 'heldout'}], False, None, None, {})
            self.assertIsNone(heldout_gate([{'role': 'development'}], False, None, None, {}))

        def test_pinned_incremental_byte_oracle(self):
            read_table, oracle = short.token_helpers()
            table = read_table(TABLE)
            saved = json.loads((ROOT / 'reports/bf16-delay240-short-paired-v1.json').read_text())
            ids = saved['passes']['original_bf16']['token_ids']
            decoder = IncrementalText(table)
            deltas = [(i, text, False) for i, token in enumerate(ids) if (text := decoder.push(token))]
            if tail := decoder.finish():
                deltas.append((len(ids), tail, True))
            expected = oracle(ids, table)
            self.assertEqual(deltas, expected['deltas'])
            self.assertEqual(decoder.text, expected['text'])

        def test_heldout_binds_external_and_entire_reference_identity(self):
            import tempfile
            with tempfile.TemporaryDirectory() as folder:
                release = Path(folder) / 'release.json'
                release.write_text('{"weights":"frozen-example","runtime":"immutable-example"}')
                identity = dict(release_candidate=release_identity(release, digest(release)), schedule='fixed')
                reference = Path(folder) / 'reference.json'
                reference.write_text(json.dumps(dict(format='audio8-torch-stream-candidate-v1', status='frozen',
                    identity=identity, identity_sha256=canonical_sha(identity))))
                rows = [{'role': 'heldout'}]
                self.assertIsNotNone(heldout_gate(rows, True, reference, digest(reference), identity))
                with self.assertRaisesRegex(ValueError, 'differs from frozen'):
                    heldout_gate(rows, True, reference, digest(reference), identity | {'schedule': 'changed'})
                release.write_text('{"weights":"changed"}')
                with self.assertRaisesRegex(ValueError, 'release candidate SHA'):
                    heldout_gate(rows, True, reference, digest(reference), identity)

    result = unittest.TextTestRunner(verbosity=2).run(unittest.defaultTestLoader.loadTestsFromTestCase(Tests))
    require(result.wasSuccessful(), 'self-tests failed')


def main():
    if len(sys.argv) == 3 and sys.argv[1] == '--_worker':
        worker(Path(sys.argv[2]))
        return
    if sys.argv[1:] == ['--self-test']:
        self_test()
        return
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--manifest', type=Path, required=True)
    parser.add_argument('--case', action='append', default=[])
    parser.add_argument('--weights', choices=('original', 'reconstructed'), required=True)
    parser.add_argument('--packed-bundle', type=Path)
    parser.add_argument('--startup', choices=('hf', 'realtime18'), required=True)
    parser.add_argument('--delay-ms', type=int, choices=(80, 160, 240, 480), required=True)
    parser.add_argument('--output', type=Path)
    parser.add_argument('--freeze-candidate', type=Path)
    parser.add_argument('--allow-heldout', action='store_true')
    parser.add_argument('--candidate-record', type=Path)
    parser.add_argument('--candidate-record-sha256')
    parser.add_argument('--release-candidate-record', type=Path)
    parser.add_argument('--release-candidate-record-sha256')
    parser.add_argument('--expected-short-report', type=Path)
    parser.add_argument('--timeout-seconds', type=int, default=3600)
    args = parser.parse_args()
    args.manifest = args.manifest.resolve(strict=True)
    identity = candidate_identity(args)
    if args.freeze_candidate:
        require(args.output is None and not args.case and not args.allow_heldout, 'freeze-only must not select cases/output')
        frozen = dict(format='audio8-torch-stream-candidate-v1', status='frozen', identity=identity,
                      identity_sha256=canonical_sha(identity),
                      frozen_at_unix=time.time(), contract='No candidate selection or fitting on heldout outputs')
        with args.freeze_candidate.open('x') as stream:
            json.dump(frozen, stream, indent=2)
            stream.write('\n')
        print(json.dumps(dict(candidate_record=str(args.freeze_candidate.resolve()), sha256=digest(args.freeze_candidate))))
        return
    require(args.output is not None, '--output required')
    require(60 <= args.timeout_seconds <= 7200, 'supervisor timeout must be 60..7200 seconds')
    cases = select_cases(args.manifest, args.case)
    authorization = heldout_gate(cases, args.allow_heldout, args.candidate_record, args.candidate_record_sha256, identity)
    expected_path = args.expected_short_report.resolve(strict=True) if args.expected_short_report else None
    if expected_path:
        require(args.weights == 'original' and args.startup == 'hf' and args.delay_ms == 240 and len(cases) == 1,
                'short parity comparison requires original HF240 one case')
        _, short = source_helpers()
        expected = json.loads(expected_path.read_text())
        require(cases[0]['sha256'] == short.INPUT_SHA and cases[0]['samples'] == short.SAMPLES
                and expected.get('status') == 'complete' and expected['input']['sha256'] == short.INPUT_SHA
                and expected['clocks']['transcription_delay_ms'] == 240
                and len(expected['passes']['original_bf16']['token_ids']) == short.EMISSIONS,
                'saved short/input comparison identity mismatch before model load')
    output = args.output.resolve()
    require(output.suffix == '.json', 'report filename must end in .json')
    paths = dict(report=output, events=output.with_suffix('.jsonl'), stdout=output.with_suffix('.stdout.txt'),
                 stderr=output.with_suffix('.stderr.txt'))
    require(all(not p.exists() and not p.is_symlink() for p in paths.values()), 'output exists; preserve old evidence')
    output.parent.mkdir(parents=True, exist_ok=True)
    report = dict(status='starting', started_at_unix=time.time(), paths={k: str(v) for k, v in paths.items()},
                  manifest=str(args.manifest), candidate_identity=identity, selected_cases=cases,
                  heldout_authorization=authorization, startup=args.startup, delay_ms=args.delay_ms,
                  expected_short_report=str(expected_path) if expected_path else None,
                  expected_short_report_sha256=digest(expected_path) if expected_path else None,
                  limits=dict(total_audio_seconds=360, timeout_seconds=args.timeout_seconds,
                              mps_driver_bytes=18_000_000_000, release_mps_cache_above_bytes=12_000_000_000),
                  parent_allocator_environment={k: os.environ.get(k) for k in ALLOCATOR_ENV},
                  python=sys.version, platform=platform.platform(), quality_approved=False, speed_claim=False,
                  limitations=['Original or exactly reconstructed weights in a dense BF16 model; no low-memory runtime claim.',
                               'Unpaced source dependencies are not measured audio-to-text latency.',
                               'Realtime18 matches audited window/prompt/EOF geometry, not the vLLM runtime numerics or feature ledger.',
                               'One-window frontend differs from upstream batch-precomputed convolution; saved short parity is an explicit gate.',
                               'Per-step MPS checks and a sampled peak are not an atomic total-host-memory cap.',
                               'Development/heldout roles constrain this experiment; unknown upstream corpus overlap remains.'])
    with output.open('x') as stream:
        json.dump(report, stream, indent=2)
    env = os.environ.copy()
    env.update(ALLOCATOR_ENV)
    env['TOKENIZERS_PARALLELISM'] = 'false'
    command = [sys.executable, str(Path(__file__).resolve()), '--_worker', str(output)]
    child = None
    try:
        with paths['stdout'].open('x') as stdout, paths['stderr'].open('x') as stderr:
            child = subprocess.Popen(command, cwd=ROOT, env=env, stdout=stdout, stderr=stderr, start_new_session=True)
            code = child.wait(timeout=args.timeout_seconds)
        require(code == 0 and json.loads(output.read_text())['status'] == 'complete', 'reference child failed; inspect retained logs')
    except BaseException as error:
        if child and child.poll() is None:
            os.killpg(child.pid, signal.SIGTERM)
            try:
                child.wait(timeout=5)
            except subprocess.TimeoutExpired:
                os.killpg(child.pid, signal.SIGKILL)
                child.wait(timeout=5)
        from baseline_probe import atomic_json
        failed = json.loads(output.read_text())
        failed.update(status='failed', supervisor_error=f'{type(error).__name__}: {error}')
        atomic_json(output, failed)
        raise
    finally:
        from baseline_probe import atomic_json
        final = json.loads(output.read_text())
        final.update(supervisor_command=command, logs={str(p): digest(p) for p in paths.values() if p != output and p.exists()})
        atomic_json(output, final)
    print(json.dumps(dict(status='complete', report=str(output), speed_claim=False)))


if __name__ == '__main__':
    main()