Spaces:
Paused
Paused
| # coding=utf-8 | |
| # Copyright 2019-present, the HuggingFace Inc. team. | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| """ | |
| Preprocessing script before training the distilled model. | |
| Specific to RoBERTa -> DistilRoBERTa and GPT2 -> DistilGPT2. | |
| """ | |
| import argparse | |
| import torch | |
| from transformers import GPT2LMHeadModel, RobertaForMaskedLM | |
| if __name__ == "__main__": | |
| parser = argparse.ArgumentParser( | |
| description=( | |
| "Extraction some layers of the full RobertaForMaskedLM or GPT2LMHeadModel for Transfer Learned" | |
| " Distillation" | |
| ) | |
| ) | |
| parser.add_argument("--model_type", default="roberta", choices=["roberta", "gpt2"]) | |
| parser.add_argument("--model_name", default="roberta-large", type=str) | |
| parser.add_argument("--dump_checkpoint", default="serialization_dir/tf_roberta_048131723.pth", type=str) | |
| parser.add_argument("--vocab_transform", action="store_true") | |
| args = parser.parse_args() | |
| if args.model_type == "roberta": | |
| model = RobertaForMaskedLM.from_pretrained(args.model_name) | |
| prefix = "roberta" | |
| elif args.model_type == "gpt2": | |
| model = GPT2LMHeadModel.from_pretrained(args.model_name) | |
| prefix = "transformer" | |
| state_dict = model.state_dict() | |
| compressed_sd = {} | |
| # Embeddings # | |
| if args.model_type == "gpt2": | |
| for param_name in ["wte.weight", "wpe.weight"]: | |
| compressed_sd[f"{prefix}.{param_name}"] = state_dict[f"{prefix}.{param_name}"] | |
| else: | |
| for w in ["word_embeddings", "position_embeddings", "token_type_embeddings"]: | |
| param_name = f"{prefix}.embeddings.{w}.weight" | |
| compressed_sd[param_name] = state_dict[param_name] | |
| for w in ["weight", "bias"]: | |
| param_name = f"{prefix}.embeddings.LayerNorm.{w}" | |
| compressed_sd[param_name] = state_dict[param_name] | |
| # Transformer Blocks # | |
| std_idx = 0 | |
| for teacher_idx in [0, 2, 4, 7, 9, 11]: | |
| if args.model_type == "gpt2": | |
| for layer in ["ln_1", "attn.c_attn", "attn.c_proj", "ln_2", "mlp.c_fc", "mlp.c_proj"]: | |
| for w in ["weight", "bias"]: | |
| compressed_sd[f"{prefix}.h.{std_idx}.{layer}.{w}"] = state_dict[ | |
| f"{prefix}.h.{teacher_idx}.{layer}.{w}" | |
| ] | |
| compressed_sd[f"{prefix}.h.{std_idx}.attn.bias"] = state_dict[f"{prefix}.h.{teacher_idx}.attn.bias"] | |
| else: | |
| for layer in [ | |
| "attention.self.query", | |
| "attention.self.key", | |
| "attention.self.value", | |
| "attention.output.dense", | |
| "attention.output.LayerNorm", | |
| "intermediate.dense", | |
| "output.dense", | |
| "output.LayerNorm", | |
| ]: | |
| for w in ["weight", "bias"]: | |
| compressed_sd[f"{prefix}.encoder.layer.{std_idx}.{layer}.{w}"] = state_dict[ | |
| f"{prefix}.encoder.layer.{teacher_idx}.{layer}.{w}" | |
| ] | |
| std_idx += 1 | |
| # Language Modeling Head ###s | |
| if args.model_type == "roberta": | |
| for layer in ["lm_head.decoder.weight", "lm_head.bias"]: | |
| compressed_sd[f"{layer}"] = state_dict[f"{layer}"] | |
| if args.vocab_transform: | |
| for w in ["weight", "bias"]: | |
| compressed_sd[f"lm_head.dense.{w}"] = state_dict[f"lm_head.dense.{w}"] | |
| compressed_sd[f"lm_head.layer_norm.{w}"] = state_dict[f"lm_head.layer_norm.{w}"] | |
| elif args.model_type == "gpt2": | |
| for w in ["weight", "bias"]: | |
| compressed_sd[f"{prefix}.ln_f.{w}"] = state_dict[f"{prefix}.ln_f.{w}"] | |
| compressed_sd["lm_head.weight"] = state_dict["lm_head.weight"] | |
| print(f"N layers selected for distillation: {std_idx}") | |
| print(f"Number of params transferred for distillation: {len(compressed_sd.keys())}") | |
| print(f"Save transferred checkpoint to {args.dump_checkpoint}.") | |
| torch.save(compressed_sd, args.dump_checkpoint) | |