Add content-only token_emb to UnifiedMultimodalEmbeddingModel output

pyhealth/models/embedding/unified.py

@@ -402,9 +402,12 @@ def forward(
             A dict with keys:
 
             * ``"sequence"``, ``(B, S_total, E')``  temporally-sorted events
+              (content + time + type embeddings)
             * ``"time"``    , ``(B, S_total)``       timestamps (hours)
             * ``"mask"``    , ``(B, S_total)``       1=real event, 0=padding
             * ``"type_ids"``, ``(B, S_total)``       modality index per event
+            * ``"token_emb"``, ``(B, S_total, E')`` content-only event embedding
+              (before time/type are added); the target for masked modeling.
         """
         all_embeddings: list[torch.Tensor] = []
         all_times:      list[torch.Tensor] = []
@@ -493,4 +496,10 @@ def forward(
             "time":     cat_time,   # (B, S_total)
             "mask":     cat_mask,   # (B, S_total)
             "type_ids": cat_types,  # (B, S_total)
+            # Per-event content embedding BEFORE time/type are added (same sort
+            # order as ``sequence``).  Masked-modeling pretrainers should
+            # reconstruct THIS rather than ``sequence``: the time/type
+            # components are largely recoverable from event position, so
+            # including them in the target dilutes the content signal.
+            "token_emb": cat_emb,   # (B, S_total, E')
         }

tests/test_unified_multimodal.py

@@ -223,6 +223,33 @@ def test_unified_model_code_only():
     assert E == 64
 
 
+def test_unified_model_token_emb_is_content_only():
+    """``token_emb`` is the content embedding BEFORE time/type are added, in the
+    same temporally-sorted order as ``sequence``; i.e.
+    ``sequence == token_emb + time_embed(time) + type_embedding(type_ids)``."""
+    from pyhealth.models.embedding import UnifiedMultimodalEmbeddingModel
+
+    processors, inputs = _make_code_processors_and_inputs()
+    model = UnifiedMultimodalEmbeddingModel(processors=processors, embedding_dim=64)
+    model.eval()
+
+    out = model(inputs)
+
+    assert "token_emb" in out
+    assert out["token_emb"].shape == out["sequence"].shape
+
+    # Reconstruct sequence from the content-only token_emb + the position-derived
+    # time and type embeddings; they must match exactly.
+    recomposed = (
+        out["token_emb"]
+        + model.time_embed(out["time"])
+        + model.type_embedding(out["type_ids"])
+    )
+    assert torch.allclose(out["sequence"], recomposed, atol=1e-5)
+    # And token_emb is genuinely distinct from sequence (time/type were added).
+    assert not torch.allclose(out["token_emb"], out["sequence"])
+
+
 def test_unified_model_rejects_non_temporal():
     from pyhealth.models.embedding import UnifiedMultimodalEmbeddingModel
     from pyhealth.processors import SequenceProcessor