#!/usr/bin/env python | |
from transformers import AutoTokenizer, TFAutoModelForCausalLM, AutoModelForCausalLM, GPT2Config | |
import tensorflow as tf | |
task_specific_params = { | |
"text-generation": { | |
"do_sample": False, | |
"max_length": 50 | |
} | |
} | |
config = GPT2Config.from_pretrained("Sentdex/GPyT", _name_or_path='prophetikai/code-gpt', use_cache=True, task_specific_params=task_specific_params) | |
tokenizer = AutoTokenizer.from_pretrained("Sentdex/GPyT") | |
tf_model = TFAutoModelForCausalLM.from_pretrained("Sentdex/GPyT", config=config) | |
pytorch_model = AutoModelForCausalLM.from_pretrained("Sentdex/GPyT", config=config) | |
config.save_pretrained('./') | |
tokenizer.save_pretrained(save_directory='./') | |
tf_model.save_pretrained(save_directory='./', saved_model=True, version='sentdex') | |
pytorch_model.save_pretrained(save_directory='./') |