import os, shutil, time, random, torch
from transformers import (
    VisionEncoderDecoderModel,
    VisionEncoderDecoderConfig,
    AutoConfig
)

torch.manual_seed(42)
random.seed(42)

BATCH_SIZE = 64
NUM_WORKERS = 12
NUM_TOKENS = 1024
MAX_SEQ_LEN = 256
PAD_ID = 0
START_ID = 1
END_ID = 2


# set up image-to-text model
def get_model():
    config = VisionEncoderDecoderConfig.from_encoder_decoder_configs(
        encoder_config=AutoConfig.for_model("vit"),  # vit encoder
        decoder_config=AutoConfig.for_model("gpt2")  # gpt2 decoder
    )
    config.decoder.vocab_size = NUM_TOKENS
    config.decoder.use_cache = False
    config.decoder_start_token_id = START_ID
    config.pad_token_id = PAD_ID
    config.eos_token_id = END_ID
    config.max_length = MAX_SEQ_LEN

    model = VisionEncoderDecoderModel(config=config)

    # remove unused pooler
    model.encoder.pooler = None

    # uncomment to specify the loss function
    # from transformers.loss.loss_utils import ForCausalLMLoss
    # model.loss_function = ForCausalLMLoss
    return model