Add time attention mask handling and baseline class time computation to DeepHealth model

models.py

@@ -141,16 +141,42 @@ class DeepHealth(nn.Module):
         summary = baseline_summary.to(device=h_disease.device, dtype=h_disease.dtype)
         return torch.where(checkup_mask.unsqueeze(-1), summary[:, None, :], h_disease)
 
+    def _baseline_cls_time(
+        self,
+        event_seq: torch.Tensor,
+        time_seq: torch.Tensor,
+        padding_mask: torch.Tensor,
+    ) -> torch.Tensor:
+        checkup_mask = event_seq == CHECKUP_IDX
+        inf = torch.full_like(time_seq, float("inf"))
+        first_checkup = torch.where(checkup_mask, time_seq, inf).min(dim=1).values
+        has_checkup = torch.isfinite(first_checkup)
+        fallback_time = torch.where(
+            padding_mask,
+            time_seq,
+            torch.full_like(time_seq, float("-inf")),
+        ).max(dim=1).values
+        fallback_time = torch.where(
+            torch.isfinite(fallback_time),
+            fallback_time,
+            torch.zeros_like(fallback_time),
+        )
+        return torch.where(has_checkup, first_checkup, fallback_time)
+
     def _encode_other_tokens(
         self,
         other_type: torch.LongTensor,
         other_value: torch.Tensor,
         other_value_kind: torch.LongTensor,
+        other_time: torch.Tensor,
+        cls_time: torch.Tensor,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         return self.token_encoder(
             other_type=other_type,
             other_value=other_value,
             other_value_kind=other_value_kind,
+            other_time=other_time,
+            cls_time=cls_time,
         )
 
     def _forward_shared(
@@ -193,16 +219,24 @@ class DeepHealth(nn.Module):
         h_disease = self.token_embedding(event_seq)
         t_disease = time_seq
 
-        h_token, token_mask, baseline_summary = self._encode_other_tokens(
-            other_type=other_type,
-            other_value=other_value,
-            other_value_kind=other_value_kind,
-        )
         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)
+        cls_time = self._baseline_cls_time(
+            event_seq=event_seq,
+            time_seq=time_seq,
+            padding_mask=padding_mask,
+        )
+        h_token, token_mask, baseline_summary = self._encode_other_tokens(
+            other_type=other_type,
+            other_value=other_value,
+            other_value_kind=other_value_kind,
+            other_time=other_time,
+            cls_time=cls_time,
+        )
         token_time = other_time.to(device=h_token.device, dtype=time_seq.dtype)
 
         h_disease = self.cross_attention(