from .gpt2 import GPT2LM
from .gpt3 import GPT3LM

MODEL_REGISTRY = {
    "gpt2": GPT2LM,
    "gpt3": GPT3LM,
}


def get_model(model_name):
    return MODEL_REGISTRY[model_name]