+ 
@@ -21,7 +22,13 @@ ## Installation ```bash -uv sync # or: pip install -e . +pip install hypnos # or: uv add hypnos +``` + +To work on the library itself, clone the repo and install from source: + +```bash +uv sync # or: pip install -e . ``` diff --git a/pyproject.toml b/pyproject.toml index 378666c..73e5af9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -9,6 +9,15 @@ license-files = ["LICENSE"] authors = [ {name = "Jonathan Carter"} ] +keywords = ["sleep", "eeg", "ecg", "physiology", "embeddings", "foundation-model", + "time-series", "polysomnography", "edf", "deep-learning"] +classifiers = [ + "Development Status :: 4 - Beta", + "Intended Audience :: Science/Research", + "Topic :: Scientific/Engineering :: Artificial Intelligence", + "Topic :: Scientific/Engineering :: Medical Science Apps.", + "Programming Language :: Python :: 3", +] dependencies = [ "torch>=2.4", # torch.nn.attention.flex_attention (hypnos.models.attention) "numpy", @@ -19,6 +28,12 @@ dependencies = [ "fsspec", # SignalTokenizer.from_checkpoint hf:// / s3:// support ] +[project.urls] +Homepage = "https://github.com/joncarter1/hypnos" +Repository = "https://github.com/joncarter1/hypnos" +Paper = "https://arxiv.org/abs/2606.09605" +"Hugging Face" = "https://huggingface.co/joncarter/hypnos" + [dependency-groups] # Dev-only: resolving checkpoints/configs from the research stack when seeding manifests. seed = [ diff --git a/src/hypnos/data/edf.py b/src/hypnos/data/edf.py index c47d824..1f91bd6 100644 --- a/src/hypnos/data/edf.py +++ b/src/hypnos/data/edf.py @@ -20,48 +20,96 @@ # Alternative EDF labels for each canonical channel, across datasets/conventions. ALT_COLUMNS = { ECG: ( - 'EKG', 'ECG1', 'ECG L', 'ECGL', 'ECG L-ECG R', - 'ECG EKG2-EKG', 'EKG2-EKG', 'ECG LA-RA', 'LA-RA', + "EKG", + "ECG1", + "ECG L", + "ECGL", + "ECG L-ECG R", + "ECG EKG2-EKG", + "EKG2-EKG", + "ECG LA-RA", + "LA-RA", ), ABD: ( - 'Abdo', 'ABDO RES', 'ABDO EFFORT', 'Abdominal', 'ABDOMINAL', 'Abdomen', 'abdomen', - 'Resp Abdominal', 'Resp Abdomen', + "Abdo", + "ABDO RES", + "ABDO EFFORT", + "Abdominal", + "ABDOMINAL", + "Abdomen", + "abdomen", + "Resp Abdominal", + "Resp Abdomen", ), THX: ( - 'Thor', 'THOR RES', 'THOR EFFORT', 'Thoracic', 'Thorax', 'Chest', 'thorax', 'CHEST', - 'Resp Thoracic', 'Resp Chest', + "Thor", + "THOR RES", + "THOR EFFORT", + "Thoracic", + "Thorax", + "Chest", + "thorax", + "CHEST", + "Resp Thoracic", + "Resp Chest", ), # EEG channels (handle referencing conventions) - EEG_C3: ('C3-M2', 'C3-A2', 'C3_M2', 'C3_A2', 'EEG C3-M2', 'EEG C3-A2', 'EEG(sec) C3', 'EEG(sec)'), - EEG_C4: ('C4-M1', 'C4-A1', 'C4_M1', 'C4_A1', 'EEG C4-M1', 'EEG C4-A1', 'EEG(sec) C4', 'EEG', 'EEG3'), + EEG_C3: ("C3-M2", "C3-A2", "C3_M2", "C3_A2", "EEG C3-M2", "EEG C3-A2", "EEG(sec) C3", "EEG(sec)"), + EEG_C4: ("C4-M1", "C4-A1", "C4_M1", "C4_A1", "EEG C4-M1", "EEG C4-A1", "EEG(sec) C4", "EEG", "EEG3"), # EOG channels (E1/E2 AASM standard) - EOG_E1: ('E1-M2', 'E1-A2', 'EOG E1-M2', 'EOG(L)', 'LOC', 'EOG-L', 'EOGl', 'EOG LOC-M2', 'LOC-M2'), + EOG_E1: ("E1-M2", "E1-A2", "EOG E1-M2", "EOG(L)", "LOC", "EOG-L", "EOGl", "EOG LOC-M2", "LOC-M2"), EOG_E2: ( - 'E2-M1', 'E2-M2', 'E2-A1', 'EOG E2-M1', 'EOG(R)', 'ROC', 'EOG-R', 'EOGr', - 'EOG ROC-M1', 'ROC-M1', 'EEG ROC-M1', + "E2-M1", + "E2-M2", + "E2-A1", + "EOG E2-M1", + "EOG(R)", + "ROC", + "EOG-R", + "EOGr", + "EOG ROC-M1", + "ROC-M1", + "EEG ROC-M1", ), # EMG chin EMG_CHIN: ( - 'Chin1-Chin2', 'CHIN1-CHIN2', 'CHIN', 'ChinA', 'Cchin', 'Chin EMG', 'EMG', 'EMG Chin', 'chin', - 'EMG Chin1-Chin2', 'EMG CHIN1-CHIN2', 'EMG Chin2-Chin1', 'Chin2-Chin1', 'EEG Chin1-Chin2', + "Chin1-Chin2", + "CHIN1-CHIN2", + "CHIN", + "ChinA", + "Cchin", + "Chin EMG", + "EMG", + "EMG Chin", + "chin", + "EMG Chin1-Chin2", + "EMG CHIN1-CHIN2", + "EMG Chin2-Chin1", + "Chin2-Chin1", + "EEG Chin1-Chin2", ), } # Contralateral mastoid referencing (AASM): canonical channel -> required reference electrode. -CONTRALATERAL_REF: dict[str, str] = {'C3': 'M2', 'C4': 'M1', 'E1': 'M2', 'E2': 'M1'} +CONTRALATERAL_REF: dict[str, str] = {"C3": "M2", "C4": "M1", "E1": "M2", "E2": "M1"} # A1/A2 are legacy equivalents of M1/M2. -REFERENCE_ALTS: dict[str, list[str]] = {'M1': ['A1'], 'M2': ['A2']} +REFERENCE_ALTS: dict[str, list[str]] = {"M1": ["A1"], "M2": ["A2"]} _ALL_REF_NAMES: set[str] = set(REFERENCE_ALTS) | {a for alts in REFERENCE_ALTS.values() for a in alts} # Pre-computed bipolar labels (already derived in the EDF), checked before components. -BIPOLAR_LABELS: dict[str, tuple[str, ...]] = {'Chin': ('Chin1-Chin2', 'CHIN1-CHIN2')} +BIPOLAR_LABELS: dict[str, tuple[str, ...]] = {"Chin": ("Chin1-Chin2", "CHIN1-CHIN2")} # Component electrode pairs for computing bipolar derivations: (positive, negative). BIPOLAR_COMPONENTS: dict[str, list[tuple[str, str]]] = { - 'Chin': [ - ('ChinR', 'ChinL'), ('Chin1', 'Chin2'), ('cchin_r', 'cchin_l'), - ('R Chin', 'L Chin'), ('Rchin', 'Lchin'), ('EMG3', 'EMG2'), ('EMG2', 'EMG1'), + "Chin": [ + ("ChinR", "ChinL"), + ("Chin1", "Chin2"), + ("cchin_r", "cchin_l"), + ("R Chin", "L Chin"), + ("Rchin", "Lchin"), + ("EMG3", "EMG2"), + ("EMG2", "EMG1"), ], } @@ -100,12 +148,13 @@ def _find_reference_label(ref_name: str, available_labels: list[str]) -> str | N def _load_reference_signals( - f: pyedflib.EdfReader, label_to_idx: dict[str, int], + f: pyedflib.EdfReader, + label_to_idx: dict[str, int], ) -> dict[str, tuple[np.ndarray, int]]: """Load M1/M2 reference signals if available in the EDF.""" available = list(label_to_idx.keys()) refs: dict[str, tuple[np.ndarray, int]] = {} - for ref_name in ('M1', 'M2'): + for ref_name in ("M1", "M2"): label = _find_reference_label(ref_name, available) if label is not None: idx = label_to_idx[label] @@ -196,9 +245,13 @@ def _resolve_one_channel( signal = f.readSignal(idx) fs, unit, pmin, pmax = _read_signal_metadata(f, idx) return ResolvedChannel( - signal=signal, sampling_rate=fs, unit=unit, - physical_min=pmin, physical_max=pmax, - method='bipolar_pre_computed', edf_labels=[label], + signal=signal, + sampling_rate=fs, + unit=unit, + physical_min=pmin, + physical_max=pmax, + method="bipolar_pre_computed", + edf_labels=[label], ) # Try computing from component electrode pairs (e.g. ChinR - ChinL) @@ -217,18 +270,22 @@ def _resolve_one_channel( min_len = min(len(pos_signal), len(neg_signal)) bipolar = (pos_signal[:min_len] - neg_signal[:min_len]).astype(np.float64) - _logger.info(f'Computed bipolar derivation {pos_label}-{neg_label} for {ch_name}') + _logger.info(f"Computed bipolar derivation {pos_label}-{neg_label} for {ch_name}") return ResolvedChannel( - signal=bipolar, sampling_rate=pos_fs, unit=pos_unit, - physical_min=pos_pmin, physical_max=pos_pmax, - method='bipolar_derived', edf_labels=[pos_label, neg_label], + signal=bipolar, + sampling_rate=pos_fs, + unit=pos_unit, + physical_min=pos_pmin, + physical_max=pos_pmax, + method="bipolar_derived", + edf_labels=[pos_label, neg_label], ) # Fall through to standard resolution (single electrode fallback) # --- Step 2: Standard name resolution --- actual_name = get_column_match(ch_name, available) if actual_name is None: - _logger.info(f'Channel {ch_name} not found in EDF') + _logger.info(f"Channel {ch_name} not found in EDF") return None idx = label_to_idx[actual_name] @@ -238,9 +295,13 @@ def _resolve_one_channel( # --- Step 3: Check if already pre-referenced (e.g. 'C3-M2') --- if _is_pre_referenced(ch_name, actual_name): return ResolvedChannel( - signal=signal, sampling_rate=fs, unit=unit, - physical_min=pmin, physical_max=pmax, - method='pre_referenced', edf_labels=[actual_name], + signal=signal, + sampling_rate=fs, + unit=unit, + physical_min=pmin, + physical_max=pmax, + method="pre_referenced", + edf_labels=[actual_name], ) # --- Step 4: Bare electrode needing contralateral reference --- @@ -252,27 +313,39 @@ def _resolve_one_channel( min_len = min(len(signal), len(ref_resampled)) signal = (signal[:min_len] - ref_resampled[:min_len]).astype(np.float64) return ResolvedChannel( - signal=signal, sampling_rate=fs, unit=unit, - physical_min=pmin, physical_max=pmax, - method='re_referenced', edf_labels=[actual_name], + signal=signal, + sampling_rate=fs, + unit=unit, + physical_min=pmin, + physical_max=pmax, + method="re_referenced", + edf_labels=[actual_name], ) # No reference electrodes in EDF at all → hardware applied referencing if not ref_signals: - _logger.info(f'{ch_name}: no reference electrodes in EDF, assuming acquisition-referenced') + _logger.info(f"{ch_name}: no reference electrodes in EDF, assuming acquisition-referenced") return ResolvedChannel( - signal=signal, sampling_rate=fs, unit=unit, - physical_min=pmin, physical_max=pmax, - method='pre_referenced', edf_labels=[actual_name], + signal=signal, + sampling_rate=fs, + unit=unit, + physical_min=pmin, + physical_max=pmax, + method="pre_referenced", + edf_labels=[actual_name], ) # Some refs exist but not the one we need → genuinely unreferenced if drop_unreferenced: - _logger.warning(f'{ch_name}: needs {ref_name} but only {set(ref_signals)} available, dropping') + _logger.warning(f"{ch_name}: needs {ref_name} but only {set(ref_signals)} available, dropping") return None - _logger.warning(f'{ch_name}: bare electrode without {ref_name} reference, using unreferenced') + _logger.warning(f"{ch_name}: bare electrode without {ref_name} reference, using unreferenced") # --- Step 5: No referencing needed (ECG, ABD, etc.) or unreferenced fallback --- return ResolvedChannel( - signal=signal, sampling_rate=fs, unit=unit, - physical_min=pmin, physical_max=pmax, - method='direct', edf_labels=[actual_name], + signal=signal, + sampling_rate=fs, + unit=unit, + physical_min=pmin, + physical_max=pmax, + method="direct", + edf_labels=[actual_name], ) diff --git a/src/hypnos/embedding/__init__.py b/src/hypnos/embedding/__init__.py index d95997e..9f56f35 100644 --- a/src/hypnos/embedding/__init__.py +++ b/src/hypnos/embedding/__init__.py @@ -29,14 +29,14 @@ from .pipeline import preprocess_edf, tokenize __all__ = [ - 'load_model', - 'preprocess_edf', - 'tokenize', - 'embed', - 'embed_edf', - 'synthesize', - 'ModelMetadata', - 'ModalitySpec', + "load_model", + "preprocess_edf", + "tokenize", + "embed", + "embed_edf", + "synthesize", + "ModelMetadata", + "ModalitySpec", ] @@ -48,7 +48,7 @@ def embed( metadata: ModelMetadata, *, chunk_tokens: int | None = None, - device: str | torch.device = 'cpu', + device: str | torch.device = "cpu", autocast_dtype: torch.dtype | None = None, ) -> dict[str, np.ndarray]: """Generate per-modality 1 Hz embeddings from assembled tokens. @@ -59,22 +59,23 @@ def embed( time (see the README). """ ctx = temporal_context( - model, tokens, modality_mask, channel_ids, - chunk_tokens=chunk_tokens, device=device, autocast_dtype=autocast_dtype, + model, + tokens, + modality_mask, + channel_ids, + chunk_tokens=chunk_tokens, + device=device, + autocast_dtype=autocast_dtype, ) # (T, M, D) float32 present = modality_mask[0].tolist() - return { - spec.name: ctx[:, i].to(torch.float16).numpy() - for i, spec in enumerate(metadata.modalities) - if present[i] - } + return {spec.name: ctx[:, i].to(torch.float16).numpy() for i, spec in enumerate(metadata.modalities) if present[i]} def embed_edf( edf_path: str, model_repo_or_path: str = DEFAULT_REPO, *, - device: str | torch.device = 'cpu', + device: str | torch.device = "cpu", dtype: torch.dtype = torch.float32, notch_freq: float = 50.0, causal: bool = True, @@ -94,6 +95,12 @@ def embed_edf( signals = preprocess_edf(edf_path, meta, notch_freq=notch_freq, causal=causal) tokens, modality_mask, channel_ids = tokenize(tokenizers, meta, signals, device=device) return embed( - model, tokens, modality_mask, channel_ids, meta, - chunk_tokens=chunk_tokens, device=device, autocast_dtype=autocast_dtype, + model, + tokens, + modality_mask, + channel_ids, + meta, + chunk_tokens=chunk_tokens, + device=device, + autocast_dtype=autocast_dtype, ) diff --git a/src/hypnos/embedding/generate.py b/src/hypnos/embedding/generate.py index d80b230..003e8f8 100644 --- a/src/hypnos/embedding/generate.py +++ b/src/hypnos/embedding/generate.py @@ -59,7 +59,7 @@ def synthesize( targets = list(names if modalities is None else modalities) unknown = [t for t in targets if t not in names] if unknown: - raise ValueError(f'unknown modalities {unknown}; choose from {names}') + raise ValueError(f"unknown modalities {unknown}; choose from {names}") device = next(model.parameters()).device total_K = sum(m.num_quantizers for m in mods) @@ -70,12 +70,8 @@ def synthesize( spans[m.name] = (offset, offset + m.num_quantizers) offset += m.num_quantizers - modality_mask = torch.tensor( - [[m.name in targets for m in mods]], dtype=torch.bool, device=device - ) - channel_ids = torch.tensor( - [[CHANNEL_REGISTRY[m.channels[0]] for m in mods]], dtype=torch.long, device=device - ) + modality_mask = torch.tensor([[m.name in targets for m in mods]], dtype=torch.bool, device=device) + channel_ids = torch.tensor([[CHANNEL_REGISTRY[m.channels[0]] for m in mods]], dtype=torch.long, device=device) if prompt_tokens is None: prompt_tokens = torch.zeros(1, 0, total_K, dtype=torch.long, device=device) @@ -83,8 +79,13 @@ def synthesize( prompt_tokens = prompt_tokens.to(device=device, dtype=torch.long) rollout = model.generate( - prompt_tokens, channel_ids, modality_mask, - num_steps=num_steps, temperature=temperature, top_k=top_k, top_p=top_p, + prompt_tokens, + channel_ids, + modality_mask, + num_steps=num_steps, + temperature=temperature, + top_k=top_k, + top_p=top_p, generator=generator, ) diff --git a/src/hypnos/embedding/infer.py b/src/hypnos/embedding/infer.py index 09570e8..b9053b0 100644 --- a/src/hypnos/embedding/infer.py +++ b/src/hypnos/embedding/infer.py @@ -62,11 +62,15 @@ def _temporal_only_forward( # the same way the model expects. backbone_group_ids = model.temporal_transformer.sample_group_ids(B, tokens.device) mod_attn_mask = model.temporal_transformer.build_modality_attn_mask( - modality_mask, B, group_ids=backbone_group_ids, + modality_mask, + B, + group_ids=backbone_group_ids, ) # cross_attn_mask=None: this model has use_cls=False, so there is no CLS cross-attention. temporal_context, _cls = model.temporal_transformer( - temporal_input, modality_attn_mask=mod_attn_mask, cross_attn_mask=None, + temporal_input, + modality_attn_mask=mod_attn_mask, + cross_attn_mask=None, ) return temporal_context # (B, M, S, D) @@ -79,7 +83,7 @@ def temporal_context( channel_ids: torch.Tensor, *, chunk_tokens: int | None = None, - device: str | torch.device = 'cpu', + device: str | torch.device = "cpu", autocast_dtype: torch.dtype | None = None, ) -> torch.Tensor: """Run the temporal model over one record (chunked) -> per-modality 1 Hz context. @@ -103,22 +107,22 @@ def temporal_context( return torch.zeros((0, len(model.modality_configs), model.embed_dim), dtype=torch.float32) if chunk_tokens is None: - chunk_tokens = 32768 if device.type == 'cuda' else 2048 + chunk_tokens = 32768 if device.type == "cuda" else 2048 chunk = max(1, chunk_tokens) tokens_t = tokens[0].to(torch.long) # (n_tokens, K) mod_mask_t = modality_mask.to(device, torch.bool) # (1, M) ch_ids_t = channel_ids.to(device, torch.long) # (1, M) autocast_ctx = ( - torch.autocast(device_type='cuda', dtype=autocast_dtype) - if autocast_dtype is not None and device.type == 'cuda' - else torch.autocast(device_type='cpu', enabled=False) + torch.autocast(device_type="cuda", dtype=autocast_dtype) + if autocast_dtype is not None and device.type == "cuda" + else torch.autocast(device_type="cpu", enabled=False) ) # On CUDA we pad each chunk to a fixed length so torch.compile reuses one graph across # records. On CPU/MPS attention runs eagerly over a materialised (S, S) score matrix, so # padding short records up to `chunk` would OOM — process actual lengths there. - pad_to_chunk = device.type == 'cuda' + pad_to_chunk = device.type == "cuda" pad_template = torch.zeros((chunk, total_k), dtype=tokens_t.dtype) if pad_to_chunk else None frames: list[torch.Tensor] = [] @@ -134,6 +138,6 @@ def temporal_context( with autocast_ctx: ctx = _temporal_only_forward(model, win, channel_ids=ch_ids_t, modality_mask=mod_mask_t) # (1, M, S, D) ctx = ctx[:, :, :actual_len] # drop padded tail - frames.append(ctx.squeeze(0).permute(1, 0, 2).to('cpu', torch.float32)) # (S, M, D) + frames.append(ctx.squeeze(0).permute(1, 0, 2).to("cpu", torch.float32)) # (S, M, D) return torch.cat(frames, dim=0) # (n_tokens, M, D) diff --git a/src/hypnos/embedding/loader.py b/src/hypnos/embedding/loader.py index b633f80..bb7004b 100644 --- a/src/hypnos/embedding/loader.py +++ b/src/hypnos/embedding/loader.py @@ -24,8 +24,8 @@ logger = logging.getLogger(__name__) # Default HuggingFace repo and bundle filename. -DEFAULT_REPO = 'joncarter/hypnos' -BUNDLE_FILENAME = 'hypnos.safetensors' +DEFAULT_REPO = "joncarter/hypnos" +BUNDLE_FILENAME = "hypnos.safetensors" def resolve_bundle(path_or_repo: str | Path, filename: str = BUNDLE_FILENAME) -> Path: @@ -43,26 +43,24 @@ def resolve_bundle(path_or_repo: str | Path, filename: str = BUNDLE_FILENAME) -> if os.path.isdir(s): return Path(s) / filename - repo_id = s[len('hf://') :] if s.startswith('hf://') else s - if '/' not in repo_id: - raise FileNotFoundError( - f'{s!r} is not a local bundle/dir and not a valid HuggingFace repo id (owner/name).' - ) + repo_id = s[len("hf://") :] if s.startswith("hf://") else s + if "/" not in repo_id: + raise FileNotFoundError(f"{s!r} is not a local bundle/dir and not a valid HuggingFace repo id (owner/name).") try: from huggingface_hub import hf_hub_download except ImportError as e: # pragma: no cover - dependency declared in pyproject - raise ImportError('huggingface-hub is required to load a bundle from the Hub.') from e + raise ImportError("huggingface-hub is required to load a bundle from the Hub.") from e # Touch config.json so the Hub counts this as a download — its stats key on config.json, # not on the .safetensors bundle. Best-effort: a GET (or cached HEAD revalidation) is # enough to register, and loading must never fail if the file is absent or the request # errors (a HEAD/GET to a missing entry raises, which we swallow). try: - hf_hub_download(repo_id=repo_id, filename='config.json') + hf_hub_download(repo_id=repo_id, filename="config.json") except Exception: - logger.debug('config.json not fetched from %s; download count may not register', repo_id) + logger.debug("config.json not fetched from %s; download count may not register", repo_id) - logger.info('Downloading %s from HuggingFace repo %s...', filename, repo_id) + logger.info("Downloading %s from HuggingFace repo %s...", filename, repo_id) return Path(hf_hub_download(repo_id=repo_id, filename=filename)) @@ -72,26 +70,26 @@ def _read_bundle(path: Path) -> tuple[dict, dict, dict[str, dict]]: model_sd: dict = {} tokenizer_sds: dict[str, dict] = {} - with safe_open(str(path), framework='pt', device='cpu') as f: + with safe_open(str(path), framework="pt", device="cpu") as f: md = f.metadata() or {} - if 'config' not in md: + if "config" not in md: raise RuntimeError(f"{path}: no 'config' in safetensors metadata; not a Hypnos bundle.") - config = json.loads(md['config']) + config = json.loads(md["config"]) for key in f.keys(): tensor = f.get_tensor(key) - if key.startswith('model/'): - model_sd[key[len('model/') :]] = tensor - elif key.startswith('tok/'): - _, stem, param = key.split('/', 2) + if key.startswith("model/"): + model_sd[key[len("model/") :]] = tensor + elif key.startswith("tok/"): + _, stem, param = key.split("/", 2) tokenizer_sds.setdefault(stem, {})[param] = tensor else: - raise RuntimeError(f'unexpected tensor key {key!r} in bundle') + raise RuntimeError(f"unexpected tensor key {key!r} in bundle") return config, model_sd, tokenizer_sds def load_model( path_or_repo: str | Path = DEFAULT_REPO, - device: str | torch.device = 'cpu', + device: str | torch.device = "cpu", dtype: torch.dtype = torch.float32, ) -> tuple[MultiModalRQTransformer, dict[str, SignalTokenizer], ModelMetadata]: """Build the model + per-modality tokenizers from a bundle. @@ -110,15 +108,17 @@ def load_model( # modality_configs MUST be built in config order — this order defines the model's # _modality_offsets (the column layout of the token tensor) and the averaging order. modality_configs = [ - ModalityConfig(name=m.name, num_quantizers=m.num_quantizers, codebook_size=m.codebook_size, signal_type=m.signal_type) + ModalityConfig( + name=m.name, num_quantizers=m.num_quantizers, codebook_size=m.codebook_size, signal_type=m.signal_type + ) for m in meta.modalities ] model = MultiModalRQTransformer(modality_configs=modality_configs, **meta.model_kwargs) missing, unexpected = model.load_state_dict(model_sd, strict=False) if unexpected: - raise RuntimeError(f'Unexpected keys loading RQ-Transformer weights: {unexpected[:10]}...') + raise RuntimeError(f"Unexpected keys loading RQ-Transformer weights: {unexpected[:10]}...") if missing: - logger.warning('Missing keys when loading RQ-Transformer (likely tied weights): %s', missing[:10]) + logger.warning("Missing keys when loading RQ-Transformer (likely tied weights): %s", missing[:10]) model.to(device=device, dtype=dtype).eval() # Build each unique tokenizer once, then fan out to every modality that uses it. @@ -128,9 +128,9 @@ def load_model( tok = SignalTokenizer(**spec.tokenizer_kwargs) tmissing, tunexpected = tok.load_state_dict(tokenizer_sds[stem], strict=False) if tunexpected: - raise RuntimeError(f'Unexpected keys loading tokenizer {stem!r}: {tunexpected[:10]}...') + raise RuntimeError(f"Unexpected keys loading tokenizer {stem!r}: {tunexpected[:10]}...") if tmissing: - logger.warning('Missing keys loading tokenizer %r: %s', stem, tmissing[:10]) + logger.warning("Missing keys loading tokenizer %r: %s", stem, tmissing[:10]) tokenizer_instances[stem] = tok.to(device=device, dtype=dtype).eval() tokenizers_by_modality = {m.name: tokenizer_instances[m.tokenizer] for m in meta.modalities} diff --git a/src/hypnos/embedding/manifest.py b/src/hypnos/embedding/manifest.py index bb4d815..fb26135 100644 --- a/src/hypnos/embedding/manifest.py +++ b/src/hypnos/embedding/manifest.py @@ -74,45 +74,45 @@ def parse_config(config: dict) -> ModelMetadata: """Validate a bundle ``config`` dict and return :class:`ModelMetadata`.""" modalities = [ ModalitySpec( - name=m['name'], - signal_type=m['signal_type'], - channels=list(m['channels']), - tokenizer=m['tokenizer'], - num_quantizers=int(m['num_quantizers']), - codebook_size=int(m['codebook_size']), - token_duration_sec=float(m['token_duration_sec']), - sample_rate=int(m['sample_rate']), - preprocess_modality=m['preprocess_modality'], + name=m["name"], + signal_type=m["signal_type"], + channels=list(m["channels"]), + tokenizer=m["tokenizer"], + num_quantizers=int(m["num_quantizers"]), + codebook_size=int(m["codebook_size"]), + token_duration_sec=float(m["token_duration_sec"]), + sample_rate=int(m["sample_rate"]), + preprocess_modality=m["preprocess_modality"], ) - for m in config['modalities'] + for m in config["modalities"] ] if not modalities: - raise ValueError('bundle config has no modalities.') + raise ValueError("bundle config has no modalities.") # All modalities must share a single token cadence — the temporal model concatenates them # along time and averages across modalities, valid only at one token-per-interval cadence. durations = {round(m.token_duration_sec, 6) for m in modalities} if len(durations) != 1: - raise ValueError(f'All modalities must share token_duration_sec; got {sorted(durations)}.') + raise ValueError(f"All modalities must share token_duration_sec; got {sorted(durations)}.") tokenizers = { stem: TokenizerSpec( - signal_type=t['signal_type'], - num_quantizers=int(t['num_quantizers']), - codebook_size=int(t['codebook_size']), - token_duration_sec=float(t['token_duration_sec']), - sample_rate=int(t['sample_rate']), - tokenizer_kwargs=dict(t['tokenizer_kwargs']), + signal_type=t["signal_type"], + num_quantizers=int(t["num_quantizers"]), + codebook_size=int(t["codebook_size"]), + token_duration_sec=float(t["token_duration_sec"]), + sample_rate=int(t["sample_rate"]), + tokenizer_kwargs=dict(t["tokenizer_kwargs"]), ) - for stem, t in config['tokenizers'].items() + for stem, t in config["tokenizers"].items() } - model_kwargs = dict(config['model_kwargs']) - if 'modality_configs' in model_kwargs: - raise ValueError('model_kwargs must NOT contain modality_configs (built from `modalities`).') + model_kwargs = dict(config["model_kwargs"]) + if "modality_configs" in model_kwargs: + raise ValueError("model_kwargs must NOT contain modality_configs (built from `modalities`).") return ModelMetadata( - model_target=config['model_target'], + model_target=config["model_target"], model_kwargs=model_kwargs, modalities=modalities, tokenizers=tokenizers, diff --git a/src/hypnos/embedding/pipeline.py b/src/hypnos/embedding/pipeline.py index bf8b392..433f479 100644 --- a/src/hypnos/embedding/pipeline.py +++ b/src/hypnos/embedding/pipeline.py @@ -64,17 +64,23 @@ def preprocess_edf( ch = m.channels[0] rc = resolved.get(ch) if rc is None: - logger.info('Channel %r for modality %r not present in %s; skipping.', ch, m.name, edf_path) + logger.info("Channel %r for modality %r not present in %s; skipping.", ch, m.name, edf_path) continue sig = resample_signal(np.asarray(rc.signal), int(rc.sampling_rate), m.sample_rate) if causal: processed, _, _ = causal_preprocess_signal( - sig, fs=m.sample_rate, modality=m.preprocess_modality, - notch_freq=notch_freq, tau_seconds=tau_seconds, + sig, + fs=m.sample_rate, + modality=m.preprocess_modality, + notch_freq=notch_freq, + tau_seconds=tau_seconds, ) else: processed, _, _ = preprocess_signal( - sig, fs=m.sample_rate, modality=m.preprocess_modality, notch_freq=notch_freq, + sig, + fs=m.sample_rate, + modality=m.preprocess_modality, + notch_freq=notch_freq, ) signals[m.name] = np.asarray(processed, dtype=np.float32) return signals @@ -85,7 +91,7 @@ def tokenize( tokenizers: dict, metadata: ModelMetadata, signals: dict[str, np.ndarray], - device: str | torch.device = 'cpu', + device: str | torch.device = "cpu", ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Tokenize per-modality signals and assemble the token tensor. @@ -108,14 +114,14 @@ def tokenize( continue x = torch.from_numpy(np.asarray(sig, dtype=np.float32)).view(1, 1, -1).to(device) tok = tokenizers[m.name].tokenize(x) # (1, n_m, K_m) - per_modality_tokens[m.name] = tok[0].to('cpu', torch.long) + per_modality_tokens[m.name] = tok[0].to("cpu", torch.long) if not per_modality_tokens: - raise ValueError('No modalities present in the recording; cannot tokenize.') + raise ValueError("No modalities present in the recording; cannot tokenize.") n_tokens = min(t.shape[0] for t in per_modality_tokens.values()) if n_tokens == 0: - raise ValueError('Recording too short: produced 0 tokens for at least one present modality.') + raise ValueError("Recording too short: produced 0 tokens for at least one present modality.") blocks: list[torch.Tensor] = [] mask: list[bool] = [] diff --git a/src/hypnos/models/attention.py b/src/hypnos/models/attention.py index b156675..fa69b67 100644 --- a/src/hypnos/models/attention.py +++ b/src/hypnos/models/attention.py @@ -45,10 +45,11 @@ def gradient_checkpoint(fn, *args, **kwargs): not ``default_dtype``, which can otherwise cause a metadata mismatch in nested checkpoints that save dtype-dependent placeholder tensors). """ - kwargs['use_reentrant'] = False - kwargs.setdefault('context_fn', _preserve_default_dtype_context_fn) + kwargs["use_reentrant"] = False + kwargs.setdefault("context_fn", _preserve_default_dtype_context_fn) return _raw_checkpoint(fn, *args, **kwargs) + # --------------------------------------------------------------------------- # Compiled FlexAttention wrappers # --------------------------------------------------------------------------- @@ -63,7 +64,7 @@ def gradient_checkpoint(fn, *args, **kwargs): def _compiled_flex_attention(query, key, value, **kwargs): - if query.device.type == 'cpu': + if query.device.type == "cpu": return flex_attention(query, key, value, **kwargs) global _compiled_flex_attention_accel if _compiled_flex_attention_accel is None: @@ -72,14 +73,15 @@ def _compiled_flex_attention(query, key, value, **kwargs): def _compiled_create_block_mask(mask_fn, **kwargs): - device = kwargs.get('device') - if getattr(device, 'type', None) == 'cpu' or str(device) == 'cpu': + device = kwargs.get("device") + if getattr(device, "type", None) == "cpu" or str(device) == "cpu": return create_block_mask(mask_fn, **kwargs) global _compiled_create_block_mask_accel if _compiled_create_block_mask_accel is None: _compiled_create_block_mask_accel = torch.compile(create_block_mask) return _compiled_create_block_mask_accel(mask_fn, **kwargs) + # --------------------------------------------------------------------------- # Rotary Position Embeddings (RoPE) # --------------------------------------------------------------------------- @@ -103,15 +105,15 @@ class RotaryEmbedding(nn.Module): def __init__(self, dim: int, max_seq_len: int = 8192, theta: float = 10000.0): super().__init__() freqs = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim)) - self.register_buffer('freqs', freqs) + self.register_buffer("freqs", freqs) self._build_cache(max_seq_len) def _build_cache(self, seq_len: int) -> None: """Build cos/sin caches up to the given sequence length.""" t = torch.arange(seq_len, device=self.freqs.device) freqs_table = torch.outer(t, self.freqs) - self.register_buffer('cos_cached', freqs_table.cos()) - self.register_buffer('sin_cached', freqs_table.sin()) + self.register_buffer("cos_cached", freqs_table.cos()) + self.register_buffer("sin_cached", freqs_table.sin()) def forward(self, seq_len: int) -> tuple[Tensor, Tensor]: """Return cos and sin tables for the given sequence length. @@ -210,9 +212,9 @@ def get_block_mask( """ if mask_fn is None: return None - if cache.get('block_mask') is not None and cache.get('seq_len') == seq_len: - return cache['block_mask'] - if getattr(device, 'type', None) == 'mps' or str(device) == 'mps': + if cache.get("block_mask") is not None and cache.get("seq_len") == seq_len: + return cache["block_mask"] + if getattr(device, "type", None) == "mps" or str(device) == "mps": mask = _dense_attention_mask(mask_fn, seq_len, device) else: mask = _compiled_create_block_mask( @@ -223,8 +225,8 @@ def get_block_mask( KV_LEN=seq_len, device=device, ) - cache['block_mask'] = mask - cache['seq_len'] = seq_len + cache["block_mask"] = mask + cache["seq_len"] = seq_len return mask @@ -647,6 +649,7 @@ def forward(self, x: Tensor) -> Tensor: is_causal = self.causal and self.window_size is None if self.use_activation_checkpointing: + def run_layer(layer, x): x, _ = layer(x, cos, sin, block_mask=block_mask, is_causal=is_causal) return x diff --git a/src/hypnos/models/rq_transformer/__init__.py b/src/hypnos/models/rq_transformer/__init__.py index 67cf808..a3d5519 100644 --- a/src/hypnos/models/rq_transformer/__init__.py +++ b/src/hypnos/models/rq_transformer/__init__.py @@ -3,6 +3,6 @@ from .model import ModalityConfig, MultiModalRQTransformer __all__ = [ - 'ModalityConfig', - 'MultiModalRQTransformer', + "ModalityConfig", + "MultiModalRQTransformer", ] diff --git a/src/hypnos/models/rq_transformer/depth.py b/src/hypnos/models/rq_transformer/depth.py index ec65c50..20efc7a 100644 --- a/src/hypnos/models/rq_transformer/depth.py +++ b/src/hypnos/models/rq_transformer/depth.py @@ -108,7 +108,7 @@ def _sample_from_logits( if top_k and top_k > 0: k = min(top_k, logits.size(-1)) kth = logits.topk(k, dim=-1).values[..., -1, None] - logits = logits.masked_fill(logits < kth, float('-inf')) + logits = logits.masked_fill(logits < kth, float("-inf")) if 0.0 < top_p < 1.0: sorted_logits, sorted_idx = torch.sort(logits, descending=True, dim=-1) @@ -117,8 +117,8 @@ def _sample_from_logits( # Keep the top token always; drop entries once cumulative prob (excluding # the current one) already exceeds top_p. remove = (cumprobs - probs) > top_p - sorted_logits = sorted_logits.masked_fill(remove, float('-inf')) - logits = torch.full_like(logits, float('-inf')).scatter(-1, sorted_idx, sorted_logits) + sorted_logits = sorted_logits.masked_fill(remove, float("-inf")) + logits = torch.full_like(logits, float("-inf")).scatter(-1, sorted_idx, sorted_logits) probs = logits.softmax(dim=-1) return torch.multinomial(probs, num_samples=1, generator=generator).squeeze(-1) @@ -170,7 +170,7 @@ def __init__( ): super().__init__() if not (0.0 <= mod_context_dropout_p < 1.0): - raise ValueError(f'mod_context_dropout_p must be in [0, 1), got {mod_context_dropout_p}') + raise ValueError(f"mod_context_dropout_p must be in [0, 1), got {mod_context_dropout_p}") self.modality_configs = modality_configs self._weight_group = weight_group @@ -204,16 +204,18 @@ def __init__( # Shared pos_embed sized for the longest modality's K. self.pos_embed = nn.Embedding(max_K, depth_dim) - self.layers = nn.ModuleList([ - DepthTransformerLayer( - d_model=depth_dim, - nhead=num_heads, - dim_feedforward=dim_feedforward, - dropout=dropout, - swiglu=swiglu, - ) - for _ in range(depth) - ]) + self.layers = nn.ModuleList( + [ + DepthTransformerLayer( + d_model=depth_dim, + nhead=num_heads, + dim_feedforward=dim_feedforward, + dropout=dropout, + swiglu=swiglu, + ) + for _ in range(depth) + ] + ) self.norm = nn.LayerNorm(depth_dim) self.modality_embeddings = nn.Parameter(torch.randn(num_modalities, depth_dim) * 0.02) @@ -256,12 +258,7 @@ def forward( # Replace mod_ctx with the learnable null at random (B, S) positions — # forces the model to lean on cls_context + modality embedding. - if ( - self.training - and self.has_cls_context - and cls_context is not None - and self.mod_context_dropout_p > 0.0 - ): + if self.training and self.has_cls_context and cls_context is not None and self.mod_context_dropout_p > 0.0: drop_mask = torch.bernoulli( torch.full((B, S, 1), self.mod_context_dropout_p, device=mod_ctx.device, dtype=mod_ctx.dtype) ).bool() @@ -284,8 +281,7 @@ def forward( N = depth_input.size(0) if N > chunk_size: outputs = [ - gradient_checkpoint(self._run_layers, c, use_reentrant=False) - for c in depth_input.split(chunk_size) + gradient_checkpoint(self._run_layers, c, use_reentrant=False) for c in depth_input.split(chunk_size) ] x = torch.cat(outputs, dim=0) else: diff --git a/src/hypnos/models/rq_transformer/model.py b/src/hypnos/models/rq_transformer/model.py index 9b1babf..31ad8a9 100644 --- a/src/hypnos/models/rq_transformer/model.py +++ b/src/hypnos/models/rq_transformer/model.py @@ -31,7 +31,7 @@ class ModalityConfig: name: str num_quantizers: int codebook_size: int - signal_type: str = '' + signal_type: str = "" def _compute_per_level_loss( @@ -110,7 +110,7 @@ def _compute_per_level_loss( if use_activation_checkpointing and h.requires_grad: def _level_loss(h, targets, _head=output_heads[k], _tw=token_weight, _nt=n_tokens): - loss = F.cross_entropy(_head(h), targets, reduction='none') + loss = F.cross_entropy(_head(h), targets, reduction="none") if _tw is not None: return (loss * _tw).sum() / _nt return loss.mean() @@ -120,7 +120,7 @@ def _level_loss(h, targets, _head=output_heads[k], _tw=token_weight, _nt=n_token # Metric-only second forward: combined loss/accuracy/per-sample reductions. with torch.no_grad(): logits = output_heads[k](h) - loss_full = F.cross_entropy(logits, targets, reduction='none') + loss_full = F.cross_entropy(logits, targets, reduction="none") correct = (logits.argmax(dim=-1) == targets).float() if token_weight is not None: per_level_accuracy[k] = (correct * token_weight).sum() / n_tokens @@ -132,7 +132,7 @@ def _level_loss(h, targets, _head=output_heads[k], _tw=token_weight, _nt=n_token per_level_logz[k] = torch.logsumexp(logits, dim=-1).mean() else: logits = output_heads[k](h).reshape(-1, codebook_size) - loss = F.cross_entropy(logits, targets, reduction='none') + loss = F.cross_entropy(logits, targets, reduction="none") correct = (logits.argmax(dim=-1) == targets).float() if token_weight is not None: per_level_loss[k] = (loss * token_weight).sum() / n_tokens @@ -292,8 +292,8 @@ def __init__( missing = [mc.name for mc in modality_configs if not mc.signal_type] if missing: raise ValueError( - f'Setups with multiple modalities require non-empty signal_type on every ModalityConfig. ' - f'Missing signal_type for: {missing}' + f"Setups with multiple modalities require non-empty signal_type on every ModalityConfig. " + f"Missing signal_type for: {missing}" ) # Compute depth layout: flat mapping from token column index → (modality_name, quantizer_idx) @@ -627,22 +627,22 @@ def forward( # depth-conditioning signal only — in the additive fusion inside depth it # learns to occupy an orthogonal complement of mod_ctx's residual subspace, # so probing CLS alone sees a structurally partial view of the signal. - embeddings: dict[str, Tensor] = {'1s': temporal_context} # (B, M, S, D) + embeddings: dict[str, Tensor] = {"1s": temporal_context} # (B, M, S, D) return { - 'loss': loss, - 'per_level_loss': per_level_loss.detach(), - 'per_level_accuracy': per_level_accuracy.detach(), - 'per_level_logit_max': per_level_logit_max.detach(), - 'per_level_logz': per_level_logz.detach(), - 'per_modality_loss': {k: v.detach() for k, v in per_modality_loss.items()}, - 'per_modality_accuracy': {k: v.detach() for k, v in per_modality_accuracy.items()}, - 'per_modality_per_sample_loss': {k: v.detach() for k, v in per_modality_per_sample_loss.items()}, - 'per_modality_per_sample_correct': {k: v.detach() for k, v in per_modality_per_sample_correct.items()}, - 'per_modality_logit_max': {k: v.detach() for k, v in per_modality_logit_max.items()}, - 'per_modality_logz': {k: v.detach() for k, v in per_modality_logz.items()}, - 'temporal_context': temporal_context, - 'embeddings': embeddings, + "loss": loss, + "per_level_loss": per_level_loss.detach(), + "per_level_accuracy": per_level_accuracy.detach(), + "per_level_logit_max": per_level_logit_max.detach(), + "per_level_logz": per_level_logz.detach(), + "per_modality_loss": {k: v.detach() for k, v in per_modality_loss.items()}, + "per_modality_accuracy": {k: v.detach() for k, v in per_modality_accuracy.items()}, + "per_modality_per_sample_loss": {k: v.detach() for k, v in per_modality_per_sample_loss.items()}, + "per_modality_per_sample_correct": {k: v.detach() for k, v in per_modality_per_sample_correct.items()}, + "per_modality_logit_max": {k: v.detach() for k, v in per_modality_logit_max.items()}, + "per_modality_logz": {k: v.detach() for k, v in per_modality_logz.items()}, + "temporal_context": temporal_context, + "embeddings": embeddings, } @torch.no_grad() diff --git a/src/hypnos/models/rq_transformer/multimodal_temporal.py b/src/hypnos/models/rq_transformer/multimodal_temporal.py index f5fc50c..3bdf5e5 100644 --- a/src/hypnos/models/rq_transformer/multimodal_temporal.py +++ b/src/hypnos/models/rq_transformer/multimodal_temporal.py @@ -283,22 +283,22 @@ def __init__( ): super().__init__() if modality_attn_start_layer < 0: - raise ValueError(f'modality_attn_start_layer must be >= 0, got {modality_attn_start_layer}') + raise ValueError(f"modality_attn_start_layer must be >= 0, got {modality_attn_start_layer}") if modality_grouping_alpha is not None and modality_grouping_alpha <= 0: - raise ValueError(f'modality_grouping_alpha must be > 0 when set, got {modality_grouping_alpha}') + raise ValueError(f"modality_grouping_alpha must be > 0 when set, got {modality_grouping_alpha}") if random_subset_masking and modality_grouping_alpha is None: - raise ValueError('random_subset_masking=True requires modality_grouping_alpha to be set') + raise ValueError("random_subset_masking=True requires modality_grouping_alpha to be set") if disable_cross_modal_attention and random_subset_masking: - raise ValueError('disable_cross_modal_attention and random_subset_masking are mutually exclusive') + raise ValueError("disable_cross_modal_attention and random_subset_masking are mutually exclusive") if modality_dropout_p is not None: if not 0.0 <= modality_dropout_p <= 1.0: - raise ValueError(f'modality_dropout_p must be in [0, 1] when set, got {modality_dropout_p}') + raise ValueError(f"modality_dropout_p must be in [0, 1] when set, got {modality_dropout_p}") if modality_grouping_alpha is not None: - raise ValueError('modality_dropout_p and modality_grouping_alpha are mutually exclusive') + raise ValueError("modality_dropout_p and modality_grouping_alpha are mutually exclusive") if random_subset_masking: - raise ValueError('modality_dropout_p and random_subset_masking are mutually exclusive') + raise ValueError("modality_dropout_p and random_subset_masking are mutually exclusive") if disable_cross_modal_attention: - raise ValueError('modality_dropout_p and disable_cross_modal_attention are mutually exclusive') + raise ValueError("modality_dropout_p and disable_cross_modal_attention are mutually exclusive") self.num_modalities = num_modalities self.causal = causal @@ -320,7 +320,7 @@ def __init__( # M modalities. Kept as a buffer (and referenced for device lookup) so # callers that pass modality_mask=None without a CUDA tensor still work. self.register_buffer( - 'modality_attn_mask', + "modality_attn_mask", torch.ones(num_modalities, num_modalities, dtype=torch.bool), ) @@ -531,7 +531,7 @@ def build_modality_attn_mask( same_group = group_ids.unsqueeze(2) == group_ids.unsqueeze(1) # (B, M, M) bool_mask = bool_mask & same_group - return torch.where(bool_mask, 0.0, float('-inf')).unsqueeze(1) + return torch.where(bool_mask, 0.0, float("-inf")).unsqueeze(1) def build_cross_attn_mask( self, @@ -584,7 +584,7 @@ def build_cross_attn_mask( safe_row = absent_q.unsqueeze(2) & eye # (B, M, M) diagonal where absent bool_mask = (bool_mask & keep_row) | safe_row - return torch.where(bool_mask, 0.0, float('-inf')).unsqueeze(1) + return torch.where(bool_mask, 0.0, float("-inf")).unsqueeze(1) def forward( self, diff --git a/src/hypnos/models/tokenizer/__init__.py b/src/hypnos/models/tokenizer/__init__.py index 532cd9f..88df311 100644 --- a/src/hypnos/models/tokenizer/__init__.py +++ b/src/hypnos/models/tokenizer/__init__.py @@ -2,4 +2,4 @@ from .tokenizer import SignalTokenizer -__all__ = ['SignalTokenizer'] +__all__ = ["SignalTokenizer"] diff --git a/src/hypnos/models/tokenizer/attention.py b/src/hypnos/models/tokenizer/attention.py index e3634d0..6ebad7c 100644 --- a/src/hypnos/models/tokenizer/attention.py +++ b/src/hypnos/models/tokenizer/attention.py @@ -2,4 +2,4 @@ from hypnos.models.attention import IdentityAttention, RoPETransformer -__all__ = ['IdentityAttention', 'RoPETransformer'] +__all__ = ["IdentityAttention", "RoPETransformer"] diff --git a/src/hypnos/models/tokenizer/quantizer.py b/src/hypnos/models/tokenizer/quantizer.py index c93010a..a247bc3 100644 --- a/src/hypnos/models/tokenizer/quantizer.py +++ b/src/hypnos/models/tokenizer/quantizer.py @@ -148,15 +148,15 @@ def __init__( self.embedding.weight.data.uniform_(-1.0 / codebook_size, 1.0 / codebook_size) # EMA tracking buffers - self.register_buffer('ema_cluster_size', torch.ones(codebook_size)) - self.register_buffer('ema_embed_sum', torch.zeros(codebook_size, dim)) - self.register_buffer('initialized', torch.tensor(False)) + self.register_buffer("ema_cluster_size", torch.ones(codebook_size)) + self.register_buffer("ema_embed_sum", torch.zeros(codebook_size, dim)) + self.register_buffer("initialized", torch.tensor(False)) # Raw assignment counts for unbiased marginal-entropy logging. # Unlike `ema_cluster_size`, these are not touched by dead-code resets, # so they reflect actual code usage frequency. Caller resets between # measurement windows (e.g. once per validation epoch). - self.register_buffer('usage_count', torch.zeros(codebook_size, dtype=torch.long)) + self.register_buffer("usage_count", torch.zeros(codebook_size, dtype=torch.long)) # Dead code check counter (not a buffer — resets on checkpoint load, which is fine) self._next_unused_check = check_unused_every @@ -452,10 +452,14 @@ def forward( indices = torch.stack(all_indices, dim=-1) - return z_q, indices, { - 'commitment_loss': total_commitment / n_q, - 'residual_norm': residual.detach().norm(dim=-1).mean(), - } + return ( + z_q, + indices, + { + "commitment_loss": total_commitment / n_q, + "residual_norm": residual.detach().norm(dim=-1).mean(), + }, + ) def encode(self, x: Tensor, num_quantizers: int | None = None) -> Tensor: """Encode to discrete indices only (for inference). @@ -489,7 +493,13 @@ def decode(self, indices: Tensor) -> Tensor: Returns: z_q: Quantized embeddings (B, T, D). """ - z_q = torch.zeros(indices.shape[0], indices.shape[1], self.dim, device=indices.device, dtype=self.quantizers[0].embedding.weight.dtype) + z_q = torch.zeros( + indices.shape[0], + indices.shape[1], + self.dim, + device=indices.device, + dtype=self.quantizers[0].embedding.weight.dtype, + ) for i in range(indices.shape[-1]): z_q = z_q + self.quantizers[i].decode(indices[..., i]) @@ -515,7 +525,7 @@ def get_codebook_usage(self) -> dict[str, Tensor]: probs = cluster_size / cluster_size.sum() entropy = -(probs * probs.log()).sum() perplexity = entropy.exp() - stats[f'codebook_usage_q{i}'] = perplexity / q.codebook_size + stats[f"codebook_usage_q{i}"] = perplexity / q.codebook_size return stats def reset_marginal_stats(self) -> None: @@ -540,7 +550,7 @@ def get_marginal_stats(self) -> dict[str, Tensor]: for i, q in enumerate(self.quantizers): H = q.marginal_entropy() log_V = torch.log(torch.tensor(float(q.codebook_size), device=H.device)) - stats[f'marginal_entropy_q{i}'] = H - stats[f'marginal_entropy_norm_q{i}'] = H / log_V - stats[f'effective_vocab_q{i}'] = H.exp() / q.codebook_size + stats[f"marginal_entropy_q{i}"] = H + stats[f"marginal_entropy_norm_q{i}"] = H / log_V + stats[f"effective_vocab_q{i}"] = H.exp() / q.codebook_size return stats diff --git a/src/hypnos/models/tokenizer/seanet.py b/src/hypnos/models/tokenizer/seanet.py index 6eb40e3..67e44d6 100644 --- a/src/hypnos/models/tokenizer/seanet.py +++ b/src/hypnos/models/tokenizer/seanet.py @@ -40,7 +40,7 @@ def __init__( bias: bool = True, causal: bool = False, norm: str | None = None, - pad_mode: str = 'reflect', + pad_mode: str = "reflect", ): super().__init__() self.causal = causal @@ -72,7 +72,7 @@ def __init__( bias=bias, ) - if norm == 'weight': + if norm == "weight": self.conv = weight_norm(self.conv) self.norm = None else: @@ -123,7 +123,7 @@ def __init__( bias=bias, ) - if norm == 'weight': + if norm == "weight": self.conv = weight_norm(self.conv) self.norm = None else: @@ -171,11 +171,11 @@ def __init__( dim: int, kernel_sizes: tp.List[int] = [3, 1], dilations: tp.List[int] = [1, 1], - activation: str = 'gelu', - norm: str | None = 'layer', + activation: str = "gelu", + norm: str | None = "layer", causal: bool = False, compress: int = 2, - pad_mode: str = 'reflect', + pad_mode: str = "reflect", ): super().__init__() assert len(kernel_sizes) == len(dilations) == 2 @@ -244,13 +244,13 @@ def __init__( ratios: tp.List[int] = [4, 4, 4, 4], n_residual_layers: int = 1, dilation_base: int = 1, - activation: str = 'gelu', - norm: str | None = 'layer', + activation: str = "gelu", + norm: str | None = "layer", causal: bool = False, kernel_size: int = 7, last_kernel_size: int = 7, stride_kernel_multiplier: int = 2, - pad_mode: str = 'reflect', + pad_mode: str = "reflect", use_activation_checkpointing: bool = False, ): super().__init__() @@ -398,12 +398,12 @@ def __init__( ratios: tp.List[int] = [4, 4, 4, 4], n_residual_layers: int = 1, dilation_base: int = 1, - activation: str = 'gelu', - norm: str | None = 'layer', + activation: str = "gelu", + norm: str | None = "layer", causal: bool = False, kernel_size: int = 7, stride_kernel_multiplier: int = 2, - pad_mode: str = 'reflect', + pad_mode: str = "reflect", use_activation_checkpointing: bool = False, ): super().__init__() diff --git a/src/hypnos/models/tokenizer/tokenizer.py b/src/hypnos/models/tokenizer/tokenizer.py index 7bea165..774de1b 100644 --- a/src/hypnos/models/tokenizer/tokenizer.py +++ b/src/hypnos/models/tokenizer/tokenizer.py @@ -65,7 +65,7 @@ class SignalTokenizer(nn.Module): rotation_trick: Use rotation trick (Fifty et al., 2025) instead of STE for VQ gradients. """ - VALID_MODES = ('discrete', 'vae') + VALID_MODES = ("discrete", "vae") def __init__( self, @@ -80,7 +80,7 @@ def __init__( n_residual_layers: int = 1, dilation_base: int = 1, # Mode - mode: str = 'discrete', + mode: str = "discrete", # Discrete (RVQ) parameters codebook_size: int = 512, codebook_dim: int | None = None, @@ -100,11 +100,11 @@ def __init__( # Causality causal: bool = False, # Architecture - activation: str = 'gelu', - norm: str | None = 'layer', + activation: str = "gelu", + norm: str | None = "layer", last_kernel_size: int = 7, stride_kernel_multiplier: int = 2, - pad_mode: str = 'reflect', + pad_mode: str = "reflect", # Memory optimization use_activation_checkpointing: bool = False, # VQ gradient method @@ -113,7 +113,7 @@ def __init__( super().__init__() if mode not in self.VALID_MODES: - raise ValueError(f'Invalid mode {mode!r}, must be one of {self.VALID_MODES}') + raise ValueError(f"Invalid mode {mode!r}, must be one of {self.VALID_MODES}") # Store configuration self.mode = mode @@ -132,8 +132,8 @@ def __init__( if hop_length != self.samples_per_token: raise ValueError( - f'Product of ratios ({hop_length}) must equal ' - f'sample_rate * token_duration_sec ({self.samples_per_token})' + f"Product of ratios ({hop_length}) must equal " + f"sample_rate * token_duration_sec ({self.samples_per_token})" ) # Encoder @@ -169,7 +169,7 @@ def __init__( self.encoder_transformer = IdentityAttention() # Mode-specific bottleneck - if mode == 'discrete': + if mode == "discrete": codebook_dim = codebook_dim if codebook_dim is not None else embed_dim self.codebook_dim = codebook_dim self.project_in = nn.Linear(embed_dim, codebook_dim) @@ -250,7 +250,7 @@ def encode(self, x: Tensor, return_embeddings: bool = False) -> Tensor | tuple[T z = self.encoder(x) z = self.encoder_transformer(z) - if self.mode == 'vae': + if self.mode == "vae": mu = self.fc_mu(z) return self._reparameterize(mu, self.fc_logvar(z)) @@ -271,7 +271,7 @@ def decode(self, z: Tensor) -> Tensor: Returns: Reconstructed signal (B, C, T). """ - if self.mode == 'discrete': + if self.mode == "discrete": z = self.project_out(z) else: z = self.fc_decode(z) @@ -288,7 +288,7 @@ def decode_tokens(self, indices: Tensor) -> Tensor: Returns: Reconstructed signal (B, C, T). """ - if self.mode != 'discrete': + if self.mode != "discrete": raise RuntimeError('decode_tokens() requires mode="discrete"') z_q = self.quantizer.decode(indices) return self.decode(z_q) @@ -310,13 +310,13 @@ def forward(self, x: Tensor) -> dict[str, Tensor]: z = self.encoder(x) z = self.encoder_transformer(z) - if self.mode == 'discrete': + if self.mode == "discrete": z_proj = self.project_in(z) # Always quantize (EMA codebook updates + commitment loss on every batch) z_q, indices, vq_losses = self.quantizer(z_proj) - commitment_loss = vq_losses['commitment_loss'] - residual_norm = vq_losses.get('residual_norm', zero) + commitment_loss = vq_losses["commitment_loss"] + residual_norm = vq_losses.get("residual_norm", zero) # Per-sequence quantization dropout (Défossez et al., 2024): # independently bypass VQ for each sequence in the batch @@ -329,16 +329,16 @@ def forward(self, x: Tensor) -> dict[str, Tensor]: z_out = self.decoder_transformer(z_out) x_recon = self.decoder(z_out) assert x_recon.size(-1) == input_length, ( - f'Decoder output length {x_recon.size(-1)} != input length {input_length}' + f"Decoder output length {x_recon.size(-1)} != input length {input_length}" ) return { - 'reconstruction': x_recon, - 'embeddings': z_q, - 'indices': indices, - 'commitment_loss': commitment_loss, - 'kl_loss': zero, - 'residual_norm': residual_norm, + "reconstruction": x_recon, + "embeddings": z_q, + "indices": indices, + "commitment_loss": commitment_loss, + "kl_loss": zero, + "residual_norm": residual_norm, } # VAE mode @@ -352,16 +352,16 @@ def forward(self, x: Tensor) -> dict[str, Tensor]: z_out = self.decoder_transformer(z_out) x_recon = self.decoder(z_out) assert x_recon.size(-1) == input_length, ( - f'Decoder output length {x_recon.size(-1)} != input length {input_length}' + f"Decoder output length {x_recon.size(-1)} != input length {input_length}" ) return { - 'reconstruction': x_recon, - 'embeddings': z_sampled, - 'indices': None, - 'commitment_loss': zero, - 'kl_loss': kl_loss, - 'residual_norm': zero, + "reconstruction": x_recon, + "embeddings": z_sampled, + "indices": None, + "commitment_loss": zero, + "kl_loss": kl_loss, + "residual_norm": zero, } def tokenize(self, x: Tensor) -> Tensor: @@ -376,7 +376,7 @@ def tokenize(self, x: Tensor) -> Tensor: Raises: RuntimeError: If mode is not 'discrete'. """ - if self.mode != 'discrete': + if self.mode != "discrete": raise RuntimeError('tokenize() requires mode="discrete"; use encode() for VAE embeddings') with torch.no_grad(): return self.encode(x) @@ -387,24 +387,24 @@ def get_num_tokens(self, signal_length: int) -> int: def get_codebook_usage(self) -> dict[str, Tensor]: """Get codebook usage statistics (discrete mode only).""" - if self.mode != 'discrete': + if self.mode != "discrete": return {} return self.quantizer.get_codebook_usage() def reset_marginal_stats(self) -> None: """Reset raw assignment counters before a measurement window.""" - if self.mode != 'discrete': + if self.mode != "discrete": return self.quantizer.reset_marginal_stats() def get_marginal_stats(self) -> dict[str, Tensor]: """Per-quantizer marginal entropy stats from raw counts (discrete mode only).""" - if self.mode != 'discrete': + if self.mode != "discrete": return {} return self.quantizer.get_marginal_stats() @classmethod - def from_checkpoint(cls, checkpoint_path: str | Path, **kwargs) -> 'SignalTokenizer': + def from_checkpoint(cls, checkpoint_path: str | Path, **kwargs) -> "SignalTokenizer": """Load a SignalTokenizer from a training checkpoint. Constructs the model from the provided kwargs, then loads the tokenizer weights from @@ -420,28 +420,28 @@ def from_checkpoint(cls, checkpoint_path: str | Path, **kwargs) -> 'SignalTokeni """ model = cls(**kwargs) checkpoint_str = str(checkpoint_path) - if checkpoint_str.startswith('hf://') or checkpoint_str.startswith('s3://'): + if checkpoint_str.startswith("hf://") or checkpoint_str.startswith("s3://"): import fsspec - with fsspec.open(checkpoint_str, 'rb') as f: - checkpoint = torch.load(f, map_location='cpu', weights_only=False) + with fsspec.open(checkpoint_str, "rb") as f: + checkpoint = torch.load(f, map_location="cpu", weights_only=False) else: - checkpoint = torch.load(checkpoint_path, map_location='cpu', weights_only=False) + checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False) # Extract model weights (strip 'model.' prefix from Lightning module state_dict, # and '_orig_mod.' prefix from torch.compile'd models) state_dict = {} - for key, value in checkpoint['state_dict'].items(): - if not key.startswith('model.'): + for key, value in checkpoint["state_dict"].items(): + if not key.startswith("model."): continue - clean_key = key.removeprefix('model.').replace('_orig_mod.', '') + clean_key = key.removeprefix("model.").replace("_orig_mod.", "") state_dict[clean_key] = value missing, unexpected = model.load_state_dict(state_dict, strict=False) if missing: - logger.warning(f'Missing keys when loading checkpoint: {missing}') + logger.warning(f"Missing keys when loading checkpoint: {missing}") if unexpected: - logger.warning(f'Unexpected keys when loading checkpoint: {unexpected}') + logger.warning(f"Unexpected keys when loading checkpoint: {unexpected}") model.eval() return model diff --git a/src/hypnos/models/utils.py b/src/hypnos/models/utils.py index 1cb2319..b2cd189 100644 --- a/src/hypnos/models/utils.py +++ b/src/hypnos/models/utils.py @@ -63,8 +63,8 @@ def __init__( self.weight = nn.Parameter(torch.ones(1, num_features, 1)) self.bias = nn.Parameter(torch.zeros(1, num_features, 1)) else: - self.register_parameter('weight', None) - self.register_parameter('bias', None) + self.register_parameter("weight", None) + self.register_parameter("bias", None) def forward(self, x: Tensor) -> Tensor: if self.window_size is None: @@ -124,37 +124,37 @@ def forward_windowed(self, x: Tensor) -> Tensor: def get_activation(name: str, **kwargs): """Return an activation function from its name.""" - if name == 'relu': + if name == "relu": return nn.ReLU(**kwargs) - elif name == 'leaky': + elif name == "leaky": return nn.LeakyReLU(**kwargs) - elif name == 'gelu': + elif name == "gelu": return nn.GELU(**kwargs) - elif name == 'elu': + elif name == "elu": return nn.ELU(**kwargs) - elif name == 'silu' or name == 'swish': + elif name == "silu" or name == "swish": return nn.SiLU(**kwargs) - elif name == 'linear': + elif name == "linear": return nn.Identity() else: - raise ValueError(f'{name=} is unsupported.') + raise ValueError(f"{name=} is unsupported.") -def get_norm(name: str | None = 'batch', causal: bool = False, *args, **kwargs) -> nn.Module: - if name == 'batch': +def get_norm(name: str | None = "batch", causal: bool = False, *args, **kwargs) -> nn.Module: + if name == "batch": return nn.BatchNorm1d(*args, **kwargs) - elif name == 'layer': + elif name == "layer": return ConvLayerNorm(*args, **kwargs) - elif name == 'rms': + elif name == "rms": return ConvRMSNorm(*args, **kwargs) elif name is None: return nn.Identity() - elif name == 'instance': # and not causal: + elif name == "instance": # and not causal: return nn.InstanceNorm1d(*args, **kwargs) - elif name == 'instance' and causal: # IGNORE FOR NOW + elif name == "instance" and causal: # IGNORE FOR NOW return CausalInstanceNorm1d(*args, **kwargs) - elif name.startswith('instance') and causal: - window_size = int(name.split('_')[-1]) + elif name.startswith("instance") and causal: + window_size = int(name.split("_")[-1]) return CausalInstanceNorm1d(*args, **kwargs, window_size=window_size) else: - raise ValueError(f'Normalisation with {name=} and {causal=} unknown.') + raise ValueError(f"Normalisation with {name=} and {causal=} unknown.") diff --git a/src/hypnos/settings.py b/src/hypnos/settings.py index 951fb74..385b605 100644 --- a/src/hypnos/settings.py +++ b/src/hypnos/settings.py @@ -9,14 +9,14 @@ """ # Canonical channel names used by the 8 modalities. -EEG_C3 = 'C3' -EEG_C4 = 'C4' -EOG_E1 = 'E1' -EOG_E2 = 'E2' -EMG_CHIN = 'Chin' -ECG = 'ECG' -ABD = 'ABD' -THX = 'THX' +EEG_C3 = "C3" +EEG_C4 = "C4" +EOG_E1 = "E1" +EOG_E2 = "E2" +EMG_CHIN = "Chin" +ECG = "ECG" +ABD = "ABD" +THX = "THX" # Canonical channel name -> row in the (8-row) channel-embedding table. CHANNEL_REGISTRY: dict[str, int] = { diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index a876071..d74295c 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -89,19 +89,34 @@ def build_bundle(path): continue kw = tokenizer_kwargs(k, sr, ratios) tokenizers[stem] = { - "signal_type": _st, "num_quantizers": k, "codebook_size": CB, - "token_duration_sec": 1.0, "sample_rate": sr, "tokenizer_kwargs": kw, + "signal_type": _st, + "num_quantizers": k, + "codebook_size": CB, + "token_duration_sec": 1.0, + "sample_rate": sr, + "tokenizer_kwargs": kw, } for k_, v in SignalTokenizer(**kw).state_dict().items(): tensors[f"tok/{stem}/{k_}"] = v.detach().contiguous().clone() modalities = [ - {"name": n, "signal_type": st, "channels": ch, "tokenizer": stem, "num_quantizers": k, - "codebook_size": CB, "token_duration_sec": 1.0, "sample_rate": sr, "preprocess_modality": pp} + { + "name": n, + "signal_type": st, + "channels": ch, + "tokenizer": stem, + "num_quantizers": k, + "codebook_size": CB, + "token_duration_sec": 1.0, + "sample_rate": sr, + "preprocess_modality": pp, + } for (n, st, ch, stem, k, sr, _r, pp) in MODALITIES ] config = { "model_target": "hypnos.models.rq_transformer.MultiModalRQTransformer", - "model_kwargs": MODEL_KWARGS, "modalities": modalities, "tokenizers": tokenizers, + "model_kwargs": MODEL_KWARGS, + "modalities": modalities, + "tokenizers": tokenizers, } save_file(tensors, str(path), metadata={"format_version": "1", "config": json.dumps(config)})