File size: 3,953 Bytes
8e28984 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 |
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.
import os
import logging
from pytorch_pretrained_bert.file_utils import http_get
logger = logging.getLogger(__name__)
# Note that the model size is roughly half of the GPT model because our model is saved by fp16
LSP_MODEL_URL = {
'multiref': {
'large_fs': 'https://acvrpublicycchen.blob.core.windows.net/dialogpt/multiref/large_fs.pkl',
'medium_fs': 'https://acvrpublicycchen.blob.core.windows.net/dialogpt/multiref/medium_fs.pkl',
'medium_ft': 'https://acvrpublicycchen.blob.core.windows.net/dialogpt/multiref/medium_ft.pkl',
'small_fs': 'https://acvrpublicycchen.blob.core.windows.net/dialogpt/multiref/small_fs.pkl',
'small_ft': 'https://acvrpublicycchen.blob.core.windows.net/dialogpt/multiref/small_ft.pkl'
},
'dstc': {
'medium_ft': 'https://acvrpublicycchen.blob.core.windows.net/dialogpt/DSTC/medium_ft.pkl'
}
}
# GPT model could be downloaded from huggingface repo
GPT2_PRETRAINED_MODEL_ARCHIVE_MAP = {
"small": "https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-pytorch_model.bin",
"medium": "https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-medium-pytorch_model.bin",
"large": "https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-large-pytorch_model.bin"
}
CONFIG_FILE = {
'small': 'https://acvrpublicycchen.blob.core.windows.net/dialogpt/117M/config.json',
'medium': 'https://acvrpublicycchen.blob.core.windows.net/dialogpt/345M/config.json',
'large': 'https://acvrpublicycchen.blob.core.windows.net/dialogpt/1542M/config.json'
}
VOCAB_FILE = {
'small': 'https://acvrpublicycchen.blob.core.windows.net/dialogpt/117M/vocab.json',
'medium': 'https://acvrpublicycchen.blob.core.windows.net/dialogpt/345M/vocab.json',
'large': 'https://acvrpublicycchen.blob.core.windows.net/dialogpt/1542M/vocab.json'
}
MERGE_FILE = {
'small': 'https://acvrpublicycchen.blob.core.windows.net/dialogpt/117M/merges.txt',
'medium': 'https://acvrpublicycchen.blob.core.windows.net/dialogpt/345M/merges.txt',
'large': 'https://acvrpublicycchen.blob.core.windows.net/dialogpt/1542M/merges.txt'
}
def download_file(url, folder):
if not os.path.exists(folder):
os.makedirs(folder, exist_ok=True)
file_name = os.path.basename(url)
if 'pytorch_model.bin' in file_name:
file_name = 'pytorch_model.bin'
if os.path.isfile(os.path.join(folder, file_name)):
logger.info(f'{os.path.join(folder, file_name)} exists, return!')
return
with open(os.path.join(folder, file_name), 'wb') as f:
http_get(url, f)
def download_model_folder(model_size, dataset=None, from_scratch=None, DATA_FOLDER=None):
assert DATA_FOLDER is not None, 'DATA_FOLDER cannot be None'
assert model_size in ['small', 'medium', 'large'], 'model size should be one of \'small\', \'medium\' or \'large\''
target_folder = os.path.join(DATA_FOLDER, model_size)
download_file(CONFIG_FILE[model_size], target_folder)
download_file(VOCAB_FILE[model_size], target_folder)
download_file(MERGE_FILE[model_size], target_folder)
download_file(GPT2_PRETRAINED_MODEL_ARCHIVE_MAP[model_size], target_folder)
if dataset is not None:
assert dataset in ['multiref', 'dstc'], \
'dataset has to be \'multiref\' or \'dstc\''
assert from_scratch in [True, False], 'from scratch has to be True or False'
if from_scratch:
model_train_type = model_size + '_fs'
else:
model_train_type = model_size + '_ft'
if model_train_type not in LSP_MODEL_URL[dataset]:
k = ','.join(list(LSP_MODEL_URL[dataset].keys()))
raise ValueError(f'\'{model_train_type}\' not exist for dataset \'{dataset}\', please choose from [{k}]')
download_file(LSP_MODEL_URL[dataset][model_train_type], target_folder)
return target_folder
|