(self)
| 31 | class HfDataProcessingTest(unittest.TestCase): |
| 32 | |
| 33 | def setUp(self): |
| 34 | super().setUp() |
| 35 | config = pyconfig.initialize( |
| 36 | [sys.argv[0], os.path.join(MAXTEXT_PKG_DIR, "configs", "base.yml")], |
| 37 | per_device_batch_size=1, |
| 38 | run_name="test", |
| 39 | mesh_axes=["data"], |
| 40 | logical_axis_rules=[["batch", "data"]], |
| 41 | data_sharding=["data"], |
| 42 | base_output_directory="gs://max-experiments/", |
| 43 | dataset_type="hf", |
| 44 | hf_path="parquet", |
| 45 | hf_data_dir="", |
| 46 | hf_train_files="gs://maxtext-dataset/hf/c4/c4-train-00000-of-01637.parquet", |
| 47 | tokenizer_path="google-t5/t5-large", |
| 48 | enable_checkpointing=False, |
| 49 | ) |
| 50 | self.config = config |
| 51 | self.mesh_shape_1d = (len(jax.devices()),) |
| 52 | self.mesh = Mesh(mesh_utils.create_device_mesh(self.mesh_shape_1d), self.config.mesh_axes) |
| 53 | self.process_indices = input_pipeline_interface.get_process_loading_real_data( |
| 54 | self.config.data_sharding, |
| 55 | self.config.global_batch_size_to_load, |
| 56 | self.config.global_batch_size_to_train_on, |
| 57 | self.config.max_target_length, |
| 58 | self.mesh, |
| 59 | ) |
| 60 | |
| 61 | self.train_iter = _hf_data_processing.make_hf_train_iterator(self.config, self.mesh, self.process_indices) |
| 62 | |
| 63 | def test_train_ds(self): |
| 64 | expected_shape = [jax.device_count(), self.config.max_target_length] |
nothing calls this directly
no test coverage detected