Refactor DeepHealth model to expand extra-info token handling and update output structure

models.py

@@ -322,7 +322,7 @@ class DeepHealth(nn.Module):
         other_time: torch.FloatTensor | None = None,
         return_output: bool = False,
         **unused_kwargs,
-    ) -> torch.Tensor:
+    ) -> torch.Tensor | DeepHealthOutput:
         if unused_kwargs:
             unknown = ", ".join(sorted(unused_kwargs))
             raise TypeError(f"Unexpected DeepHealth forward arguments: {unknown}")
@@ -363,11 +363,6 @@ class DeepHealth(nn.Module):
         )
         h_other = h_other.to(device=event_seq.device)
         other_mask = other_mask.to(device=event_seq.device, dtype=torch.bool)
-        h_other, other_time, other_mask = self._pool_other_by_time(
-            h_other=h_other,
-            other_time=other_time,
-            other_mask=other_mask,
-        )
 
         h_disease = torch.cat([h_disease, h_other], dim=1)
         t_disease = torch.cat([t_disease, other_time], dim=1)
@@ -433,10 +428,18 @@ class DeepHealth(nn.Module):
                 )
             return hidden
         if return_output:
+            h_event = h_disease[:, :event_len, :]
+            t_event = t_disease[:, :event_len]
+            event_mask = padding_mask[:, :event_len]
+            h_extra, t_extra, extra_mask = self._pool_other_by_time(
+                h_other=h_disease[:, event_len:, :],
+                other_time=t_disease[:, event_len:],
+                other_mask=padding_mask[:, event_len:],
+            )
             return DeepHealthOutput(
-                hidden=h_disease,
-                time_seq=t_disease,
-                padding_mask=padding_mask,
+                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 h_disease[:, :event_len, :]