Upload tgn checkpoint 20260430_135010

README.md

@@ -0,0 +1,44 @@
+# tgn
+
+JAX checkpoint for the `Tgn` temporal graph model with separate intention and outcome heads.
+
+## Checkpoint
+
+- Model class: `Tgn`
+- Framework: `jax`
+- Run timestamp: `20260430_135010`
+- Parameters: `model_params.joblib`
+- Hyperparameters: `config.json`
+- Metrics: `metrics.json`
+
+## Validation Metrics
+
+```json
+{
+  "val/loss": 1.1337319612503052,
+  "val/intention_top1_accuracy": 0.6611314415931702,
+  "val/intention_top2_accuracy": 0.8483576774597168,
+  "val/intention_top3_accuracy": 0.9177007675170898,
+  "val/intention_mrr": 0.7940667867660522,
+  "val/outcome_top1_accuracy": 0.5961679220199585,
+  "val/outcome_top2_accuracy": 0.7636861801147461,
+  "val/outcome_top3_accuracy": 0.8319343328475952,
+  "val/outcome_mrr": 0.7304578423500061
+}
+```
+
+## Loading
+
+```python
+import json
+from pathlib import Path
+
+import joblib
+
+checkpoint_path = Path("path/to/checkpoint")
+
+with open(checkpoint_path / "config.json") as f:
+    config = json.load(f)
+
+params = joblib.load(checkpoint_path / "model_params.joblib")
+```

config.json

@@ -0,0 +1,13 @@
+{
+  "prev_hidden_dims": 128,
+  "curr_hidden_dims": 512,
+  "time_embedding_dims": 8,
+  "lstm_embedding_dims": 128,
+  "dropout_rate": 0.2,
+  "num_layers": 4,
+  "num_heads": 4,
+  "batch_size": 8,
+  "learning_rate": "1e-4",
+  "weight_decay": "1e-4",
+  "grad_accum_steps": 16
+}

metrics.json

@@ -0,0 +1,15 @@
+{
+  "misc/epoch": 73,
+  "misc/patience": 9,
+  "train/loss": 1.0968098640441895,
+  "val/loss": 1.1337319612503052,
+  "val/intention_top1_accuracy": 0.6611314415931702,
+  "val/intention_top2_accuracy": 0.8483576774597168,
+  "val/intention_top3_accuracy": 0.9177007675170898,
+  "val/intention_mrr": 0.7940667867660522,
+  "val/outcome_top1_accuracy": 0.5961679220199585,
+  "val/outcome_top2_accuracy": 0.7636861801147461,
+  "val/outcome_top3_accuracy": 0.8319343328475952,
+  "val/outcome_mrr": 0.7304578423500061,
+  "misc/grad": 9.830899238586426
+}

model_info.json

@@ -0,0 +1,6 @@
+{
+  "model_class": "Tgn",
+  "framework": "jax",
+  "timestamp": "20260430_135010",
+  "model_name": "tgn"
+}

model_params.joblib

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:08a648e5a33a6784645b8bb214f4ef240e3174ad8de9dba652e07d38a4713618
+size 15387556