763 lines
28 KiB
Python
763 lines
28 KiB
Python
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_v1"
|
|
|
|
|
|
@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=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
|