Audio8-ASR-Infinite-Compressed / 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 | |
| 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') | |
| def delay(self): | |
| return self.delay_ms // 80 | |
| 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() | |