Spaces:
Configuration error
Configuration error
File size: 4,911 Bytes
a01ef8c |
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 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 |
#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# Copyright (c) 2023 Intel Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# SPDX-License-Identifier: Apache-2.0
#
import os
from pydoc import locate
from downloader.types import ModelType
class ModelDownloader():
"""
A unified model downloader class.
Can download models from TF Hub, Torchvision, and Hugging Face.
"""
def __init__(self, model_name, hub, model_dir=None, **kwargs):
"""
Class constructor for a ModelDownloader.
Args:
model_name (str): Name of the model
hub (str, optional): The catalog to download the model from; options are 'tf_hub',
'torchvision', 'pytorch_hub', 'hugging_face', and 'keras'
model_dir (str): Local destination directory of the model, if None the model hub's default cache
directory will be used
kwargs (optional): Some model hubs accept additional keyword arguments when downloading
"""
if model_dir is not None and not os.path.isdir(model_dir):
os.makedirs(model_dir)
self._model_name = model_name
self._model_dir = model_dir
self._type = ModelType.from_str(hub)
self._args = kwargs
def download(self):
"""
Download the model
Returns:
A torch.nn.Module, keras.engine.functional.Functional, or tensorflow_hub.keras_layer.KerasLayer object
"""
if self._type == ModelType.TF_HUB:
from tensorflow_hub import KerasLayer
if self._model_dir is not None:
os.environ['TFHUB_CACHE_DIR'] = self._model_dir
return KerasLayer(self._model_name, **self._args)
elif self._type == ModelType.TORCHVISION:
if self._model_dir is not None:
os.environ['TORCH_HOME'] = self._model_dir
pretrained_model_class = locate('torchvision.models.{}'.format(self._model_name))
return pretrained_model_class(**self._args)
elif self._type == ModelType.PYTORCH_HUB:
from tlt.utils.file_utils import read_json_file
from tlt import TLT_BASE_DIR
import torch
if self._model_dir is not None:
os.environ['TORCH_HOME'] = self._model_dir
config_file = os.path.join(TLT_BASE_DIR, "models/configs/pytorch_hub_image_classification_models.json")
pytorch_hub_model_map = read_json_file(config_file)
self._repo = pytorch_hub_model_map[self._model_name]["repo"]
# Some models have pretrained=True by default, which error out if passed in load()
if pytorch_hub_model_map[self._model_name]["pretrained_default"] == "True":
return torch.hub.load(self._repo, self._model_name)
else:
return torch.hub.load(self._repo, self._model_name, pretrained=True)
elif self._type == ModelType.HUGGING_FACE:
if self._model_dir is not None:
os.environ['TRANSFORMERS_CACHE'] = self._model_dir
# AutoModelForSequenceClassification is currently the only supported model type
from transformers import AutoModelForSequenceClassification
return AutoModelForSequenceClassification.from_pretrained(self._model_name, **self._args)
elif self._type == ModelType.KERAS_APPLICATIONS:
if self._model_dir is not None:
os.environ['KERAS_HOME'] = self._model_dir
try:
pretrained_model_class = locate('keras.applications.{}'.format(self._model_name))
except TypeError:
pretrained_model_class = locate('keras.applications.{}.{}'.format(self._model_name.lower(),
self._model_name))
return pretrained_model_class(**self._args)
elif self._type == ModelType.TF_BERT_HUGGINGFACE:
if self._model_dir is not None:
os.environ['TRANSFORMERS_CACHE'] = self._model_dir
from transformers import BertConfig, TFBertModel
config = BertConfig.from_pretrained(self._model_name, output_hidden_states=True)
return TFBertModel.from_pretrained(self._model_name, config=config, from_pt=True, **self._args)
|