Unify FFN and TrajMixer model architectures

models.py

@@ -6,11 +6,12 @@ import torch.nn.functional as F
 
 from backbones import (
     AgeSinusoidalEncoding,
-    GPTBlock,
     GaussianRBFTimeBasis,
     TimeRoPE,
     TokenAutoDiscretization,
+    build_backbone_block,
 )
+from model_architectures import resolve_model_architecture
 from targets import PAD_IDX
 
 
@@ -147,8 +148,7 @@ class DeepHealth(nn.Module):
         vocab_size: int,
         n_embd: int,
         n_head: int,
-        n_hist_layer: int,
-        n_tab_layer: int,
+        n_layer: int,
         n_types: int,
         n_cont_types: int,
         n_categories: int,
@@ -160,6 +160,7 @@ class DeepHealth(nn.Module):
         dist_mode: str = "exponential",  # "exponential", "weibull" or "mixed"
         extra_pool_reduce: str = "mean",
         dropout: float = 0.0,
+        model_architecture: str | None = None,
     ):
         super().__init__()
         if target_mode not in ["next_token", "all_future"]:
@@ -173,6 +174,9 @@ class DeepHealth(nn.Module):
                 "dist_mode must be either 'exponential', 'weibull' or 'mixed'")
         if extra_pool_reduce not in {"mean", "sum"}:
             raise ValueError("extra_pool_reduce must be either 'mean' or 'sum'")
+        if n_layer < 1:
+            raise ValueError(f"n_layer must be >= 1, got {n_layer}")
+        model_architecture = resolve_model_architecture(model_architecture)
         self.token_embedding = nn.Embedding(vocab_size, n_embd, padding_idx=0)
         self.gender_embedding = nn.Embedding(
             2, n_embd)  # Assuming binary gender
@@ -189,6 +193,8 @@ class DeepHealth(nn.Module):
         self.time_mode = time_mode
         self.dist_mode = dist_mode
         self.extra_pool_reduce = extra_pool_reduce
+        self.model_architecture = model_architecture
+        self.n_layer = n_layer
         self.n_embd = n_embd
         self.vocab_size = vocab_size
         nn.init.normal_(self.token_embedding.weight, mean=0.0, std=0.02)
@@ -208,26 +214,28 @@ class DeepHealth(nn.Module):
         if time_mode == "absolute":
             self.age_encoding = AgeSinusoidalEncoding(n_embd)
             self.blocks = nn.ModuleList([
-                GPTBlock(
+                build_backbone_block(
+                    model_architecture,
                     n_embd=n_embd,
                     n_head=n_head,
                     use_time_rope=False,
                     use_rbf_bias=False,
                     mlp_dropout=dropout,
-                ) for _ in range(n_hist_layer)
+                ) for _ in range(n_layer)
             ])
             self.rope = None
             self.rbf = None
         elif time_mode == "relative":
             self.age_encoding = None
             self.blocks = nn.ModuleList([
-                GPTBlock(
+                build_backbone_block(
+                    model_architecture,
                     n_embd=n_embd,
                     n_head=n_head,
                     use_time_rope=True,
                     use_rbf_bias=True,
                     mlp_dropout=dropout,
-                ) for _ in range(n_hist_layer)
+                ) for _ in range(n_layer)
             ])
             self.rope = TimeRoPE(n_embd // n_head)
             self.rbf = GaussianRBFTimeBasis(n_bases=16, max_time_diff=40.0)