Add task adapter for exposure embeddings

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"