Create and save a model for `model_arch`. Also copy the set of processors to each model (under the same model type) output folder.
(model_arch, tiny_config, output_dir, keep_model=False)
| 1020 | |
| 1021 | |
| 1022 | def build_model(model_arch, tiny_config, output_dir, keep_model=False): |
| 1023 | """Create and save a model for `model_arch`. |
| 1024 | |
| 1025 | Also copy the set of processors to each model (under the same model type) output folder. |
| 1026 | """ |
| 1027 | |
| 1028 | checkpoint_dir = get_checkpoint_dir(output_dir, model_arch) |
| 1029 | |
| 1030 | if not os.path.isdir(checkpoint_dir): |
| 1031 | os.makedirs(checkpoint_dir, exist_ok=True) |
| 1032 | |
| 1033 | processor_output_dir = os.path.join(output_dir, "processors") |
| 1034 | # copy the (same set of) processors (for a model type) to the model arch. specific folder |
| 1035 | if os.path.isdir(processor_output_dir): |
| 1036 | shutil.copytree(processor_output_dir, checkpoint_dir, dirs_exist_ok=True) |
| 1037 | |
| 1038 | tiny_config = copy.deepcopy(tiny_config) |
| 1039 | |
| 1040 | if any(model_arch.__name__.endswith(x) for x in ["ForCausalLM", "LMHeadModel"]): |
| 1041 | tiny_config.is_encoder_decoder = False |
| 1042 | tiny_config.is_decoder = True |
| 1043 | |
| 1044 | model = model_arch(config=tiny_config) |
| 1045 | |
| 1046 | with tempfile.TemporaryDirectory(dir=checkpoint_dir) as tmpdir: |
| 1047 | if keep_model: |
| 1048 | checkpoint_dir_tmp = checkpoint_dir |
| 1049 | else: |
| 1050 | checkpoint_dir_tmp = tmpdir |
| 1051 | |
| 1052 | model.save_pretrained(checkpoint_dir_tmp) |
| 1053 | |
| 1054 | # can't call from_pretrained from saved one |
| 1055 | if not tiny_config.__class__.__name__.endswith(("TimmBackboneConfig",)): |
| 1056 | model.from_pretrained(checkpoint_dir_tmp) |
| 1057 | |
| 1058 | return model |
| 1059 | |
| 1060 | |
| 1061 | def fill_result_with_error(result, error, trace, models_to_create): |
no test coverage detected