# Copyright 2021 Google LLC
#
# 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.

import transformers

MODELS = {
    "gpt2": [
        "gpt2",
        "gpt2-medium",
        "gpt2-large",
        "gpt2-xl",
    ],
    "openai-gpt": [
        "openai-gpt",
    ],
}

MODEL_NAMES = sum(MODELS.values(), [])
MODEL_TYPES = dict([(y, k) for k, v in MODELS.items() for y in v])

MODEL_CLASSES = {
    "gpt2": {
        "lm": transformers.TFGPT2LMHeadModel,
        "tokenizer": transformers.GPT2Tokenizer,
    },
    "openai-gpt": {
        "lm": transformers.TFOpenAIGPTLMHeadModel,
        "tokenizer": transformers.OpenAIGPTTokenizer,
    },
}

# ----------------------------------------------------------------------------
if __name__ == "__main__":
    for model_name in ["openai-gpt", "gpt2"]:  # , "gpt2-medium", "gpt2-large"]:
        model_type = MODEL_TYPES[model_name]
        model_class = MODEL_CLASSES[model_type]["lm"]
        tokenizer_class = MODEL_CLASSES[model_type]["tokenizer"]

        tokenizer = tokenizer_class.from_pretrained(model_name)
        model = model_class.from_pretrained(
            model_name, pad_token_id=tokenizer.eos_token_id
        )