Synapse/notebooks/image_captioning/standalone_image_captioning.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "source": [
        "Copyright (c) Microsoft Corporation. All rights reserved.\n",
        "\n",
        "Licensed under the MIT License.\n",
        "\n",
        "# This cell configures the spark session - Do not change"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 18,
      "metadata": {
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "outputs": [
        {
          "data": {
            "application/vnd.livy.statement-meta+json": {
              "execution_finish_time": "2022-04-06T08:36:46.4581395Z",
              "execution_start_time": "2022-04-06T08:36:46.4577438Z",
              "livy_statement_state": "available",
              "queued_time": "2022-04-06T08:35:55.2113151Z",
              "session_id": 195,
              "session_start_time": "2022-04-06T08:35:55.5813115Z",
              "spark_pool": null,
              "state": "finished",
              "statement_id": -1
            },
            "text/plain": [
              "StatementMeta(, 195, -1, Finished, Available)"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "%%configure -f\n",
        "{\n",
        "\"conf\": {\n",
        "     \"spark.rpc.message.maxSize\": 1024,\n",
        "     \"spark.kryoserializer.buffer.max\": \"256m\"\n",
        "   }\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "source": [
        "# These are the parameters that need to be changed to your values"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# The blob account name = [accountname]\n",
        "account_name = ''\n",
        "# The name of the primary ADLS share\n",
        "file_system_name=\"share\"\n",
        "# The blob account key [iufquq34r423r2==] - used to generate a SAS key\n",
        "account_key = ''\n",
        "# The directory folders where your files reside  \n",
        "directory_name='animal_faces/walkthrough'\n",
        "# The name of the output file\n",
        "output_filename = 'penalty_cleansed_animal_faces.csv'\n",
        "\n",
        "# AML Experiment tracking\n",
        "\n",
        "# Azure SubscriptionId\n",
        "subscription_id=\"\"\n",
        "# AzureML Workspace Resource Group\n",
        "resource_group=\"\"\n",
        "# AzureML Workspace Name\n",
        "workspace_name=\"\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "outputs": [
        {
          "data": {
            "application/vnd.livy.statement-meta+json": {
              "execution_finish_time": "2022-04-06T08:36:52.5645635Z",
              "execution_start_time": "2022-04-06T08:36:52.4220922Z",
              "livy_statement_state": "available",
              "queued_time": "2022-04-06T08:35:55.2676826Z",
              "session_id": 195,
              "session_start_time": null,
              "spark_pool": "small",
              "state": "finished",
              "statement_id": 1
            },
            "text/plain": [
              "StatementMeta(small, 195, 1, Finished, Available)"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "# The input directory where the images reside, can be nested\n",
        "input_directory = f'abfss://{file_system_name}@{account_name}.dfs.core.windows.net/{directory_name}/'\n",
        "\n",
        "# The output directory where the output file will be written to\n",
        "output_directory = f'abfss://{file_system_name}@{account_name}.dfs.core.windows.net/{directory_name}/walkthrough/'\n",
        "\n",
        "# The blob account url - https://[accountname].blob.core.windows.net\n",
        "account_url = f\"https://{account_name}.blob.core.windows.net\"\n",
        "\n",
        "# This is the model config file which can be found in the repo at /Data-Discovery-Toolset/synapse/image_captioning/med_config.json. This needs to be uploaded to your primary ADLS linked to Synapse\n",
        "model_config_file = f\"abfss://{file_system_name}@{account_name}.dfs.core.windows.net/notebooks/image_captioning/med_config.json\" \n",
        "\n",
        "# If this is set to True then the Coalesce notebook will need to be run to merge the partition files into a single file\n",
        "LOW_MEMORY_MODE = True\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Track the Experiment in Azure ML"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from azureml.core import Workspace, Experiment, Run\n",
        "import mlflow\n",
        "\n",
        "ws = Workspace(subscription_id=subscription_id, resource_group=resource_group, workspace_name=workspace_name)    \n",
        "mlflow.set_tracking_uri(ws.get_mlflow_tracking_uri())\n",
        "experiment_name = f\"({mssparkutils.runtime.context['notebookname']}_{str(mssparkutils.env.getJobId())})\"\n",
        "mlflow.set_experiment(experiment_name)\n",
        "mlflow.log_param(\"input_directory\", input_directory)\n",
        "mlflow.log_param(\"output_directory\", output_directory)\n",
        "mlflow.log_param(\"account_url\", account_url)\n",
        "mlflow.log_param(\"account_name\", account_name)\n",
        "mlflow.log_param(\"file_system_name\", file_system_name)\n",
        "mlflow.log_param(\"model_config_file\", model_config_file)\n",
        "mlflow.log_param(\"LOW_MEMORY_MODE\", LOW_MEMORY_MODE)\n",
        "params = {\n",
        "    \"sparkpool\": mssparkutils.runtime.context['sparkpool'],\n",
        "    \"workspace\": mssparkutils.runtime.context['workspace'],\n",
        "    \"notebookname\": mssparkutils.runtime.context['notebookname'],\n",
        "    \"isForPipeline\": mssparkutils.runtime.context['isForPipeline'],\n",
        "    \"pipelinejobid\": mssparkutils.runtime.context['pipelinejobid']\n",
        "}\n",
        "\n",
        "mlflow.log_params(params)\n",
        "mlflow.pyspark.ml.autolog()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 20,
      "metadata": {
        "jupyter": {
          "outputs_hidden": true,
          "source_hidden": true
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "outputs": [
        {
          "data": {
            "application/vnd.livy.statement-meta+json": {
              "execution_finish_time": "2022-04-06T08:36:54.8280993Z",
              "execution_start_time": "2022-04-06T08:36:52.8039534Z",
              "livy_statement_state": "available",
              "queued_time": "2022-04-06T08:35:55.3748318Z",
              "session_id": 195,
              "session_start_time": null,
              "spark_pool": "small",
              "state": "finished",
              "statement_id": 2
            },
            "text/plain": [
              "StatementMeta(small, 195, 2, Finished, Available)"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "from PIL import Image\n",
        "import requests\n",
        "import torch\n",
        "from torchvision import transforms\n",
        "from torchvision.transforms.functional import InterpolationMode\n",
        "from torch import nn\n",
        "from transformers import BertTokenizer\n",
        "from functools import partial\n",
        "from timm.models.vision_transformer import _cfg, PatchEmbed\n",
        "from timm.models.registry import register_model\n",
        "from timm.models.layers import trunc_normal_, DropPath\n",
        "from timm.models.helpers import named_apply, adapt_input_conv\n",
        "from transformers import BertTokenizer\n",
        "from transformers.models.bert.configuration_bert import BertConfig\n",
        "from transformers.modeling_utils import (\n",
        "    PreTrainedModel,\n",
        "    apply_chunking_to_forward,\n",
        "    find_pruneable_heads_and_indices,\n",
        "    prune_linear_layer,\n",
        ")\n",
        "from transformers.activations import ACT2FN\n",
        "from urllib.parse import urlparse\n",
        "from timm.models.hub import download_cached_file\n",
        "from fairscale.nn.checkpoint.checkpoint_activations import checkpoint_wrapper\n",
        "import ntpath\n",
        "\n",
        "'''\n",
        " * Copyright (c) 2022, salesforce.com, inc.\n",
        " * All rights reserved.\n",
        " * SPDX-License-Identifier: BSD-3-Clause\n",
        " * For full license text, see LICENSE.txt file in the repo root or https://opensource.org/licenses/BSD-3-Clause\n",
        " * By Junnan Li\n",
        " * Based on timm code base\n",
        " * https://github.com/rwightman/pytorch-image-models/tree/master/timm\n",
        "'''\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "from functools import partial\n",
        "\n",
        "from timm.models.vision_transformer import _cfg, PatchEmbed\n",
        "from timm.models.registry import register_model\n",
        "from timm.models.layers import trunc_normal_, DropPath\n",
        "from timm.models.helpers import named_apply, adapt_input_conv\n",
        "\n",
        "from fairscale.nn.checkpoint.checkpoint_activations import checkpoint_wrapper\n",
        "\n",
        "class Mlp(nn.Module):\n",
        "    \"\"\" MLP as used in Vision Transformer, MLP-Mixer and related networks\n",
        "    \"\"\"\n",
        "    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):\n",
        "        super().__init__()\n",
        "        out_features = out_features or in_features\n",
        "        hidden_features = hidden_features or in_features\n",
        "        self.fc1 = nn.Linear(in_features, hidden_features)\n",
        "        self.act = act_layer()\n",
        "        self.fc2 = nn.Linear(hidden_features, out_features)\n",
        "        self.drop = nn.Dropout(drop)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.fc1(x)\n",
        "        x = self.act(x)\n",
        "        x = self.drop(x)\n",
        "        x = self.fc2(x)\n",
        "        x = self.drop(x)\n",
        "        return x\n",
        "\n",
        "\n",
        "class Attention(nn.Module):\n",
        "    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):\n",
        "        super().__init__()\n",
        "        self.num_heads = num_heads\n",
        "        head_dim = dim // num_heads\n",
        "        # NOTE scale factor was wrong in my original version, can set manually to be compat with prev weights\n",
        "        self.scale = qk_scale or head_dim ** -0.5\n",
        "        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)\n",
        "        self.attn_drop = nn.Dropout(attn_drop)\n",
        "        self.proj = nn.Linear(dim, dim)\n",
        "        self.proj_drop = nn.Dropout(proj_drop)\n",
        "        self.attn_gradients = None\n",
        "        self.attention_map = None\n",
        "        \n",
        "    def save_attn_gradients(self, attn_gradients):\n",
        "        self.attn_gradients = attn_gradients\n",
        "        \n",
        "    def get_attn_gradients(self):\n",
        "        return self.attn_gradients\n",
        "    \n",
        "    def save_attention_map(self, attention_map):\n",
        "        self.attention_map = attention_map\n",
        "        \n",
        "    def get_attention_map(self):\n",
        "        return self.attention_map\n",
        "    \n",
        "    def forward(self, x, register_hook=False):\n",
        "        B, N, C = x.shape\n",
        "        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)\n",
        "        q, k, v = qkv[0], qkv[1], qkv[2]   # make torchscript happy (cannot use tensor as tuple)\n",
        "\n",
        "        attn = (q @ k.transpose(-2, -1)) * self.scale\n",
        "        attn = attn.softmax(dim=-1)\n",
        "        attn = self.attn_drop(attn)\n",
        "                \n",
        "        if register_hook:\n",
        "            self.save_attention_map(attn)\n",
        "            attn.register_hook(self.save_attn_gradients)        \n",
        "\n",
        "        x = (attn @ v).transpose(1, 2).reshape(B, N, C)\n",
        "        x = self.proj(x)\n",
        "        x = self.proj_drop(x)\n",
        "        return x\n",
        "\n",
        "\n",
        "class Block(nn.Module):\n",
        "\n",
        "    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,\n",
        "                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, use_grad_checkpointing=False):\n",
        "        super().__init__()\n",
        "        self.norm1 = norm_layer(dim)\n",
        "        self.attn = Attention(\n",
        "            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)\n",
        "        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here\n",
        "        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()\n",
        "        self.norm2 = norm_layer(dim)\n",
        "        mlp_hidden_dim = int(dim * mlp_ratio)\n",
        "        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)\n",
        "\n",
        "        if use_grad_checkpointing:\n",
        "            self.attn = checkpoint_wrapper(self.attn)\n",
        "            self.mlp = checkpoint_wrapper(self.mlp)\n",
        "\n",
        "    def forward(self, x, register_hook=False):\n",
        "        x = x + self.drop_path(self.attn(self.norm1(x), register_hook=register_hook))\n",
        "        x = x + self.drop_path(self.mlp(self.norm2(x)))\n",
        "        return x\n",
        "\n",
        "    \n",
        "class VisionTransformer(nn.Module):\n",
        "    \"\"\" Vision Transformer\n",
        "    A PyTorch impl of : `An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale`  -\n",
        "        https://arxiv.org/abs/2010.11929\n",
        "    \"\"\"\n",
        "    def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, embed_dim=768, depth=12,\n",
        "                 num_heads=12, mlp_ratio=4., qkv_bias=True, qk_scale=None, representation_size=None,\n",
        "                 drop_rate=0., attn_drop_rate=0., drop_path_rate=0., norm_layer=None, \n",
        "                 use_grad_checkpointing=False, ckpt_layer=0):\n",
        "        \"\"\"\n",
        "        Args:\n",
        "            img_size (int, tuple): input image size\n",
        "            patch_size (int, tuple): patch size\n",
        "            in_chans (int): number of input channels\n",
        "            num_classes (int): number of classes for classification head\n",
        "            embed_dim (int): embedding dimension\n",
        "            depth (int): depth of transformer\n",
        "            num_heads (int): number of attention heads\n",
        "            mlp_ratio (int): ratio of mlp hidden dim to embedding dim\n",
        "            qkv_bias (bool): enable bias for qkv if True\n",
        "            qk_scale (float): override default qk scale of head_dim ** -0.5 if set\n",
        "            representation_size (Optional[int]): enable and set representation layer (pre-logits) to this value if set\n",
        "            drop_rate (float): dropout rate\n",
        "            attn_drop_rate (float): attention dropout rate\n",
        "            drop_path_rate (float): stochastic depth rate\n",
        "            norm_layer: (nn.Module): normalization layer\n",
        "        \"\"\"\n",
        "        super().__init__()\n",
        "        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models\n",
        "        norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)\n",
        "\n",
        "        self.patch_embed = PatchEmbed(\n",
        "            img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)\n",
        "\n",
        "        num_patches = self.patch_embed.num_patches\n",
        "\n",
        "        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))\n",
        "        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))\n",
        "        self.pos_drop = nn.Dropout(p=drop_rate)\n",
        "\n",
        "        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule\n",
        "        self.blocks = nn.ModuleList([\n",
        "            Block(\n",
        "                dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,\n",
        "                drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer,\n",
        "                use_grad_checkpointing=(use_grad_checkpointing and i>=depth-ckpt_layer)\n",
        "            )\n",
        "            for i in range(depth)])\n",
        "        self.norm = norm_layer(embed_dim)\n",
        "\n",
        "        trunc_normal_(self.pos_embed, std=.02)\n",
        "        trunc_normal_(self.cls_token, std=.02)\n",
        "        self.apply(self._init_weights)\n",
        "\n",
        "    def _init_weights(self, m):\n",
        "        if isinstance(m, nn.Linear):\n",
        "            trunc_normal_(m.weight, std=.02)\n",
        "            if isinstance(m, nn.Linear) and m.bias is not None:\n",
        "                nn.init.constant_(m.bias, 0)\n",
        "        elif isinstance(m, nn.LayerNorm):\n",
        "            nn.init.constant_(m.bias, 0)\n",
        "            nn.init.constant_(m.weight, 1.0)\n",
        "\n",
        "    @torch.jit.ignore\n",
        "    def no_weight_decay(self):\n",
        "        return {'pos_embed', 'cls_token'}\n",
        "\n",
        "    def forward(self, x, register_blk=-1):\n",
        "        B = x.shape[0]\n",
        "        x = self.patch_embed(x)\n",
        "\n",
        "        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks\n",
        "        x = torch.cat((cls_tokens, x), dim=1)\n",
        "  \n",
        "        x = x + self.pos_embed[:,:x.size(1),:]\n",
        "        x = self.pos_drop(x)\n",
        "\n",
        "        for i,blk in enumerate(self.blocks):\n",
        "            x = blk(x, register_blk==i)\n",
        "        x = self.norm(x)\n",
        "        \n",
        "        return x\n",
        "\n",
        "    @torch.jit.ignore()\n",
        "    def load_pretrained(self, checkpoint_path, prefix=''):\n",
        "        _load_weights(self, checkpoint_path, prefix)\n",
        "        \n",
        "\n",
        "@torch.no_grad()\n",
        "def _load_weights(model: VisionTransformer, checkpoint_path: str, prefix: str = ''):\n",
        "    \"\"\" Load weights from .npz checkpoints for official Google Brain Flax implementation\n",
        "    \"\"\"\n",
        "    import numpy as np\n",
        "\n",
        "    def _n2p(w, t=True):\n",
        "        if w.ndim == 4 and w.shape[0] == w.shape[1] == w.shape[2] == 1:\n",
        "            w = w.flatten()\n",
        "        if t:\n",
        "            if w.ndim == 4:\n",
        "                w = w.transpose([3, 2, 0, 1])\n",
        "            elif w.ndim == 3:\n",
        "                w = w.transpose([2, 0, 1])\n",
        "            elif w.ndim == 2:\n",
        "                w = w.transpose([1, 0])\n",
        "        return torch.from_numpy(w)\n",
        "\n",
        "    w = np.load(checkpoint_path)\n",
        "    if not prefix and 'opt/target/embedding/kernel' in w:\n",
        "        prefix = 'opt/target/'\n",
        "\n",
        "    if hasattr(model.patch_embed, 'backbone'):\n",
        "        # hybrid\n",
        "        backbone = model.patch_embed.backbone\n",
        "        stem_only = not hasattr(backbone, 'stem')\n",
        "        stem = backbone if stem_only else backbone.stem\n",
        "        stem.conv.weight.copy_(adapt_input_conv(stem.conv.weight.shape[1], _n2p(w[f'{prefix}conv_root/kernel'])))\n",
        "        stem.norm.weight.copy_(_n2p(w[f'{prefix}gn_root/scale']))\n",
        "        stem.norm.bias.copy_(_n2p(w[f'{prefix}gn_root/bias']))\n",
        "        if not stem_only:\n",
        "            for i, stage in enumerate(backbone.stages):\n",
        "                for j, block in enumerate(stage.blocks):\n",
        "                    bp = f'{prefix}block{i + 1}/unit{j + 1}/'\n",
        "                    for r in range(3):\n",
        "                        getattr(block, f'conv{r + 1}').weight.copy_(_n2p(w[f'{bp}conv{r + 1}/kernel']))\n",
        "                        getattr(block, f'norm{r + 1}').weight.copy_(_n2p(w[f'{bp}gn{r + 1}/scale']))\n",
        "                        getattr(block, f'norm{r + 1}').bias.copy_(_n2p(w[f'{bp}gn{r + 1}/bias']))\n",
        "                    if block.downsample is not None:\n",
        "                        block.downsample.conv.weight.copy_(_n2p(w[f'{bp}conv_proj/kernel']))\n",
        "                        block.downsample.norm.weight.copy_(_n2p(w[f'{bp}gn_proj/scale']))\n",
        "                        block.downsample.norm.bias.copy_(_n2p(w[f'{bp}gn_proj/bias']))\n",
        "        embed_conv_w = _n2p(w[f'{prefix}embedding/kernel'])\n",
        "    else:\n",
        "        embed_conv_w = adapt_input_conv(\n",
        "            model.patch_embed.proj.weight.shape[1], _n2p(w[f'{prefix}embedding/kernel']))\n",
        "    model.patch_embed.proj.weight.copy_(embed_conv_w)\n",
        "    model.patch_embed.proj.bias.copy_(_n2p(w[f'{prefix}embedding/bias']))\n",
        "    model.cls_token.copy_(_n2p(w[f'{prefix}cls'], t=False))\n",
        "    pos_embed_w = _n2p(w[f'{prefix}Transformer/posembed_input/pos_embedding'], t=False)\n",
        "    if pos_embed_w.shape != model.pos_embed.shape:\n",
        "        pos_embed_w = resize_pos_embed(  # resize pos embedding when different size from pretrained weights\n",
        "            pos_embed_w, model.pos_embed, getattr(model, 'num_tokens', 1), model.patch_embed.grid_size)\n",
        "    model.pos_embed.copy_(pos_embed_w)\n",
        "    model.norm.weight.copy_(_n2p(w[f'{prefix}Transformer/encoder_norm/scale']))\n",
        "    model.norm.bias.copy_(_n2p(w[f'{prefix}Transformer/encoder_norm/bias']))\n",
        "#     if isinstance(model.head, nn.Linear) and model.head.bias.shape[0] == w[f'{prefix}head/bias'].shape[-1]:\n",
        "#         model.head.weight.copy_(_n2p(w[f'{prefix}head/kernel']))\n",
        "#         model.head.bias.copy_(_n2p(w[f'{prefix}head/bias']))\n",
        "#     if isinstance(getattr(model.pre_logits, 'fc', None), nn.Linear) and f'{prefix}pre_logits/bias' in w:\n",
        "#         model.pre_logits.fc.weight.copy_(_n2p(w[f'{prefix}pre_logits/kernel']))\n",
        "#         model.pre_logits.fc.bias.copy_(_n2p(w[f'{prefix}pre_logits/bias']))\n",
        "    for i, block in enumerate(model.blocks.children()):\n",
        "        block_prefix = f'{prefix}Transformer/encoderblock_{i}/'\n",
        "        mha_prefix = block_prefix + 'MultiHeadDotProductAttention_1/'\n",
        "        block.norm1.weight.copy_(_n2p(w[f'{block_prefix}LayerNorm_0/scale']))\n",
        "        block.norm1.bias.copy_(_n2p(w[f'{block_prefix}LayerNorm_0/bias']))\n",
        "        block.attn.qkv.weight.copy_(torch.cat([\n",
        "            _n2p(w[f'{mha_prefix}{n}/kernel'], t=False).flatten(1).T for n in ('query', 'key', 'value')]))\n",
        "        block.attn.qkv.bias.copy_(torch.cat([\n",
        "            _n2p(w[f'{mha_prefix}{n}/bias'], t=False).reshape(-1) for n in ('query', 'key', 'value')]))\n",
        "        block.attn.proj.weight.copy_(_n2p(w[f'{mha_prefix}out/kernel']).flatten(1))\n",
        "        block.attn.proj.bias.copy_(_n2p(w[f'{mha_prefix}out/bias']))\n",
        "        for r in range(2):\n",
        "            getattr(block.mlp, f'fc{r + 1}').weight.copy_(_n2p(w[f'{block_prefix}MlpBlock_3/Dense_{r}/kernel']))\n",
        "            getattr(block.mlp, f'fc{r + 1}').bias.copy_(_n2p(w[f'{block_prefix}MlpBlock_3/Dense_{r}/bias']))\n",
        "        block.norm2.weight.copy_(_n2p(w[f'{block_prefix}LayerNorm_2/scale']))\n",
        "        block.norm2.bias.copy_(_n2p(w[f'{block_prefix}LayerNorm_2/bias']))\n",
        "\n",
        "            \n",
        "def interpolate_pos_embed(pos_embed_checkpoint, visual_encoder):        \n",
        "    # interpolate position embedding\n",
        "    embedding_size = pos_embed_checkpoint.shape[-1]\n",
        "    num_patches = visual_encoder.patch_embed.num_patches\n",
        "    num_extra_tokens = visual_encoder.pos_embed.shape[-2] - num_patches\n",
        "    # height (== width) for the checkpoint position embedding\n",
        "    orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5)\n",
        "    # height (== width) for the new position embedding\n",
        "    new_size = int(num_patches ** 0.5)\n",
        "\n",
        "    if orig_size!=new_size:\n",
        "        # class_token and dist_token are kept unchanged\n",
        "        extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]\n",
        "        # only the position tokens are interpolated\n",
        "        pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]\n",
        "        pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, embedding_size).permute(0, 3, 1, 2)\n",
        "        pos_tokens = torch.nn.functional.interpolate(\n",
        "            pos_tokens, size=(new_size, new_size), mode='bicubic', align_corners=False)\n",
        "        pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)\n",
        "        new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)\n",
        "        print('reshape position embedding from %d to %d'%(orig_size ** 2,new_size ** 2))\n",
        "        \n",
        "        return new_pos_embed    \n",
        "    else:\n",
        "        return pos_embed_checkpoint\n",
        "\n",
        "\n",
        "'''\n",
        " * Copyright (c) 2022, salesforce.com, inc.\n",
        " * All rights reserved.\n",
        " * SPDX-License-Identifier: BSD-3-Clause\n",
        " * For full license text, see LICENSE.txt file in the repo root or https://opensource.org/licenses/BSD-3-Clause\n",
        " * By Junnan Li\n",
        " * Based on huggingface code base\n",
        " * https://github.com/huggingface/transformers/blob/v4.15.0/src/transformers/models/bert\n",
        "'''\n",
        "\n",
        "import math\n",
        "import os\n",
        "import warnings\n",
        "from dataclasses import dataclass\n",
        "from typing import Optional, Tuple\n",
        "\n",
        "import torch\n",
        "from torch import Tensor, device, dtype, nn\n",
        "import torch.utils.checkpoint\n",
        "from torch import nn\n",
        "from torch.nn import CrossEntropyLoss\n",
        "import torch.nn.functional as F\n",
        "\n",
        "from transformers.activations import ACT2FN\n",
        "from transformers.file_utils import (\n",
        "    ModelOutput,\n",
        ")\n",
        "from transformers.modeling_outputs import (\n",
        "    BaseModelOutputWithPastAndCrossAttentions,\n",
        "    BaseModelOutputWithPoolingAndCrossAttentions,\n",
        "    CausalLMOutputWithCrossAttentions,\n",
        "    MaskedLMOutput,\n",
        "    MultipleChoiceModelOutput,\n",
        "    NextSentencePredictorOutput,\n",
        "    QuestionAnsweringModelOutput,\n",
        "    SequenceClassifierOutput,\n",
        "    TokenClassifierOutput,\n",
        ")\n",
        "from transformers.modeling_utils import (\n",
        "    PreTrainedModel,\n",
        "    apply_chunking_to_forward,\n",
        "    find_pruneable_heads_and_indices,\n",
        "    prune_linear_layer,\n",
        ")\n",
        "from transformers.utils import logging\n",
        "from transformers.models.bert.configuration_bert import BertConfig\n",
        "\n",
        "\n",
        "logger = logging.get_logger(__name__)\n",
        "\n",
        "\n",
        "class BertEmbeddings(nn.Module):\n",
        "    \"\"\"Construct the embeddings from word and position embeddings.\"\"\"\n",
        "\n",
        "    def __init__(self, config):\n",
        "        super().__init__()\n",
        "        self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)\n",
        "        self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)\n",
        "\n",
        "        # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load\n",
        "        # any TensorFlow checkpoint file\n",
        "        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)\n",
        "        self.dropout = nn.Dropout(config.hidden_dropout_prob)\n",
        "\n",
        "        # position_ids (1, len position emb) is contiguous in memory and exported when serialized\n",
        "        self.register_buffer(\"position_ids\", torch.arange(config.max_position_embeddings).expand((1, -1)))\n",
        "        self.position_embedding_type = getattr(config, \"position_embedding_type\", \"absolute\")\n",
        "        \n",
        "        self.config = config\n",
        "\n",
        "    def forward(\n",
        "        self, input_ids=None, position_ids=None, inputs_embeds=None, past_key_values_length=0\n",
        "    ):\n",
        "        if input_ids is not None:\n",
        "            input_shape = input_ids.size()\n",
        "        else:\n",
        "            input_shape = inputs_embeds.size()[:-1]\n",
        "\n",
        "        seq_length = input_shape[1]\n",
        "\n",
        "        if position_ids is None:\n",
        "            position_ids = self.position_ids[:, past_key_values_length : seq_length + past_key_values_length]\n",
        "\n",
        "        if inputs_embeds is None:\n",
        "            inputs_embeds = self.word_embeddings(input_ids)\n",
        "\n",
        "        embeddings = inputs_embeds\n",
        "\n",
        "        if self.position_embedding_type == \"absolute\":\n",
        "            position_embeddings = self.position_embeddings(position_ids)\n",
        "            embeddings += position_embeddings\n",
        "        embeddings = self.LayerNorm(embeddings)\n",
        "        embeddings = self.dropout(embeddings)\n",
        "        return embeddings\n",
        "\n",
        "\n",
        "class BertSelfAttention(nn.Module):\n",
        "    def __init__(self, config, is_cross_attention):\n",
        "        super().__init__()\n",
        "        self.config = config\n",
        "        if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, \"embedding_size\"):\n",
        "            raise ValueError(\n",
        "                \"The hidden size (%d) is not a multiple of the number of attention \"\n",
        "                \"heads (%d)\" % (config.hidden_size, config.num_attention_heads)\n",
        "            )\n",
        "        \n",
        "        self.num_attention_heads = config.num_attention_heads\n",
        "        self.attention_head_size = int(config.hidden_size / config.num_attention_heads)\n",
        "        self.all_head_size = self.num_attention_heads * self.attention_head_size\n",
        "\n",
        "        self.query = nn.Linear(config.hidden_size, self.all_head_size)\n",
        "        if is_cross_attention:\n",
        "            self.key = nn.Linear(config.encoder_width, self.all_head_size)\n",
        "            self.value = nn.Linear(config.encoder_width, self.all_head_size)\n",
        "        else:\n",
        "            self.key = nn.Linear(config.hidden_size, self.all_head_size)\n",
        "            self.value = nn.Linear(config.hidden_size, self.all_head_size)\n",
        "\n",
        "        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)\n",
        "        self.position_embedding_type = getattr(config, \"position_embedding_type\", \"absolute\")\n",
        "        if self.position_embedding_type == \"relative_key\" or self.position_embedding_type == \"relative_key_query\":\n",
        "            self.max_position_embeddings = config.max_position_embeddings\n",
        "            self.distance_embedding = nn.Embedding(2 * config.max_position_embeddings - 1, self.attention_head_size)\n",
        "        self.save_attention = False   \n",
        "            \n",
        "    def save_attn_gradients(self, attn_gradients):\n",
        "        self.attn_gradients = attn_gradients\n",
        "        \n",
        "    def get_attn_gradients(self):\n",
        "        return self.attn_gradients\n",
        "    \n",
        "    def save_attention_map(self, attention_map):\n",
        "        self.attention_map = attention_map\n",
        "        \n",
        "    def get_attention_map(self):\n",
        "        return self.attention_map\n",
        "    \n",
        "    def transpose_for_scores(self, x):\n",
        "        new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)\n",
        "        x = x.view(*new_x_shape)\n",
        "        return x.permute(0, 2, 1, 3)\n",
        "\n",
        "    def forward(\n",
        "        self,\n",
        "        hidden_states,\n",
        "        attention_mask=None,\n",
        "        head_mask=None,\n",
        "        encoder_hidden_states=None,\n",
        "        encoder_attention_mask=None,\n",
        "        past_key_value=None,\n",
        "        output_attentions=False,\n",
        "    ):\n",
        "        mixed_query_layer = self.query(hidden_states)\n",
        "\n",
        "        # If this is instantiated as a cross-attention module, the keys\n",
        "        # and values come from an encoder; the attention mask needs to be\n",
        "        # such that the encoder's padding tokens are not attended to.\n",
        "        is_cross_attention = encoder_hidden_states is not None\n",
        "\n",
        "        if is_cross_attention:\n",
        "            key_layer = self.transpose_for_scores(self.key(encoder_hidden_states))\n",
        "            value_layer = self.transpose_for_scores(self.value(encoder_hidden_states))\n",
        "            attention_mask = encoder_attention_mask\n",
        "        elif past_key_value is not None:\n",
        "            key_layer = self.transpose_for_scores(self.key(hidden_states))\n",
        "            value_layer = self.transpose_for_scores(self.value(hidden_states))\n",
        "            key_layer = torch.cat([past_key_value[0], key_layer], dim=2)\n",
        "            value_layer = torch.cat([past_key_value[1], value_layer], dim=2)\n",
        "        else:\n",
        "            key_layer = self.transpose_for_scores(self.key(hidden_states))\n",
        "            value_layer = self.transpose_for_scores(self.value(hidden_states))\n",
        "\n",
        "        query_layer = self.transpose_for_scores(mixed_query_layer)\n",
        "\n",
        "        past_key_value = (key_layer, value_layer)\n",
        "\n",
        "        # Take the dot product between \"query\" and \"key\" to get the raw attention scores.\n",
        "        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))\n",
        "\n",
        "        if self.position_embedding_type == \"relative_key\" or self.position_embedding_type == \"relative_key_query\":\n",
        "            seq_length = hidden_states.size()[1]\n",
        "            position_ids_l = torch.arange(seq_length, dtype=torch.long, device=hidden_states.device).view(-1, 1)\n",
        "            position_ids_r = torch.arange(seq_length, dtype=torch.long, device=hidden_states.device).view(1, -1)\n",
        "            distance = position_ids_l - position_ids_r\n",
        "            positional_embedding = self.distance_embedding(distance + self.max_position_embeddings - 1)\n",
        "            positional_embedding = positional_embedding.to(dtype=query_layer.dtype)  # fp16 compatibility\n",
        "\n",
        "            if self.position_embedding_type == \"relative_key\":\n",
        "                relative_position_scores = torch.einsum(\"bhld,lrd->bhlr\", query_layer, positional_embedding)\n",
        "                attention_scores = attention_scores + relative_position_scores\n",
        "            elif self.position_embedding_type == \"relative_key_query\":\n",
        "                relative_position_scores_query = torch.einsum(\"bhld,lrd->bhlr\", query_layer, positional_embedding)\n",
        "                relative_position_scores_key = torch.einsum(\"bhrd,lrd->bhlr\", key_layer, positional_embedding)\n",
        "                attention_scores = attention_scores + relative_position_scores_query + relative_position_scores_key\n",
        "\n",
        "        attention_scores = attention_scores / math.sqrt(self.attention_head_size)\n",
        "        if attention_mask is not None:\n",
        "            # Apply the attention mask is (precomputed for all layers in BertModel forward() function)\n",
        "            attention_scores = attention_scores + attention_mask\n",
        "\n",
        "        # Normalize the attention scores to probabilities.\n",
        "        attention_probs = nn.Softmax(dim=-1)(attention_scores)\n",
        "        \n",
        "        if is_cross_attention and self.save_attention:\n",
        "            self.save_attention_map(attention_probs)\n",
        "            attention_probs.register_hook(self.save_attn_gradients)         \n",
        "\n",
        "        # This is actually dropping out entire tokens to attend to, which might\n",
        "        # seem a bit unusual, but is taken from the original Transformer paper.\n",
        "        attention_probs_dropped = self.dropout(attention_probs)\n",
        "\n",
        "        # Mask heads if we want to\n",
        "        if head_mask is not None:\n",
        "            attention_probs_dropped = attention_probs_dropped * head_mask\n",
        "\n",
        "        context_layer = torch.matmul(attention_probs_dropped, value_layer)\n",
        "\n",
        "        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()\n",
        "        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)\n",
        "        context_layer = context_layer.view(*new_context_layer_shape)\n",
        "\n",
        "        outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)\n",
        "\n",
        "        outputs = outputs + (past_key_value,)\n",
        "        return outputs\n",
        "\n",
        "\n",
        "class BertSelfOutput(nn.Module):\n",
        "    def __init__(self, config):\n",
        "        super().__init__()\n",
        "        self.dense = nn.Linear(config.hidden_size, config.hidden_size)\n",
        "        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)\n",
        "        self.dropout = nn.Dropout(config.hidden_dropout_prob)\n",
        "\n",
        "    def forward(self, hidden_states, input_tensor):\n",
        "        hidden_states = self.dense(hidden_states)\n",
        "        hidden_states = self.dropout(hidden_states)\n",
        "        hidden_states = self.LayerNorm(hidden_states + input_tensor)\n",
        "        return hidden_states\n",
        "\n",
        "\n",
        "class BertAttention(nn.Module):\n",
        "    def __init__(self, config, is_cross_attention=False):\n",
        "        super().__init__()\n",
        "        self.self = BertSelfAttention(config, is_cross_attention)\n",
        "        self.output = BertSelfOutput(config)\n",
        "        self.pruned_heads = set()\n",
        "\n",
        "    def prune_heads(self, heads):\n",
        "        if len(heads) == 0:\n",
        "            return\n",
        "        heads, index = find_pruneable_heads_and_indices(\n",
        "            heads, self.self.num_attention_heads, self.self.attention_head_size, self.pruned_heads\n",
        "        )\n",
        "\n",
        "        # Prune linear layers\n",
        "        self.self.query = prune_linear_layer(self.self.query, index)\n",
        "        self.self.key = prune_linear_layer(self.self.key, index)\n",
        "        self.self.value = prune_linear_layer(self.self.value, index)\n",
        "        self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)\n",
        "\n",
        "        # Update hyper params and store pruned heads\n",
        "        self.self.num_attention_heads = self.self.num_attention_heads - len(heads)\n",
        "        self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads\n",
        "        self.pruned_heads = self.pruned_heads.union(heads)\n",
        "\n",
        "    def forward(\n",
        "        self,\n",
        "        hidden_states,\n",
        "        attention_mask=None,\n",
        "        head_mask=None,\n",
        "        encoder_hidden_states=None,\n",
        "        encoder_attention_mask=None,\n",
        "        past_key_value=None,\n",
        "        output_attentions=False,\n",
        "    ):\n",
        "        self_outputs = self.self(\n",
        "            hidden_states,\n",
        "            attention_mask,\n",
        "            head_mask,\n",
        "            encoder_hidden_states,\n",
        "            encoder_attention_mask,\n",
        "            past_key_value,\n",
        "            output_attentions,\n",
        "        )\n",
        "        attention_output = self.output(self_outputs[0], hidden_states)\n",
        "        outputs = (attention_output,) + self_outputs[1:]  # add attentions if we output them\n",
        "        return outputs\n",
        "\n",
        "\n",
        "class BertIntermediate(nn.Module):\n",
        "    def __init__(self, config):\n",
        "        super().__init__()\n",
        "        self.dense = nn.Linear(config.hidden_size, config.intermediate_size)\n",
        "        if isinstance(config.hidden_act, str):\n",
        "            self.intermediate_act_fn = ACT2FN[config.hidden_act]\n",
        "        else:\n",
        "            self.intermediate_act_fn = config.hidden_act\n",
        "\n",
        "    def forward(self, hidden_states):\n",
        "        hidden_states = self.dense(hidden_states)\n",
        "        hidden_states = self.intermediate_act_fn(hidden_states)\n",
        "        return hidden_states\n",
        "\n",
        "\n",
        "class BertOutput(nn.Module):\n",
        "    def __init__(self, config):\n",
        "        super().__init__()\n",
        "        self.dense = nn.Linear(config.intermediate_size, config.hidden_size)\n",
        "        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)\n",
        "        self.dropout = nn.Dropout(config.hidden_dropout_prob)\n",
        "\n",
        "    def forward(self, hidden_states, input_tensor):\n",
        "        hidden_states = self.dense(hidden_states)\n",
        "        hidden_states = self.dropout(hidden_states)\n",
        "        hidden_states = self.LayerNorm(hidden_states + input_tensor)\n",
        "        return hidden_states\n",
        "\n",
        "\n",
        "class BertLayer(nn.Module):\n",
        "    def __init__(self, config, layer_num):\n",
        "        super().__init__()\n",
        "        self.config = config\n",
        "        self.chunk_size_feed_forward = config.chunk_size_feed_forward\n",
        "        self.seq_len_dim = 1\n",
        "        self.attention = BertAttention(config)      \n",
        "        self.layer_num = layer_num          \n",
        "        if self.config.add_cross_attention:\n",
        "            self.crossattention = BertAttention(config, is_cross_attention=self.config.add_cross_attention)\n",
        "        self.intermediate = BertIntermediate(config)\n",
        "        self.output = BertOutput(config)\n",
        "\n",
        "    def forward(\n",
        "        self,\n",
        "        hidden_states,\n",
        "        attention_mask=None,\n",
        "        head_mask=None,\n",
        "        encoder_hidden_states=None,\n",
        "        encoder_attention_mask=None,\n",
        "        past_key_value=None,\n",
        "        output_attentions=False,\n",
        "        mode=None,\n",
        "    ):\n",
        "        # decoder uni-directional self-attention cached key/values tuple is at positions 1,2\n",
        "        self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None\n",
        "        self_attention_outputs = self.attention(\n",
        "            hidden_states,\n",
        "            attention_mask,\n",
        "            head_mask,\n",
        "            output_attentions=output_attentions,\n",
        "            past_key_value=self_attn_past_key_value,\n",
        "        )\n",
        "        attention_output = self_attention_outputs[0]\n",
        "\n",
        "        outputs = self_attention_outputs[1:-1]\n",
        "        present_key_value = self_attention_outputs[-1]\n",
        "\n",
        "        if mode=='multimodal':\n",
        "            assert encoder_hidden_states is not None, \"encoder_hidden_states must be given for cross-attention layers\"\n",
        "\n",
        "            cross_attention_outputs = self.crossattention(\n",
        "                attention_output,\n",
        "                attention_mask,\n",
        "                head_mask,\n",
        "                encoder_hidden_states,\n",
        "                encoder_attention_mask,\n",
        "                output_attentions=output_attentions,\n",
        "            )\n",
        "            attention_output = cross_attention_outputs[0]\n",
        "            outputs = outputs + cross_attention_outputs[1:-1]  # add cross attentions if we output attention weights                               \n",
        "        layer_output = apply_chunking_to_forward(\n",
        "            self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output\n",
        "        )\n",
        "        outputs = (layer_output,) + outputs\n",
        "\n",
        "        outputs = outputs + (present_key_value,)\n",
        "\n",
        "        return outputs\n",
        "\n",
        "    def feed_forward_chunk(self, attention_output):\n",
        "        intermediate_output = self.intermediate(attention_output)\n",
        "        layer_output = self.output(intermediate_output, attention_output)\n",
        "        return layer_output\n",
        "\n",
        "\n",
        "class BertEncoder(nn.Module):\n",
        "    def __init__(self, config):\n",
        "        super().__init__()\n",
        "        self.config = config\n",
        "        self.layer = nn.ModuleList([BertLayer(config,i) for i in range(config.num_hidden_layers)])\n",
        "        self.gradient_checkpointing = False\n",
        "\n",
        "    def forward(\n",
        "        self,\n",
        "        hidden_states,\n",
        "        attention_mask=None,\n",
        "        head_mask=None,\n",
        "        encoder_hidden_states=None,\n",
        "        encoder_attention_mask=None,\n",
        "        past_key_values=None,\n",
        "        use_cache=None,\n",
        "        output_attentions=False,\n",
        "        output_hidden_states=False,\n",
        "        return_dict=True,\n",
        "        mode='multimodal',\n",
        "    ):\n",
        "        all_hidden_states = () if output_hidden_states else None\n",
        "        all_self_attentions = () if output_attentions else None\n",
        "        all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None\n",
        "\n",
        "        next_decoder_cache = () if use_cache else None\n",
        "               \n",
        "        for i in range(self.config.num_hidden_layers):\n",
        "            layer_module = self.layer[i]\n",
        "            if output_hidden_states:\n",
        "                all_hidden_states = all_hidden_states + (hidden_states,)\n",
        "\n",
        "            layer_head_mask = head_mask[i] if head_mask is not None else None\n",
        "            past_key_value = past_key_values[i] if past_key_values is not None else None\n",
        "\n",
        "            if self.gradient_checkpointing and self.training:\n",
        "\n",
        "                if use_cache:\n",
        "                    logger.warn(\n",
        "                        \"`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...\"\n",
        "                    )\n",
        "                    use_cache = False\n",
        "\n",
        "                def create_custom_forward(module):\n",
        "                    def custom_forward(*inputs):\n",
        "                        return module(*inputs, past_key_value, output_attentions)\n",
        "\n",
        "                    return custom_forward\n",
        "\n",
        "                layer_outputs = torch.utils.checkpoint.checkpoint(\n",
        "                    create_custom_forward(layer_module),\n",
        "                    hidden_states,\n",
        "                    attention_mask,\n",
        "                    layer_head_mask,\n",
        "                    encoder_hidden_states,\n",
        "                    encoder_attention_mask,\n",
        "                    mode=mode,\n",
        "                )\n",
        "            else:\n",
        "                layer_outputs = layer_module(\n",
        "                    hidden_states,\n",
        "                    attention_mask,\n",
        "                    layer_head_mask,\n",
        "                    encoder_hidden_states,\n",
        "                    encoder_attention_mask,\n",
        "                    past_key_value,\n",
        "                    output_attentions,\n",
        "                    mode=mode,\n",
        "                )\n",
        "\n",
        "            hidden_states = layer_outputs[0]\n",
        "            if use_cache:\n",
        "                next_decoder_cache += (layer_outputs[-1],)\n",
        "            if output_attentions:\n",
        "                all_self_attentions = all_self_attentions + (layer_outputs[1],)\n",
        "\n",
        "        if output_hidden_states:\n",
        "            all_hidden_states = all_hidden_states + (hidden_states,)\n",
        "\n",
        "        if not return_dict:\n",
        "            return tuple(\n",
        "                v\n",
        "                for v in [\n",
        "                    hidden_states,\n",
        "                    next_decoder_cache,\n",
        "                    all_hidden_states,\n",
        "                    all_self_attentions,\n",
        "                    all_cross_attentions,\n",
        "                ]\n",
        "                if v is not None\n",
        "            )\n",
        "        return BaseModelOutputWithPastAndCrossAttentions(\n",
        "            last_hidden_state=hidden_states,\n",
        "            past_key_values=next_decoder_cache,\n",
        "            hidden_states=all_hidden_states,\n",
        "            attentions=all_self_attentions,\n",
        "            cross_attentions=all_cross_attentions,\n",
        "        )\n",
        "\n",
        "\n",
        "class BertPooler(nn.Module):\n",
        "    def __init__(self, config):\n",
        "        super().__init__()\n",
        "        self.dense = nn.Linear(config.hidden_size, config.hidden_size)\n",
        "        self.activation = nn.Tanh()\n",
        "\n",
        "    def forward(self, hidden_states):\n",
        "        # We \"pool\" the model by simply taking the hidden state corresponding\n",
        "        # to the first token.\n",
        "        first_token_tensor = hidden_states[:, 0]\n",
        "        pooled_output = self.dense(first_token_tensor)\n",
        "        pooled_output = self.activation(pooled_output)\n",
        "        return pooled_output\n",
        "\n",
        "\n",
        "class BertPredictionHeadTransform(nn.Module):\n",
        "    def __init__(self, config):\n",
        "        super().__init__()\n",
        "        self.dense = nn.Linear(config.hidden_size, config.hidden_size)\n",
        "        if isinstance(config.hidden_act, str):\n",
        "            self.transform_act_fn = ACT2FN[config.hidden_act]\n",
        "        else:\n",
        "            self.transform_act_fn = config.hidden_act\n",
        "        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)\n",
        "\n",
        "    def forward(self, hidden_states):\n",
        "        hidden_states = self.dense(hidden_states)\n",
        "        hidden_states = self.transform_act_fn(hidden_states)\n",
        "        hidden_states = self.LayerNorm(hidden_states)\n",
        "        return hidden_states\n",
        "\n",
        "\n",
        "class BertLMPredictionHead(nn.Module):\n",
        "    def __init__(self, config):\n",
        "        super().__init__()\n",
        "        self.transform = BertPredictionHeadTransform(config)\n",
        "\n",
        "        # The output weights are the same as the input embeddings, but there is\n",
        "        # an output-only bias for each token.\n",
        "        self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=False)\n",
        "\n",
        "        self.bias = nn.Parameter(torch.zeros(config.vocab_size))\n",
        "\n",
        "        # Need a link between the two variables so that the bias is correctly resized with `resize_token_embeddings`\n",
        "        self.decoder.bias = self.bias\n",
        "\n",
        "    def forward(self, hidden_states):\n",
        "        hidden_states = self.transform(hidden_states)\n",
        "        hidden_states = self.decoder(hidden_states)\n",
        "        return hidden_states\n",
        "\n",
        "\n",
        "class BertOnlyMLMHead(nn.Module):\n",
        "    def __init__(self, config):\n",
        "        super().__init__()\n",
        "        self.predictions = BertLMPredictionHead(config)\n",
        "\n",
        "    def forward(self, sequence_output):\n",
        "        prediction_scores = self.predictions(sequence_output)\n",
        "        return prediction_scores\n",
        "\n",
        "\n",
        "class BertPreTrainedModel(PreTrainedModel):\n",
        "    \"\"\"\n",
        "    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained\n",
        "    models.\n",
        "    \"\"\"\n",
        "\n",
        "    config_class = BertConfig\n",
        "    base_model_prefix = \"bert\"\n",
        "    _keys_to_ignore_on_load_missing = [r\"position_ids\"]\n",
        "\n",
        "    def _init_weights(self, module):\n",
        "        \"\"\" Initialize the weights \"\"\"\n",
        "        if isinstance(module, (nn.Linear, nn.Embedding)):\n",
        "            # Slightly different from the TF version which uses truncated_normal for initialization\n",
        "            # cf https://github.com/pytorch/pytorch/pull/5617\n",
        "            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)\n",
        "        elif isinstance(module, nn.LayerNorm):\n",
        "            module.bias.data.zero_()\n",
        "            module.weight.data.fill_(1.0)\n",
        "        if isinstance(module, nn.Linear) and module.bias is not None:\n",
        "            module.bias.data.zero_()\n",
        "\n",
        "\n",
        "class BertModel(BertPreTrainedModel):\n",
        "    \"\"\"\n",
        "    The model can behave as an encoder (with only self-attention) as well as a decoder, in which case a layer of\n",
        "    cross-attention is added between the self-attention layers, following the architecture described in `Attention is\n",
        "    all you need <https://arxiv.org/abs/1706.03762>`__ by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit,\n",
        "    Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin.\n",
        "    argument and :obj:`add_cross_attention` set to :obj:`True`; an :obj:`encoder_hidden_states` is then expected as an\n",
        "    input to the forward pass.\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, config, add_pooling_layer=True):\n",
        "        super().__init__(config)\n",
        "        self.config = config\n",
        "\n",
        "        self.embeddings = BertEmbeddings(config)\n",
        "        \n",
        "        self.encoder = BertEncoder(config)\n",
        "\n",
        "        self.pooler = BertPooler(config) if add_pooling_layer else None\n",
        "\n",
        "        self.init_weights()\n",
        " \n",
        "\n",
        "    def get_input_embeddings(self):\n",
        "        return self.embeddings.word_embeddings\n",
        "\n",
        "    def set_input_embeddings(self, value):\n",
        "        self.embeddings.word_embeddings = value\n",
        "\n",
        "    def _prune_heads(self, heads_to_prune):\n",
        "        \"\"\"\n",
        "        Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base\n",
        "        class PreTrainedModel\n",
        "        \"\"\"\n",
        "        for layer, heads in heads_to_prune.items():\n",
        "            self.encoder.layer[layer].attention.prune_heads(heads)\n",
        "\n",
        "    \n",
        "    def get_extended_attention_mask(self, attention_mask: Tensor, input_shape: Tuple[int], device: device, is_decoder: bool) -> Tensor:\n",
        "        \"\"\"\n",
        "        Makes broadcastable attention and causal masks so that future and masked tokens are ignored.\n",
        "\n",
        "        Arguments:\n",
        "            attention_mask (:obj:`torch.Tensor`):\n",
        "                Mask with ones indicating tokens to attend to, zeros for tokens to ignore.\n",
        "            input_shape (:obj:`Tuple[int]`):\n",
        "                The shape of the input to the model.\n",
        "            device: (:obj:`torch.device`):\n",
        "                The device of the input to the model.\n",
        "\n",
        "        Returns:\n",
        "            :obj:`torch.Tensor` The extended attention mask, with a the same dtype as :obj:`attention_mask.dtype`.\n",
        "        \"\"\"\n",
        "        # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]\n",
        "        # ourselves in which case we just need to make it broadcastable to all heads.\n",
        "        if attention_mask.dim() == 3:\n",
        "            extended_attention_mask = attention_mask[:, None, :, :]\n",
        "        elif attention_mask.dim() == 2:\n",
        "            # Provided a padding mask of dimensions [batch_size, seq_length]\n",
        "            # - if the model is a decoder, apply a causal mask in addition to the padding mask\n",
        "            # - if the model is an encoder, make the mask broadcastable to [batch_size, num_heads, seq_length, seq_length]\n",
        "            if is_decoder:\n",
        "                batch_size, seq_length = input_shape\n",
        "\n",
        "                seq_ids = torch.arange(seq_length, device=device)\n",
        "                causal_mask = seq_ids[None, None, :].repeat(batch_size, seq_length, 1) <= seq_ids[None, :, None]\n",
        "                # in case past_key_values are used we need to add a prefix ones mask to the causal mask\n",
        "                # causal and attention masks must have same type with pytorch version < 1.3\n",
        "                causal_mask = causal_mask.to(attention_mask.dtype)\n",
        "   \n",
        "                if causal_mask.shape[1] < attention_mask.shape[1]:\n",
        "                    prefix_seq_len = attention_mask.shape[1] - causal_mask.shape[1]\n",
        "                    causal_mask = torch.cat(\n",
        "                        [\n",
        "                            torch.ones((batch_size, seq_length, prefix_seq_len), device=device, dtype=causal_mask.dtype),\n",
        "                            causal_mask,\n",
        "                        ],\n",
        "                        axis=-1,\n",
        "                    )                     \n",
        "\n",
        "                extended_attention_mask = causal_mask[:, None, :, :] * attention_mask[:, None, None, :]\n",
        "            else:\n",
        "                extended_attention_mask = attention_mask[:, None, None, :]\n",
        "        else:\n",
        "            raise ValueError(\n",
        "                \"Wrong shape for input_ids (shape {}) or attention_mask (shape {})\".format(\n",
        "                    input_shape, attention_mask.shape\n",
        "                )\n",
        "            )\n",
        "\n",
        "        # Since attention_mask is 1.0 for positions we want to attend and 0.0 for\n",
        "        # masked positions, this operation will create a tensor which is 0.0 for\n",
        "        # positions we want to attend and -10000.0 for masked positions.\n",
        "        # Since we are adding it to the raw scores before the softmax, this is\n",
        "        # effectively the same as removing these entirely.\n",
        "        extended_attention_mask = extended_attention_mask.to(dtype=self.dtype)  # fp16 compatibility\n",
        "        extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0\n",
        "        return extended_attention_mask\n",
        "    \n",
        "    def forward(\n",
        "        self,\n",
        "        input_ids=None,\n",
        "        attention_mask=None,\n",
        "        position_ids=None,\n",
        "        head_mask=None,\n",
        "        inputs_embeds=None,\n",
        "        encoder_embeds=None,\n",
        "        encoder_hidden_states=None,\n",
        "        encoder_attention_mask=None,\n",
        "        past_key_values=None,\n",
        "        use_cache=None,\n",
        "        output_attentions=None,\n",
        "        output_hidden_states=None,\n",
        "        return_dict=None,\n",
        "        is_decoder=False,\n",
        "        mode='multimodal',\n",
        "    ):\n",
        "        r\"\"\"\n",
        "        encoder_hidden_states  (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`):\n",
        "            Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if\n",
        "            the model is configured as a decoder.\n",
        "        encoder_attention_mask (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):\n",
        "            Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in\n",
        "            the cross-attention if the model is configured as a decoder. Mask values selected in ``[0, 1]``:\n",
        "            - 1 for tokens that are **not masked**,\n",
        "            - 0 for tokens that are **masked**.\n",
        "        past_key_values (:obj:`tuple(tuple(torch.FloatTensor))` of length :obj:`config.n_layers` with each tuple having 4 tensors of shape :obj:`(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):\n",
        "            Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.\n",
        "            If :obj:`past_key_values` are used, the user can optionally input only the last :obj:`decoder_input_ids`\n",
        "            (those that don't have their past key value states given to this model) of shape :obj:`(batch_size, 1)`\n",
        "            instead of all :obj:`decoder_input_ids` of shape :obj:`(batch_size, sequence_length)`.\n",
        "        use_cache (:obj:`bool`, `optional`):\n",
        "            If set to :obj:`True`, :obj:`past_key_values` key value states are returned and can be used to speed up\n",
        "            decoding (see :obj:`past_key_values`).\n",
        "        \"\"\"\n",
        "        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions\n",
        "        output_hidden_states = (\n",
        "            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states\n",
        "        )\n",
        "        return_dict = return_dict if return_dict is not None else self.config.use_return_dict\n",
        "\n",
        "        if is_decoder:\n",
        "            use_cache = use_cache if use_cache is not None else self.config.use_cache\n",
        "        else:\n",
        "            use_cache = False\n",
        "\n",
        "        if input_ids is not None and inputs_embeds is not None:\n",
        "            raise ValueError(\"You cannot specify both input_ids and inputs_embeds at the same time\")\n",
        "        elif input_ids is not None:\n",
        "            input_shape = input_ids.size()\n",
        "            batch_size, seq_length = input_shape\n",
        "            device = input_ids.device\n",
        "        elif inputs_embeds is not None:\n",
        "            input_shape = inputs_embeds.size()[:-1]\n",
        "            batch_size, seq_length = input_shape\n",
        "            device = inputs_embeds.device\n",
        "        elif encoder_embeds is not None:    \n",
        "            input_shape = encoder_embeds.size()[:-1]\n",
        "            batch_size, seq_length = input_shape \n",
        "            device = encoder_embeds.device\n",
        "        else:\n",
        "            raise ValueError(\"You have to specify either input_ids or inputs_embeds or encoder_embeds\")\n",
        "\n",
        "        # past_key_values_length\n",
        "        past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0\n",
        "\n",
        "        if attention_mask is None:\n",
        "            attention_mask = torch.ones(((batch_size, seq_length + past_key_values_length)), device=device)\n",
        "            \n",
        "        # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]\n",
        "        # ourselves in which case we just need to make it broadcastable to all heads.\n",
        "        extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(attention_mask, input_shape, \n",
        "                                                                                 device, is_decoder)\n",
        "\n",
        "        # If a 2D or 3D attention mask is provided for the cross-attention\n",
        "        # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]\n",
        "        if encoder_hidden_states is not None:\n",
        "            if type(encoder_hidden_states) == list:\n",
        "                encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states[0].size()\n",
        "            else:\n",
        "                encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()\n",
        "            encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)\n",
        "            \n",
        "            if type(encoder_attention_mask) == list:\n",
        "                encoder_extended_attention_mask = [self.invert_attention_mask(mask) for mask in encoder_attention_mask]\n",
        "            elif encoder_attention_mask is None:\n",
        "                encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)\n",
        "                encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)\n",
        "            else:    \n",
        "                encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)\n",
        "        else:\n",
        "            encoder_extended_attention_mask = None\n",
        "\n",
        "        # Prepare head mask if needed\n",
        "        # 1.0 in head_mask indicate we keep the head\n",
        "        # attention_probs has shape bsz x n_heads x N x N\n",
        "        # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]\n",
        "        # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]\n",
        "        head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)\n",
        "        \n",
        "        if encoder_embeds is None:\n",
        "            embedding_output = self.embeddings(\n",
        "                input_ids=input_ids,\n",
        "                position_ids=position_ids,\n",
        "                inputs_embeds=inputs_embeds,\n",
        "                past_key_values_length=past_key_values_length,\n",
        "            )\n",
        "        else:\n",
        "            embedding_output = encoder_embeds\n",
        "            \n",
        "        encoder_outputs = self.encoder(\n",
        "            embedding_output,\n",
        "            attention_mask=extended_attention_mask,\n",
        "            head_mask=head_mask,\n",
        "            encoder_hidden_states=encoder_hidden_states,\n",
        "            encoder_attention_mask=encoder_extended_attention_mask,\n",
        "            past_key_values=past_key_values,\n",
        "            use_cache=use_cache,\n",
        "            output_attentions=output_attentions,\n",
        "            output_hidden_states=output_hidden_states,\n",
        "            return_dict=return_dict,\n",
        "            mode=mode,\n",
        "        )\n",
        "        sequence_output = encoder_outputs[0]\n",
        "        pooled_output = self.pooler(sequence_output) if self.pooler is not None else None\n",
        "\n",
        "        if not return_dict:\n",
        "            return (sequence_output, pooled_output) + encoder_outputs[1:]\n",
        "\n",
        "        return BaseModelOutputWithPoolingAndCrossAttentions(\n",
        "            last_hidden_state=sequence_output,\n",
        "            pooler_output=pooled_output,\n",
        "            past_key_values=encoder_outputs.past_key_values,\n",
        "            hidden_states=encoder_outputs.hidden_states,\n",
        "            attentions=encoder_outputs.attentions,\n",
        "            cross_attentions=encoder_outputs.cross_attentions,\n",
        "        )\n",
        "\n",
        "\n",
        "\n",
        "class BertLMHeadModel(BertPreTrainedModel):\n",
        "\n",
        "    _keys_to_ignore_on_load_unexpected = [r\"pooler\"]\n",
        "    _keys_to_ignore_on_load_missing = [r\"position_ids\", r\"predictions.decoder.bias\"]\n",
        "\n",
        "    def __init__(self, config):\n",
        "        super().__init__(config)\n",
        "\n",
        "        self.bert = BertModel(config, add_pooling_layer=False)\n",
        "        self.cls = BertOnlyMLMHead(config)\n",
        "\n",
        "        self.init_weights()\n",
        "\n",
        "    def get_output_embeddings(self):\n",
        "        return self.cls.predictions.decoder\n",
        "\n",
        "    def set_output_embeddings(self, new_embeddings):\n",
        "        self.cls.predictions.decoder = new_embeddings\n",
        "\n",
        "    def forward(\n",
        "        self,\n",
        "        input_ids=None,\n",
        "        attention_mask=None,\n",
        "        position_ids=None,\n",
        "        head_mask=None,\n",
        "        inputs_embeds=None,\n",
        "        encoder_hidden_states=None,\n",
        "        encoder_attention_mask=None,\n",
        "        labels=None,\n",
        "        past_key_values=None,\n",
        "        use_cache=None,\n",
        "        output_attentions=None,\n",
        "        output_hidden_states=None,\n",
        "        return_dict=None,\n",
        "        return_logits=False,            \n",
        "        is_decoder=True,\n",
        "        reduction='mean',\n",
        "        mode='multimodal', \n",
        "    ):\n",
        "        r\"\"\"\n",
        "        encoder_hidden_states  (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`):\n",
        "            Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if\n",
        "            the model is configured as a decoder.\n",
        "        encoder_attention_mask (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):\n",
        "            Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in\n",
        "            the cross-attention if the model is configured as a decoder. Mask values selected in ``[0, 1]``:\n",
        "            - 1 for tokens that are **not masked**,\n",
        "            - 0 for tokens that are **masked**.\n",
        "        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):\n",
        "            Labels for computing the left-to-right language modeling loss (next word prediction). Indices should be in\n",
        "            ``[-100, 0, ..., config.vocab_size]`` (see ``input_ids`` docstring) Tokens with indices set to ``-100`` are\n",
        "            ignored (masked), the loss is only computed for the tokens with labels n ``[0, ..., config.vocab_size]``\n",
        "        past_key_values (:obj:`tuple(tuple(torch.FloatTensor))` of length :obj:`config.n_layers` with each tuple having 4 tensors of shape :obj:`(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):\n",
        "            Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.\n",
        "            If :obj:`past_key_values` are used, the user can optionally input only the last :obj:`decoder_input_ids`\n",
        "            (those that don't have their past key value states given to this model) of shape :obj:`(batch_size, 1)`\n",
        "            instead of all :obj:`decoder_input_ids` of shape :obj:`(batch_size, sequence_length)`.\n",
        "        use_cache (:obj:`bool`, `optional`):\n",
        "            If set to :obj:`True`, :obj:`past_key_values` key value states are returned and can be used to speed up\n",
        "            decoding (see :obj:`past_key_values`).\n",
        "        Returns:\n",
        "        Example::\n",
        "            >>> from transformers import BertTokenizer, BertLMHeadModel, BertConfig\n",
        "            >>> import torch\n",
        "            >>> tokenizer = BertTokenizer.from_pretrained('bert-base-cased')\n",
        "            >>> config = BertConfig.from_pretrained(\"bert-base-cased\")\n",
        "            >>> model = BertLMHeadModel.from_pretrained('bert-base-cased', config=config)\n",
        "            >>> inputs = tokenizer(\"Hello, my dog is cute\", return_tensors=\"pt\")\n",
        "            >>> outputs = model(**inputs)\n",
        "            >>> prediction_logits = outputs.logits\n",
        "        \"\"\"\n",
        "        return_dict = return_dict if return_dict is not None else self.config.use_return_dict\n",
        "        if labels is not None:\n",
        "            use_cache = False\n",
        "\n",
        "        outputs = self.bert(\n",
        "            input_ids,\n",
        "            attention_mask=attention_mask,\n",
        "            position_ids=position_ids,\n",
        "            head_mask=head_mask,\n",
        "            inputs_embeds=inputs_embeds,\n",
        "            encoder_hidden_states=encoder_hidden_states,\n",
        "            encoder_attention_mask=encoder_attention_mask,\n",
        "            past_key_values=past_key_values,\n",
        "            use_cache=use_cache,\n",
        "            output_attentions=output_attentions,\n",
        "            output_hidden_states=output_hidden_states,\n",
        "            return_dict=return_dict,\n",
        "            is_decoder=is_decoder,\n",
        "            mode=mode,\n",
        "        )\n",
        "        \n",
        "        sequence_output = outputs[0]\n",
        "        prediction_scores = self.cls(sequence_output)\n",
        "        \n",
        "        if return_logits:\n",
        "            return prediction_scores[:, :-1, :].contiguous()  \n",
        "\n",
        "        lm_loss = None\n",
        "        if labels is not None:\n",
        "            # we are doing next-token prediction; shift prediction scores and input ids by one\n",
        "            shifted_prediction_scores = prediction_scores[:, :-1, :].contiguous()\n",
        "            labels = labels[:, 1:].contiguous()\n",
        "            loss_fct = CrossEntropyLoss(reduction=reduction, label_smoothing=0.1) \n",
        "            lm_loss = loss_fct(shifted_prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))\n",
        "            if reduction=='none':\n",
        "                lm_loss = lm_loss.view(prediction_scores.size(0),-1).sum(1)               \n",
        "\n",
        "        if not return_dict:\n",
        "            output = (prediction_scores,) + outputs[2:]\n",
        "            return ((lm_loss,) + output) if lm_loss is not None else output\n",
        "\n",
        "        return CausalLMOutputWithCrossAttentions(\n",
        "            loss=lm_loss,\n",
        "            logits=prediction_scores,\n",
        "            past_key_values=outputs.past_key_values,\n",
        "            hidden_states=outputs.hidden_states,\n",
        "            attentions=outputs.attentions,\n",
        "            cross_attentions=outputs.cross_attentions,\n",
        "        )\n",
        "\n",
        "    def prepare_inputs_for_generation(self, input_ids, past=None, attention_mask=None, **model_kwargs):\n",
        "        input_shape = input_ids.shape\n",
        "        # if model is used as a decoder in encoder-decoder model, the decoder attention mask is created on the fly\n",
        "        if attention_mask is None:\n",
        "            attention_mask = input_ids.new_ones(input_shape)\n",
        "\n",
        "        # cut decoder_input_ids if past is used\n",
        "        if past is not None:\n",
        "            input_ids = input_ids[:, -1:]\n",
        "\n",
        "        return {\n",
        "            \"input_ids\": input_ids, \n",
        "            \"attention_mask\": attention_mask, \n",
        "            \"past_key_values\": past,\n",
        "            \"encoder_hidden_states\": model_kwargs.get(\"encoder_hidden_states\", None),\n",
        "            \"encoder_attention_mask\": model_kwargs.get(\"encoder_attention_mask\", None),\n",
        "            \"is_decoder\": True,\n",
        "        }\n",
        "\n",
        "    def _reorder_cache(self, past, beam_idx):\n",
        "        reordered_past = ()\n",
        "        for layer_past in past:\n",
        "            reordered_past += (tuple(past_state.index_select(0, beam_idx) for past_state in layer_past),)\n",
        "        return reordered_past\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 21,
      "metadata": {
        "jupyter": {
          "outputs_hidden": true,
          "source_hidden": true
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "outputs": [
        {
          "data": {
            "application/vnd.livy.statement-meta+json": {
              "execution_finish_time": "2022-04-06T08:36:55.0981173Z",
              "execution_start_time": "2022-04-06T08:36:54.9353048Z",
              "livy_statement_state": "available",
              "queued_time": "2022-04-06T08:35:55.4733366Z",
              "session_id": 195,
              "session_start_time": null,
              "spark_pool": "small",
              "state": "finished",
              "statement_id": 3
            },
            "text/plain": [
              "StatementMeta(small, 195, 3, Finished, Available)"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "import json\n",
        "def read_batch_config(batch_root: str):\n",
        "    \"\"\"\n",
        "    We read the config file using the Java File System API as we do not need to let multiple nodes read individual lines and join it\n",
        "    all back together again\n",
        "    \"\"\"\n",
        "    # Change our file system from 'synapse' to 'input'\n",
        "    sc._jsc.hadoopConfiguration().set(\"fs.defaultFS\", f\"abfss://{file_system_name}@{account_name}.dfs.core.windows.net\")\n",
        "\n",
        "    fs = sc._jvm.org.apache.hadoop.fs.FileSystem.get(sc._jsc.hadoopConfiguration())\n",
        "    config_path = sc._jvm.org.apache.hadoop.fs.Path(f'{batch_root}')\n",
        "\n",
        "    # If we don't have a batch config, copy the global one.\n",
        "    if fs.exists(config_path) != True:\n",
        "        logger.error(f'{config_path} not found.')\n",
        "\n",
        "    # Open our file directly rather than through spark\n",
        "    input_stream = fs.open(config_path)  # FSDataInputStream\n",
        "\n",
        "    config_string = sc._jvm.java.io.BufferedReader(\n",
        "        sc._jvm.java.io.InputStreamReader(input_stream, sc._jvm.java.nio.charset.StandardCharsets.UTF_8)\n",
        "        ).lines().collect(sc._jvm.java.util.stream.Collectors.joining(\"\\n\"))\n",
        "\n",
        "    # Load it into json    \n",
        "    #return json.loads(''.join(config_string), object_hook=lambda dictionary: SimpleNamespace(**dictionary))\n",
        "    return json.loads(config_string)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 22,
      "metadata": {
        "jupyter": {
          "outputs_hidden": true,
          "source_hidden": true
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "outputs": [
        {
          "data": {
            "application/vnd.livy.statement-meta+json": {
              "execution_finish_time": "2022-04-06T08:37:00.7000337Z",
              "execution_start_time": "2022-04-06T08:36:55.2401301Z",
              "livy_statement_state": "available",
              "queued_time": "2022-04-06T08:35:55.563476Z",
              "session_id": 195,
              "session_start_time": null,
              "spark_pool": "small",
              "state": "finished",
              "statement_id": 4
            },
            "text/plain": [
              "StatementMeta(small, 195, 4, Finished, Available)"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "from transformers import BertTokenizer\n",
        "\n",
        "import torch\n",
        "from torch import nn\n",
        "import torch.nn.functional as F\n",
        "\n",
        "\n",
        "import os\n",
        "from urllib.parse import urlparse\n",
        "from timm.models.hub import download_cached_file\n",
        "\n",
        "class BLIP_Base(nn.Module):\n",
        "    def __init__(self,                 \n",
        "                 med_config = read_batch_config(model_config_file),\n",
        "                 image_size = 224,\n",
        "                 vit = 'base',\n",
        "                 vit_grad_ckpt = False,\n",
        "                 vit_ckpt_layer = 0,                 \n",
        "                 ):\n",
        "        \"\"\"\n",
        "        Args:\n",
        "            med_config (str): path for the mixture of encoder-decoder model's configuration file\n",
        "            image_size (int): input image size\n",
        "            vit (str): model size of vision transformer\n",
        "        \"\"\"               \n",
        "        super().__init__()\n",
        "        \n",
        "        self.visual_encoder, vision_width = create_vit(vit,image_size, vit_grad_ckpt, vit_ckpt_layer)\n",
        "        self.tokenizer = init_tokenizer()   \n",
        "        med_config = BertConfig.from_dict(med_config)\n",
        "        #med_config = BertConfig.from_json_file(med_config)\n",
        "        #med_config = json.loads(''.join(med_config), object_hook=lambda dictionary: SimpleNamespace(**dictionary))\n",
        "        med_config.encoder_width = vision_width\n",
        "        self.text_encoder = BertModel(config=med_config, add_pooling_layer=False)  \n",
        "\n",
        "        \n",
        "    def forward(self, image, caption, mode):\n",
        "        \n",
        "        assert mode in ['image', 'text', 'multimodal'], \"mode parameter must be image, text, or multimodal\"\n",
        "        text = self.tokenizer(caption, return_tensors=\"pt\").to(image.device) \n",
        "        \n",
        "        if mode=='image':    \n",
        "            # return image features\n",
        "            image_embeds = self.visual_encoder(image)             \n",
        "            return image_embeds\n",
        "        \n",
        "        elif mode=='text':\n",
        "            # return text features\n",
        "            text_output = self.text_encoder(text.input_ids, attention_mask = text.attention_mask,                      \n",
        "                                            return_dict = True, mode = 'text')  \n",
        "            return text_output.last_hidden_state\n",
        "        \n",
        "        elif mode=='multimodal':\n",
        "            # return multimodel features\n",
        "            image_embeds = self.visual_encoder(image)    \n",
        "            image_atts = torch.ones(image_embeds.size()[:-1],dtype=torch.long).to(image.device)      \n",
        "            \n",
        "            text.input_ids[:,0] = self.tokenizer.enc_token_id\n",
        "            output = self.text_encoder(text.input_ids,\n",
        "                                       attention_mask = text.attention_mask,\n",
        "                                       encoder_hidden_states = image_embeds,\n",
        "                                       encoder_attention_mask = image_atts,      \n",
        "                                       return_dict = True,\n",
        "                                      )              \n",
        "            return output.last_hidden_state\n",
        "        \n",
        "        \n",
        "        \n",
        "class BLIP_Decoder(nn.Module):\n",
        "    def __init__(self,                 \n",
        "                 med_config = read_batch_config(model_config_file),\n",
        "                 image_size = 384,\n",
        "                 vit = 'base',\n",
        "                 vit_grad_ckpt = False,\n",
        "                 vit_ckpt_layer = 0,\n",
        "                 prompt = 'a picture of ',\n",
        "                 ):\n",
        "        \"\"\"\n",
        "        Args:\n",
        "            med_config (str): path for the mixture of encoder-decoder model's configuration file\n",
        "            image_size (int): input image size\n",
        "            vit (str): model size of vision transformer\n",
        "        \"\"\"            \n",
        "        super().__init__()\n",
        "        \n",
        "        self.visual_encoder, vision_width = create_vit(vit,image_size, vit_grad_ckpt, vit_ckpt_layer)\n",
        "        self.tokenizer = init_tokenizer()  \n",
        "\n",
        "        med_config = BertConfig.from_dict(med_config)\n",
        "        med_config.encoder_width = vision_width\n",
        "        self.text_decoder = BertLMHeadModel(config=med_config)    \n",
        "        \n",
        "        self.prompt = prompt\n",
        "        self.prompt_length = len(self.tokenizer(self.prompt).input_ids)-1\n",
        "\n",
        "        \n",
        "    def forward(self, image, caption):\n",
        "        \n",
        "        image_embeds = self.visual_encoder(image) \n",
        "        image_atts = torch.ones(image_embeds.size()[:-1],dtype=torch.long).to(image.device)\n",
        "        \n",
        "        text = self.tokenizer(caption, padding='longest', truncation=True, max_length=40, return_tensors=\"pt\").to(image.device) \n",
        "        \n",
        "        text.input_ids[:,0] = self.tokenizer.bos_token_id\n",
        "        \n",
        "        decoder_targets = text.input_ids.masked_fill(text.input_ids == self.tokenizer.pad_token_id, -100)         \n",
        "        decoder_targets[:,:self.prompt_length] = -100\n",
        "     \n",
        "        decoder_output = self.text_decoder(text.input_ids, \n",
        "                                           attention_mask = text.attention_mask, \n",
        "                                           encoder_hidden_states = image_embeds,\n",
        "                                           encoder_attention_mask = image_atts,                  \n",
        "                                           labels = decoder_targets,\n",
        "                                           return_dict = True,   \n",
        "                                          )   \n",
        "        loss_lm = decoder_output.loss\n",
        "        \n",
        "        return loss_lm\n",
        "\n",
        "def blip_decoder(pretrained='',**kwargs):\n",
        "    model = BLIP_Decoder(**kwargs)\n",
        "    if pretrained:\n",
        "        model,msg = load_checkpoint(model,pretrained)\n",
        "        assert(len(msg.missing_keys)==0)\n",
        "    return model    \n",
        "    \n",
        "def blip_feature_extractor(pretrained='',**kwargs):\n",
        "    model = BLIP_Base(**kwargs)\n",
        "    if pretrained:\n",
        "        model,msg = load_checkpoint(model,pretrained)\n",
        "        assert(len(msg.missing_keys)==0)\n",
        "    return model        \n",
        "\n",
        "def init_tokenizer():\n",
        "    tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')\n",
        "    tokenizer.add_special_tokens({'bos_token':'[DEC]'})\n",
        "    tokenizer.add_special_tokens({'additional_special_tokens':['[ENC]']})       \n",
        "    tokenizer.enc_token_id = tokenizer.additional_special_tokens_ids[0]  \n",
        "    return tokenizer\n",
        "\n",
        "\n",
        "def create_vit(vit, image_size, use_grad_checkpointing=False, ckpt_layer=0, drop_path_rate=0):\n",
        "        \n",
        "    assert vit in ['base', 'large'], \"vit parameter must be base or large\"\n",
        "    if vit=='base':\n",
        "        vision_width = 768\n",
        "        visual_encoder = VisionTransformer(img_size=image_size, patch_size=16, embed_dim=vision_width, depth=12, \n",
        "                                           num_heads=12, use_grad_checkpointing=use_grad_checkpointing, ckpt_layer=ckpt_layer,\n",
        "                                           drop_path_rate=0 or drop_path_rate\n",
        "                                          )   \n",
        "    elif vit=='large':\n",
        "        vision_width = 1024\n",
        "        visual_encoder = VisionTransformer(img_size=image_size, patch_size=16, embed_dim=vision_width, depth=24, \n",
        "                                           num_heads=16, use_grad_checkpointing=use_grad_checkpointing, ckpt_layer=ckpt_layer,\n",
        "                                           drop_path_rate=0.1 or drop_path_rate\n",
        "                                          )   \n",
        "    return visual_encoder, vision_width\n",
        "\n",
        "def is_url(url_or_filename):\n",
        "    parsed = urlparse(url_or_filename)\n",
        "    return parsed.scheme in (\"http\", \"https\")\n",
        "\n",
        "def load_checkpoint(model,url_or_filename):\n",
        "    if is_url(url_or_filename):\n",
        "        cached_file = download_cached_file(url_or_filename, check_hash=False, progress=True)\n",
        "        checkpoint = torch.load(cached_file, map_location='cpu') \n",
        "    elif os.path.isfile(url_or_filename):        \n",
        "        checkpoint = torch.load(url_or_filename, map_location='cpu') \n",
        "    else:\n",
        "        raise RuntimeError('checkpoint url or path is invalid')\n",
        "        \n",
        "    state_dict = checkpoint['model']\n",
        "    \n",
        "    state_dict['visual_encoder.pos_embed'] = interpolate_pos_embed(state_dict['visual_encoder.pos_embed'],model.visual_encoder) \n",
        "    if 'visual_encoder_m.pos_embed' in model.state_dict().keys():\n",
        "        state_dict['visual_encoder_m.pos_embed'] = interpolate_pos_embed(state_dict['visual_encoder_m.pos_embed'],\n",
        "                                                                         model.visual_encoder_m)    \n",
        "    for key in model.state_dict().keys():\n",
        "        if key in state_dict.keys():\n",
        "            if state_dict[key].shape!=model.state_dict()[key].shape:\n",
        "                del state_dict[key]\n",
        "    \n",
        "    msg = model.load_state_dict(state_dict,strict=False)\n",
        "    print('load checkpoint from %s'%url_or_filename)  \n",
        "    return model,msg"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 23,
      "metadata": {
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": true
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "outputs": [
        {
          "data": {
            "application/vnd.livy.statement-meta+json": {
              "execution_finish_time": "2022-04-06T08:37:03.6225878Z",
              "execution_start_time": "2022-04-06T08:37:00.8144056Z",
              "livy_statement_state": "available",
              "queued_time": "2022-04-06T08:35:55.6295939Z",
              "session_id": 195,
              "session_start_time": null,
              "spark_pool": "small",
              "state": "finished",
              "statement_id": 5
            },
            "text/plain": [
              "StatementMeta(small, 195, 5, Finished, Available)"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "from pyspark.sql.types import StringType\n",
        "from pyspark.sql.functions import col, pandas_udf, lit, struct, PandasUDFType, udf\n",
        "import pyspark.sql.types as Types\n",
        "from pyspark.ml.linalg import Vectors, VectorUDT\n",
        "import io     \n",
        "\n",
        "# Load all images\n",
        "images = spark.read.format(\"binaryFile\") \\\n",
        "  .option(\"recursiveFileLookup\", \"true\") \\\n",
        "  .option(\"pathGlobFilter\", \"*.jpg\") \\\n",
        "  .load(input_directory)\n",
        "\n",
        "# Convert to spark df\n",
        "images_df = images.select(\n",
        "  col(\"path\"),\n",
        "  col(\"content\"))\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 24,
      "metadata": {
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": true
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "outputs": [
        {
          "data": {
            "application/vnd.livy.statement-meta+json": {
              "execution_finish_time": "2022-04-06T08:37:04.2679439Z",
              "execution_start_time": "2022-04-06T08:37:03.7462305Z",
              "livy_statement_state": "available",
              "queued_time": "2022-04-06T08:35:55.6994653Z",
              "session_id": 195,
              "session_start_time": null,
              "spark_pool": "small",
              "state": "finished",
              "statement_id": 6
            },
            "text/plain": [
              "StatementMeta(small, 195, 6, Finished, Available)"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "import numpy as np\n",
        "import pandas as pd\n",
        "from pyspark.sql.types import DoubleType, ArrayType, FloatType, DoubleType\n",
        "from PIL import Image\n",
        "\n",
        "def transformer(img_data):\n",
        "    image = \"\"\n",
        "    image_size = 384\n",
        "    raw_image = Image.open(io.BytesIO(img_data)).convert('RGB')   \n",
        "\n",
        "    w,h = raw_image.size\n",
        "    \n",
        "    transform = transforms.Compose([\n",
        "        transforms.Resize((image_size,image_size),interpolation=InterpolationMode.BICUBIC),\n",
        "        transforms.ToTensor(),\n",
        "        transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))\n",
        "        ]) \n",
        "\n",
        "    image = transform(raw_image).unsqueeze(0)\n",
        "    return image\n",
        "\n",
        "\n",
        "def blip_model_udf(model_fn):\n",
        "    def generate(image_batch_iter):\n",
        "\n",
        "        num_beams=3\n",
        "        max_length=40\n",
        "        min_length=20\n",
        "        top_p=0.9\n",
        "        repetition_penalty=1.0\n",
        "        \n",
        "\n",
        "        with torch.no_grad():\n",
        "        \n",
        "            for img_series in image_batch_iter:\n",
        "                processed_images = [transformer(img) for img in img_series]\n",
        "                for img_array in processed_images:\n",
        "\n",
        "                    try:\n",
        "\n",
        "                        image_embeds = model.visual_encoder(img_array)\n",
        "                        image_embeds = image_embeds.repeat_interleave(num_beams,dim=0)\n",
        "                        \n",
        "                        image_atts = torch.ones(image_embeds.size()[:-1],dtype=torch.long).to(img_array)\n",
        "                        model_kwargs = {\"encoder_hidden_states\": image_embeds, \"encoder_attention_mask\":image_atts}\n",
        "                        \n",
        "                        prompt = [model.prompt] * img_array.size(0)\n",
        "                        input_ids = model.tokenizer(prompt, return_tensors=\"pt\").input_ids.to(img_array) \n",
        "                        input_ids[:,0] = model.tokenizer.bos_token_id\n",
        "                        input_ids = input_ids[:, :-1] \n",
        "                        # Fix\n",
        "                        input_ids = input_ids.type(torch.LongTensor)\n",
        "                        outputs = model.text_decoder.generate(input_ids=input_ids,\n",
        "                                                                max_length=max_length,\n",
        "                                                                min_length=min_length,\n",
        "                                                                num_beams=num_beams,\n",
        "                                                                eos_token_id=model.tokenizer.sep_token_id,\n",
        "                                                                pad_token_id=model.tokenizer.pad_token_id,     \n",
        "                                                                repetition_penalty=repetition_penalty,\n",
        "                                                                **model_kwargs)\n",
        "\n",
        "                    except Exception as ImageError:\n",
        "                        print(f\"ImageError: {ImageError} {image_batch_iter} {img_series}\")\n",
        "                        outputs = [0,0,0]     \n",
        "\n",
        "                    yield pd.DataFrame([outputs.tolist()])\n",
        "\n",
        "    return_type = \"captions: array<float>\"\n",
        "    return pandas_udf(return_type, PandasUDFType.SCALAR_ITER)(generate)\n",
        "        \n",
        "        \n",
        "def blip_decoder_fn():\n",
        "    model_url = 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_base_capfilt_large.pth'\n",
        "    image_size = 384\n",
        "    model = blip_decoder(pretrained=model_url, image_size=image_size, vit='base')\n",
        "    model.eval()\n",
        "    return model\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 25,
      "metadata": {
        "jupyter": {
          "outputs_hidden": true,
          "source_hidden": true
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "outputs": [
        {
          "data": {
            "application/vnd.livy.statement-meta+json": {
              "execution_finish_time": "2022-04-06T08:37:39.3575899Z",
              "execution_start_time": "2022-04-06T08:37:04.3800028Z",
              "livy_statement_state": "available",
              "queued_time": "2022-04-06T08:35:55.8017607Z",
              "session_id": 195,
              "session_start_time": null,
              "spark_pool": "small",
              "state": "finished",
              "statement_id": 7
            },
            "text/plain": [
              "StatementMeta(small, 195, 7, Finished, Available)"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Downloading: 100%|##########| 226k/226k [00:00<00:00, 30.7MB/s]\n",
            "Downloading: 100%|##########| 28.0/28.0 [00:00<00:00, 35.2kB/s]\n",
            "Downloading: 100%|##########| 570/570 [00:00<00:00, 582kB/s]\n",
            "100%|##########| 1.97G/1.97G [00:11<00:00, 192MB/s]\n",
            "reshape position embedding from 196 to 576\n",
            "load checkpoint from https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_base_capfilt_large.pth\n",
            "/opt/spark/python/lib/pyspark.zip/pyspark/sql/pandas/functions.py:389: UserWarning: In Python 3.6+ and Spark 3.0+, it is preferred to specify type hints for pandas UDF instead of specifying pandas UDF type which will be deprecated in the future releases. See SPARK-28264 for more details."
          ]
        }
      ],
      "source": [
        "images_df_small = images_df#.limit(1)\n",
        "model = blip_decoder_fn()\n",
        "bc_model_state = sc.broadcast(model.state_dict())\n",
        "blip_udf = blip_model_udf(blip_decoder_fn)\n",
        "predictions = images_df_small.withColumn(\"text\", blip_udf(col(\"content\")))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 26,
      "metadata": {
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": true
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "outputs": [
        {
          "data": {
            "application/vnd.livy.statement-meta+json": {
              "execution_finish_time": "2022-04-06T08:37:41.0256806Z",
              "execution_start_time": "2022-04-06T08:37:39.9177591Z",
              "livy_statement_state": "available",
              "queued_time": "2022-04-06T08:35:55.9200004Z",
              "session_id": 195,
              "session_start_time": null,
              "spark_pool": "small",
              "state": "finished",
              "statement_id": 8
            },
            "text/plain": [
              "StatementMeta(small, 195, 8, Finished, Available)"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "from pyspark.ml.linalg import Vectors, VectorUDT\n",
        "from pyspark.sql.functions import udf\n",
        "from transformers import AutoTokenizer, AutoModelForSequenceClassification\n",
        "import torch\n",
        "from tqdm import tqdm\n",
        "\n",
        "\n",
        "list_to_vector_udf = udf(lambda l: Vectors.dense(l), VectorUDT())\n",
        "df_with_vectors = predictions.select(\n",
        "    predictions[\"path\"],\n",
        "    predictions[\"content\"],\n",
        "    predictions[\"text.captions\"]\n",
        "    #list_to_vector_udf(predictions[\"captions.captions\"]).alias(\"tokens\")\n",
        ")\n",
        "\n",
        "from azure.identity import DefaultAzureCredential\n",
        "from datetime import datetime, timedelta\n",
        "from azure.storage.blob import BlobServiceClient, generate_container_sas, BlobSasPermissions\n",
        "token_credential = DefaultAzureCredential()\n",
        "\n",
        "blob_service_client = BlobServiceClient(\n",
        "    account_url=account_url,\n",
        "    credential=token_credential\n",
        ")\n",
        "\n",
        "from azure.storage.filedatalake import DataLakeServiceClient, generate_directory_sas\n",
        "SAS = generate_directory_sas(\n",
        "        account_name=account_name,\n",
        "        file_system_name=file_system_name,\n",
        "        directory_name=directory_name,\n",
        "        credential=account_key,\n",
        "        permission=BlobSasPermissions(read=True),\n",
        "        expiry=datetime.utcnow() + timedelta(days=100))\n",
        "\n",
        "SAS_key  = \"?\" + SAS\n",
        "\n",
        "storage_path = os.path.join(account_url, file_system_name, directory_name)\n",
        "SAS_path = []\n",
        "\n",
        "def build_sas_path(row):\n",
        "    file_name = ntpath.basename(row)\n",
        "    return account_url + \"/\" + file_system_name + \"/\" + directory_name + \"/\" + file_name + SAS_key\n",
        "\n",
        "udf_build_sas_path = udf(build_sas_path, StringType())\n",
        "\n",
        "df_coords = df_with_vectors.withColumn(\"blob_path\", udf_build_sas_path(df_with_vectors.path))\n",
        "\n",
        "def decode(indices):\n",
        "        \n",
        "    tokens = torch.Tensor(indices)\n",
        "    tokens = tokens.type(torch.LongTensor)\n",
        "    caption = btokenizer.decode(tokens, skip_special_tokens=True)\n",
        "\n",
        "    text = caption.split()\n",
        "    # Let's remove the duplicate repeats as incurring a penalty increased run time\n",
        "    caption = \" \".join([text[i] for i in range(len(text)) if (i==0) or text[i] != text[i-1]])\n",
        "   \n",
        "    return caption[len('a picture of '):]\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 27,
      "metadata": {
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": true
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "outputs": [
        {
          "data": {
            "application/vnd.livy.statement-meta+json": {
              "execution_finish_time": "2022-04-06T08:37:41.6527224Z",
              "execution_start_time": "2022-04-06T08:37:41.1299481Z",
              "livy_statement_state": "available",
              "queued_time": "2022-04-06T08:35:56.0110118Z",
              "session_id": 195,
              "session_start_time": null,
              "spark_pool": "small",
              "state": "finished",
              "statement_id": 9
            },
            "text/plain": [
              "StatementMeta(small, 195, 9, Finished, Available)"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\rDownloading:   0%|          | 0.00/455k [00:00<?, ?B/s]\rDownloading: 100%|##########| 455k/455k [00:00<00:00, 43.4MB/s]"
          ]
        }
      ],
      "source": [
        "btokenizer = AutoTokenizer.from_pretrained('bert-base-uncased')\n",
        "btokenizer.add_special_tokens({'bos_token':'[DEC]'})\n",
        "btokenizer.add_special_tokens({'additional_special_tokens':['[ENC]']})       \n",
        "btokenizer.enc_token_id = btokenizer.additional_special_tokens_ids[0]\n",
        "\n",
        "udf_decode = udf(decode, StringType())\n",
        "df = df_coords.withColumn(\"text\", udf_decode(df_coords.captions)) \n",
        "\n",
        "df = df.drop(df.captions)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 28,
      "metadata": {
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      },
      "outputs": [
        {
          "data": {
            "application/vnd.livy.statement-meta+json": {
              "execution_finish_time": "2022-04-06T08:48:27.822094Z",
              "execution_start_time": "2022-04-06T08:37:41.7610214Z",
              "livy_statement_state": "available",
              "queued_time": "2022-04-06T08:35:56.0621485Z",
              "session_id": 195,
              "session_start_time": null,
              "spark_pool": "small",
              "state": "finished",
              "statement_id": 10
            },
            "text/plain": [
              "StatementMeta(small, 195, 10, Finished, Available)"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "if LOW_MEMORY_MODE:\n",
        "    df.write.mode('overwrite').options(header='true').csv(os.path.join(output_directory, output_filename))\n",
        "else:\n",
        "    df.coalesce(1).write.mode('overwrite').options(header='true').csv(os.path.join(output_directory, output_filename))\n",
        "\n",
        "mlflow.pyspark.ml.mlflow.end_run()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Add Azure Cognitive Search"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Add Search Parameters"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Azure Search Admin Key\n",
        "search_admin_key = \"\"\n",
        "# The name of the search service\n",
        "search_service_name = \"\"\n",
        "# The Azure Search Query Key\n",
        "search_query_key = \"\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from synapse.ml.cognitive import *\n",
        "from pyspark.sql.functions import monotonically_increasing_id, lit\n",
        "\n",
        "df = df.drop(\"_c0\")\n",
        "\n",
        "(\n",
        "    df.withColumn(\"key\", monotonically_increasing_id().cast(\"string\"))\n",
        "    .withColumn(\"SearchAction\", lit(\"upload\"))\n",
        "    .writeToAzureSearch(\n",
        "        subscriptionKey=search_admin_key,\n",
        "        actionCol=\"SearchAction\",\n",
        "        serviceName=search_service_name,\n",
        "        indexName=experiment_name,  # Defaults to the notebook name\n",
        "        keyCol=\"key\",\n",
        "    )\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Search the generated Azure Search Index"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import requests\n",
        "from IPython.display import display, Markdown\n",
        "\n",
        "term_to_search_for = \"covid\"\n",
        "\n",
        "url = \"https://{}.search.windows.net/indexes/{}/docs/search?api-version=2019-05-06\".format(\n",
        "    search_service_name, experiment_name\n",
        "\n",
        ")\n",
        "jdata = requests.post(url, json={\"search\": term_to_search_for}, headers={\"api-key\": search_query_key}).json()\n",
        "\n",
        "for doc in jdata['value']:\n",
        "    display(Markdown(f'**Search Score {doc[\"@search.score\"]}** Document {doc[\"filename\"]}'))\n",
        "    display(Markdown(f'{doc[\"text\"]}'))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Implement Semantic Search"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "1) [Enable Semantic Search](https://docs.microsoft.com/en-us/azure/search/semantic-search-overview#enable-semantic-search) on your search instance\n",
        "\n",
        "2) [Configure Semantic Search](https://docs.microsoft.com/en-us/azure/search/semantic-how-to-query-request?tabs=semanticConfiguration%2Cportal#create-a-semantic-configuration)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "term_to_search_for = \"Whose thumb was fractured?\"\n",
        "\n",
        "url = \"https://{}.search.windows.net/indexes/{}/docs/search?api-version=2021-04-30-Preview\".format(\n",
        "    search_service_name, experiment_name\n",
        ")\n",
        "jdata = requests.post(url, json={\"search\": term_to_search_for, \"queryType\": \"semantic\", \"semanticConfiguration\": \"config\", \"queryLanguage\": \"en-us\", \"answers\": \"extractive|count-3\",\n",
        "\"captions\": \"extractive|highlight-true\",  \"highlightPreTag\": \"<mark>\",\"highlightPostTag\": \"</mark>\"}, headers={\"api-key\": search_query_key}).json()\n",
        "\n",
        "for doc in jdata['value']:\n",
        "    display(Markdown(f'**Search Score {doc[\"@search.score\"]}** **Search rerankerScore Score {doc[\"@search.rerankerScore\"]}** Document {doc[\"filename\"]}'))\n",
        "    display(Markdown(f'@search.captions {doc[\"@search.captions\"]}'))\n"
      ]
    }
  ],
  "metadata": {
    "description": null,
    "kernel_info": {
      "name": "synapse_pyspark"
    },
    "kernelspec": {
      "display_name": "Python 3.8.2 64-bit",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "name": "python",
      "version": "3.8.2"
    },
    "save_output": true,
    "synapse_widget": {
      "state": {},
      "version": "0.1"
    },
    "vscode": {
      "interpreter": {
        "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 2
}