JiaruiLi/DeepHealth
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add nano and tiny model presets

Browse Source
This commit is contained in:
Jiarui Li
2026-07-23 15:01:33 +08:00
parent 06f29c0f0a
commit 6f7b5be405
3 changed files with 38 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files

33
test_event_trajectory_backbone.py
View File

@@ -58,7 +58,8 @@ def model_inputs() -> dict[str, torch.Tensor]:
class EventTrajectoryBackboneTest(unittest.TestCase):
def test_model_size_presets(self) -> None:
expected = {
"nano": (256, 8, 32, 32),
"nano": (120, 6, 20, 24),
"tiny": (256, 8, 32, 32),
"small": (512, 8, 64, 32),
"medium": (768, 12, 64, 48),
"huge": (1024, 16, 64, 64),
@@ -186,9 +187,9 @@ class EventTrajectoryBackboneTest(unittest.TestCase):
inputs = model_inputs()
next_model = build_test_model(target_mode="next_token")
next_hidden = next_model(**inputs)
self.assertEqual(tuple(next_hidden.shape), (2, 4, 256))
self.assertEqual(tuple(next_hidden.shape), (2, 4, 120))
next_output = next_model(**inputs, return_output=True)
self.assertEqual(tuple(next_output.hidden.shape), (2, 4, 256))
self.assertEqual(tuple(next_output.hidden.shape), (2, 4, 120))
self.assertEqual(tuple(next_output.padding_mask.shape), (2, 4))
future_model = build_test_model(target_mode="all_future")
@@ -196,7 +197,7 @@ class EventTrajectoryBackboneTest(unittest.TestCase):
**inputs,
t_query=torch.tensor([5.0, 3.0]),
)
self.assertEqual(tuple(future_hidden.shape), (2, 256))
self.assertEqual(tuple(future_hidden.shape), (2, 120))
def test_model_contains_one_shared_core_and_no_block_stack(self) -> None:
model = build_test_model(n_reasoning_rounds=12)
@@ -261,10 +262,10 @@ class EventTrajectoryBackboneTest(unittest.TestCase):
{
"model_architecture": EVENT_TRAJECTORY_ARCHITECTURE,
"model_size": "nano",
"d_model": 256,
"n_trajectory": 8,
"trajectory_dim": 32,
"traj_hidden": 32,
"d_model": 120,
"n_trajectory": 6,
"trajectory_dim": 20,
"traj_hidden": 24,
"n_reasoning_rounds": 3,
}
)
@@ -272,15 +273,19 @@ class EventTrajectoryBackboneTest(unittest.TestCase):
validate_event_trajectory_config(
{"model_architecture": "traj_mixer_v2"}
)
with self.assertRaisesRegex(ValueError, "only accepts models trained"):
validate_event_trajectory_config(
{"model_architecture": "event_trajectory_shared_v1"}
)
with self.assertRaisesRegex(ValueError, "trajectory_dim"):
validate_event_trajectory_config(
{
"model_architecture": EVENT_TRAJECTORY_ARCHITECTURE,
"model_size": "nano",
"d_model": 256,
"n_trajectory": 8,
"trajectory_dim": 16,
"traj_hidden": 32,
"d_model": 120,
"n_trajectory": 6,
"trajectory_dim": 10,
"traj_hidden": 24,
"n_reasoning_rounds": 3,
}
)
@@ -289,8 +294,8 @@ class EventTrajectoryBackboneTest(unittest.TestCase):
state_dict = model.state_dict()
validate_event_trajectory_state_dict(
state_dict,
expected_d_model=256,
expected_n_trajectory=8,
expected_d_model=120,
expected_n_trajectory=6,
expected_n_reasoning_rounds=3,
)
with self.assertRaisesRegex(