Upload FIMMJP

config.json

@@ -7,57 +7,52 @@
     "AutoModel": "mjp.FIMMJP"
   },
   "initial_distribution_decoder": {
-    "dropout": 0.1,
+    "dropout": 0,
+    "hidden_act": {
+      "name": "torch.nn.SELU"
+    },
     "hidden_layers": [
-      64,
-      64
+      128,
+      128
     ],
-    "name": "fim.models.blocks.MLP"
+    "in_features": 2049,
+    "initialization_scheme": "lecun_normal",
+    "name": "fim.models.blocks.base.MLP"
   },
   "intensity_matrix_decoder": {
-    "dropout": 0.1,
+    "dropout": 0,
+    "hidden_act": {
+      "name": "torch.nn.SELU"
+    },
     "hidden_layers": [
-      64,
-      64
+      128,
+      128
     ],
-    "name": "fim.models.blocks.MLP"
+    "in_features": 2049,
+    "initialization_scheme": "lecun_normal",
+    "name": "fim.models.blocks.base.MLP"
   },
   "model_type": "fimmjp",
   "n_states": 6,
   "path_attention": {
-    "batch_first": true,
-    "embed_dim": 64,
-    "name": "torch.nn.MultiheadAttention",
-    "num_heads": 8
+    "embed_dim": 512,
+    "kv_dim": 128,
+    "n_heads": 1,
+    "n_queries": 16,
+    "name": "fim.models.blocks.MultiHeadLearnableQueryAttention"
   },
   "pos_encodings": {
-    "name": "fim.models.blocks.SineTimeEncoding",
-    "out_features": 64
+    "name": "fim.models.blocks.positional_encodings.DeltaTimeEncoding"
   },
   "torch_dtype": "float32",
   "transformers_version": "4.46.0",
   "ts_encoder": {
-    "embed_dim": 64,
-    "name": "fim.models.blocks.base.TransformerEncoder",
-    "num_layers": 4,
-    "transformer_block": {
-      "activation": {
-        "name": "torch.nn.ReLU"
-      },
-      "attention_head": {
-        "batch_first": true,
-        "embed_dim": 64,
-        "name": "torch.nn.MultiheadAttention",
-        "num_heads": 8
-      },
-      "dropout": 0.1,
-      "ff_dim": 256,
-      "in_features": 64,
-      "name": "fim.models.blocks.TransformerBlock",
-      "normalization": {
-        "name": "torch.nn.LayerNorm",
-        "normalized_shape": 64
-      }
+    "name": "fim.models.blocks.base.RNNEncoder",
+    "rnn": {
+      "batch_first": true,
+      "bidirectional": true,
+      "hidden_size": 256,
+      "name": "torch.nn.LSTM"
     }
   },
   "use_adjacency_matrix": false

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:97b87e5a321a9c57addcee2d03d9cd7f996f67bdc91f67023eaf495b2d3a885a
-size 1025616
+oid sha256:932ce52b1ca140a94b92bfe2ec7dea9fa2625c857e146970aa4689a4f361892f
+size 4979384