"""Service layer wiring UI callbacks to domain logic.""" from __future__ import annotations from contextlib import contextmanager from dataclasses import dataclass import time from typing import Dict, Iterable, Iterator, List, Sequence, Tuple import numpy as np from .audio import load_audio_file, play_until_boundary, get_patch_info from .constraints import ( ConstraintDefinition, default_thresholds_for_metric, get_constraint_definition, guard_range_for_metric, iter_constraint_definitions, ) from .data_loader import PrecomputedRepository from .patching import apply_patching, compute_patch_statistics from .visualization import create_visualization MetricName = str # "Entropy" or "Perplexity" def normalize_metric_name(metric: MetricName) -> str: key = metric.lower() if key not in {"entropy", "perplexity"}: raise ValueError(f"Unsupported metric '{metric}'") return key @dataclass class PatchService: """Stateful helper for the Gradio UI callbacks.""" repository: PrecomputedRepository @classmethod def from_data_root(cls) -> "PatchService": repo = PrecomputedRepository() repo.load() return cls(repository=repo) # ------------------------------------------------------------------ # # Data lookups # ------------------------------------------------------------------ # def audio_ids(self) -> Sequence[str]: return list(self.repository.ids()) def get_data(self, file_id: str) -> Dict[str, object]: return self.repository.get(file_id) def file_info_markdown(self, file_id: str) -> str: if not file_id or file_id not in self.repository.records: return "No file selected" data = self.repository.get(file_id) return ( "### Audio Information\n\n" f"**File ID:** `{data['file_id']}`\n\n" f"**Duration:** {data['audio_duration_sec']:.2f}s\n\n" f"**Tokens:** {data['num_tokens']}\n\n" "**Entropy Statistics:**\n" f"- Average: {data['avg_entropy']:.2f} bits\n" f"- Min: {data['min_entropy']:.2f} bits\n" f"- Max: {data['max_entropy']:.2f} bits\n\n" "**Perplexity Statistics:**\n" f"- Average: {data['avg_perplexity']:.2f}\n\n" "**Transcript:**\n" f"> {data['transcript']}\n" ) # ------------------------------------------------------------------ # # Constraint helpers # ------------------------------------------------------------------ # def constraint_definitions(self) -> Iterable[ConstraintDefinition]: return iter_constraint_definitions() def guard_slider_config(self, metric: MetricName): return guard_range_for_metric(metric) def threshold_defaults(self, metric: MetricName) -> Dict[str, float]: return dict(default_thresholds_for_metric(metric)) # ------------------------------------------------------------------ # # Patching pipeline # ------------------------------------------------------------------ # def run_patching( self, file_id: str, metric_name: MetricName, selected_constraints: Sequence[str], guard_threshold: float, thresholds: Dict[str, float], ) -> Tuple[object, str, Tuple[int, np.ndarray], Tuple[int, np.ndarray], List[List[int]]]: """Apply patching constraints and return visualization artefacts.""" metric_key = normalize_metric_name(metric_name) data = self.repository.get(file_id) if not selected_constraints: raise ValueError("At least one constraint must be selected") metric = np.array( data["entropies"] if metric_key == "entropy" else data["perplexities"] ) constraints = {name: thresholds[name] for name in selected_constraints} constraints["min_metric"] = float(guard_threshold) with log_phase("apply_patching()"): patches = apply_patching(metric, constraints=constraints) with log_phase("compute_patch_statistics()"): stats = compute_patch_statistics(patches) display_threshold = thresholds[selected_constraints[0]] with log_phase("create_visualization()"): figure = create_visualization( data=data, patches=patches, metric_name=metric_key, threshold=display_threshold, ) stats_text = self._format_stats( metric_key=metric_key, guard_threshold=guard_threshold, selected_constraints=selected_constraints, thresholds=thresholds, stats=stats, data=data, ) with log_phase("load_audio_file(original)"): original_audio = load_audio_file(data["audio_path"]) with log_phase("load_audio_file(reconstructed)"): reconstructed_audio = load_audio_file(data["reconstructed_audio_path"]) return figure, stats_text, original_audio, reconstructed_audio, patches def _format_stats( self, metric_key: str, guard_threshold: float, selected_constraints: Sequence[str], thresholds: Dict[str, float], stats: Dict[str, float], data: Dict[str, object], ) -> str: guard = guard_range_for_metric(metric_key) lines = [] for name in selected_constraints: definition = get_constraint_definition(name) lines.append( f" - **{definition.label}**: {thresholds[name]:.2f} ({definition.description})" ) guard_label = guard.label or "Minimum Metric for Boundary" guard_value = f"{guard_threshold:.2f}" if metric_key == "entropy" else f"{guard_threshold:.1f}" lines.append(f" - **{guard_label}**: {guard_value} ({guard.description})") constraints_display = "\n".join(lines) return f""" ### Patch Statistics **Active Constraints** (OR logic): {constraints_display} - **Total Patches:** {stats['num_patches']} - **Average Size:** {stats['avg_size']:.2f} tokens ({stats['avg_size']/40:.3f}s) - **Min Size:** {stats['min_size']} tokens - **Max Size:** {stats['max_size']} tokens - **Median Size:** {stats['median_size']:.1f} tokens - **Std Dev:** {stats['std_size']:.2f} tokens **Inference Efficiency:** - Entropy model steps: {data['num_tokens']} - BLT global transformer steps: {stats['num_patches']} - **Speedup:** {data['num_tokens'] / stats['num_patches']:.2f}x """ # ------------------------------------------------------------------ # # Audio utilities # ------------------------------------------------------------------ # def play_until_boundary( self, file_id: str, current_token: int, patches: Sequence[Sequence[int]] ) -> Tuple[Tuple[int, np.ndarray], str, int]: if not patches: raise ValueError("Patches must be computed before playback.") data = self.repository.get(file_id) with log_phase("play_until_boundary()"): sr_audio = play_until_boundary(data["audio_path"], patches, current_token) patch_info = None next_token = 0 for idx, patch in enumerate(patches): if current_token in patch: patch_info = get_patch_info(patches, idx) if idx + 1 < len(patches): next_token = patches[idx + 1][0] else: next_token = patches[0][0] break if not patch_info: return sr_audio, f"Token {current_token} not found in any patch", next_token info_text = ( f"**Playing Patch #{patch_info['patch_index']} / {len(patches)}**\n\n" f"- Tokens: {patch_info['start_token']} → {patch_info['end_token']} " f"({patch_info['num_tokens']} tokens)\n" f"- Time: {patch_info['start_time_sec']:.2f}s → {patch_info['end_time_sec']:.2f}s\n" f"- Duration: {patch_info['duration_sec']:.2f}s\n" ) return sr_audio, info_text, next_token @contextmanager def log_phase(label: str) -> Iterator[None]: """Profile a phase and emit duration in seconds.""" start = time.perf_counter() yield elapsed = time.perf_counter() - start print(f" {label}: {elapsed:.3f}s")