def remove_projection_matrices(model_file: str): """ Remove all projection matrices used for distillation from the model and re-save it. """ print(f'Creating a backup copy of the original model at {model_file}._orig.') PathManager.copy(model_file, f'{model_file}._orig') print(f"Loading {model_file}.") with PathManager.open(model_file, 'rb') as f: states = torch.load(f, map_location=lambda cpu, _: cpu, pickle_module=pickle) print('Deleting projection matrices.') orig_num_keys = len(states['model']) states['model'] = { key: val for key, val in states['model'].items() if key.split('.')[0] not in ['encoder_proj_layer', 'embedding_proj_layers', 'hidden_proj_layers'] } new_num_keys = len(states['model']) print(f'{orig_num_keys-new_num_keys:d} model keys removed.') print(f"Saving to {model_file}.") atomic_save(states, model_file)
def save(self, dir_name: str, file_name: str): """ Save appropriate files. :param dir_name: directory to save. :param file_name: file to save. """ out_json_path = os.path.join(dir_name, file_name + "-vocab.json") out_merge_path = os.path.join(dir_name, file_name + "-merges.txt") # Possibly bad assumption: if the destination file already exists, # we don't need to copy it over again. if not PathManager.exists(out_json_path): logging.info(f"Copying {self.json_path} to {out_json_path}") PathManager.copy(self.json_path, out_json_path) if not PathManager.exists(out_merge_path): logging.info(f"Copying {self.merge_path} to {out_merge_path}") PathManager.copy(self.merge_path, out_merge_path)