JiaruiLi/DeepHealth
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Remove legacy event and mixed distribution paths

Browse Source
This commit is contained in:
Jiarui Li
2026-08-01 14:23:18 +08:00
parent dfb22adf2d
commit de6f9b75b9
22 changed files with 370 additions and 463 deletions
Download Patch File Download Diff File Expand all files Collapse all files

60
models.py
View File

@@ -37,7 +37,6 @@ class OtherInfoTokenizer(nn.Module):
cont_type_ids: list[int],
n_value_kinds: int = 3,
n_bins: int = 16,
continuous_value_scaling: str = "none",
continuous_value_center: torch.Tensor | list[float] | None = None,
continuous_value_scale: torch.Tensor | list[float] | None = None,
):
@@ -57,13 +56,6 @@ class OtherInfoTokenizer(nn.Module):
raise ValueError(
f"n_value_kinds must be > {self.CATE_KIND}, got {n_value_kinds}"
)
continuous_value_scaling = str(continuous_value_scaling).lower()
if continuous_value_scaling not in {"none", "robust"}:
raise ValueError(
"continuous_value_scaling must be either 'none' or 'robust', "
f"got {continuous_value_scaling!r}"
)
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 = (
@@ -80,18 +72,20 @@ class OtherInfoTokenizer(nn.Module):
n_embd,
padding_idx=0,
)
self.continuous_value_scaling = continuous_value_scaling
if continuous_value_scaling == "robust" and n_cont_types > 0:
if n_cont_types > 0:
if continuous_value_center is None or continuous_value_scale is None:
raise ValueError(
"Continuous values require train-split RobustScale center "
"and scale statistics"
)
center = self._coerce_scaler_buffer(
continuous_value_center,
n_cont_types=n_cont_types,
default=0.0,
name="continuous_value_center",
)
scale = self._coerce_scaler_buffer(
continuous_value_scale,
n_cont_types=n_cont_types,
default=1.0,
name="continuous_value_scale",
)
if not torch.isfinite(center).all():
@@ -130,11 +124,10 @@ class OtherInfoTokenizer(nn.Module):
value: torch.Tensor | list[float] | None,
*,
n_cont_types: int,
default: float,
name: str,
) -> torch.Tensor:
if value is None:
return torch.full((n_cont_types,), float(default), dtype=torch.float32)
raise ValueError(f"{name} is required")
tensor = torch.as_tensor(value, dtype=torch.float32).detach().clone()
if tensor.shape != (n_cont_types,):
raise ValueError(
@@ -184,15 +177,14 @@ class OtherInfoTokenizer(nn.Module):
"cont_type_ids"
)
cont_value = other_value[cont_pos].to(type_emb.dtype)
if self.continuous_value_scaling == "robust":
if (
self.continuous_value_center is None
or self.continuous_value_scale is None
):
raise RuntimeError("Robust continuous-value scaler buffers are missing")
center = self.continuous_value_center[cont_idx].to(type_emb.dtype)
scale = self.continuous_value_scale[cont_idx].to(type_emb.dtype)
cont_value = (cont_value - center) / scale
if (
self.continuous_value_center is None
or self.continuous_value_scale is None
):
raise RuntimeError("RobustScale buffers are missing")
center = self.continuous_value_center[cont_idx].to(type_emb.dtype)
scale = self.continuous_value_scale[cont_idx].to(type_emb.dtype)
cont_value = (cont_value - center) / scale
value_emb[cont_pos] = self.cont_value_encoder(
cont_type_idx=cont_idx,
value=cont_value,
@@ -221,12 +213,11 @@ class DeepHealth(nn.Module):
cont_type_ids: list[int],
n_value_kinds: int = 3,
n_bins: int = 16,
continuous_value_scaling: str = "none",
continuous_value_center: torch.Tensor | list[float] | None = None,
continuous_value_scale: torch.Tensor | list[float] | None = None,
target_mode: str = "next_token", # "next_token" or "all_future"
time_mode: str = "absolute", # next_token requires absolute
dist_mode: str = "exponential", # "exponential", "weibull" or "mixed"
dist_mode: str = "exponential", # "exponential" or "weibull"
extra_pool_reduce: str = "mean",
dropout: float = 0.0,
model_architecture: str | None = None,
@@ -243,9 +234,9 @@ class DeepHealth(nn.Module):
"next_token is reserved for Delphi2M reproduction and "
"requires time_mode='absolute'"
)
if dist_mode not in ["exponential", "weibull", "mixed"]:
if dist_mode not in ["exponential", "weibull"]:
raise ValueError(
"dist_mode must be either 'exponential', 'weibull' or 'mixed'")
"dist_mode must be either 'exponential' or 'weibull'")
if extra_pool_reduce not in {"mean", "sum"}:
raise ValueError("extra_pool_reduce must be either 'mean' or 'sum'")
if n_layer < 1:
@@ -262,7 +253,6 @@ class DeepHealth(nn.Module):
cont_type_ids=cont_type_ids,
n_value_kinds=n_value_kinds,
n_bins=n_bins,
continuous_value_scaling=continuous_value_scaling,
continuous_value_center=continuous_value_center,
continuous_value_scale=continuous_value_scale,
)
@@ -270,7 +260,6 @@ class DeepHealth(nn.Module):
self.time_mode = time_mode
self.dist_mode = dist_mode
self.extra_pool_reduce = extra_pool_reduce
self.continuous_value_scaling = str(continuous_value_scaling).lower()
self.model_architecture = model_architecture
self.n_layer = n_layer
self.n_embd = n_embd
@@ -283,12 +272,6 @@ class DeepHealth(nn.Module):
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(n_embd, 1)
nn.init.zeros_(self.rho_death_head.weight)
nn.init.constant_(self.rho_death_head.bias, 0.5413)
if time_mode == "absolute":
self.age_encoding = AgeSinusoidalEncoding(n_embd)
self.blocks = nn.ModuleList([
@@ -551,10 +534,3 @@ class DeepHealth(nn.Module):
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