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

Revert "Add nano and tiny model presets"

Browse Source 
This reverts commit 6f7b5be405.
This commit is contained in:
Jiarui Li
2026-07-23 16:03:29 +08:00
parent 6f7b5be405
commit 22faee7c51
3 changed files with 29 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files

33
test_event_trajectory_backbone.py
View File

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