Divide params into with-weight-decay and without-weight-decay groups. Layernorms and biases will have no weight decay but the rest will.

Reduce a tensor of losses across all GPUs.

Adds an adapter layer to `self` at the specified location

Builds labels from input embeddings. Masks out the labels with -100 in positions up to the seq length of the embeddings, so loss is only com

Loads clip's image encoder module, discarding the lm component. If the variant is a resnet model, we also remove the attention pooling.

Configures the different parameter groups in the model for training. If a separate learning rate for the image prefix is provided, we separat

Embeds a list of tensors In the correct format to input into the LM (b, s, d). For each tensor, if it's 2d assume it's text and use w

Loads a model checkpoint from disk / downloads from url if not present

Generates captions for a batch of embeddings. :param model: The model to use for generation. :param embeddings: The embeddings to genera

Generates captions for a batch of embeddings.

Loads GPTJ language model from HF

Loads image encoder module

Gets tokenizer for LM

Loads a model from disk and returns the global step to resume from if loading was successful, otherwise returns 0

Loads nfresnet50 model, removing the pooling layer and replacing it with an adaptive pooling layer.

Expects a list of strings and instances of ImageInput Converts them into a list of tensors and then optionally runs self.embed over i

Iterator over data from a jsonl file

Top K sampling

Nucleus sampling

Builds a dataset directory in our standard format. ds_iterator should return data of the form image_path, {"captions": [...], "metadata": {..

Counts the number of trainable parameters in a model

Yields image_path, image_caption from cc jsonl files