| from transformers import GPTNeoConfig | |
| import os | |
| # export DATASET="${HOME}/data/nedd_wiki_news/nedd_wiki_news.py" # Name of the dataset in the Huggingface Hub | |
| # export DATASET_CONFIG="nedd_nl" # Config of the dataset in the Huggingface Hub | |
| # export DATASET_SPLIT="train" # Split to use for training tokenizer and model | |
| # export VOCAB_SIZE="50257" | |
| # export MODEL_PATH="${HOME}/data/${HF_PROJECT}" # Path to the model, e.g. here inside the mount | |
| # export CONFIG_TYPE="gpt2-medium" # Config that our model will use | |
| config_type = os.environ.get("CONFIG_TYPE") | |
| dataset_name = os.environ.get("DATASET") | |
| dataset_config = os.environ.get("DATASET_CONFIG") | |
| dataset_split = os.environ.get("DATASET_SPLIT") | |
| vocab_size = int(os.environ.get("VOCAB_SIZE")) | |
| model_path = os.environ.get("MODEL_PATH") | |
| config = GPTNeoConfig.from_pretrained(config_type, embed_dropout=0.0, attention_dropout=0.0, vocab_size=vocab_size) | |
| config.save_pretrained(model_path) | |