Refactor DeepHealth model and related components

- Removed BaselineEncoder and CrossAttention classes from models.py.
- Introduced OtherInfoTokenizer for handling additional token types.
- Updated DeepHealth class to integrate OtherInfoTokenizer and manage extra pooling logic.
- Added support for extra_pool_reduce parameter to control pooling behavior.
- Modified forward methods to return structured output using DeepHealthOutput dataclass.
- Updated training scripts to accommodate changes in model architecture and output handling.
- Enhanced error handling and validation for input shapes and types.

train_all_future.py

@@ -64,6 +64,8 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--n_hist_layer", type=int, default=12)
     parser.add_argument("--n_tab_layer", type=int, default=4)
     parser.add_argument("--n_bins", type=int, default=16)
+    parser.add_argument("--extra_pool_reduce", type=str, default="mean",
+                        choices=["mean", "sum"])
     parser.add_argument("--time_mode", type=str, default="relative",
                         choices=["relative", "absolute"])
     parser.add_argument("--dist_mode", type=str, default="exponential",
@@ -129,6 +131,7 @@ def build_model(args: argparse.Namespace, dataset: AllFutureHealthDataset) -> De
         n_categories=dataset.n_categories,
         cont_type_ids=dataset.cont_type_ids,
         n_bins=args.n_bins,
+        extra_pool_reduce=args.extra_pool_reduce,
         target_mode="all_future",
         time_mode=args.time_mode,
         dist_mode=args.dist_mode,