from collections.abc import Mapping from dataclasses import dataclass import torch import torch.nn as nn import torch.nn.functional as F from backbones import ( AgeSinusoidalEncoding, GaussianRBFTimeBasis, SharedEventTrajectoryCore, TimeRoPE, TokenAutoDiscretization, ) from targets import PAD_IDX EVENT_TRAJECTORY_ARCHITECTURE = "event_trajectory_shared_v2" @dataclass(frozen=True) class EventTrajectoryModelSize: d_model: int n_trajectory: int @property def trajectory_dim(self) -> int: return self.d_model // self.n_trajectory @property def traj_hidden(self) -> int: return 4 * self.n_trajectory MODEL_SIZE_PRESETS = { "nano": EventTrajectoryModelSize(d_model=120, n_trajectory=6), "tiny": EventTrajectoryModelSize(d_model=256, n_trajectory=8), "small": EventTrajectoryModelSize(d_model=512, n_trajectory=8), "medium": EventTrajectoryModelSize(d_model=768, n_trajectory=12), "huge": EventTrajectoryModelSize(d_model=1024, n_trajectory=16), } MODEL_SIZE_NAMES = tuple(MODEL_SIZE_PRESETS) def resolve_model_size(model_size: str) -> EventTrajectoryModelSize: if not isinstance(model_size, str): raise ValueError( f"model_size must be a string, got {type(model_size).__name__}" ) normalized = model_size.strip().lower() try: return MODEL_SIZE_PRESETS[normalized] except KeyError as exc: choices = ", ".join(MODEL_SIZE_NAMES) raise ValueError( f"Unknown model_size {model_size!r}; expected one of: {choices}" ) from exc def _required_config_int( config: Mapping[str, object], key: str, ) -> int: raw_value = config.get(key) if isinstance(raw_value, bool): raise ValueError(f"Config field {key!r} must be an integer") try: value = int(raw_value) except (TypeError, ValueError) as exc: raise ValueError( f"Config field {key!r} must be present and integer-valued; " f"got {raw_value!r}" ) from exc if isinstance(raw_value, float) and not raw_value.is_integer(): raise ValueError(f"Config field {key!r} must be an integer") return value def validate_event_trajectory_config(config: Mapping[str, object]) -> None: actual = config.get("model_architecture") if actual != EVENT_TRAJECTORY_ARCHITECTURE: raise ValueError( "This branch only accepts models trained with the shared " "event-trajectory architecture marker " f"{EVENT_TRAJECTORY_ARCHITECTURE!r}; got {actual!r}." ) raw_model_size = config.get("model_size") if not isinstance(raw_model_size, str): raise ValueError( "Config field 'model_size' must be one of: " + ", ".join(MODEL_SIZE_NAMES) ) model_size = raw_model_size.strip().lower() preset = resolve_model_size(model_size) d_model = _required_config_int(config, "d_model") n_trajectory = _required_config_int(config, "n_trajectory") n_reasoning_rounds = _required_config_int( config, "n_reasoning_rounds", ) trajectory_dim = _required_config_int(config, "trajectory_dim") traj_hidden = _required_config_int(config, "traj_hidden") if n_reasoning_rounds <= 0: raise ValueError( "n_reasoning_rounds must be positive" ) expected_values = { "d_model": preset.d_model, "n_trajectory": preset.n_trajectory, "trajectory_dim": preset.trajectory_dim, "traj_hidden": preset.traj_hidden, } actual_values = { "d_model": d_model, "n_trajectory": n_trajectory, "trajectory_dim": trajectory_dim, "traj_hidden": traj_hidden, } mismatches = [ f"{key}: expected {expected}, got {actual_values[key]}" for key, expected in expected_values.items() if actual_values[key] != expected ] if mismatches: raise ValueError( f"Config does not match model_size={model_size!r}: " + "; ".join(mismatches) ) def _checkpoint_scalar_int( state_dict: Mapping[str, object], key: str, ) -> int: value = state_dict[key] if not isinstance(value, torch.Tensor) or value.numel() != 1: raise ValueError( f"Checkpoint architecture field {key!r} must be a scalar tensor" ) return int(value.detach().cpu().item()) def validate_event_trajectory_state_dict( state_dict: Mapping[str, object], *, expected_d_model: int | None = None, expected_n_trajectory: int | None = None, expected_n_reasoning_rounds: int | None = None, ) -> None: required_keys = { "architecture_d_model", "architecture_n_trajectory", "architecture_n_reasoning_rounds", "event_projection.weight", "trajectory_prototypes", "query_projection.weight", "reasoning_core.cross_attention.q_proj.weight", "reasoning_core.cross_attention.k_proj.weight", "reasoning_core.cross_attention.v_proj.weight", "reasoning_core.traj_mixer.gate_proj", "reasoning_core.traj_mixer.value_proj", "reasoning_core.traj_mixer.output_proj", "reasoning_core.attn_scale", "reasoning_core.mixer_scale", } missing = sorted(required_keys.difference(state_dict)) if missing: raise ValueError( "Checkpoint is not a shared event-trajectory checkpoint; " "missing required " f"parameters: {', '.join(missing)}" ) checkpoint_values = { "d_model": _checkpoint_scalar_int( state_dict, "architecture_d_model", ), "n_trajectory": _checkpoint_scalar_int( state_dict, "architecture_n_trajectory", ), "n_reasoning_rounds": _checkpoint_scalar_int( state_dict, "architecture_n_reasoning_rounds", ), } expected_values = { "d_model": expected_d_model, "n_trajectory": expected_n_trajectory, "n_reasoning_rounds": expected_n_reasoning_rounds, } mismatches = [ f"{name}: checkpoint={checkpoint_values[name]}, expected={expected}" for name, expected in expected_values.items() if expected is not None and checkpoint_values[name] != expected ] if mismatches: raise ValueError( "Checkpoint architecture does not match the constructed model: " + "; ".join(mismatches) ) @dataclass class DeepHealthOutput: hidden: torch.Tensor time_seq: torch.Tensor padding_mask: torch.Tensor event_len: int class OtherInfoTokenizer(nn.Module): PAD_KIND = 0 CONT_KIND = 1 CATE_KIND = 2 def __init__( self, n_embd: int, n_types: int, n_cont_types: int, n_categories: int, cont_type_ids: list[int], n_value_kinds: int = 3, n_bins: int = 16, ): super().__init__() if len(cont_type_ids) != n_cont_types: raise ValueError( "cont_type_ids length must match n_cont_types, got " f"{len(cont_type_ids)} vs {n_cont_types}" ) if n_types <= 0: raise ValueError(f"n_types must include PAD and be > 0, got {n_types}") if n_categories <= 0: raise ValueError( f"n_categories must include PAD and be > 0, got {n_categories}" ) if n_value_kinds <= self.CATE_KIND: raise ValueError( f"n_value_kinds must be > {self.CATE_KIND}, got {n_value_kinds}" ) self.type_emb = nn.Embedding(n_types, n_embd, padding_idx=0) self.kind_emb = nn.Embedding(n_value_kinds, n_embd, padding_idx=0) self.cont_value_encoder = ( TokenAutoDiscretization( n_cont_types=n_cont_types, n_bins=n_bins, n_embd=n_embd, ) if n_cont_types > 0 else None ) self.cate_value_emb = nn.Embedding( n_categories, n_embd, padding_idx=0, ) cont_type_index = torch.full((n_types,), -1, dtype=torch.long) for idx, type_id in enumerate(cont_type_ids): if type_id <= 0 or type_id >= n_types: raise ValueError( f"continuous type id {type_id} must be in [1, {n_types})" ) cont_type_index[type_id] = idx self.register_buffer( "cont_type_index", cont_type_index, persistent=False, ) self.reset_parameters() def reset_parameters(self) -> None: nn.init.normal_(self.type_emb.weight, mean=0.0, std=0.02) nn.init.zeros_(self.type_emb.weight[0]) nn.init.normal_(self.kind_emb.weight, mean=0.0, std=0.02) nn.init.zeros_(self.kind_emb.weight[0]) nn.init.normal_(self.cate_value_emb.weight, mean=0.0, std=0.02) nn.init.zeros_(self.cate_value_emb.weight[0]) def forward( self, other_type: torch.LongTensor, other_value: torch.Tensor, other_value_kind: torch.LongTensor, ) -> tuple[torch.Tensor, torch.Tensor]: if other_type.shape != other_value.shape: raise ValueError( "other_type and other_value must have the same shape, got " f"{tuple(other_type.shape)} vs {tuple(other_value.shape)}" ) if other_type.shape != other_value_kind.shape: raise ValueError( "other_type and other_value_kind must have the same shape, got " f"{tuple(other_type.shape)} vs {tuple(other_value_kind.shape)}" ) other_valid = other_type > 0 type_emb = self.type_emb(other_type) kind_emb = self.kind_emb(other_value_kind) value_emb = torch.zeros_like(type_emb) cont_pos = other_valid & (other_value_kind == self.CONT_KIND) if cont_pos.any(): if self.cont_value_encoder is None: raise ValueError("continuous tokens found but n_cont_types is 0") cont_idx = self.cont_type_index[other_type[cont_pos]] if (cont_idx < 0).any(): bad_type = other_type[cont_pos][cont_idx < 0][0].item() raise ValueError( f"type_id={bad_type} is marked continuous but is not in " "cont_type_ids" ) value_emb[cont_pos] = self.cont_value_encoder( cont_type_idx=cont_idx, value=other_value[cont_pos].to(type_emb.dtype), ) cate_pos = other_valid & (other_value_kind == self.CATE_KIND) if cate_pos.any(): cate_id = other_value[cate_pos].long() value_emb[cate_pos] = self.cate_value_emb(cate_id) out = type_emb + kind_emb + value_emb out = out * other_valid.unsqueeze(-1).to(out.dtype) return out, other_valid class DeepHealth(nn.Module): def __init__( self, vocab_size: int, model_size: str, n_reasoning_rounds: int, n_types: int, n_cont_types: int, n_categories: int, cont_type_ids: list[int], n_value_kinds: int = 3, n_bins: int = 16, target_mode: str = "next_token", # "next_token" or "all_future" time_mode: str = "relative", # "relative" or "absolute" dist_mode: str = "exponential", # "exponential", "weibull" or "mixed" extra_pool_reduce: str = "mean", dropout: float = 0.0, ): super().__init__() if target_mode not in ["next_token", "all_future"]: raise ValueError( "target_mode must be either 'next_token' or 'all_future'") if time_mode not in ["relative", "absolute"]: raise ValueError( "time_mode must be either 'relative' or 'absolute'") if dist_mode not in ["exponential", "weibull", "mixed"]: raise ValueError( "dist_mode must be either 'exponential', 'weibull' or 'mixed'") if extra_pool_reduce not in {"mean", "sum"}: raise ValueError("extra_pool_reduce must be either 'mean' or 'sum'") if n_reasoning_rounds <= 0: raise ValueError( "n_reasoning_rounds must be positive, got " f"{n_reasoning_rounds}" ) size_config = resolve_model_size(model_size) normalized_model_size = model_size.strip().lower() d_model = size_config.d_model n_trajectory = size_config.n_trajectory self.token_embedding = nn.Embedding(vocab_size, d_model, padding_idx=0) self.gender_embedding = nn.Embedding( 2, d_model) # Assuming binary gender self.tokenizer = OtherInfoTokenizer( n_embd=d_model, n_types=n_types, n_cont_types=n_cont_types, n_categories=n_categories, cont_type_ids=cont_type_ids, n_value_kinds=n_value_kinds, n_bins=n_bins, ) self.target_mode = target_mode self.time_mode = time_mode self.dist_mode = dist_mode self.extra_pool_reduce = extra_pool_reduce self.model_size = normalized_model_size self.d_model = d_model self.n_trajectory = n_trajectory self.trajectory_dim = d_model // n_trajectory self.traj_hidden = 4 * n_trajectory self.n_reasoning_rounds = n_reasoning_rounds self.vocab_size = vocab_size self.register_buffer( "architecture_d_model", torch.tensor(d_model, dtype=torch.int64), ) self.register_buffer( "architecture_n_trajectory", torch.tensor(n_trajectory, dtype=torch.int64), ) self.register_buffer( "architecture_n_reasoning_rounds", torch.tensor(n_reasoning_rounds, dtype=torch.int64), ) nn.init.normal_(self.token_embedding.weight, mean=0.0, std=0.02) nn.init.zeros_(self.token_embedding.weight[0]) nn.init.normal_(self.gender_embedding.weight, mean=0.0, std=0.02) if dist_mode == "weibull": self.rho_head = nn.Linear(d_model, vocab_size) nn.init.zeros_(self.rho_head.weight) nn.init.constant_(self.rho_head.bias, 0.5413) if dist_mode == "mixed": self.death_idx = vocab_size - 1 self.rho_death_head = nn.Linear(d_model, 1) nn.init.zeros_(self.rho_death_head.weight) nn.init.constant_(self.rho_death_head.bias, 0.5413) # Event and query time are encoded once before shared reasoning. In # relative mode, cross-attention additionally uses TimeRoPE and RBF. self.age_encoding = AgeSinusoidalEncoding(d_model) self.event_projection = nn.Linear(d_model, d_model, bias=False) self.event_norm = nn.LayerNorm(d_model) self.query_projection = nn.Linear(d_model, d_model, bias=False) self.trajectory_prototypes = nn.Parameter( torch.empty(n_trajectory, self.trajectory_dim) ) self.query_token = nn.Parameter(torch.empty(d_model)) nn.init.normal_(self.event_projection.weight, mean=0.0, std=0.02) nn.init.normal_(self.query_projection.weight, mean=0.0, std=0.02) nn.init.normal_(self.trajectory_prototypes, mean=0.0, std=0.02) nn.init.normal_(self.query_token, mean=0.0, std=0.02) use_relative_time = time_mode == "relative" self.reasoning_core = SharedEventTrajectoryCore( d_model=d_model, n_trajectory=n_trajectory, n_reasoning_rounds=n_reasoning_rounds, dropout=dropout, n_rbf_bases=16, use_time_rope=use_relative_time, use_rbf_bias=use_relative_time, ) if use_relative_time: self.rope = TimeRoPE(self.trajectory_dim) self.rbf = GaussianRBFTimeBasis( n_bases=16, max_time_diff=40.0, ) else: self.rope = None self.rbf = None self.final_ln = nn.LayerNorm(d_model) self.risk_head = nn.Linear(d_model, vocab_size, bias=False) if target_mode == "next_token": self.risk_head.weight = self.token_embedding.weight def _make_event_invalid_mask( self, event_valid_mask: torch.Tensor, event_time: torch.Tensor, query_time: torch.Tensor, query_position: torch.Tensor | None = None, ) -> torch.Tensor: valid_key = event_valid_mask[:, None, :] key_time = event_time[:, None, :] query_time = query_time[:, :, None] if query_position is None: visible_by_time = key_time <= query_time else: key_position = torch.arange( event_time.size(1), device=event_time.device, ).view(1, 1, -1) visible_by_time = (key_time < query_time) | ( (key_time == query_time) & (key_position <= query_position[:, :, None]) ) return ~(valid_key & visible_by_time) def _pool_other_by_time( self, h_other: torch.Tensor, other_time: torch.Tensor, other_mask: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: batch_size, n_other, n_embd = h_other.shape if n_other == 0: empty_h = h_other.new_zeros(batch_size, 0, n_embd) empty_t = other_time.new_zeros(batch_size, 0) empty_m = torch.zeros(batch_size, 0, dtype=torch.bool, device=h_other.device) return empty_h, empty_t, empty_m masked_time = other_time.masked_fill(~other_mask, float("inf")) _sorted_time_with_pad, order = masked_time.sort(dim=1) sorted_time = other_time.gather(1, order) sorted_mask = other_mask.gather(1, order) sorted_h = h_other.gather(1, order.unsqueeze(-1).expand(-1, -1, n_embd)) group_start = torch.zeros_like(sorted_mask) group_start[:, 0] = sorted_mask[:, 0] group_start[:, 1:] = sorted_mask[:, 1:] & ( sorted_time[:, 1:] != sorted_time[:, :-1] ) group_id = group_start.long().cumsum(dim=1) - 1 max_groups = int(group_start.sum(dim=1).max().item()) pooled_h = h_other.new_zeros(batch_size, max_groups, n_embd) pooled_time = other_time.new_zeros(batch_size, max_groups) pooled_mask = torch.zeros( batch_size, max_groups, dtype=torch.bool, device=h_other.device, ) if max_groups == 0: return pooled_h, pooled_time, pooled_mask safe_group_id = group_id.clamp_min(0) pooled_h.scatter_add_( 1, safe_group_id.unsqueeze(-1).expand_as(sorted_h), sorted_h * sorted_mask.unsqueeze(-1).to(sorted_h.dtype), ) if self.extra_pool_reduce == "mean": counts = h_other.new_zeros(batch_size, max_groups, 1) counts.scatter_add_( 1, safe_group_id.unsqueeze(-1), sorted_mask.unsqueeze(-1).to(h_other.dtype), ) pooled_h = pooled_h / counts.clamp_min(1.0) pooled_time.scatter_add_( 1, safe_group_id, sorted_time * group_start.to(sorted_time.dtype), ) group_count = group_start.sum(dim=1) arange_groups = torch.arange(max_groups, device=h_other.device) pooled_mask = arange_groups.unsqueeze(0) < group_count.unsqueeze(1) return pooled_h, pooled_time, pooled_mask def _forward_shared( self, event_seq: torch.LongTensor, time_seq: torch.FloatTensor, sex: torch.LongTensor, mode: str, padding_mask: torch.Tensor | None = None, t_query: torch.FloatTensor | None = None, other_type: torch.LongTensor | None = None, other_value: torch.Tensor | None = None, other_value_kind: torch.LongTensor | None = None, other_time: torch.FloatTensor | None = None, return_output: bool = False, **unused_kwargs, ) -> torch.Tensor | DeepHealthOutput: if unused_kwargs: unknown = ", ".join(sorted(unused_kwargs)) raise TypeError(f"Unexpected DeepHealth forward arguments: {unknown}") if mode not in {"next_token", "all_future"}: raise ValueError("mode must be either 'next_token' or 'all_future'") if mode == "all_future" and t_query is None: raise ValueError("t_query is required when mode='all_future'") if ( other_type is None or other_value is None or other_value_kind is None or other_time is None ): raise ValueError( "DeepHealth expects other_type, other_value, " "other_value_kind, and other_time." ) if padding_mask is None: padding_mask = event_seq > PAD_IDX else: padding_mask = padding_mask.to(device=event_seq.device, dtype=torch.bool) event_len = event_seq.size(1) event_features = self.token_embedding(event_seq) event_time = time_seq if other_time.shape != other_type.shape: raise ValueError( "other_time must have the same shape as other_type, got " f"{tuple(other_time.shape)} vs {tuple(other_type.shape)}" ) other_time = other_time.to(device=event_seq.device, dtype=time_seq.dtype) other_features, other_mask = self.tokenizer( other_type=other_type, other_value=other_value, other_value_kind=other_value_kind, ) other_features = other_features.to(device=event_seq.device) other_mask = other_mask.to(device=event_seq.device, dtype=torch.bool) event_features = torch.cat([event_features, other_features], dim=1) event_time = torch.cat([event_time, other_time], dim=1) event_valid_mask = torch.cat([padding_mask, other_mask], dim=1) batch_size = event_seq.size(0) sex_context = self.gender_embedding(sex)[:, None, :] event_features = ( event_features + sex_context + self.age_encoding(event_time) ) event_features = event_features * event_valid_mask.unsqueeze(-1).to( event_features.dtype ) event_memory = self.event_norm( self.event_projection(event_features) ) event_memory = event_memory * event_valid_mask.unsqueeze(-1).to( event_memory.dtype ) if mode == "all_future": query_time = t_query[:, None] query_position = None query_features = ( self.query_token.view(1, 1, -1) + sex_context + self.age_encoding(query_time) ) query_valid_mask = torch.ones( batch_size, 1, dtype=torch.bool, device=event_seq.device, ) else: # Each event position is an independent parallel query. Including # its event feature preserves token-level next-step semantics. # Equal-time memory is additionally position-causal so a token # cannot read a later token that may be its Delphi2M target. query_time = event_time query_position = torch.arange( event_time.size(1), device=event_time.device, ).view(1, -1).expand(batch_size, -1) query_features = event_features query_valid_mask = event_valid_mask n_query = query_time.size(1) query_context = self.query_projection(query_features).reshape( batch_size, n_query, self.n_trajectory, self.trajectory_dim, ) trajectory_state = ( self.trajectory_prototypes.view( 1, 1, self.n_trajectory, self.trajectory_dim, ) + query_context ) event_invalid_mask = self._make_event_invalid_mask( event_valid_mask=event_valid_mask, event_time=event_time, query_time=query_time, query_position=query_position, ) event_rope_cache = None query_rope_cache = None rbf_cache = None if self.time_mode == "relative": if self.rope is None or self.rbf is None: raise RuntimeError("Relative-time modules are not initialized") event_rope_cache = self.rope.precompute_cache(event_time) query_rope_cache = self.rope.precompute_cache(query_time) rbf_cache = self.rbf.precompute_cross_cache( query_time, event_time, ) event_key_value = self.reasoning_core.project_event_memory( event_memory, event_rope_cache=event_rope_cache, ) for _ in range(self.n_reasoning_rounds): trajectory_state = self.reasoning_core( trajectory_state=trajectory_state, event_key_value=event_key_value, event_invalid_mask=event_invalid_mask, query_rope_cache=query_rope_cache, rbf_cache=rbf_cache, ) hidden_sequence = self.final_ln( trajectory_state.reshape(batch_size, n_query, self.d_model) ) hidden_sequence = hidden_sequence * query_valid_mask.unsqueeze(-1).to( hidden_sequence.dtype ) if mode == "all_future": hidden = hidden_sequence[:, 0, :] if return_output: return DeepHealthOutput( hidden=hidden, time_seq=t_query[:, None], padding_mask=torch.ones( hidden.size(0), 1, dtype=torch.bool, device=hidden.device, ), event_len=event_len, ) return hidden if return_output: h_event = hidden_sequence[:, :event_len, :] t_event = event_time[:, :event_len] event_mask = event_valid_mask[:, :event_len] h_extra, t_extra, extra_mask = self._pool_other_by_time( h_other=hidden_sequence[:, event_len:, :], other_time=event_time[:, event_len:], other_mask=event_valid_mask[:, event_len:], ) return DeepHealthOutput( hidden=torch.cat([h_event, h_extra], dim=1), time_seq=torch.cat([t_event, t_extra], dim=1), padding_mask=torch.cat([event_mask, extra_mask], dim=1), event_len=event_len, ) return hidden_sequence[:, :event_len, :] def forward_next_token(self, **kwargs) -> torch.Tensor: return self._forward_shared(mode="next_token", **kwargs) def forward_all_future(self, **kwargs) -> torch.Tensor: return self._forward_shared(mode="all_future", **kwargs) def forward(self, target_mode: str | None = None, **kwargs) -> torch.Tensor: mode = self.target_mode if target_mode is None else target_mode return self._forward_shared(mode=mode, **kwargs) def calc_risk(self, x: torch.Tensor) -> torch.Tensor: return self.risk_head(x) def calc_weibull_rho(self, x: torch.Tensor) -> torch.Tensor: if self.dist_mode != "weibull": raise RuntimeError( f"calc_weibull_rho called with dist_mode={self.dist_mode!r}" ) return F.softplus(self.rho_head(x)) + 1e-6 def calc_death_rho(self, x: torch.Tensor) -> torch.Tensor: if self.dist_mode != "mixed": raise RuntimeError( f"calc_death_rho called with dist_mode={self.dist_mode!r}" ) return F.softplus(self.rho_death_head(x)).squeeze(-1) + 1e-6