init pyconfig
(self, **kwargs)
| 46 | self.rng = jax.random.PRNGKey(0) |
| 47 | |
| 48 | def init_pyconfig(self, **kwargs): |
| 49 | """init pyconfig""" |
| 50 | init_kwargs = { |
| 51 | "per_device_batch_size": 1.0, |
| 52 | "run_name": "test", |
| 53 | "enable_checkpointing": False, |
| 54 | "base_num_decoder_layers": 2, |
| 55 | "attention": "dot_product", |
| 56 | "max_target_length": 16, |
| 57 | "base_emb_dim": 256, |
| 58 | "base_num_query_heads": 2, |
| 59 | "base_num_kv_heads": 2, |
| 60 | "max_prefill_predict_length": 4, |
| 61 | "return_log_prob": True, |
| 62 | } | kwargs |
| 63 | config = pyconfig.initialize( |
| 64 | [sys.argv[0], os.path.join(MAXTEXT_PKG_DIR, "configs", "base.yml")], |
| 65 | **init_kwargs, |
| 66 | ) |
| 67 | return config |
| 68 | |
| 69 | def get_data(self): |
| 70 | s = (self.cfg.global_batch_size_to_train_on, self.cfg.max_target_length) |
no test coverage detected