|
|
|
|
|
|
|
|
|
|
|
import os |
|
|
import logging |
|
|
|
|
|
from pytorch_pretrained_bert.file_utils import http_get |
|
|
|
|
|
|
|
|
logger = logging.getLogger(__name__) |
|
|
|
|
|
|
|
|
|
|
|
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' |
|
|
} |
|
|
} |
|
|
|
|
|
|
|
|
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 |
|
|
|
|
|
|
