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)
