(self)
| 35 | class SFTHooksTest(unittest.TestCase): |
| 36 | |
| 37 | def setUp(self): |
| 38 | super().setUp() |
| 39 | self.config = pyconfig.initialize( |
| 40 | [os.path.join(MAXTEXT_PKG_DIR, "sft.hooks"), os.path.join(MAXTEXT_PKG_DIR, "configs", "sft.yml")], |
| 41 | per_device_batch_size=1, |
| 42 | run_name="test", |
| 43 | base_output_directory="test", |
| 44 | skip_jax_distributed_system=True, |
| 45 | ) |
| 46 | self.mesh = Mesh(create_device_mesh(self.config), self.config.mesh_axes) |
| 47 | learning_rate_schedule = maxtext_utils.create_learning_rate_schedule(self.config) |
| 48 | |
| 49 | self.training_hooks = hooks.SFTTrainingHooks(self.config, self.mesh, learning_rate_schedule, goodput_recorder=None) |
| 50 | self.training_hooks.metric_logger = MagicMock() |
| 51 | |
| 52 | expected_shape = [jax.device_count(), self.config.max_target_length] |
| 53 | self.expected_batch = { |
| 54 | "inputs": np.zeros(expected_shape, dtype=int), |
| 55 | "targets_segmentation": np.ones(expected_shape, dtype=int), |
| 56 | } |
| 57 | self.mock_data_iterator = MagicMock() |
| 58 | self.mock_data_iterator.__next__.return_value = self.expected_batch |
| 59 | |
| 60 | self.mock_train_ctx = MagicMock() |
| 61 | |
| 62 | @patch("MaxText.sft.hooks.create_data_iterator") |
| 63 | def test_data_hooks_load_next_train_batch(self, mock_create_data_iterator): |
nothing calls this directly
no test coverage detected