From 0339f1a2782bffe7664b85417117e9aefe8f85c7 Mon Sep 17 00:00:00 2001 From: Jiarui Li Date: Sun, 12 Jul 2026 10:48:52 +0800 Subject: [PATCH] Add task adapter for exposure embeddings --- models.py | 21 ++++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/models.py b/models.py index 3b79f45..4780ab3 100644 --- a/models.py +++ b/models.py @@ -49,6 +49,21 @@ class DeepHealth(nn.Module): 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 self.use_exposure_embeddings: + self.exposure_adapter = nn.Sequential( + nn.LayerNorm(n_embd), + nn.Linear(n_embd, n_embd), + nn.GELU(), + nn.Linear(n_embd, n_embd), + ) + self.exposure_gate = nn.Parameter(torch.tensor(-2.0)) + for module in self.exposure_adapter: + if isinstance(module, nn.Linear): + nn.init.normal_(module.weight, mean=0.0, std=0.02) + nn.init.zeros_(module.bias) + else: + self.exposure_adapter = None + self.register_parameter("exposure_gate", None) if dist_mode == "weibull": self.rho_head = nn.Linear(n_embd, vocab_size) nn.init.zeros_(self.rho_head.weight) @@ -132,9 +147,13 @@ class DeepHealth(nn.Module): "exposure_embedding must have shape " f"{tuple(h_disease.shape)}, got {tuple(exposure_embedding.shape)}" ) - h_disease = h_disease + exposure_embedding.to( + exposure = exposure_embedding.to( device=h_disease.device, dtype=h_disease.dtype ) + if self.exposure_adapter is None or self.exposure_gate is None: + raise RuntimeError("Exposure adapter is not initialized") + exposure = self.exposure_adapter(exposure) + h_disease = h_disease + torch.sigmoid(self.exposure_gate) * exposure elif exposure_embedding is not None: raise ValueError( "exposure_embedding provided but use_exposure_embeddings=False"