(cls)
| 32 | |
| 33 | @classmethod |
| 34 | def setUpClass(cls): |
| 35 | super().setUpClass() |
| 36 | gin_file = [ |
| 37 | 'base_htrans.gin', |
| 38 | 'size/medium_150M.gin', |
| 39 | 'options/positions_t5.gin', |
| 40 | 'options/lr_cosine_decay.gin', |
| 41 | 'options/seq_1024_nocache.gin', |
| 42 | 'geometry_150M_generate.gin', |
| 43 | ] |
| 44 | |
| 45 | gin_param = [ |
| 46 | 'DecoderOnlyLanguageModelGenerate.output_token_losses=True', |
| 47 | 'TransformerTaskConfig.batch_size=2', |
| 48 | 'TransformerTaskConfig.sequence_length=128', |
| 49 | 'Trainer.restore_state_variables=False', |
| 50 | ] |
| 51 | |
| 52 | gin_search_paths = [ |
| 53 | os.path.join(_MELIAD_PATH.value, 'transformer/configs'), |
| 54 | os.getcwd(), |
| 55 | ] |
| 56 | |
| 57 | vocab_path = os.path.join(_DATA_PATH.value, 'geometry.757.model') |
| 58 | |
| 59 | lm.parse_gin_configuration(gin_file, gin_param, gin_paths=gin_search_paths) |
| 60 | |
| 61 | cls.loaded_lm = lm.LanguageModelInference( |
| 62 | vocab_path, _DATA_PATH.value, mode='beam_search' |
| 63 | ) |
| 64 | |
| 65 | def test_lm_decode(self): |
| 66 | outputs = LmInferenceTest.loaded_lm.beam_decode( |
nothing calls this directly
no outgoing calls
no test coverage detected