File size: 1,093 Bytes
2ea9ced |
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 |
from typing import * # pylint: disable=wildcard-import,unused-wildcard-import
import itertools
from .abc.base import LMScorer
from .gpt2 import GPT2LMScorer
class AutoLMScorer:
MODEL_CLASSES = [GPT2LMScorer]
def __init__(self):
raise EnvironmentError(
"AutoLMscorer is designed to be instantiated "
"using the `AutoLMscorer.from_pretrained(model_name)`"
"method"
)
@classmethod
def from_pretrained(cls, model_name: str, **kwargs: Any) -> LMScorer:
for model_class in cls.MODEL_CLASSES:
if model_name not in model_class.supported_model_names():
continue
return model_class(model_name, **kwargs)
raise ValueError(
"Unrecognized model name."
"Can be one of: %s" % ", ".join(cls.supported_model_names()),
)
@classmethod
def supported_model_names(cls) -> Iterable[str]:
classes = cls.MODEL_CLASSES
models = map(lambda c: c.supported_model_names(), classes)
return itertools.chain.from_iterable(models)
|