PythonCopilot / model.py
TRM-coding's picture
Upload 7 files
f9cc63a verified
raw
history blame contribute delete
564 Bytes
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)