| from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer | |
| model_ckpt = "codeparrot" | |
| org = "transformersbook" | |
| def model_size(model): | |
| return sum(t.numel() for t in model.parameters()) | |
| tokenizer = AutoTokenizer.from_pretrained(org + "/" + model_ckpt) | |
| config_small = AutoConfig.from_pretrained("gpt2", vocab_size=len(tokenizer)) | |
| model_small = AutoModelForCausalLM.from_config(config_small) | |
| print(f'GPT-2 size: {model_size(model_small)/1000**2:.1f}M parameters') | |
| model_small.save_pretrained("models/" + model_ckpt + "-small", push_to_hub=True) |