clean

Image_gen.ipynb

@@ -1,301 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "markdown",
-   "id": "399f2fcf-9241-4910-a30d-6ca19880d0ad",
-   "metadata": {},
-   "source": [
-    "## Import"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 15,
-   "id": "97e68340-0096-475e-8ed8-22f5d627e3ad",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import torch\n",
-    "import numpy as np\n",
-    "from fairseq import utils, tasks\n",
-    "from fairseq import checkpoint_utils\n",
-    "from utils.eval_utils import eval_step\n",
-    "from tasks.mm_tasks import ImageGenTask\n",
-    "from models.unival import UnIVALModel\n",
-    "from PIL import Image\n",
-    "from torchvision import transforms\n",
-    "import time\n",
-    "\n",
-    "\n",
-    "# turn on cuda if GPU is available\n",
-    "use_cuda = torch.cuda.is_available()\n",
-    "# use fp16 only when GPU is available\n",
-    "use_fp16 = True if use_cuda else False"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 16,
-   "id": "719cef65-c00c-4c9c-90b2-e660b386c3d5",
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "<function fairseq.tasks.register_task.<locals>.register_task_cls(cls)>"
-      ]
-     },
-     "execution_count": 16,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "# Register caption task\n",
-    "tasks.register_task('image_gen', ImageGenTask)\n"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "cc9c1d7b-898b-4ac4-adf3-832891d9e4be",
-   "metadata": {},
-   "source": [
-    "### Load model "
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 12,
-   "id": "568bb6ea-eef9-4024-98e6-35e74b5ffeec",
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "self.sample_patch_num 784\n",
-      "self.sample_audio_patch_num None\n",
-      "self.sample_video_patch_num None\n",
-      "self.with_cls False\n",
-      "Frozen image bn <class 'models.ofa.frozen_bn.FrozenBatchNorm2d'>\n",
-      "Loading:  all_resnext101\n",
-      "use bn: <class 'torch.nn.modules.batchnorm.BatchNorm3d'>\n",
-      "load pretrained_model /data/mshukor/logs/ofa/best_models/resnext-101-kinetics.pth\n",
-      "_IncompatibleKeys(missing_keys=[], unexpected_keys=['fc.weight', 'fc.bias'])\n",
-      "load resnet /data/mshukor/logs/ofa/best_models/resnet101-5d3b4d8f.pth\n",
-      "<All keys matched successfully>\n",
-      "RAM memory % used: 10.5\n",
-      "RAM Used (GB): 19.574349824\n",
-      "encoder\n",
-      "RAM memory % used: 10.5\n",
-      "decoder\n",
-      "RAM memory % used: 10.5\n",
-      "ofa\n",
-      "Working with z of shape (1, 256, 32, 32) = 262144 dimensions.\n"
-     ]
-    }
-   ],
-   "source": [
-    "# Load pretrained ckpt & config\n",
-    "clip_model_path='/data/mshukor/data/ofa/clip/ViT-B-16.pt'\n",
-    "vqgan_model_path='/data/mshukor/data/ofa/vqgan/last.ckpt'\n",
-    "vqgan_config_path='/data/mshukor/data/ofa/vqgan/model.yaml'\n",
-    "\n",
-    "# checkpoint_path = '/data/mshukor/logs/ofa/best_models/image_gen_ofa_stage_1_base_s2_hsep1_long/checkpoint_best.pt'\n",
-    "# checkpoint_path = '/data/mshukor/logs/ofa/best_models/image_gen_ofaplus_stage_1_base_s2_long/checkpoint_best.pt'\n",
-    "# checkpoint_path = '/data/mshukor/logs/ofa/best_models/image_gen_base_best.pt'\n",
-    "# checkpoint_path = '/data/mshukor/logs/ofa/best_models/image_gen_large_best.pt'\n",
-    "\n",
-    "# checkpoint_path = '/data/mshukor/logs/ofa/best_models/image_gen_ofaplus_stage_1_base_s2_hsep1_long/checkpoint_best.pt'\n",
-    "checkpoint_path = '/data/mshukor/logs/ofa/best_models/image_gen_ofaplus_stage_2_base_s2_hsep1_long/checkpoint_best.pt'\n",
-    "\n",
-    "\n",
-    "\n",
-    "video_model_path = '/data/mshukor/logs/ofa/best_models/resnext-101-kinetics.pth'\n",
-    "resnet_model_path = '/data/mshukor/logs/ofa/best_models/resnet101-5d3b4d8f.pth'\n",
-    "\n",
-    "gen_images_path='results/image_gen/'\n",
-    "\n",
-    "overrides = {\"bpe_dir\": \"utils/BPE\",\n",
-    "             \"eval_cider\": False,\n",
-    "             \"beam\": 24,\n",
-    "             \"max_len_b\": 1024,\n",
-    "             \"max_len_a\": 0,\n",
-    "             \"min_len\": 1024,\n",
-    "             \"sampling_topk\": 256,\n",
-    "             \"constraint_range\": \"50265,58457\",\n",
-    "             \"clip_model_path\": clip_model_path,\n",
-    "             \"vqgan_model_path\": vqgan_model_path,\n",
-    "             \"vqgan_config_path\": vqgan_config_path,\n",
-    "             \"seed\": 42,\n",
-    "             \"video_model_path\": video_model_path, \n",
-    "             \"resnet_model_path\": resnet_model_path,\n",
-    "             \"gen_images_path\":gen_images_path,\n",
-    "             \"patch_image_size\": 256,\n",
-    "             \"temperature\": 1.5,\n",
-    "            }\n",
-    "\n",
-    "models, cfg, task = checkpoint_utils.load_model_ensemble_and_task(\n",
-    "    utils.split_paths(checkpoint_path),\n",
-    "    arg_overrides=overrides\n",
-    ")\n",
-    "\n",
-    "task.cfg.sampling_times = 2\n",
-    "# Move models to GPU\n",
-    "for model in models:\n",
-    "    model.eval()\n",
-    "    if use_fp16:\n",
-    "        model.half()\n",
-    "    if use_cuda and not cfg.distributed_training.pipeline_model_parallel:\n",
-    "        model.cuda()\n",
-    "    model.prepare_for_inference_(cfg)\n",
-    "\n",
-    "# Initialize generator\n",
-    "generator = task.build_generator(models, cfg.generation)\n",
-    "\n",
-    "# Text preprocess\n",
-    "bos_item = torch.LongTensor([task.src_dict.bos()])\n",
-    "eos_item = torch.LongTensor([task.src_dict.eos()])\n",
-    "pad_idx = task.src_dict.pad()"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "5e4a45ec-bce1-495b-8033-3b574367b360",
-   "metadata": {},
-   "source": [
-    "### Preprocess"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 13,
-   "id": "9f2e7e32-c9a0-43b3-bf86-2419d9f7dfe0",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def encode_text(text, length=None, append_bos=False, append_eos=False):\n",
-    "    s = task.tgt_dict.encode_line(\n",
-    "        line=task.bpe.encode(text),\n",
-    "        add_if_not_exist=False,\n",
-    "        append_eos=False\n",
-    "    ).long()\n",
-    "    if length is not None:\n",
-    "        s = s[:length]\n",
-    "    if append_bos:\n",
-    "        s = torch.cat([bos_item, s])\n",
-    "    if append_eos:\n",
-    "        s = torch.cat([s, eos_item])\n",
-    "    return s\n",
-    "\n",
-    "\n",
-    "# Construct input for image generation task\n",
-    "def construct_sample(query: str):\n",
-    "    code_mask = torch.tensor([True])\n",
-    "    src_text = encode_text(\" what is the complete image? caption: {}\".format(query), append_bos=True,\n",
-    "                           append_eos=True).unsqueeze(0)\n",
-    "    src_length = torch.LongTensor([s.ne(pad_idx).long().sum() for s in src_text])\n",
-    "    sample = {\n",
-    "        \"id\": np.array(['42']),\n",
-    "        \"net_input\": {\n",
-    "            \"src_tokens\": src_text,\n",
-    "            \"src_lengths\": src_length,\n",
-    "            \"code_masks\": code_mask\n",
-    "        }\n",
-    "    }\n",
-    "    return sample\n",
-    "\n",
-    "\n",
-    "# Function to turn FP32 to FP16\n",
-    "def apply_half(t):\n",
-    "    if t.dtype is torch.float32:\n",
-    "        return t.to(dtype=torch.half)\n",
-    "    return t\n",
-    "\n",
-    "\n",
-    "# Function for image generation\n",
-    "def image_generation(caption):\n",
-    "    sample = construct_sample(caption)\n",
-    "    sample = utils.move_to_cuda(sample) if use_cuda else sample\n",
-    "    sample = utils.apply_to_sample(apply_half, sample) if use_fp16 else sample\n",
-    "    print('|Start|', time.strftime(\"%Y-%m-%d %H:%M:%S\", time.localtime()), caption)\n",
-    "    with torch.no_grad():\n",
-    "        result, scores = eval_step(task, generator, models, sample)\n",
-    "\n",
-    "    # return top-4 results (ranked by clip)\n",
-    "    images = [result[i]['image'] for i in range(4)]\n",
-    "    pic_size = 256\n",
-    "    retImage = Image.new('RGB', (pic_size * 2, pic_size * 2))\n",
-    "    print('|FINISHED|', time.strftime(\"%Y-%m-%d %H:%M:%S\", time.localtime()), caption)\n",
-    "    for i in range(4):\n",
-    "        loc = ((i % 2) * pic_size, int(i / 2) * pic_size)\n",
-    "        retImage.paste(images[i], loc)\n",
-    "    return retImage"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "44dec799-c5c2-4d22-8b08-7a7ca2cdf3c9",
-   "metadata": {},
-   "source": [
-    "### Inference"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 14,
-   "id": "02d5cd7a-8d63-4fa4-9da1-d4b79ec01445",
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "|Start| 2023-06-29 12:57:39 A brown horse in the street\n",
-      "|FINISHED| 2023-06-29 12:59:03 A brown horse in the street\n"
-     ]
-    }
-   ],
-   "source": [
-    "query = \"A brown horse in the street\"\n",
-    "# query = \"Cattle grazing on grass near a lake surrounded by mountain.\"\n",
-    "# query = 'A street scene with a double-decker bus on the road.'\n",
-    "# query = 'A path.'\n",
-    "\n",
-    "\n",
-    "retImage = image_generation(query)\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "1a8a1654-1f17-41c7-b410-c7491a96dcee",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "retImage.save(f'{query}.png')"
-   ]
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "ofa",
-   "language": "python",
-   "name": "ofa"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.7.4"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 5
-}

README_EncouragingLoss.md

@@ -1,34 +0,0 @@
-# Finetuning with Encouraging Loss (EL)
-Below we provide methods for finetuning with label smoothed encouraging loss proposed in [_Well-classified Examples are Underestimated in Classification with Deep Neural Networks_](https://arxiv.org/pdf/2110.06537.pdf) on different downstream tasks.
-The implementation is in [label_smoothed_encouraging_loss.py](criterions/label_smoothed_encouraging_loss.py).
-You can set the `--criterion` to `adjust_label_smoothed_encouraging_loss` to use it. This criterion has a hyper-parameter `--log-end`. 
-`--log-end < 1` results in a approximated and conservative version of the full encouraging loss. 
-A high log_end will more strongly weaken the gradient vanishing, enhance the modeling of the data, and increase the growth rate of the margin, but it will also bring a larger gradient norm, which will bring challenges to the existing optimization system.
-We recommend higher log_end for cases with higher performance, and 0.75 or 0.5 as your first try.
-## Image Captioning
-We provide procedures for image captioning with EL below. The preprocessing is identical to default setting.
-
-<details>
-    <summary><b>Finetuning</b></summary>
-    <p>
-        We propose two scripts for stage1. </b>
-    </p>
-<pre>
-cd run_scripts/caption
-nohup sh train_caption_stage1_el.sh > train_stage1_el.out &  # stage 1, train with encouraging loss, expected cider 1.403
-nohup sh train_caption_stage1_el_db.sh > train_stage1_el.out &  # stage 1, train with encouraging loss, and drop best examples, expected cider 1.404
-</pre>
-</details>
-
-## Referring Expression Comprehension 
-We provide procedures for image captioning with EL below. The preprocessing is identical to default setting.
-<details>
-    <summary><b>Finetuning</b></summary>
-<pre>
-cd run_scripts/refcoco
-nohup sh train_refcoco_el.sh > train_refcoco_el.out &  # finetune for refcoco
-nohup sh train_refcocoplus_el.sh > train_refcocoplus_el.out &  # finetune for refcoco+
-nohup sh train_refcocog_el.sh > train_refcocog_el.out &  # finetune for refcocog
-</pre>
-</details>
-Evaluation is also the same as the default setting.

checkpoints.md

@@ -1,36 +0,0 @@
-# Checkpoints
-
-We provide links for you to download our checkpoints, including pretrained and finetuned models on different tasks. If you would like to use OFA with Transformers, please download checkpoints at [https://huggingface.co/OFA-Sys](https://huggingface.co/OFA-Sys), and check the code in the branch `feature/add_transformers`. 
-
-## Pretraining
-* <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/checkpoints/ofa_huge.pt"> Pre-trained checkpoint (OFA-Huge) </a> (~930M parameters)
-* <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/checkpoints/ofa_large.pt"> Pre-trained checkpoint (OFA-Large) </a> (~470M parameters)
-* <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/checkpoints/ofa_base.pt"> Pre-trained checkpoint (OFA-Base) </a> (~180M parameters)
-* <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/checkpoints/ofa_medium.pt"> Pre-trained checkpoint (OFA-Medium) </a> (~93M parameters)
-* <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/checkpoints/ofa_tiny.pt"> Pre-trained checkpoint (OFA-Tiny) </a> (~33M parameters)
-
-## Finetuning (OFA-Huge)
-* <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/checkpoints/caption_huge_best.pt"> Finetuned checkpoint for Caption on COCO </a>
-
-## Finetuning (OFA-Large)
-
-* <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/checkpoints/caption_large_best_clean.pt"> Finetuned checkpoint for Caption on COCO </a>
-* <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/checkpoints/caption_stage1_best.pt"> Finetuned checkpoint for Caption on COCO During Stage1 Finetuning </a>
-* <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/checkpoints/refcoco_large_best.pt"> Finetuned checkpoint for RefCOCO </a>
-* <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/checkpoints/refcocoplus_large_best.pt"> Finetuned checkpoint for RefCOCO+ </a>
-* <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/checkpoints/refcocog_large_best.pt"> Finetuned checkpoint for RefCOCOg </a>
-* <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/checkpoints/vqa_large_best.pt"> Finetuned checkpoint for VQAv2 </a>
-* <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/checkpoints/snli_ve_large_best.pt"> Finetuned checkpoint for SNLI-VE </a>
-* <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/checkpoints/image_gen_large_best.zip"> Finetuned checkpoint for Text-to-Image Generation on COCO && CLIP checkpoint && VQGAN checkpoint </a>
-* <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/checkpoints/imagenet_1k_large_best.pt"> Finetuned checkpoint for ImageNet-1K </a>
-* <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/checkpoints/gigaword_large_best.pt"> Finetuned checkpoint for Gigaword </a>
-
-
-## Finetuning (OFA-Base)
-* <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/checkpoints/caption_base_best.pt"> Finetuned base checkpoint for Caption on COCO </a>
-* <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/checkpoints/refcoco_base_best.pt"> Finetuned base checkpoint for RefCOCO </a>
-* <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/checkpoints/refcocoplus_base_best.pt"> Finetuned base checkpoint for RefCOCO+ </a>
-* <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/checkpoints/refcocog_base_best.pt"> Finetuned base checkpoint for RefCOCOg </a>
-* <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/checkpoints/vqa_base_best.pt"> Finetuned base checkpoint for VQAv2 </a>
-* <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/checkpoints/snli_ve_base_best.pt"> Finetuned base checkpoint for SNLI-VE </a>
-* <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/checkpoints/image_gen_base_best.pt"> Finetuned base checkpoint for Text-to-Image Generation on COCO  </a>

checkpoints_cn.md

@@ -1,82 +0,0 @@
-# Checkpoints (OFA-CN)
-
-We provide checkpoints of OFA-CN, which is the Chinese version of OFA. We provide Base-size and Large-size models, including pretrained and finetuned models on image captioning and referring expression comprehension. Note that we translated the texts in the RefCOCO(-/+/g) datasets and finetuned OFA-CN on them. We plan to release the related new datasets in the near future. 
-<br>
-
-## Checkpoints
-Below we provide the links for downloading the Chinese OFA checkpoints.
-
-### Pretraining
-* <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/checkpoints/ofa_cn_large.pt"> Pretrained checkpoint (OFA-CN-Large) </a> (~443M parameters)
-* <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/checkpoints/ofa_cn_base.pt "> Pretrained checkpoint (OFA-CN-Base) </a> (~160M parameters)
-
-### Finetuning (OFA-Large)
-* <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/checkpoints/caption_cn_large.pt"> Finetuned checkpoint for MUGE Caption (Stage 1) </a>
-* <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/checkpoints/refcoco_cn_large.pt"> Finetuned checkpoint for RefCOCO-CN </a>
-* <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/checkpoints/refcocoplus_cn_large.pt"> Finetuned checkpoint for RefCOCO+-CN </a>
-* <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/checkpoints/refcocog_cn_large.pt"> Finetuned checkpoint for RefCOCOg-CN </a>
-
-### Finetuning (OFA-Base)
-* <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/checkpoints/caption_cn_base.pt"> Finetuned checkpoint for MUGE Caption (Stage 1) </a>
-* <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/checkpoints/refcoco_cn_base.pt"> Finetuned checkpoint for RefCOCO-CN </a>
-* <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/checkpoints/refcocoplus_cn_base.pt"> Finetuned checkpoint for RefCOCO+-CN </a>
-* <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/checkpoints/refcocog_cn_base.pt"> Finetuned checkpoint for RefCOCOg-CN </a>
-<br>
-
-## Model Card
-Below we provide the basic information of the base-size and large-size OFA-CN. 
-
-<table border="1" width="100%">
-    <tr align="center">
-        <th>Model</th><th>#Params</th><th>Backbone</th><th>Hidden Size</th><th>Intermediate Size</th><th>#Heads</th><th>#Enc. Layers</th><th>#Dec. Layers</th>
-    </tr>
-    <tr align="center">
-        <td>OFA<sub>Base</sub><td>160M</td><td>ResNet101</td><td>768</td></td><td>3072</td><td>12</td><td>6</td><td>6</td>
-    </tr>
-    <tr align="center">
-        <td>OFA<sub>Large</sub></td><td>443M</td><td>ResNet152</td><td>1024</td></td><td>4096</td><td>16</td><td>12</td><td>12</td>
-    </tr>
-    </tr>
-</table>
-<br>
-
-## Results
-Below we provide the results of OFA-CN and the baselines for comparison. 
-
-### [MUGE Caption]("https://tianchi.aliyun.com/muge")
-<table border="1" width="100%">
-    <tr align="center">
-        <td>Model</td><td>BLEU@4</td><td>ROUGE-L</td><td>CIDEr-D</td>
-    </tr>
-    <tr align="center">
-        <td>Trm </td><td>7.33</td><td>51.51</td><td>11.00</td>
-    </tr>
-    <tr align="center">
-        <td>M6</td><td>16.19</td><td>55.06</td><td>30.75</td>
-    </tr>
-    <tr align="center">
-        <td>OFA<sub>Base</sub></td><td>26.23</td><td>58.95</td><td>50.70</td>
-    </tr>
-    <tr align="center">
-        <td>OFA<sub>Large</sub></td><td><b>27.32</b></td><td><b>59.20</b></td><td><b>53.51</b></td>
-    </tr>
-</table>
-
-### RefCOCO-CN Series
-<table border="1" width="100%">
-    <tr align="center">
-        <td>Model</td><td>RefCOCO(val/testA/testB)</td><td>RefCOCO+(val/testA/testB)</td><td>RefCOCOg(val/test-u)</td>
-    </tr>
-    <tr align="center">
-        <td>OFA<sub>Base</sub>(random-init)</td><td>30.13/35.07/25.03</td><td>17.89/20.90/15.83</td><td>20.30/20.45</td>
-    </tr>
-    <tr align="center">
-        <td>OFA<sub>Base</sub></td><td>82.18/86.07/<b>76.68</b></td><td>69.38/77.26/60.14</td><td><b>73.57/72.53</b></td>
-    </tr>
-    <tr align="center">
-        <td>OFA<sub>Large</sub></td><td><b>82.84/86.54</b>/76.50</td><td><b>71.30/78.56/61.85</b></td><td>71.96/71.30</td>
-    </tr>
-</table>
-<br>
-
-

colab.md

@@ -1,9 +0,0 @@
-# Colab Notebooks
-
-We provide Colab notebooks of different downstream tasks for you guys to enjoy OFA. See below.
-
-* [Image Captioning in Huggingface Transformers](https://colab.research.google.com/drive/1Ho81RBV8jysZ7e0FhsSCk_v938QeDuy3?usp=sharing)
-* [Generic Interface](https://colab.research.google.com/drive/1jogyZ-2rdHU3XxZOf3TBfhex1XHqX-1m?usp=sharing#scrollTo=s9Vni6YUZOpC) (using different instructions to perform various tasks with just one model.)
-* [Image Captioning](https://colab.research.google.com/drive/1Q4eNhhhLcgOP4hHqwZwU1ijOlabgve1W?usp=sharing)
-* [Referring Expression Comprehension](https://colab.research.google.com/drive/1AHQNRdaUpRTgr3XySHSlba8aXwBAjwPB?usp=sharing)
-* [Open-Domain Visual Question Answering](https://colab.research.google.com/drive/14v6OQe_MxV_HMnsiKfnEeMR1UMqhzZNb?usp=sharing)

datasets.md

@@ -1,44 +0,0 @@
-# Datasets
-
-We provide links to download our preprocessed dataset. If you would like to process the data on your own, we will soon provide scripts for you to do so. 
-
-## Pretraining
- * <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/datasets/pretrain_data/pretrain_data_examples.zip"> A small subset of the pretraining data </a>
-
- The pretraining datasets used in OFA are all publicly available. Here we provide the public links to these data, it is recommended that you download the data from the links first, and then process the downloaded dataset into a similar format as the examples we provided.
--   _CC12M_:  https://github.com/google-research-datasets/conceptual-12m
--   _CC3M_: https://github.com/google-research-datasets/conceptual-captions
--   _SBU_: https://www.cs.virginia.edu/~vicente/sbucaptions
--   _COCO_: https://cocodataset.org/#home
--   _VG_: https://visualgenome.org/
--   _VQAv2_: https://visualqa.org/
-- _GQA_: https://cs.stanford.edu/people/dorarad/gqa/about.html
-- _RefCOCO_/_RefCOCO+_/RefCOCOg: https://github.com/lichengunc/refer
--   _OpenImages_: https://storage.googleapis.com/openimages/web/index.html
--   _Object365_: https://www.objects365.org/overview.html
--   _YFCC100M (subset)_: https://github.com/openai/CLIP/blob/main/data/yfcc100m.md
--   _ImageNet-21K_: https://image-net.org/index.php
--   _Pile_: https://pile.eleuther.ai
-
-## Vision & Language Tasks
- * <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/datasets/caption_data/caption_data.zip"> Dataset for Caption </a>
- * <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/datasets/refcoco_data/refcoco_data.zip"> Dataset for RefCOCO </a>
- * <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/datasets/refcocoplus_data/refcocoplus_data.zip"> Dataset for RefCOCO+ </a>
- * <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/datasets/refcocog_data/refcocog_data.zip"> Dataset for RefCOCOg </a>
- * <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/datasets/vqa_data/vqa_data.zip"> Dataset for VQAv2 </a> (we have also provided chunked parts of the dataset files for more convenient downloading, please refer to <a href="https://github.com/OFA-Sys/OFA/issues/68#issuecomment-1096837349">issue #68</a>)
- * <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/datasets/snli_ve_data/snli_ve_data.zip"> Dataset for SNLI-VE </a>
- * <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/datasets/coco_image_gen_data/coco_image_gen.zip"> Dataset for Text-to-Image Genearion </a>
- * <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/datasets/coco_image_gen_data/coco_image_gen_origin_id.zip"> Dataset for Text-to-Image Genearion (with original id) </a>
-
-## Vision Tasks
- * <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/datasets/imagenet_1k_data/imagenet_1k_data.zip"> Dataset for ImageNet-1K </a>
-
-## Language Tasks
- * <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/datasets/glue_data/cola_data.zip"> Dataset for COLA </a>
- * <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/datasets/glue_data/mnli_data.zip"> Dataset for MNLI </a>
- * <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/datasets/glue_data/mrpc_data.zip"> Dataset for MRPC </a>
- * <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/datasets/glue_data/qnli_data.zip"> Dataset for QNLI </a>
- * <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/datasets/glue_data/qqp_data.zip"> Dataset for QQP </a>
- * <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/datasets/glue_data/rte_data.zip"> Dataset for RTE </a>
- * <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/datasets/glue_data/sst2_data.zip"> Dataset for SST2 </a>
- * <a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/datasets/gigaword_data/gigaword_data.zip"> Dataset for Gigaword </a>

evaluate.py

@@ -1,239 +0,0 @@
-#!/usr/bin/env python3 -u
-# Copyright 2022 The OFA-Sys Team. 
-# All rights reserved.
-# This source code is licensed under the Apache 2.0 license 
-# found in the LICENSE file in the root directory.
-
-import logging
-import os
-import sys
-
-import numpy as np
-import torch
-from fairseq import distributed_utils, options, tasks, utils
-from fairseq.dataclass.utils import convert_namespace_to_omegaconf
-from fairseq.logging import progress_bar
-from fairseq.utils import reset_logging
-from omegaconf import DictConfig
-
-from utils import checkpoint_utils
-from utils.eval_utils import eval_step, merge_results
-from utils.zero_shot_utils import zero_shot_step
-
-logging.basicConfig(
-    format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
-    datefmt="%Y-%m-%d %H:%M:%S",
-    level=os.environ.get("LOGLEVEL", "INFO").upper(),
-    stream=sys.stdout,
-)
-logger = logging.getLogger("ofa.evaluate")
-
-from utils.utils import print_trainable_params_percentage, setup_for_distributed
-
-def apply_half(t):
-    if t.dtype is torch.float32:
-        return t.to(dtype=torch.half)
-    return t
-
-
-def main(cfg: DictConfig, **kwargs):
-    utils.import_user_module(cfg.common)
-
-    setup_for_distributed(distributed_utils.is_master(cfg.distributed_training))
-
-    reset_logging()
-    # logger.info(cfg)
-
-    assert (
-            cfg.dataset.max_tokens is not None or cfg.dataset.batch_size is not None
-    ), "Must specify batch size either with --max-tokens or --batch-size"
-
-    # Fix seed for stochastic decoding
-    if cfg.common.seed is not None and not cfg.generation.no_seed_provided:
-        np.random.seed(cfg.common.seed)
-        utils.set_torch_seed(cfg.common.seed)
-
-    use_fp16 = cfg.common.fp16
-    use_cuda = torch.cuda.is_available() and not cfg.common.cpu
-
-    if use_cuda:
-        torch.cuda.set_device(cfg.distributed_training.device_id)
-
-    # Load ensemble
-    overrides = eval(cfg.common_eval.model_overrides)
-    # Deal with beam-search / all-candidate VQA eval
-    if cfg.task._name == "vqa_gen":
-        overrides['val_inference_type'] = "beamsearch" if kwargs['beam_search_vqa_eval'] else "allcand"
-
-    logger.info("loading model(s) from {}".format(cfg.common_eval.path))
-
-    # print("cfg", cfg)
-    # print(kwargs)
-    # cfg.model.num_frames = kwargs["num_frames"]
-    # cfg.model.patch_frame_size = kwargs["patch_frame_size"]
-    # print("cfg.model", cfg.model)
-    # strict = getattr(kwargs, 'strict', True)
-    strict = kwargs['strict']
-    logger.info('load checkpoint, strict:{}'.format(strict))
-
-    if kwargs["zero_shot"]:
-        for arg_name, arg_val in overrides.items():
-            cfg.task[arg_name] = arg_val
-        # print("Zero-shot eval", cfg.task, cfg)
-
-        if hasattr(cfg.task, "add_caption"):
-            cfg.task.add_caption = False
-        print("cfg.task", cfg.task)
-        task = tasks.setup_task(cfg.task)
-        # cfg.criterion.sample_patch_num = 776
-            
-        
-        models, saved_cfg = checkpoint_utils.load_model_ensemble(
-            utils.split_paths(cfg.common_eval.path),
-            arg_overrides=overrides,
-            task=task,
-            suffix=cfg.checkpoint.checkpoint_suffix,
-            strict=((cfg.checkpoint.checkpoint_shard_count == 1) and strict),
-            num_shards=cfg.checkpoint.checkpoint_shard_count,
-        )
-        for m in models:
-            m.encoder.sample_patch_num = 776
-        saved_cfg.task = cfg.task
-        # print("saved_cfg", saved_cfg)
-    else:
-        models, saved_cfg, task = checkpoint_utils.load_model_ensemble_and_task(
-            utils.split_paths(cfg.common_eval.path),
-            arg_overrides=overrides,
-            suffix=cfg.checkpoint.checkpoint_suffix,
-            strict=((cfg.checkpoint.checkpoint_shard_count == 1) and strict),
-            num_shards=cfg.checkpoint.checkpoint_shard_count,
-        )
-
-
-
-    # task.cfg['evaluate_cfg'] = cfg.task
-    # print(task.cfg)
-    kwargs['evaluate_cfg'] = cfg.task
-    # print(kwargs)
-    # loading the dataset should happen after the checkpoint has been loaded so we can give it the saved task config
-    task.load_dataset(cfg.dataset.gen_subset, task_cfg=saved_cfg.task)
-
-    # Move models to GPU
-    for model, ckpt_path in zip(models, utils.split_paths(cfg.common_eval.path)):
-        if kwargs['ema_eval']:
-            logger.info("loading EMA weights from {}".format(ckpt_path))
-            model.load_state_dict(checkpoint_utils.load_ema_from_checkpoint(ckpt_path)['model'])
-        model.eval()
-        print("use fp16", use_fp16)
-        if use_fp16:
-            
-            model.half()
-        if use_cuda and not cfg.distributed_training.pipeline_model_parallel:
-            model.cuda()
-        model.prepare_for_inference_(cfg)
-
-    # Load dataset (possibly sharded)
-    itr = task.get_batch_iterator(
-        dataset=task.dataset(cfg.dataset.gen_subset),
-        max_tokens=cfg.dataset.max_tokens,
-        max_sentences=cfg.dataset.batch_size,
-        max_positions=utils.resolve_max_positions(
-            task.max_positions(), *[m.max_positions() for m in models]
-        ),
-        ignore_invalid_inputs=cfg.dataset.skip_invalid_size_inputs_valid_test,
-        required_batch_size_multiple=cfg.dataset.required_batch_size_multiple,
-        seed=cfg.common.seed,
-        num_shards=cfg.distributed_training.distributed_world_size,
-        shard_id=cfg.distributed_training.distributed_rank,
-        num_workers=cfg.dataset.num_workers,
-        data_buffer_size=cfg.dataset.data_buffer_size,
-    ).next_epoch_itr(shuffle=False)
-    progress = progress_bar.progress_bar(
-        itr,
-        log_format=cfg.common.log_format,
-        log_interval=cfg.common.log_interval,
-        default_log_format=("tqdm" if not cfg.common.no_progress_bar else "simple"),
-    )
-
-    # Initialize generator
-    generator = task.build_generator(models, cfg.generation)
-
-    results = []
-    score_sum = torch.FloatTensor([0]).cuda()
-    score_cnt = torch.FloatTensor([0]).cuda()
-
-    score_sum_list = []
-    score_cnt_list = []
-    for sample in progress:
-        if "net_input" not in sample:
-            continue
-        sample = utils.move_to_cuda(sample) if use_cuda else sample
-        sample = utils.apply_to_sample(apply_half, sample) if cfg.common.fp16 else sample
-        with torch.no_grad():
-            if kwargs["zero_shot"] and kwargs['noconstraints']:
-                result, scores = zero_shot_step(task, generator, models, sample)
-            else:
-                result, scores = eval_step(task, generator, models, sample, **kwargs)
-        ### else refcoco res, score, other_scores
-
-        # print(scores)
-        scalar = False
-        if isinstance(scores, list):
-            if not isinstance(scores[0], list):
-                try:
-                    tmp = sum(scores[0])
-                    scalar=False
-                except:
-                    scalar=True
-        # print(scalar)
-        # print(sum(scores[0]))
-        if isinstance(scores, list) and not scalar:
-            names = result[0]
-            result = result[1]
-            if len(score_sum_list) == 0:
-                score_sum_list = [torch.FloatTensor([0]).cuda() for i in range(len(scores))]
-                score_cnt_list = [torch.FloatTensor([0]).cuda() for i in range(len(scores))]
-                                
-                for i in range(len(scores)):
-                    
-
-                    score_sum_list[i] += sum(scores[i]) if scores[i] is not None else 0
-                    score_cnt_list[i] += len(scores[i]) if scores[i] is not None else 0
-            else:
-                for i in range(len(scores)):
-                    score_sum_list[i] += sum(scores[i]) if scores[i] is not None else 0
-                    score_cnt_list[i] += len(scores[i]) if scores[i] is not None else 0
-        else:
-            score_sum += sum(scores) if scores is not None else 0
-            score_cnt += len(scores) if scores is not None else 0
-        results += result
-        progress.log({"sentences": sample["nsentences"]})
-        
-
-    ### merge per metric 
-    if len(score_sum_list) > 0:
-        print(names, len(score_sum_list))
-        for i in range(len(score_sum_list)):
-            print(names[i])
-            merge_results(task, cfg, logger, score_cnt_list[i], score_sum_list[i], results)
-    else:
-        merge_results(task, cfg, logger, score_cnt, score_sum, results)
-
-
-def cli_main():
-    parser = options.get_generation_parser()
-    parser.add_argument("--ema-eval", action='store_true', help="Use EMA weights to make evaluation.")
-    parser.add_argument("--beam-search-vqa-eval", action='store_true', help="Use beam search for vqa evaluation (faster inference speed but sub-optimal result), if not specified, we compute scores for each answer in the candidate set, which is slower but can obtain best result.")
-    parser.add_argument("--zero-shot", action='store_true')
-    parser.add_argument("--strict", action='store_false')
-    parser.add_argument("--noconstraints", action='store_true')
-    args = options.parse_args_and_arch(parser)
-    cfg = convert_namespace_to_omegaconf(args)
-    distributed_utils.call_main(
-        cfg, main, ema_eval=args.ema_eval, beam_search_vqa_eval=args.beam_search_vqa_eval, 
-        zero_shot=args.zero_shot, strict=args.strict, noconstraints=args.noconstraints
-    )
-
-
-if __name__ == "__main__":
-    cli_main()

modelscope.md

@@ -1,23 +0,0 @@
-# ModelScope
-
-ModelScope is a new platform that provides "Model-As-A-Service", where users can use state-of-the-art models with the lowest costs of efforts as possible. We have released:
-* The pretrained and finetuned **OFA** models
-* **Chinese CLIP** (the CLIP pretrained Chinese data, which was previously released in our organization)
-
-on the platform, including the English and Chinese ones. Feel free to check these models and use them with ModelScope, and also feel free to send us feedbacks to help us improve the product.
-
-## Chinese
-* Chinese CLIP \[[Base](https://www.modelscope.cn/#/models/damo/multi-modal_clip-vit-base-patch16_zh/summary) | [Large](https://www.modelscope.cn/#/models/damo/multi-modal_clip-vit-large-patch14_zh/summary)\]
-* Finetuned OFA on Visual Grounding (RefCOCO) \[[Large](https://www.modelscope.cn/#/models/damo/ofa_visual-grounding_refcoco_large_zh/summary)\]
-
-## English
-* Finetuned OFA on Image Captioning \[[Large](https://www.modelscope.cn/#/models/damo/ofa_image-caption_coco_large_en/summary) | [Distill](https://modelscope.cn/#/models/damo/ofa_image-caption_coco_distilled_en/summary)\]
-* Finetuned OFA on Text-to-Image Generation \[[Large](https://www.modelscope.cn/#/models/damo/ofa_text-to-image-synthesis_coco_large_en/summary) | [Distill](https://modelscope.cn/#/models/damo/ofa_visual-grounding_refcoco_distilled_en/summary)\]
-* Finetuned OFA on Visual Question Answering \[[Large](https://www.modelscope.cn/#/models/damo/ofa_visual-question-answering_pretrain_large_en/summary)\]
-* Finetuned OFA on Visual Grounding (RefCOCO) \[[Large](https://www.modelscope.cn/#/models/damo/ofa_visual-grounding_refcoco_large_en/summary)\]
-* Finetuned OFA on Visual Entailment \[[Large](https://www.modelscope.cn/#/models/damo/ofa_visual-entailment_snli-ve_large_en/summary) | [Distill](https://modelscope.cn/#/models/damo/ofa_visual-entailment_snli-ve_distilled_v2_en/summary)\]
-* Finetuned OFA on Summarization (Gigaword) \[[Large](https://www.modelscope.cn/#/models/damo/ofa_summarization_gigaword_large_en/summary)\]
-* Finetuned OFA on Natural Language Entailment (MNLI, can be used to finetune on the GLUE benchmark) \[[Large](https://modelscope.cn/#/models/damo/ofa_text-classification_mnli_large_en/summary)\]
-* Finetuned OFA on Image Classification (ImageNet-1k) \[[Large](https://modelscope.cn/#/models/damo/ofa_image-classification_imagenet_large_en/summary)\]
-
-

ofa_test.ipynb

@@ -1,2499 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# Import"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 1,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "%load_ext autoreload\n",
-    "%autoreload 2"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 2,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import os\n",
-    "import json \n",
-    "import torch\n",
-    "# import clip\n",
-    "from PIL import Image\n",
-    "# import sng_parser\n",
-    "from tqdm import tqdm \n",
-    "import codecs\n",
-    "import numpy as np\n",
-    "import csv\n",
-    "import sys\n",
-    "\n",
-    "from io import BytesIO\n",
-    "import base64"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# Data"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Explore"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 16,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "100it [00:00, 14325.30it/s]\n"
-     ]
-    }
-   ],
-   "source": [
-    "csv.field_size_limit(sys.maxsize)\n",
-    "\n",
-    "# path_data =  '/data/mshukor/data/ofa/pretrain_example/vision_language_examples.tsv'\n",
-    "# selected_cols='0,1,2,3,4,5,6,7'\n",
-    "\n",
-    "# path_data =  '/data/mshukor/data/ofa/pretrain_example/detection_examples.tsv'\n",
-    "# selected_cols='0,1,2'\n",
-    "\n",
-    "# path_data =  '/data/mshukor/data/ofa/pretrain_example/image_examples.tsv'\n",
-    "# selected_cols='0,1,2'\n",
-    "\n",
-    "path_data =  '/data/mshukor/data/ofa/pretrain_example/text_examples.tsv'\n",
-    "selected_cols='0,1'\n",
-    "\n",
-    "data_example = []\n",
-    "\n",
-    "selected_col_ids = [int(col_id) for col_id in selected_cols.split(\",\")]\n",
-    "\n",
-    "with open(path_data) as file:\n",
-    "    tsv_file = csv.reader(file, delimiter='\\t')\n",
-    "    for line in tqdm(tsv_file):\n",
-    "\n",
-    "        d = [line[i] for i in selected_col_ids]\n",
-    "#         print(d)\n",
-    "        data_example.append(d)\n",
-    "    \n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 17,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "['100',\n",
-       " '...please depart this field clean unless you might be answering the question. do not ask questions you already know the answer to. thanks.retrieved from \" \" ad blocker interference detected! wikia is a single-to-usefulness web site that makes cash from promoting. we\\'ve a experience for viewers using ad blockers wikia shouldn\\'t be if youve made further modifications. take away the custom ad blocker (s) and the page leave timber as expected. categories : un-answered questionsadd class cancelsave per the reddit twine, flac files will be synced to an ios gadget via icloud impel, then accessed through thenew information utility , which will allow for local playback of the excessive-high quality audio files straight by the side of the device. if , it could stamp the first time that apple has offered help for the favored flac format an ios gadget.']"
-      ]
-     },
-     "execution_count": 17,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "line"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 12,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "['7',\n",
-       " 'perhaps the clearest indication of who won and lost came quickly on the heels of the event itself: the democratic post-debate message was that joe biden scored a clear win; the republican message was that joe biden was too mean to paul ryan. the former is a boast of success; the latter is an excuse for failure. in the larger context, it’s hard to overstate how much democrats needed a shot in the arm like this. the surface-level goals of any vice presidential debate is for the candidates to demonstrate a capacity to step up in the event of a crisis, while defending their ticket’s agenda and knocking their rivals’ agenda. but for biden, the overarching benefit was about the basic morale of his party with less than four weeks to go until election day: he wanted to give democratic voters something to feel good about, and he did. who the hell am i! i’m a liberal that is extreme in some ways and not in others. i support president obama and make no apologies for it. i think he has done a phenomenal job, especially when you consider that he inherited a huge mess and has faced unprecedented opposition from a lazy & desperate republican party. i’m a film producer/director/editor, adjunct professor, technician, media critic and photographer when i’m not reading left wing blogs and typing on this one. – on twitter @extremeliberal or email at liberalforreal (at) gmail.com own an important part of american history! cicely tyson narrates this award winning documentary that tells the story of african american migration from the old south to the prosperous north. winner of 5 awards including \"best film\" at the astoria international film festival, the \"paul robeson award\" at the newark black film festival and \"best film relating to the black experience\" at the xxv international black cinema berlin/germany!']"
-      ]
-     },
-     "execution_count": 12,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "data_[6]"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "tasks = set()\n",
-    "datasets = set()\n",
-    "\n",
-    "for d in data:\n",
-    "    tasks.add(d[-1])\n",
-    "    datasets.add(d[-2])"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "int(data[10][0])"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# len(data[0][2:][0].split(' '))\n",
-    "# len(data[0][1])\n",
-    "text = data[10][1]\n",
-    "print(len(text.split(' ')))\n",
-    "print(len(text))\n",
-    "from nltk.tokenize import word_tokenize\n",
-    "len(word_tokenize(text))"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "text"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from nltk.tokenize.treebank import TreebankWordDetokenizer\n",
-    "TreebankWordDetokenizer().detokenize(word_tokenize(text))"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "key = 'refcoco_train'\n",
-    "index = -2\n",
-    "for d in data:\n",
-    "    if d[index] == key:\n",
-    "        print(d[2:])\n",
-    "#         break"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "d[4].split(',')\n",
-    "str([287.0, 127.0, 340.0, 162.0])\n",
-    "'{:.2f},{:.2f},{:.2f},{:.2f}'.format(287.0, 127.0, 340.0, 162.0)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "print(len(data))\n",
-    "data[0]"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "all_captions_path = '/data/mshukor/data/ofa/pretrain_example/negative_sample/all_captions.txt'\n",
-    "all_objects_path = '/data/mshukor/data/ofa/pretrain_example/negative_sample/object.txt'\n",
-    "\n",
-    "all_object_list = [\n",
-    "    row.strip() for row in open(all_objects_path) if row.strip() != ''\n",
-    "]\n",
-    "all_caption_list = [\n",
-    "    row.strip() for row in open(all_captions_path) if row.strip() != ''\n",
-    "]\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "len(all_object_list)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "len(all_caption_list)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "all_object_list[:10]"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "all_caption_list[:10]"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "json_path = '/data/mshukor/data/ofa/pretrain_example/negative_sample/type2ans.json'\n",
-    "type2ans = json.load(open(json_path,'r'))"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "type2ans.keys()\n",
-    "# type2ans['what color is the']"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "### Our data"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 3,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "181767it [00:02, 70482.72it/s]\n"
-     ]
-    }
-   ],
-   "source": [
-    "# path_data =  '/data/mshukor/data/ofa/pretrain_ours/text_mini.tsv'\n",
-    "# selected_cols='0,1'\n",
-    "\n",
-    "path_data =  '/data/mshukor/data/ofa/pretrain_ours/detection_mini.tsv'\n",
-    "selected_cols='0,1,2'\n",
-    "\n",
-    "# path_data =  '/data/mshukor/data/ofa/pretrain_ours/vision_language_mini.tsv'\n",
-    "# selected_cols='0,1,2,3,4,5,6,7'\n",
-    "\n",
-    "# path_data =  '/data/mshukor/data/ofa/pretrain_ours/image_mini.tsv'\n",
-    "# selected_cols='0,1,2'\n",
-    "\n",
-    "data = []\n",
-    "\n",
-    "selected_col_ids = [int(col_id) for col_id in selected_cols.split(\",\")]\n",
-    "\n",
-    "with open(path_data) as file:\n",
-    "    tsv_file = csv.reader(file, delimiter='\\t')\n",
-    "    for line in tqdm(tsv_file):\n",
-    "\n",
-    "        d = [line[i] for i in selected_col_ids]\n",
-    "#         print(d)\n",
-    "        data.append(d)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 21,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# new_data = []\n",
-    "# for d in data:\n",
-    "#     label_list = d[2].strip().split('&&')\n",
-    "#     new_label_list = []\n",
-    "#     for label in label_list:\n",
-    "#         lab = label.strip().split(',', 5)[:4] # x0, y0, x1, y1, cat_id, cat\n",
-    "        \n",
-    "#         if any([\"&\" in l for l in lab]):\n",
-    "#             lab = [remove_special(l) for l in lab]\n",
-    "            \n",
-    "#             print(lab)\n",
-    "#         lab_ = lab + label.strip().split(',', 5)[4:]\n",
-    "#         lab_ = ','.join(lab_)\n",
-    "#         new_label_list.append(lab_)\n",
-    "#     new_label_list = ['&&'.join(new_label_list)]\n",
-    "#     new_data.append(d[:2]+new_label_list)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 24,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "['&40.000', '155.000', '44.000', '164.000']\n"
-     ]
-    }
-   ],
-   "source": [
-    "for d in data:\n",
-    "    label_list = d[2].strip().split('&&')\n",
-    "    new_label_list = []\n",
-    "    for label in label_list:\n",
-    "        lab = label.strip().split(',', 5)[:4] # x0, y0, x1, y1, cat_id, cat\n",
-    "        \n",
-    "        if any([\"&\" in l for l in lab]):\n",
-    "            print(lab)\n",
-    "            # lab = [remove_special(l) for l in lab]"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 27,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "['0',\n",
-       " 'coco/train2014/COCO_train2014_000000057870.jpg',\n",
-       " '1.020,279.960,534.110,480.000,67,dining table&&90.670,271.490,262.510,480.000,62,chair&&233.290,270.450,403.610,473.810,62,chair&&367.820,264.270,506.970,480.000,62,chair&&476.760,261.030,596.490,462.740,62,chair&&263.030,174.370,417.670,299.400,64,potted plant&&539.330,290.160,640.000,469.210,62,chair&&10.790,260.030,125.120,384.070,62,chair&&560.800,413.950,639.090,479.200,67,dining table&&20.540,376.760,103.780,431.890,62,chair&&1.080,373.210,32.360,480.000,62,chair&&298.200,235.170,381.210,269.250,86,vase&&152.170,256.670,230.580,285.780,62,chair&&364.400,256.570,417.060,283.210,62,chair&&296.780,277.790,329.260,289.780,84,book&&292.800,289.310,314.210,300.650,84,book&&285.800,257.460,299.770,273.600,62,chair']"
-      ]
-     },
-     "execution_count": 27,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "new_data[0]"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 8,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "['1.020,279.960,534.110,480.000,67,dining table',\n",
-       " '90.670,271.490,262.510,480.000,62,chair',\n",
-       " '233.290,270.450,403.610,473.810,62,chair',\n",
-       " '367.820,264.270,506.970,480.000,62,chair',\n",
-       " '476.760,261.030,596.490,462.740,62,chair',\n",
-       " '263.030,174.370,417.670,299.400,64,potted plant',\n",
-       " '539.330,290.160,640.000,469.210,62,chair',\n",
-       " '10.790,260.030,125.120,384.070,62,chair',\n",
-       " '560.800,413.950,639.090,479.200,67,dining table',\n",
-       " '20.540,376.760,103.780,431.890,62,chair',\n",
-       " '1.080,373.210,32.360,480.000,62,chair',\n",
-       " '298.200,235.170,381.210,269.250,86,vase',\n",
-       " '152.170,256.670,230.580,285.780,62,chair',\n",
-       " '364.400,256.570,417.060,283.210,62,chair',\n",
-       " '296.780,277.790,329.260,289.780,84,book',\n",
-       " '292.800,289.310,314.210,300.650,84,book',\n",
-       " '285.800,257.460,299.770,273.600,62,chair']"
-      ]
-     },
-     "execution_count": 8,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "label_list"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 12,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def remove_special(input_string):\n",
-    "    final_string = \"\"\n",
-    "    for character in input_string:\n",
-    "        if  character == \" \":\n",
-    "            final_string = final_string + character\n",
-    "        else:\n",
-    "            if(character.isalnum()):\n",
-    "                final_string = final_string + character\n",
-    "    return final_string"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 60,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "100%|█| 5593207/5593207 [00:35<00:0\n"
-     ]
-    }
-   ],
-   "source": [
-    "for d in tqdm(data):\n",
-    "    label = d[2]\n",
-    "    d[2] = remove_special(caption)\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 4,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "['0',\n",
-       " 'coco/train2014/COCO_train2014_000000057870.jpg',\n",
-       " '1.020,279.960,534.110,480.000,67,dining table&&90.670,271.490,262.510,480.000,62,chair&&233.290,270.450,403.610,473.810,62,chair&&367.820,264.270,506.970,480.000,62,chair&&476.760,261.030,596.490,462.740,62,chair&&263.030,174.370,417.670,299.400,64,potted plant&&539.330,290.160,640.000,469.210,62,chair&&10.790,260.030,125.120,384.070,62,chair&&560.800,413.950,639.090,479.200,67,dining table&&20.540,376.760,103.780,431.890,62,chair&&1.080,373.210,32.360,480.000,62,chair&&298.200,235.170,381.210,269.250,86,vase&&152.170,256.670,230.580,285.780,62,chair&&364.400,256.570,417.060,283.210,62,chair&&296.780,277.790,329.260,289.780,84,book&&292.800,289.310,314.210,300.650,84,book&&285.800,257.460,299.770,273.600,62,chair']"
-      ]
-     },
-     "execution_count": 4,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "data[0]"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 78,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "100%|█| 181767/181767 [00:00<00:00,\n"
-     ]
-    }
-   ],
-   "source": [
-    "for d in tqdm(data):\n",
-    "    d[2] = d[2].replace('\\\"', '')\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 49,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "data_  = []\n",
-    "with open(path_data) as file:\n",
-    "    for i in tqdm(range(6458670)):\n",
-    "        column_l = file.readline().rstrip(\"\\n\").split(\"\\t\")\n",
-    "        data_.append(column_l)\n",
-    "        if len(column_l) < 2:\n",
-    "            break"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 64,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "5593207it [00:03, 1463300.52it/s]\n"
-     ]
-    }
-   ],
-   "source": [
-    "\n",
-    "data_example  = []\n",
-    "fp = open('/data/mshukor/data/ofa/pretrain_ours/vision_language_mini_.tsv', \"r\")\n",
-    "data_example  = []\n",
-    "for line in tqdm(fp):\n",
-    "    data_example.append(line)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 74,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "2796604\tcc3m/train/8/2d0d96e4ecb8e2e959a3bf10d59b9d05ac114aea.jpg\tthe residential development under construction in district\t\t\t\tcc3m\tcaption\n",
-      "\n"
-     ]
-    }
-   ],
-   "source": [
-    "print(data_example[2796604])"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 73,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "/val2014/COCO_val2014_000000329789.jpg\tA young man is eating a slice of pizza in his room\t\t\t\tcoco_karp\tcaption\n",
-      "\n"
-     ]
-    }
-   ],
-   "source": [
-    "data_example[2796604]\n",
-    "fp.seek(2796604)\n",
-    "for l in fp:\n",
-    "    print(l)\n",
-    "    break"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 46,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "2317 2317\n",
-      "2510 2514 2510\n"
-     ]
-    }
-   ],
-   "source": [
-    "len(data_[10].rstrip(\"\\n\").split(\"\\t\")[1])# len(line)\n",
-    "# len(data_[10].rstrip(\"\\n\").split(\"\\t\")[1].encode('utf-8'))\n",
-    "print(len(data_example[10]), len(data_example[10].encode('utf-8')))\n",
-    "print(len(data_[10]), len(data_[10].encode('utf-8')), len(data_[10].encode('utf-8').decode(\"utf-8\")))\n",
-    "\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "print(data_[10].encode('utf-8'))\n",
-    "print(data_[10])\n",
-    "\n",
-    "print(data_example[10].encode('utf-8'))\n",
-    "print(data_example[10])"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 4,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "6458670it [01:45, 61129.36it/s] \n"
-     ]
-    }
-   ],
-   "source": [
-    "output_path = '/data/mshukor/data/ofa/pretrain_ours/text_mini.tsv'\n",
-    "\n",
-    "fp = open(output_path, \"r\")\n",
-    "data_  = []\n",
-    "for line in tqdm(fp):\n",
-    "    data_.append(line)\n",
-    "    "
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 12,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "6458670it [04:08, 25941.37it/s]\n"
-     ]
-    }
-   ],
-   "source": [
-    "output_path = '/data/mshukor/data/ofa/pretrain_ours/text_mini.tsv'\n",
-    "\n",
-    "start_id = 0 \n",
-    "num_max_characters = 2500\n",
-    "\n",
-    "with open(output_path, 'w', newline='\\n') as f_output:\n",
-    "    csv_output = csv.writer(f_output, delimiter='\\t')\n",
-    "\n",
-    "    for i, t in tqdm(enumerate(data)):\n",
-    "        text = t[1]\n",
-    "        item = [start_id, text]\n",
-    "        csv_output.writerow(item)\n",
-    "        start_id+=1"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 28,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "100%|█████████████████████████████████████████████████████████| 181767/181767 [00:03<00:00, 51934.09it/s]\n"
-     ]
-    }
-   ],
-   "source": [
-    "output_path = '/data/mshukor/data/ofa/pretrain_ours/detection_mini.tsv'\n",
-    "\n",
-    "with open(output_path, 'w', newline='\\n') as f_output:\n",
-    "    csv_output = csv.writer(f_output, delimiter='\\t')\n",
-    "\n",
-    "    for t in tqdm(data):\n",
-    "        csv_output.writerow(t)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 63,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "['5593206',\n",
-       " 'train2014/COCO_train2014_000000524286.jpg',\n",
-       " '',\n",
-       " 'Is that a laptop?',\n",
-       " '1.0|!+yes',\n",
-       " '',\n",
-       " 'vqa_train',\n",
-       " 'qa']"
-      ]
-     },
-     "execution_count": 63,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "t"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Create data tsv"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def convert_img_to_str(file_name):\n",
-    "    img = Image.open(file_name) # path to file\n",
-    "    img_buffer = BytesIO()\n",
-    "    img.save(img_buffer, format=img.format)\n",
-    "    byte_data = img_buffer.getvalue()\n",
-    "    base64_str = base64.b64encode(byte_data) # bytes\n",
-    "    base64_str = base64_str.decode(\"utf-8\") # str\n",
-    "    return base64_str"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "### Create VL tsv"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "#### Caption"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "original_data_path = '/data/mshukor/data/our_albef_data/json_pretrain/sbu.json'\n",
-    "original_data = json.load(open(original_data_path,'r'))"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from preprocess.utils import get_tsv_data_from_jsons\n",
-    "    \n",
-    "datasets = [\n",
-    "            '/data/mshukor/data/our_albef_data/json_pretrain/coco_karp.json',\n",
-    "            '/data/mshukor/data/our_albef_data/json_pretrain/vg_albef.json',\n",
-    "            '/data/mshukor/data/our_albef_data/json_pretrain/sbu.json',\n",
-    "            '/data/mshukor/data/our_albef_data/json_pretrain/cc3m.json', \n",
-    "    \n",
-    "            ['/data/mshukor/data/refcoco/refcoco+/refs(unc).p', '/data/mshukor/data/refcoco/refcoco+/instances.json'],\n",
-    "            \n",
-    "            '/data/mshukor/data/our_albef_data/data/vqa_train.json',\n",
-    "]\n",
-    "\n",
-    "start_id = 0\n",
-    "task_types = ['caption',\n",
-    "             'caption',\n",
-    "             'caption',\n",
-    "             'caption',\n",
-    "             'visual_grounding',\n",
-    "             'qa',]\n",
-    "\n",
-    "tsvs = get_tsv_data_from_jsons(datasets, start_id, task_types, convert_images=False)\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "len(tsvs)\n",
-    "# tsvs[-10000]\n",
-    "tsvs[-1000000]"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import csv\n",
-    "from io import StringIO\n",
-    "\n",
-    "output_path = '/data/mshukor/data/ofa/pretrain_ours/vision_language_mini.tsv'\n",
-    "\n",
-    "with open(output_path, 'w', newline='') as f_output:\n",
-    "    csv_output = csv.writer(f_output, delimiter='\\t')\n",
-    "\n",
-    "    for t in tqdm(tsvs):\n",
-    "        csv_output.writerow(t)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "csv.field_size_limit(sys.maxsize)\n",
-    "\n",
-    "\n",
-    "out_data = []\n",
-    "selected_cols='0,1,2,3,4,5,6,7'\n",
-    "\n",
-    "selected_col_ids = [int(col_id) for col_id in selected_cols.split(\",\")]\n",
-    "\n",
-    "with open(output_path) as file:\n",
-    "    tsv_file = csv.reader(file, delimiter='\\t')\n",
-    "    for line in tqdm(tsv_file):\n",
-    "        d = [line[i] for i in selected_col_ids]\n",
-    "        out_data.append(d)\n",
-    "    "
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "out_data[-1]"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "#### VQA"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "original_data_path = '/data/mshukor/data/our_albef_data/data/vqa_train.json'\n",
-    "original_data = json.load(open(original_data_path,'r'))\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "original_data[100]"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# 1.0|!+horizontal&&0.3|!+south&&0.3|!+straight&&0.3|!+vertical\n",
-    "\n",
-    "from preprocess.utils import get_tsv_vqa_data_from_json\n",
-    "\n",
-    "\n",
-    "start_id = 0\n",
-    "dataset_name = 'vqav2'\n",
-    "task_type = 'qa'\n",
-    "\n",
-    "image_root = '/data/mshukor/data/coco'\n",
-    "tmp = get_tsv_vqa_data_from_json(original_data, start_id, dataset_name, task_type, image_root=image_root, convert_images=False)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "tmp[10]"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "#### Visual Grounding "
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "original_data_path = '/data/mshukor/data/our_albef_data/data/refcoco+_train.json'\n",
-    "original_data = json.load(open(original_data_path,'r'))\n",
-    "\n",
-    "original_data_path = '/data/mshukor/data/our_albef_data/data/refcoco+/dets.json'\n",
-    "det_file = json.load(open(original_data_path,'r'))\n",
-    "\n",
-    "original_data_path = '/data/mshukor/data/our_albef_data/data/refcoco+/cocos.json'\n",
-    "coco_file = json.load(open(original_data_path,'r'))"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "list(det_file.keys())[:10]"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "original_data_path = '/data/mshukor/data/refcoco/refcoco+/instances.json'\n",
-    "original_data = json.load(open(original_data_path,'r'))"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import pickle\n",
-    "\n",
-    "ref_path = '/data/mshukor/data/refcoco/refcoco+/refs(unc).p'\n",
-    "refs = pickle.load(open(ref_path, 'rb'))"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "for i, ref in tqdm(enumerate(refs)):\n",
-    "    \n",
-    "    "
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "len(refs)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "refs[500]"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "id_to_annot = {}\n",
-    "for annot in original_data['annotations']:\n",
-    "    id_to_annot[annot['id']] = annot\n",
-    "    \n",
-    "    "
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "id_to_images = {}\n",
-    "for annot in tqdm(original_data['images']):\n",
-    "    id_to_images[annot['id']] = annot"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "id_to_images[576457]"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "list(id_to_annot.keys())[:10]\n",
-    "id_to_annot[1640859]['bbox']\n",
-    "for r in tqdm(id_to_annot.values()):\n",
-    "    if r['bbox'][0] > 0:\n",
-    "        print(r['bbox'])"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "list(original_data.keys())[:10]"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "ref_path = '/data/mshukor/data/refcoco/refcoco+/refs(unc).p'\n",
-    "instances_path = '/data/mshukor/data/refcoco/refcoco+/instances.json'\n",
-    "start_id = 0\n",
-    "dataset_name='refcoco_train'\n",
-    "task_type='visual_grounding'\n",
-    "convert_images=False\n",
-    "split='train'\n",
-    "\n",
-    "tmp = get_tsv_from_refcoco(ref_path, instances_path, start_id, dataset_name=dataset_name, task_type=task_type, convert_images=convert_images, split=split)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "tmp[-1]"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "Image.open('/data/mshukor/data/coco/train2014/COCO_train2014_000000000072.jpg').convert('RGB')"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "original_data['images'][:10]"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# ['third book starting from left', '', '29.1,11.72,66.81,343.41', '', 'refcoco_train', 'visual_grounding']\n",
-    "\n",
-    "original_data['categories']"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "### Imagenet"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# image-id and image base64 string .txt file \n",
-    "# id, image, code in tsv final \n",
-    "\n",
-    "from preprocesss.utils import create_imagenet_txt_files\n",
-    "\n",
-    "\n",
-    "path_data = '/data/mshukor/data/imagenet/val'\n",
-    "output_path = '/data/mshukor/data/ofa/pretrain_ours/imagenet_val.txt'\n",
-    "\n",
-    "\n",
-    "create_imagenet_txt_files(path_data, output_path)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "start_id\n",
-    "len(data)\n",
-    "data[0]"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "\n",
-    "code_path = '/data/mshukor/data/ofa/pretrain_ours/imagenet_train_codes.tsv'\n",
-    "output_path = '/data/mshukor/data/ofa/pretrain_ours/image_mini.tsv'\n",
-    "\n",
-    "def save_image_only_tsv_from_code_files(code_path, output_path, start_id=0):\n",
-    "    selected_col_ids = [0,1]\n",
-    "    out_data = []\n",
-    "    with open(code_path) as file:\n",
-    "        tsv_file = csv.reader(file, delimiter='\\t')\n",
-    "        for line in tqdm(tsv_file):\n",
-    "            d = [line[i] for i in selected_col_ids]\n",
-    "            d = [start_id]+d\n",
-    "            out_data.append(d)\n",
-    "\n",
-    "\n",
-    "    with open(output_path, 'w', newline='') as f_output:\n",
-    "        csv_output = csv.writer(f_output, delimiter='\\t')\n",
-    "\n",
-    "        for t in tqdm(out_data):\n",
-    "            csv_output.writerow(t)\n",
-    "\n",
-    "save_image_only_tsv_from_code_files(code_path, output_path, start_id=0)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "selected_col_ids = [0,1,2]\n",
-    "out_data = []\n",
-    "with open(output_path) as file:\n",
-    "    tsv_file = csv.reader(file, delimiter='\\t')\n",
-    "    for line in tqdm(tsv_file):\n",
-    "        d = [line[i] for i in selected_col_ids]\n",
-    "        out_data.append(d)\n",
-    "        break"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "len(out_data[0][2].split(' '))"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "#### Fix image paths"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 33,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "1281167it [00:16, 79250.80it/s]\n"
-     ]
-    }
-   ],
-   "source": [
-    "\n",
-    "path_data =  '/data/mshukor/data/ofa/pretrain_ours/image_mini.tsv'\n",
-    "selected_cols='0,1,2'\n",
-    "\n",
-    "data = []\n",
-    "\n",
-    "selected_col_ids = [int(col_id) for col_id in selected_cols.split(\",\")]\n",
-    "\n",
-    "with open(path_data) as file:\n",
-    "    tsv_file = csv.reader(file, delimiter='\\t')\n",
-    "    for line in tqdm(tsv_file):\n",
-    "\n",
-    "        d = [line[i] for i in selected_col_ids]\n",
-    "#         print(d)\n",
-    "        data.append(d)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 44,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "1281167it [00:16, 76760.12it/s]\n",
-      "1281167it [00:01, 671149.72it/s]\n",
-      "100%|█████| 1281167/1281167 [00:01<00:00, 947543.73it/s]\n"
-     ]
-    }
-   ],
-   "source": [
-    "# from imge-id img-path to \n",
-    "def replace_image_id_by_path(input_tsv, output_tsv, mapping_file):\n",
-    "    selected_cols='0,1,2'\n",
-    "    data = []\n",
-    "    selected_col_ids = [int(col_id) for col_id in selected_cols.split(\",\")]\n",
-    "    with open(input_tsv) as file:\n",
-    "        tsv_file = csv.reader(file, delimiter='\\t')\n",
-    "        for line in tqdm(tsv_file):\n",
-    "            d = [line[i] for i in selected_col_ids]\n",
-    "            data.append(d)\n",
-    "            \n",
-    "    im_id_to_path = {}\n",
-    "    with open(mapping_file) as file:\n",
-    "        tsv_file = csv.reader(file, delimiter='\\t')\n",
-    "        for line in tqdm(tsv_file):\n",
-    "            d = [line[i] for i in [0, 1]]\n",
-    "            im_id_to_path[d[0]] = d[1]\n",
-    "            \n",
-    "    for d in tqdm(data):\n",
-    "        im_id = d[1].split('/')[-1].split('.')[0]\n",
-    "        im_path = im_id_to_path[im_id]\n",
-    "        d[1] = im_path\n",
-    "        \n",
-    "    with open(output_tsv, 'w', newline='') as f_output:\n",
-    "        csv_output = csv.writer(f_output, delimiter='\\t')\n",
-    "\n",
-    "        for t in tqdm(data):\n",
-    "            csv_output.writerow(t)\n",
-    "        \n",
-    "    return data\n",
-    "\n",
-    "input_tsv = '/data/mshukor/data/ofa/pretrain_ours/image_mini.tsv'\n",
-    "output_tsv = '/data/mshukor/data/ofa/pretrain_ours/image_mini.tsv'\n",
-    "mapping_file = '/data/mshukor/data/ofa/pretrain_ours/imagenet_train.txt'\n",
-    "\n",
-    "tmp = replace_image_id_by_path(input_tsv, output_tsv, mapping_file)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 45,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "['0',\n",
-       " 'RawImages/train/n03146219/n03146219_8050.JPEG',\n",
-       " '7442 662 7977 1652 6320 650 4376 992 1596 7734 1925 5335 3935 5604 5697 4504 5114 4050 144 215 144 6691 5321 7769 4755 3346 4691 3469 4175 1351 6907 9 6948 7749 7166 215 1026 931 970 4168 2675 6874 6248 2306 6138 8052 2970 6302 5550 2491 6931 969 6574 8014 6588 6639 389 1882 688 4691 4266 675 6248 3938 2387 4365 5999 261 2966 3499 651 5290 970 3526 5583 516 167 2103 1513 198 6657 7442 1118 7207 7307 1792 2078 388 4285 3417 5450 6959 6999 1306 1649 4556 2533 1103 6869 7681 8051 1916 7160 7743 2704 8063 2726 4860 2383 1635 8061 3497 7327 5915 7836 5697 1719 2136 96 970 7184 5167 2250 404 7007 7565 2742 33 7076 5250 7790 1838 1298 2847 3250 1204 1934 5550 4360 5688 1791 3465 634 4663 2991 5352 4066 4157 946 1596 3504 5855 5629 5411 7695 3627 3942 5631 2736 2883 5059 1423 2009 2643 1873 4960 1661 545 1396 3450 3145 211 6869 2226 6780 2724 4606 3702 3667 891 6236 6419 3531 7032 5277 3381 3031 7878 725 1652 1813 5037 949 3087 405 7884 3784 5432 633 4256 235 3182 3686 5450 2419 1593 7948 5741 6237 7233 20 7470 7071 182 1584 6780 7913 2691 7207 5094 5199 4502 5030 2360 448 5129 2713 1094 1678 1934 2458 2970 2133 867 3332 6138 294 3260 5495 4189 5732 3940 5629 4139 7335 7607 3248 4981 2109 3660 4364 7763 3964 7163 6702 691']"
-      ]
-     },
-     "execution_count": 45,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "tmp[0]"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 36,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "100%|█████| 1281167/1281167 [00:03<00:00, 336250.44it/s]\n"
-     ]
-    }
-   ],
-   "source": [
-    "# imgage_dir = 'imagenet/RawImages/train/'\n",
-    "# for d in tqdm(data):\n",
-    "#     im_id = d[1]\n",
-    "#     im_dir = im_id.split('_')[0]\n",
-    "#     im_path = os.path.join(im_dir, im_id+'.JPEG')\n",
-    "#     d[1] = os.path.join(imgage_dir, im_path)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 39,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "['0',\n",
-       " 'imagenet/RawImages/train/n03146219/n03146219_8050.JPEG',\n",
-       " '7442 662 7977 1652 6320 650 4376 992 1596 7734 1925 5335 3935 5604 5697 4504 5114 4050 144 215 144 6691 5321 7769 4755 3346 4691 3469 4175 1351 6907 9 6948 7749 7166 215 1026 931 970 4168 2675 6874 6248 2306 6138 8052 2970 6302 5550 2491 6931 969 6574 8014 6588 6639 389 1882 688 4691 4266 675 6248 3938 2387 4365 5999 261 2966 3499 651 5290 970 3526 5583 516 167 2103 1513 198 6657 7442 1118 7207 7307 1792 2078 388 4285 3417 5450 6959 6999 1306 1649 4556 2533 1103 6869 7681 8051 1916 7160 7743 2704 8063 2726 4860 2383 1635 8061 3497 7327 5915 7836 5697 1719 2136 96 970 7184 5167 2250 404 7007 7565 2742 33 7076 5250 7790 1838 1298 2847 3250 1204 1934 5550 4360 5688 1791 3465 634 4663 2991 5352 4066 4157 946 1596 3504 5855 5629 5411 7695 3627 3942 5631 2736 2883 5059 1423 2009 2643 1873 4960 1661 545 1396 3450 3145 211 6869 2226 6780 2724 4606 3702 3667 891 6236 6419 3531 7032 5277 3381 3031 7878 725 1652 1813 5037 949 3087 405 7884 3784 5432 633 4256 235 3182 3686 5450 2419 1593 7948 5741 6237 7233 20 7470 7071 182 1584 6780 7913 2691 7207 5094 5199 4502 5030 2360 448 5129 2713 1094 1678 1934 2458 2970 2133 867 3332 6138 294 3260 5495 4189 5732 3940 5629 4139 7335 7607 3248 4981 2109 3660 4364 7763 3964 7163 6702 691']"
-      ]
-     },
-     "execution_count": 39,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "data[0]"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 40,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "100%|██████| 1281167/1281167 [00:27<00:00, 46704.02it/s]\n"
-     ]
-    }
-   ],
-   "source": [
-    "output_path = '/data/mshukor/data/ofa/pretrain_ours/image_mini.tsv'\n",
-    "with open(output_path, 'w', newline='') as f_output:\n",
-    "    csv_output = csv.writer(f_output, delimiter='\\t')\n",
-    "\n",
-    "    for t in tqdm(data):\n",
-    "        csv_output.writerow(t)\n"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "### Object detection"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "#### COCO"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# '505.856,189.994,799.744,450.016,/m/07j7r,tree&&753.664,384.00600000000003,827.392,446.572,/m/0c9ph5,flower'\n",
-    "\n",
-    "path_json = '/data/mshukor/data/coco/annotations/instances_train2014.json'\n",
-    "\n",
-    "data = json.load(open(path_json,'r'))"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def get_tsv_from_coco_detection(instances_path, start_id, convert_images=True, split='train'):\n",
-    "\n",
-    "    instances = json.load(open(instances_path,'r'))\n",
-    "    imgid_to_annot = {}\n",
-    "    for annot in tqdm(instances['annotations']):\n",
-    "        if annot['image_id'] not in imgid_to_annot:\n",
-    "            imgid_to_annot[annot['image_id']] = [annot]\n",
-    "        else:\n",
-    "            imgid_to_annot[annot['image_id']].append(annot)\n",
-    "\n",
-    "    id_to_category = {}\n",
-    "    for annot in tqdm(instances['categories']):\n",
-    "        id_to_category[annot['id']] = annot['name']\n",
-    "\n",
-    "    tsv_data = []\n",
-    "    missied = []\n",
-    "    for ref in tqdm(instances['images']):\n",
-    "        ref_split = split\n",
-    "        image_id = ref['id']\n",
-    "        file_name = ref['file_name']\n",
-    "\n",
-    "        if ref_split == 'train':\n",
-    "            file_name = os.path.join('coco/train2014', file_name)\n",
-    "\n",
-    "        if convert_images:\n",
-    "            img_path = os.path.join('/data/mshukor/data/', file_name)\n",
-    "            img = convert_img_to_str(img_path)\n",
-    "        else:\n",
-    "            img_path = file_name.replace('/data/mshukor/data/', '')\n",
-    "            img = img_path\n",
-    "\n",
-    "        # ann_id = ref['id']\n",
-    "        # annot = id_to_annot[ann_id]\n",
-    "        if image_id not in imgid_to_annot:\n",
-    "            missied.append(image_id)\n",
-    "            continue\n",
-    "        annots = imgid_to_annot[image_id]\n",
-    "        detections = []\n",
-    "        areas = []\n",
-    "        for annot in annots:\n",
-    "            bbox = annot['bbox'] # x,y,w,h bottom left\n",
-    "            area = bbox[2]*bbox[3]\n",
-    "            x1, y1, x2, y2 = bbox[0], bbox[1], bbox[0] + bbox[2], bbox[1] + bbox[3]   # top left, bottom right \n",
-    "            # box = '{:.3f},{:.3f},{:.3f},{:.3f}'.format(x1, y1, x2, y2)\n",
-    "\n",
-    "            object_id = annot['category_id']\n",
-    "            category = id_to_category[object_id]\n",
-    "\n",
-    "            tmp = '{:.3f},{:.3f},{:.3f},{:.3f},{},{}'.format(x1, y1, x2, y2, object_id, category)\n",
-    "            areas.append(area)\n",
-    "            detections.append(tmp)\n",
-    "\n",
-    "        sorted_indices = sorted(range(len(areas)), key=lambda k: areas[k], reverse=True)\n",
-    "        detections = [detections[k] for k in sorted_indices]\n",
-    "        detections = '&&'.join(detections)\n",
-    "        t = [start_id, img, detections]\n",
-    "\n",
-    "        tsv_data.append(t)\n",
-    "        start_id+=1\n",
-    "\n",
-    "    return tsv_data\n",
-    "\n",
-    "instances_path = '/data/mshukor/data/coco/annotations/instances_train2014.json'\n",
-    "start_id = 0\n",
-    "tmp = get_tsv_from_coco_detection(instances_path, start_id, convert_images=False, split='train')"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "list(imgid_to_annot.keys())[:10]\n",
-    "len(missied)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "#### VG"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def get_tsv_from_vg_detection(instances_path, path_images, start_id, convert_images=True, split='train'):\n",
-    "    \n",
-    "    instances = json.load(open(instances_path,'r'))\n",
-    "    \n",
-    "    id_to_objects = {}\n",
-    "    for d in instances:\n",
-    "        id_to_objects[d['id']] = d\n",
-    "\n",
-    "\n",
-    "    \n",
-    "    id_to_image_path = {}\n",
-    "    for root, dirs, files, in os.walk(path_images):\n",
-    "        for d in dirs:\n",
-    "            dir_path = os.path.join(root, d)\n",
-    "            for _, _, dir_files in os.walk(dir_path):\n",
-    "                for f in dir_files:\n",
-    "                    file_path = os.path.join(dir_path, f)\n",
-    "                    file_path = '/'.join(file_path.split('/')[-4:])\n",
-    "                    image_id = f.split('.')[0]\n",
-    "                    id_to_image_path[image_id] = file_path\n",
-    "\n",
-    "                    \n",
-    "\n",
-    "\n",
-    "    tsv_data = []\n",
-    "    missied = []\n",
-    "    negs = []\n",
-    "    for ref in tqdm(id_to_image_path.keys()):\n",
-    "        ref_split = split\n",
-    "        \n",
-    "        image_id = ref\n",
-    "        \n",
-    "        file_name = id_to_image_path[image_id]\n",
-    "        if convert_images:\n",
-    "            img_path = os.path.join('/data/mshukor/data/', file_name)\n",
-    "            img = convert_img_to_str(img_path)\n",
-    "        else:\n",
-    "            img_path = file_name.replace('/data/mshukor/data/', '')\n",
-    "            img = img_path\n",
-    "\n",
-    "            \n",
-    "        if int(image_id) in id_to_objects:\n",
-    "            objects = id_to_objects[int(image_id)]['objects']\n",
-    "        else:\n",
-    "            missied.append(image_id)\n",
-    "            continue\n",
-    "        \n",
-    "        if len(objects) == 0:\n",
-    "            missied.append(image_id)\n",
-    "            continue\n",
-    "            \n",
-    "        \n",
-    "        areas = []\n",
-    "        detections = []\n",
-    "        for annot in objects:\n",
-    "            x,y,w,h = annot['x'], annot['y'], annot['w'], annot['h'] # x,y,w,h bottom left\n",
-    "            \n",
-    "            area = w*h\n",
-    "            \n",
-    "            x1, y1, x2, y2 = x, y, x + w, y + h  # top left, bottom right \n",
-    "            \n",
-    "            if x1 < 0 or x2 < 0:\n",
-    "                negs.append(annot)\n",
-    "            x1 = max(0, x1)\n",
-    "            x2 = max(0, x2)\n",
-    "            \n",
-    "            \n",
-    "            category = ','.join(annot['names']).replace('\\x00','')\n",
-    "            object_id = annot['id']\n",
-    "            \n",
-    "            \n",
-    "            tmp = '{:.3f},{:.3f},{:.3f},{:.3f},{},{}'.format(x1, y1, x2, y2, object_id, category)\n",
-    "            detections.append(tmp)\n",
-    "            areas.append(area)\n",
-    "\n",
-    "        sorted_indices = sorted(range(len(areas)), key=lambda k: areas[k], reverse=True)\n",
-    "        detections = [detections[k] for k in sorted_indices]\n",
-    "        \n",
-    "        detections = '&&'.join(detections)\n",
-    "        t = [start_id, img, detections]\n",
-    "\n",
-    "        tsv_data.append(t)\n",
-    "        start_id+=1\n",
-    "    print('missed images:', len(missied), 'negs', len(negs))\n",
-    "    return tsv_data\n",
-    "\n",
-    "\n",
-    "instances_path = '/data/mshukor/data/visual_genome/annotations/objects.json'\n",
-    "path_images = '/data/mshukor/data/visual_genome/images'\n",
-    "start_id = 0\n",
-    "\n",
-    "tmp = get_tsv_from_vg_detection(instances_path, path_images, start_id, convert_images=False, split='train')"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "image_root = '/data/mshukor/data/'\n",
-    "\n",
-    "Image.open(image_root+id_to_image_path['1087']).convert('RGB')"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "#### OpenImagesV5"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# data_path = '/data/mshukor/data/OpenImagesV5/train-annotations-bbox.csv'\n",
-    "# data_path = '/data/mshukor/data/OpenImagesV5/train-images-boxable.csv'\n",
-    "# data_path = '/data/mshukor/data/OpenImagesV5/train-images-boxable-with-rotation.csv'\n",
-    "data_path = '/data/mshukor/data/OpenImagesV5/class-descriptions-boxable.csv'\n",
-    "\n",
-    "\n",
-    "\n",
-    "\n",
-    "selected_col_ids = [0,1,2]\n",
-    "out_data = []\n",
-    "with open(data_path) as file:\n",
-    "    tsv_file = csv.reader(file, delimiter='\\t')\n",
-    "    for i, line in tqdm(enumerate(tsv_file)):\n",
-    "        # d = [line[i] for i in selected_col_ids]\n",
-    "        out_data.append(line)\n",
-    "#         print(line)\n",
-    "#         if i > 2:\n",
-    "#             break\n",
-    "        "
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def get_tsv_from_openimages_detection(instances_path, path_images, start_id, convert_images=False, split='train')\n",
-    "\n",
-    "    id_to_image_path = {}\n",
-    "    for root, dirs, files, in os.walk(path_images):\n",
-    "        for d in dirs:\n",
-    "            dir_path = os.path.join(root, d)\n",
-    "            for _, _, dir_files in os.walk(dir_path):\n",
-    "                for f in dir_files:\n",
-    "                    file_path = os.path.join(dir_path, f)\n",
-    "                    file_path = '/'.join(file_path.split('/')[-4:])\n",
-    "                    image_id = f.split('.')[0]\n",
-    "                    id_to_image_path[image_id] = file_path\n",
-    "\n",
-    "    image_root = '/gpfsdswork/dataset'\n",
-    "\n",
-    "    def imagepath_to_image_size(path):\n",
-    "        w, h = Image.open(path).size\n",
-    "\n",
-    "    id_to_annot = {}\n",
-    "    with open(instances_path) as file:\n",
-    "        tsv_file = csv.reader(file, delimiter='\\t')\n",
-    "        for i, line in tqdm(enumerate(tsv_file)):\n",
-    "            img_id = line[0].split(',')[0]\n",
-    "            if img_id in id_to_annot:\n",
-    "                id_to_annot[img_id].append(line)\n",
-    "            else:\n",
-    "                id_to_annot[img_id] = [line]\n",
-    "\n",
-    "    classid_to_class = {}\n",
-    "\n",
-    "    with open(class_path) as file:\n",
-    "        tsv_file = csv.reader(file, delimiter=',')\n",
-    "        for i, line in tqdm(enumerate(tsv_file)):\n",
-    "            classid_to_class[line[0]] = line[1]\n",
-    "\n",
-    "    tsv_data = []\n",
-    "    for img_id in id_to_annot.keys():\n",
-    "        annots = id_to_annot[img_id]\n",
-    "        img_path = id_to_image_path[img_id]\n",
-    "        orig_img_path = os.path.join(image_root, img_path)\n",
-    "        w, h = imagepath_to_image_size(path)\n",
-    "\n",
-    "        if convert_images:\n",
-    "            img = convert_img_to_str(orig_img_path)\n",
-    "        else:\n",
-    "            img = img_path\n",
-    "\n",
-    "        areas = []\n",
-    "        detections = []\n",
-    "        for d in annots:\n",
-    "            d = d[0].split(',')\n",
-    "\n",
-    "            x1, x2, y1, y2 = d[4:8]\n",
-    "            x1, x2, y1, y2 = x1*w, x2*w, y1*h, y2*h\n",
-    "            box_w, box_h = x2 - x1, y2 - y1\n",
-    "            area = box_w*box_h\n",
-    "            areas.append(area)\n",
-    "\n",
-    "            object_id = d[2]\n",
-    "            category = classid_to_class[object_id]\n",
-    "\n",
-    "            tmp = '{:.3f},{:.3f},{:.3f},{:.3f},{},{}'.format(x1, y1, x2, y2, object_id, category)\n",
-    "            detections.append(tmp)\n",
-    "\n",
-    "\n",
-    "        sorted_indices = sorted(range(len(areas)), key=lambda k: areas[k], reverse=True)\n",
-    "        detections = [detections[k] for k in sorted_indices]\n",
-    "\n",
-    "        detections = '&&'.join(detections)\n",
-    "        t = [start_id, img, detections]\n",
-    "\n",
-    "        tsv_data.append(t)\n",
-    "        start_id+=1\n",
-    "        \n",
-    "    return tsv_data\n",
-    "\n",
-    "    "
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "e39871fd9fd74f55"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "### Text"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "#### En Wikipedia"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from datasets import load_dataset"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "%env http_proxy='http://192.168.0.100:3128' \n",
-    "%env https_proxy='http://192.168.0.100:3128'\n",
-    "\n",
-    "%env HF_DATASETS_CACHE=\"/data/mshukor/data/.cache\"\n",
-    "%env HF_DATASETS_OFFLINE=1"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "tmp = load_dataset(\"wikipedia\", \"20220301.en\", cache_dir=\"/data/mshukor/data/.cache\")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "len(tmp['train'][0]['text'])\n",
-    "tmp['train'][0]['text'][:512]"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def remove_special(input_string):\n",
-    "    final_string = \"\"\n",
-    "    for character in input_string:\n",
-    "        if  character == \" \":\n",
-    "            final_string = final_string + character\n",
-    "        else:\n",
-    "            if(character.isalnum()):\n",
-    "                final_string = final_string + character\n",
-    "    return final_string\n",
-    "\n",
-    "def get_tsv_from_text_data(data_name=\"wikipedia\", data_subname=\"20220301.en\", \n",
-    "                           output_path, cache_dir=\"/data/mshukor/data/.cache\", start_id=0, num_max_characters=2500):\n",
-    "    from datasets import load_dataset\n",
-    "    tmp = load_dataset(data_name, data_subname, cache_dir=cache_dir)\n",
-    "\n",
-    "    with open(output_path, 'w', newline='') as f_output:\n",
-    "        csv_output = csv.writer(f_output, delimiter='\\t')\n",
-    "\n",
-    "        for i, t in tqdm(enumerate(tmp['train'])):\n",
-    "            text = t['text'][:num_max_characters].replace('\\t', ' ').replace(\"\\n\", ' ').replace('\\\"', '')\n",
-    "            text = remove_special(text)\n",
-    "            item = [start_id, text]\n",
-    "            csv_output.writerow(item)\n",
-    "            start_id+=1\n",
-    "\n",
-    "    "
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import csv\n",
-    "from io import StringIO\n",
-    "\n",
-    "output_path = '/data/mshukor/data/ofa/pretrain_ours/text_mini.tsv'\n",
-    "\n",
-    "start_id = 0 \n",
-    "num_max_characters = 2500\n",
-    "\n",
-    "with open(output_path, 'w', newline='') as f_output:\n",
-    "    csv_output = csv.writer(f_output, delimiter='\\t')\n",
-    "\n",
-    "    for i, t in tqdm(enumerate(tmp['train'])):\n",
-    "        text = t['text'][:num_max_characters]\n",
-    "        item = [start_id, text]\n",
-    "        csv_output.writerow(item)\n",
-    "        start_id+=1"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "out_data = []\n",
-    "selected_cols='0,1'\n",
-    "\n",
-    "selected_col_ids = [int(col_id) for col_id in selected_cols.split(\",\")]\n",
-    "\n",
-    "with open(output_path) as file:\n",
-    "    tsv_file = csv.reader(file, delimiter='\\t')\n",
-    "    for line in tqdm(tsv_file):\n",
-    "        d = [line[i] for i in selected_col_ids]\n",
-    "        out_data.append(d)\n",
-    "    "
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "out_data"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "### Create from finetuned data"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "read from tsv and write to tsv directly \n",
-    "same for vqa v2\n",
-    "then create ofa_mini 4m, vqa and refcoco for pretraining "
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# Convert weights"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 3,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from transformers import T5Tokenizer, T5ForConditionalGeneration\n",
-    "from models import ofa_base_architecture, OFAModel\n",
-    "from transformers.tokenization_utils_base import BatchEncoding"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "### Explore ofa"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 4,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "2022-11-15 08:52:08 | INFO | tasks.ofa_task | source dictionary: 59457 types\n",
-      "2022-11-15 08:52:08 | INFO | tasks.ofa_task | target dictionary: 59457 types\n"
-     ]
-    }
-   ],
-   "source": [
-    "import torch\n",
-    "import numpy as np\n",
-    "from fairseq import utils, tasks\n",
-    "from fairseq import checkpoint_utils\n",
-    "from utils.eval_utils import eval_step\n",
-    "from tasks.mm_tasks.caption import CaptionTask\n",
-    "from models.ofa import OFAModel\n",
-    "from PIL import Image\n",
-    "\n",
-    "# Register refcoco task\n",
-    "tasks.register_task('caption', CaptionTask)\n",
-    "\n",
-    "# turn on cuda if GPU is available\n",
-    "use_cuda = torch.cuda.is_available()\n",
-    "# use fp16 only when GPU is available\n",
-    "use_fp16 = False\n",
-    "\n",
-    "# Load pretrained ckpt & config\n",
-    "overrides={\"eval_cider\":False, \"beam\":5, \"max_len_b\":16, \"no_repeat_ngram_size\":3, \"seed\":7}\n",
-    "models, cfg, task = checkpoint_utils.load_model_ensemble_and_task(\n",
-    "        utils.split_paths('/data/mshukor/logs/ofa/checkpoints/caption/ofa_caption_stage_1/5_0.06_6000/checkpoint_best.pt'),\n",
-    "        arg_overrides=overrides\n",
-    "    )\n",
-    "\n",
-    "# Move models to GPU\n",
-    "for model in models:\n",
-    "    model.eval()\n",
-    "    if use_fp16:\n",
-    "        model.half()\n",
-    "    if use_cuda and not cfg.distributed_training.pipeline_model_parallel:\n",
-    "        model.cuda()\n",
-    "    model.prepare_for_inference_(cfg)\n",
-    "\n",
-    "# Initialize generator\n",
-    "generator = task.build_generator(models, cfg.generation)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 5,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "model_ofa = models[0]\n",
-    "ofa_state = model_ofa.state_dict()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def get_state_given_key(state, key, excluded_keys=None):\n",
-    "    new_state = {}\n",
-    "    for k, v in state.items():\n",
-    "        if key in k:\n",
-    "            if excluded_keys is not None:\n",
-    "                if not any([ek in k for ek in excluded_keys]):\n",
-    "                    new_state[k] = v\n",
-    "            else:\n",
-    "                new_state[k] = v\n",
-    "    return new_state\n",
-    "\n",
-    "key = 'encoder.layers.0'\n",
-    "excluded_keys = ['embed', 'image']\n",
-    "ofa_tmp = get_state_given_key(ofa_state, key, excluded_keys=excluded_keys)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# def get_ofa_args_large(args):\n",
-    "#     args['encoder_embed_path'] = getattr(args, \"encoder_embed_path\", None)\n",
-    "#     args['encoder_embed_dim'] = getattr(args, \"encoder_embed_dim\", 1024)\n",
-    "#     args['encoder_ffn_embed_dim'] = getattr(args, \"encoder_ffn_embed_dim\", 4 * 1024)\n",
-    "#     args['encoder_layers'] = getattr(args, \"encoder_layers\", 12)\n",
-    "#     args['encoder_attention_heads'] = getattr(args, \"encoder_attention_heads\", 16)\n",
-    "#     args['encoder_normalize_before'] = getattr(args, \"encoder_normalize_before\", True)\n",
-    "#     args['encoder_learned_pos'] = getattr(args, \"encoder_learned_pos\", True)\n",
-    "#     args['decoder_embed_path'] = getattr(args, \"decoder_embed_path\", None)\n",
-    "#     args['decoder_embed_dim'] = getattr(args, \"decoder_embed_dim\", args['encoder_embed_dim'])\n",
-    "#     args['decoder_ffn_embed_dim'] = getattr(\n",
-    "#         args, \"decoder_ffn_embed_dim\", args['encoder_ffn_embed_dim']\n",
-    "#     )\n",
-    "#     args['decoder_layers'] = getattr(args, \"decoder_layers\", 12)\n",
-    "#     args['decoder_attention_heads'] = getattr(args, \"decoder_attention_heads\", 16)\n",
-    "#     args['decoder_normalize_before'] = getattr(args, \"decoder_normalize_before\", True)\n",
-    "#     args['decoder_learned_pos'] = getattr(args, \"decoder_learned_pos\", True)\n",
-    "#     args['attention_dropout'] = getattr(args, \"attention_dropout\", 0.0)\n",
-    "#     args['relu_dropout'] = getattr(args, \"relu_dropout\", 0.0)\n",
-    "#     args['dropout'] = getattr(args, \"dropout\", 0.0)\n",
-    "#     args['max_target_positions'] = getattr(args, \"max_target_positions\", 1024)\n",
-    "#     args['max_source_positions'] = getattr(args, \"max_source_positions\", 1024)\n",
-    "#     args['adaptive_softmax_cutoff'] = getattr(args, \"adaptive_softmax_cutoff\", None)\n",
-    "#     args['adaptive_softmax_dropout'] = getattr(args, \"adaptive_softmax_dropout\", 0)\n",
-    "#     args['share_decoder_input_output_embed'] = getattr(\n",
-    "#         args, \"share_decoder_input_output_embed\", True\n",
-    "#     )\n",
-    "#     args['share_all_embeddings'] = getattr(args, \"share_all_embeddings\", True)\n",
-    "\n",
-    "#     args['decoder_output_dim'] = getattr(\n",
-    "#         args, \"decoder_output_dim\", args['decoder_embed_dim']\n",
-    "#     )\n",
-    "#     args['decoder_input_dim'] = getattr(args, \"decoder_input_dim\", args['decoder_embed_dim'])\n",
-    "\n",
-    "#     args['no_scale_embedding'] = getattr(args, \"no_scale_embedding\", True)\n",
-    "#     args['layernorm_embedding'] = getattr(args, \"layernorm_embedding\", True)\n",
-    "\n",
-    "#     args['activation_fn'] = getattr(args, \"activation_fn\", \"gelu\")\n",
-    "#     args['pooler_activation_fn'] = getattr(args, \"pooler_activation_fn\", \"tanh\")\n",
-    "#     args['pooler_dropout'] = getattr(args, \"pooler_dropout\", 0.0)\n",
-    "#     args['pooler_classifier'] = getattr(args, \"pooler_classifier\", \"mlp\")\n",
-    "\n",
-    "#     args['resnet_drop_path_rate'] = getattr(args, \"resnet_drop_path_rate\", 0.0)\n",
-    "#     args['encoder_drop_path_rate'] = getattr(args, \"encoder_drop_path_rate\", 0.0)\n",
-    "#     args['decoder_drop_path_rate'] = getattr(args, \"decoder_drop_path_rate\", 0.0)\n",
-    "\n",
-    "#     args['resnet_type'] = getattr(args, \"resnet_type\", \"resnet152\")\n",
-    "#     args['token_bucket_size'] = getattr(args, \"token_bucket_size\", 256)\n",
-    "#     args['image_bucket_size'] = getattr(args, \"image_bucket_size\", 42)\n",
-    "\n",
-    "#     args['freeze_encoder_embedding'] = getattr(args, \"freeze_encoder_embedding\", False)\n",
-    "#     args['freeze_decoder_embedding'] = getattr(args, \"freeze_decoder_embedding\", False)\n",
-    "#     args['add_type_embedding'] = getattr(args, \"add_type_embedding\", True)\n",
-    "#     args['attn_scale_factor'] = getattr(args, \"attn_scale_factor\", 2)\n",
-    "\n",
-    "#     args['code_image_size'] = getattr(args, \"code_image_size\", 128)\n",
-    "#     args['patch_layernorm_embedding'] = getattr(args, \"patch_layernorm_embedding\", True)\n",
-    "#     args['code_layernorm_embedding'] = getattr(args, \"code_layernorm_embedding\", True)\n",
-    "#     args['entangle_position_embedding'] = getattr(args, \"entangle_position_embedding\", False)\n",
-    "#     args['disable_entangle'] = getattr(args, \"disable_entangle\", False)\n",
-    "#     args['sync_bn'] = getattr(args, \"sync_bn\", False)\n",
-    "\n",
-    "#     args['scale_attn'] = getattr(args, \"scale_attn\", False)\n",
-    "#     args['scale_fc'] = getattr(args, \"scale_fc\", False)\n",
-    "#     args['scale_heads'] = getattr(args, \"scale_heads\", False)\n",
-    "#     args['scale_resids'] = getattr(args, \"scale_resids\", False)\n",
-    "\n",
-    "#     args['orig_patch_image_size'] = getattr(args, \"orig_patch_image_size\", 256)\n",
-    "\n",
-    "#     return args"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# args = {}\n",
-    "# args = get_ofa_args_large(args)\n",
-    "# args = BatchEncoding(args)\n",
-    "# ofa_base_architecture(args)\n",
-    "# data_dir = '/data/mshukor/data/ofa/pretrain_example'\n",
-    "\n",
-    "# cfg.task.neg_sample_dir = data_dir+'/negative_sample'"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "### convert t5 weights"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 6,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "model_t5 = T5ForConditionalGeneration.from_pretrained(\"t5-base\")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "model_t5"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 7,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "t5_state = model_t5.state_dict()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 56,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import re\n",
-    "# line = re.sub(r\"</?\\[\\d+>\", \"\", line)\n",
-    "\n",
-    "mapping_dict = {\n",
-    "    ## encoder\n",
-    "    'block': 'layers', \n",
-    "    'layer.[0-9]+.SelfAttention': 'self_attn', \n",
-    "    '.q.': '.q_proj.', \n",
-    "    '.k.weight': '.k_proj.weight', \n",
-    "    '.v.': '.v_proj.', \n",
-    "    # '.o.weight': '.out_proj.weight',    \n",
-    "    'layer.0.layer_norm.': 'self_attn_layer_norm.', \n",
-    "    'layer.[0-9]+.DenseReluDense.': '', \n",
-    "    '.wi.': '.fc1.', \n",
-    "    '.wo.': '.fc2.', \n",
-    "    \n",
-    "    \n",
-    "    # decoder\n",
-    "    'layer.[0-9]+.EncDecAttention': 'encoder_attn', \n",
-    "    # 'layer.1.layer_norm.': 'encoder_attn_layer_norm.', \n",
-    "    \n",
-    "    \n",
-    "}\n",
-    "\n",
-    "encoder_mapping = {\n",
-    "    'layer.1.layer_norm.': 'final_layer_norm.',   \n",
-    "}\n",
-    "\n",
-    "decoder_mapping = {\n",
-    "    'layer.1.layer_norm.': 'encoder_attn_layer_norm.', \n",
-    "    'layer.2.layer_norm.': 'final_layer_norm.', \n",
-    "}\n",
-    "\n",
-    "\n",
-    "simple_replace_mapping = {\n",
-    "    \n",
-    " '.o.weight': '.out_proj.weight',    \n",
-    "}\n",
-    "def modify_state(state, mapping_dict, encoder_mapping, decoder_mapping, simple_replace_mapping):\n",
-    "    # orig_keys = ['block', 'layer.[0-9]+.SelfAttention', '.q.', '.k.', '.v.', '.o.', '0.layer_norm.', '.DenseReluDense.wi.', '.DenseReluDense.wo.', '.1.layer_norm.']\n",
-    "    # new_keys = ['layers', 'layer.self_attn', '.q_proj.', '.k_proj.', '.v_proj.', '.out_proj.', '.self_attn_layer_norm.', '.fc1.', '.fc2.', '.final_layer_norm.']\n",
-    "    \n",
-    "    new_state = state.copy()\n",
-    "    old_keys = []\n",
-    "    for k, v in state.items():\n",
-    "        \n",
-    "        new_key = '%s' % k  \n",
-    "        for old, new in simple_replace_mapping.items():\n",
-    "            new_key = new_key.replace(old, new)\n",
-    "            \n",
-    "        for old, new in mapping_dict.items():\n",
-    "            new_key = re.sub(r\"{}\".format(old), new, new_key)\n",
-    "            \n",
-    "        if 'encoder' in new_key:\n",
-    "            for old, new in encoder_mapping.items():\n",
-    "                new_key = re.sub(r\"{}\".format(old), new, new_key)\n",
-    "            \n",
-    "        if 'decoder' in new_key:\n",
-    "            for old, new in decoder_mapping.items():\n",
-    "                new_key = re.sub(r\"{}\".format(old), new, new_key)\n",
-    "        \n",
-    "        new_state[new_key] = v\n",
-    "        old_keys.append(k)\n",
-    "        \n",
-    "        \n",
-    "    \n",
-    "    \n",
-    "    for k in old_keys:\n",
-    "        del new_state[k]\n",
-    "        \n",
-    "    final_state = {}\n",
-    "    final_state['model'] = new_state\n",
-    "    return final_state\n",
-    "    \n",
-    "new_state = modify_state(t5_state, mapping_dict, encoder_mapping, decoder_mapping, simple_replace_mapping)\n",
-    "\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "new_state['model'].keys()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def compare_states(state1, state2):\n",
-    "    different = []\n",
-    "    for k1 in state1.keys():\n",
-    "        if k1 not in state2:\n",
-    "            different.append(k1)\n",
-    "    return different\n",
-    "            \n",
-    "tmp = compare_states(new_state, ofa_state)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 35,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "output_path = '/data/mshukor/logs/ofa/pretrained_models/t5_base.pt'\n",
-    "torch.save(new_state, output_path)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 51,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "output_path = '/data/mshukor/logs/ofa/pretrained_models/t5_base.pt'\n",
-    "\n",
-    "tmp_state = torch.load(output_path)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "\n",
-    "model_ofa.load_state_dict(tmp_state['model'],  strict=False)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "tmp_state['model'].keys()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 18,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "tmp_state = torch.load('/data/mshukor/logs/ofa/pretrained_models/ofa_base.pt')\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 19,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "dict_keys(['args', 'cfg', 'model', 'criterion', 'optimizer_history', 'task_state', 'extra_state', 'last_optimizer_state'])"
-      ]
-     },
-     "execution_count": 19,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "tmp_state.keys()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "model_t5.encoder.block[0]"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "model_ofa.encoder.layers[0]"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "### convert BART weights"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 7,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "weights_path = '/data/mshukor/logs/ofa/pretrained_models/bart.base/model.pt'\n",
-    "bart_state = torch.load(weights_path, map_location=torch.device('cpu'))"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 13,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "<All keys matched successfully>"
-      ]
-     },
-     "execution_count": 13,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "model_ofa.load_state_dict(bart_state['model'],  strict=True)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 9,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "odict_keys(['encoder.version', 'encoder.embed_tokens.weight', 'encoder.embed_positions.weight', 'encoder.layers.0.self_attn.k_proj.weight', 'encoder.layers.0.self_attn.k_proj.bias', 'encoder.layers.0.self_attn.v_proj.weight', 'encoder.layers.0.self_attn.v_proj.bias', 'encoder.layers.0.self_attn.q_proj.weight', 'encoder.layers.0.self_attn.q_proj.bias', 'encoder.layers.0.self_attn.out_proj.weight', 'encoder.layers.0.self_attn.out_proj.bias', 'encoder.layers.0.self_attn_layer_norm.weight', 'encoder.layers.0.self_attn_layer_norm.bias', 'encoder.layers.0.fc1.weight', 'encoder.layers.0.fc1.bias', 'encoder.layers.0.fc2.weight', 'encoder.layers.0.fc2.bias', 'encoder.layers.0.final_layer_norm.weight', 'encoder.layers.0.final_layer_norm.bias', 'encoder.layers.1.self_attn.k_proj.weight', 'encoder.layers.1.self_attn.k_proj.bias', 'encoder.layers.1.self_attn.v_proj.weight', 'encoder.layers.1.self_attn.v_proj.bias', 'encoder.layers.1.self_attn.q_proj.weight', 'encoder.layers.1.self_attn.q_proj.bias', 'encoder.layers.1.self_attn.out_proj.weight', 'encoder.layers.1.self_attn.out_proj.bias', 'encoder.layers.1.self_attn_layer_norm.weight', 'encoder.layers.1.self_attn_layer_norm.bias', 'encoder.layers.1.fc1.weight', 'encoder.layers.1.fc1.bias', 'encoder.layers.1.fc2.weight', 'encoder.layers.1.fc2.bias', 'encoder.layers.1.final_layer_norm.weight', 'encoder.layers.1.final_layer_norm.bias', 'encoder.layers.2.self_attn.k_proj.weight', 'encoder.layers.2.self_attn.k_proj.bias', 'encoder.layers.2.self_attn.v_proj.weight', 'encoder.layers.2.self_attn.v_proj.bias', 'encoder.layers.2.self_attn.q_proj.weight', 'encoder.layers.2.self_attn.q_proj.bias', 'encoder.layers.2.self_attn.out_proj.weight', 'encoder.layers.2.self_attn.out_proj.bias', 'encoder.layers.2.self_attn_layer_norm.weight', 'encoder.layers.2.self_attn_layer_norm.bias', 'encoder.layers.2.fc1.weight', 'encoder.layers.2.fc1.bias', 'encoder.layers.2.fc2.weight', 'encoder.layers.2.fc2.bias', 'encoder.layers.2.final_layer_norm.weight', 'encoder.layers.2.final_layer_norm.bias', 'encoder.layers.3.self_attn.k_proj.weight', 'encoder.layers.3.self_attn.k_proj.bias', 'encoder.layers.3.self_attn.v_proj.weight', 'encoder.layers.3.self_attn.v_proj.bias', 'encoder.layers.3.self_attn.q_proj.weight', 'encoder.layers.3.self_attn.q_proj.bias', 'encoder.layers.3.self_attn.out_proj.weight', 'encoder.layers.3.self_attn.out_proj.bias', 'encoder.layers.3.self_attn_layer_norm.weight', 'encoder.layers.3.self_attn_layer_norm.bias', 'encoder.layers.3.fc1.weight', 'encoder.layers.3.fc1.bias', 'encoder.layers.3.fc2.weight', 'encoder.layers.3.fc2.bias', 'encoder.layers.3.final_layer_norm.weight', 'encoder.layers.3.final_layer_norm.bias', 'encoder.layers.4.self_attn.k_proj.weight', 'encoder.layers.4.self_attn.k_proj.bias', 'encoder.layers.4.self_attn.v_proj.weight', 'encoder.layers.4.self_attn.v_proj.bias', 'encoder.layers.4.self_attn.q_proj.weight', 'encoder.layers.4.self_attn.q_proj.bias', 'encoder.layers.4.self_attn.out_proj.weight', 'encoder.layers.4.self_attn.out_proj.bias', 'encoder.layers.4.self_attn_layer_norm.weight', 'encoder.layers.4.self_attn_layer_norm.bias', 'encoder.layers.4.fc1.weight', 'encoder.layers.4.fc1.bias', 'encoder.layers.4.fc2.weight', 'encoder.layers.4.fc2.bias', 'encoder.layers.4.final_layer_norm.weight', 'encoder.layers.4.final_layer_norm.bias', 'encoder.layers.5.self_attn.k_proj.weight', 'encoder.layers.5.self_attn.k_proj.bias', 'encoder.layers.5.self_attn.v_proj.weight', 'encoder.layers.5.self_attn.v_proj.bias', 'encoder.layers.5.self_attn.q_proj.weight', 'encoder.layers.5.self_attn.q_proj.bias', 'encoder.layers.5.self_attn.out_proj.weight', 'encoder.layers.5.self_attn.out_proj.bias', 'encoder.layers.5.self_attn_layer_norm.weight', 'encoder.layers.5.self_attn_layer_norm.bias', 'encoder.layers.5.fc1.weight', 'encoder.layers.5.fc1.bias', 'encoder.layers.5.fc2.weight', 'encoder.layers.5.fc2.bias', 'encoder.layers.5.final_layer_norm.weight', 'encoder.layers.5.final_layer_norm.bias', 'encoder.layernorm_embedding.weight', 'encoder.layernorm_embedding.bias', 'decoder.version', 'decoder.embed_tokens.weight', 'decoder.embed_positions.weight', 'decoder.layers.0.self_attn.k_proj.weight', 'decoder.layers.0.self_attn.k_proj.bias', 'decoder.layers.0.self_attn.v_proj.weight', 'decoder.layers.0.self_attn.v_proj.bias', 'decoder.layers.0.self_attn.q_proj.weight', 'decoder.layers.0.self_attn.q_proj.bias', 'decoder.layers.0.self_attn.out_proj.weight', 'decoder.layers.0.self_attn.out_proj.bias', 'decoder.layers.0.self_attn_layer_norm.weight', 'decoder.layers.0.self_attn_layer_norm.bias', 'decoder.layers.0.encoder_attn.k_proj.weight', 'decoder.layers.0.encoder_attn.k_proj.bias', 'decoder.layers.0.encoder_attn.v_proj.weight', 'decoder.layers.0.encoder_attn.v_proj.bias', 'decoder.layers.0.encoder_attn.q_proj.weight', 'decoder.layers.0.encoder_attn.q_proj.bias', 'decoder.layers.0.encoder_attn.out_proj.weight', 'decoder.layers.0.encoder_attn.out_proj.bias', 'decoder.layers.0.encoder_attn_layer_norm.weight', 'decoder.layers.0.encoder_attn_layer_norm.bias', 'decoder.layers.0.fc1.weight', 'decoder.layers.0.fc1.bias', 'decoder.layers.0.fc2.weight', 'decoder.layers.0.fc2.bias', 'decoder.layers.0.final_layer_norm.weight', 'decoder.layers.0.final_layer_norm.bias', 'decoder.layers.1.self_attn.k_proj.weight', 'decoder.layers.1.self_attn.k_proj.bias', 'decoder.layers.1.self_attn.v_proj.weight', 'decoder.layers.1.self_attn.v_proj.bias', 'decoder.layers.1.self_attn.q_proj.weight', 'decoder.layers.1.self_attn.q_proj.bias', 'decoder.layers.1.self_attn.out_proj.weight', 'decoder.layers.1.self_attn.out_proj.bias', 'decoder.layers.1.self_attn_layer_norm.weight', 'decoder.layers.1.self_attn_layer_norm.bias', 'decoder.layers.1.encoder_attn.k_proj.weight', 'decoder.layers.1.encoder_attn.k_proj.bias', 'decoder.layers.1.encoder_attn.v_proj.weight', 'decoder.layers.1.encoder_attn.v_proj.bias', 'decoder.layers.1.encoder_attn.q_proj.weight', 'decoder.layers.1.encoder_attn.q_proj.bias', 'decoder.layers.1.encoder_attn.out_proj.weight', 'decoder.layers.1.encoder_attn.out_proj.bias', 'decoder.layers.1.encoder_attn_layer_norm.weight', 'decoder.layers.1.encoder_attn_layer_norm.bias', 'decoder.layers.1.fc1.weight', 'decoder.layers.1.fc1.bias', 'decoder.layers.1.fc2.weight', 'decoder.layers.1.fc2.bias', 'decoder.layers.1.final_layer_norm.weight', 'decoder.layers.1.final_layer_norm.bias', 'decoder.layers.2.self_attn.k_proj.weight', 'decoder.layers.2.self_attn.k_proj.bias', 'decoder.layers.2.self_attn.v_proj.weight', 'decoder.layers.2.self_attn.v_proj.bias', 'decoder.layers.2.self_attn.q_proj.weight', 'decoder.layers.2.self_attn.q_proj.bias', 'decoder.layers.2.self_attn.out_proj.weight', 'decoder.layers.2.self_attn.out_proj.bias', 'decoder.layers.2.self_attn_layer_norm.weight', 'decoder.layers.2.self_attn_layer_norm.bias', 'decoder.layers.2.encoder_attn.k_proj.weight', 'decoder.layers.2.encoder_attn.k_proj.bias', 'decoder.layers.2.encoder_attn.v_proj.weight', 'decoder.layers.2.encoder_attn.v_proj.bias', 'decoder.layers.2.encoder_attn.q_proj.weight', 'decoder.layers.2.encoder_attn.q_proj.bias', 'decoder.layers.2.encoder_attn.out_proj.weight', 'decoder.layers.2.encoder_attn.out_proj.bias', 'decoder.layers.2.encoder_attn_layer_norm.weight', 'decoder.layers.2.encoder_attn_layer_norm.bias', 'decoder.layers.2.fc1.weight', 'decoder.layers.2.fc1.bias', 'decoder.layers.2.fc2.weight', 'decoder.layers.2.fc2.bias', 'decoder.layers.2.final_layer_norm.weight', 'decoder.layers.2.final_layer_norm.bias', 'decoder.layers.3.self_attn.k_proj.weight', 'decoder.layers.3.self_attn.k_proj.bias', 'decoder.layers.3.self_attn.v_proj.weight', 'decoder.layers.3.self_attn.v_proj.bias', 'decoder.layers.3.self_attn.q_proj.weight', 'decoder.layers.3.self_attn.q_proj.bias', 'decoder.layers.3.self_attn.out_proj.weight', 'decoder.layers.3.self_attn.out_proj.bias', 'decoder.layers.3.self_attn_layer_norm.weight', 'decoder.layers.3.self_attn_layer_norm.bias', 'decoder.layers.3.encoder_attn.k_proj.weight', 'decoder.layers.3.encoder_attn.k_proj.bias', 'decoder.layers.3.encoder_attn.v_proj.weight', 'decoder.layers.3.encoder_attn.v_proj.bias', 'decoder.layers.3.encoder_attn.q_proj.weight', 'decoder.layers.3.encoder_attn.q_proj.bias', 'decoder.layers.3.encoder_attn.out_proj.weight', 'decoder.layers.3.encoder_attn.out_proj.bias', 'decoder.layers.3.encoder_attn_layer_norm.weight', 'decoder.layers.3.encoder_attn_layer_norm.bias', 'decoder.layers.3.fc1.weight', 'decoder.layers.3.fc1.bias', 'decoder.layers.3.fc2.weight', 'decoder.layers.3.fc2.bias', 'decoder.layers.3.final_layer_norm.weight', 'decoder.layers.3.final_layer_norm.bias', 'decoder.layers.4.self_attn.k_proj.weight', 'decoder.layers.4.self_attn.k_proj.bias', 'decoder.layers.4.self_attn.v_proj.weight', 'decoder.layers.4.self_attn.v_proj.bias', 'decoder.layers.4.self_attn.q_proj.weight', 'decoder.layers.4.self_attn.q_proj.bias', 'decoder.layers.4.self_attn.out_proj.weight', 'decoder.layers.4.self_attn.out_proj.bias', 'decoder.layers.4.self_attn_layer_norm.weight', 'decoder.layers.4.self_attn_layer_norm.bias', 'decoder.layers.4.encoder_attn.k_proj.weight', 'decoder.layers.4.encoder_attn.k_proj.bias', 'decoder.layers.4.encoder_attn.v_proj.weight', 'decoder.layers.4.encoder_attn.v_proj.bias', 'decoder.layers.4.encoder_attn.q_proj.weight', 'decoder.layers.4.encoder_attn.q_proj.bias', 'decoder.layers.4.encoder_attn.out_proj.weight', 'decoder.layers.4.encoder_attn.out_proj.bias', 'decoder.layers.4.encoder_attn_layer_norm.weight', 'decoder.layers.4.encoder_attn_layer_norm.bias', 'decoder.layers.4.fc1.weight', 'decoder.layers.4.fc1.bias', 'decoder.layers.4.fc2.weight', 'decoder.layers.4.fc2.bias', 'decoder.layers.4.final_layer_norm.weight', 'decoder.layers.4.final_layer_norm.bias', 'decoder.layers.5.self_attn.k_proj.weight', 'decoder.layers.5.self_attn.k_proj.bias', 'decoder.layers.5.self_attn.v_proj.weight', 'decoder.layers.5.self_attn.v_proj.bias', 'decoder.layers.5.self_attn.q_proj.weight', 'decoder.layers.5.self_attn.q_proj.bias', 'decoder.layers.5.self_attn.out_proj.weight', 'decoder.layers.5.self_attn.out_proj.bias', 'decoder.layers.5.self_attn_layer_norm.weight', 'decoder.layers.5.self_attn_layer_norm.bias', 'decoder.layers.5.encoder_attn.k_proj.weight', 'decoder.layers.5.encoder_attn.k_proj.bias', 'decoder.layers.5.encoder_attn.v_proj.weight', 'decoder.layers.5.encoder_attn.v_proj.bias', 'decoder.layers.5.encoder_attn.q_proj.weight', 'decoder.layers.5.encoder_attn.q_proj.bias', 'decoder.layers.5.encoder_attn.out_proj.weight', 'decoder.layers.5.encoder_attn.out_proj.bias', 'decoder.layers.5.encoder_attn_layer_norm.weight', 'decoder.layers.5.encoder_attn_layer_norm.bias', 'decoder.layers.5.fc1.weight', 'decoder.layers.5.fc1.bias', 'decoder.layers.5.fc2.weight', 'decoder.layers.5.fc2.bias', 'decoder.layers.5.final_layer_norm.weight', 'decoder.layers.5.final_layer_norm.bias', 'decoder.layernorm_embedding.weight', 'decoder.layernorm_embedding.bias'])"
-      ]
-     },
-     "execution_count": 9,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "bart_state['model'].keys()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 12,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "tensor([[ 0.0125,  0.0014, -0.0096,  ...,  0.0022,  0.1057,  0.0103],\n",
-       "        [-0.0114, -0.0169, -0.0184,  ..., -0.0131, -0.0043, -0.0053],\n",
-       "        [ 0.0842, -0.0389,  0.0096,  ...,  0.0583,  0.0082,  0.0357],\n",
-       "        ...,\n",
-       "        [-0.0032, -0.0313, -0.1026,  ...,  0.0138,  0.0056, -0.0023],\n",
-       "        [ 0.0104, -0.0045,  0.0263,  ...,  0.0158,  0.0324, -0.0111],\n",
-       "        [-0.0038, -0.0532, -0.0147,  ...,  0.0067,  0.0256,  0.0009]])"
-      ]
-     },
-     "execution_count": 12,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "ofa_state.keys()\n",
-    "ofa_state['encoder.embed_tokens.weight']"
-   ]
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "ofa",
-   "language": "python",
-   "name": "ofa"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.7.4"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 4
-}

prompt_tuning.md

@@ -1,66 +0,0 @@
-<!---
-Copyright 2022 The OFA-Sys Team. 
-All rights reserved.
-This source code is licensed under the Apache 2.0 license found in the LICENSE file in the root directory.
--->
-
-## Prompt Tuning for Generative Multimodal Pretrained Models
-
-### Overview
-This is the code for **"Prompt Tuning for Generative Multimodal Pretrained Models"**, [Check our paper on ArXiv](https://arxiv.org/abs/2208.02532). This paper explores prompt tuning for generative multimodal pretrained models, instead of the constrastive learning models. We specifically focuses on the unified sequence-to-sequence learning framework and implement on our OFA models. 
-<br>
-
-### Requirements
-* python 3.7.4
-* pytorch 1.8.1
-* torchvision 0.9.1
-* JAVA 1.8 (for COCO evaluation)
-<br></br>
-
-### Installation
-```bash
-pip install -r requirements.txt
-```
-<br>
-
-### Datasets and Checkpoints
-See [datasets.md](datasets.md) and [checkpoints.md](checkpoints.md).
-<br>
-
-### Training
-We provide a demo script (`run_scripts/refcoco/train_refcoco_prefix.sh`) that has all the required parts for training.
-
-```sh
-sh ./run_scripts/refcoco/train_refcoco_prefix.sh
-```
-A few options of note:
-*   `--encoder-prompt` :: whether to insert prompts to the encoder
-*   `--decoder-prompt` :: whether to insert prompts to the decoder
-*   `--encoder-prompt-length` :: encoder prompt length
-*   `--decoder-prompt-length` :: decoder prompt length
-*   `--bitfit` :: whether to use bitfit
-*   `--adapter` :: whether to use adapter
-*   `--adapter-dim` :: adapter projection dim
-
-We recommend that your workspace directory should be organized like this: 
-```
-OFA/
-├── checkpoints/
-│   ├── ofa_base.pt
-│   ├── ofa_large.pt
-│   └── ...
-├── criterions/
-├── data/
-├── dataset/
-│   ├── caption_data/
-│   ├── refcoco_data/
-│   └── ...
-├── fairseq/
-├── models/
-├── run_scripts/
-├── tasks/
-├── train.py
-├── trainer.py
-└── utils/
-```
-<br>

spaces.md

@@ -1,8 +0,0 @@
-# Spaces
-To provide better experience, we plan to build demos for our OFA models on Huggingface Spaces. Below we provide links to the demos. Have fun!
-
-* Generic Interface: [![Hugging Face Spaces](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue)](https://huggingface.co/spaces/OFA-Sys/OFA-Generic_Interface)
-* Text-to-Image Generation: [![Hugging Face Spaces](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue)](https://huggingface.co/spaces/OFA-Sys/OFA-Text2Image_Generation)
-* Image Captioning: [![Hugging Face Spaces](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue)](https://huggingface.co/spaces/OFA-Sys/OFA-Image_Caption)
-* Referring Expression Comprehension: [![Hugging Face Spaces](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue)](https://huggingface.co/spaces/OFA-Sys/OFA-Visual_Grounding)
-* Visual Question Answering: [![Hugging Face Spaces](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue)](https://huggingface.co/spaces/OFA-Sys/OFA-Visual_Question_Answering)

test.py

@@ -1,101 +0,0 @@
-from data.audio_utils import get_audio_features, int16_to_float32, float32_to_int16, AUDIO_CFG
-import soundfile as sf
-import io 
-import torch 
-import numpy as np
-
-AUDIO_CFG =  {
-        "audio_length": 1024,
-        "clip_samples": 480000,
-        "mel_bins": 64,
-        "sample_rate": 48000,
-        "window_size": 1024,
-        "hop_size": 480,
-        "fmin": 50,
-        "fmax": 14000,
-        "class_num": 527,
-    }
-    
-
-
-audio_cfg = AUDIO_CFG
-max_len = 480000
-data_path = '/work/NAT/gda2204/mshukor/data/audiocaps/train/--CHY2qO5zc.wav'
-
-audio_data, orig_sr = sf.read(data_path)
-# import librosa
-# audio_data, orig_sr = librosa.load(data_path, sr=48000)
-
-print(orig_sr)
-if audio_data.ndim>1:
-  audio_data = np.mean(audio_data,axis=1)
-
- 
-print(audio_data.shape, audio_data)
-
-audio_data = int16_to_float32(float32_to_int16(audio_data))
-audio_data = torch.tensor(audio_data).float()
-print(audio_data.dtype)
-print(audio_data.shape, audio_data)
-# the 'fusion' truncate mode can be changed to 'rand_trunc' if run in unfusion mode
-sample = {}
-
-sample = get_audio_features(
-    sample, audio_data, max_len, 
-    data_truncating='fusion', 
-    data_filling='repeatpad',
-    audio_cfg=audio_cfg,
-)
-
-patch_audio = sample['waveform'] #.half()
-print(patch_audio.shape, patch_audio.min(), patch_audio.max(), patch_audio)
-
-patch_audio = torch.zeros(480000)
-print(patch_audio.shape)
-
-
-from torchlibrosa.stft import Spectrogram, LogmelFilterBank
-
-AUDIO_CFG =  {
-        "sample_rate": 48000,
-        "audio_length": 1024,
-        "clip_samples": 480000,
-        "mel_bins": 64,
-        "sample_rate": 48000,
-        "window_size": 1024,
-        "hop_size": 480,
-        "fmin": 50,
-        "fmax": 14000,
-        "class_num": 527,
-    }
-
-window = 'hann'
-center = True
-pad_mode = 'reflect'
-ref = 1.0
-amin = 1e-10
-top_db = None
-
-spectrogram_extractor = Spectrogram(n_fft=AUDIO_CFG['window_size'], hop_length=AUDIO_CFG['hop_size'],
-            win_length=AUDIO_CFG['window_size'], window=window, center=center, pad_mode=pad_mode,
-            freeze_parameters=True)
-
-
-logmel_extractor = LogmelFilterBank(sr=AUDIO_CFG['sample_rate'], n_fft=AUDIO_CFG['window_size'],
-            n_mels=AUDIO_CFG['mel_bins'], fmin=AUDIO_CFG['fmin'], fmax=AUDIO_CFG['fmax'], 
-            ref=ref, amin=amin, top_db=top_db,
-            freeze_parameters=True)#.half()
-
-
-patch_audio = patch_audio[None, :]
-print(patch_audio.shape)
-spectro = spectrogram_extractor(patch_audio)
-
-print(spectro.shape)
-print(spectro)
-
-
-mel = logmel_extractor(spectro)
-
-print(mel.shape)
-print(mel)

train.py

@@ -1,729 +0,0 @@
-#!/usr/bin/env python3 -u
-# Copyright 2022 The OFA-Sys Team. 
-# All rights reserved.
-# This source code is licensed under the Apache 2.0 license 
-# found in the LICENSE file in the root directory.
-
-"""
-Train a new model on one or across multiple GPUs.
-"""
-
-import argparse
-import logging
-import math
-import os
-import sys
-from typing import Dict, Optional, Any, List, Tuple, Callable
- 
-# We need to setup root logger before importing any fairseq libraries.
-logging.basicConfig(
-    format='%(asctime)s - %(filename)s[line:%(lineno)d] - %(levelname)s: %(message)s',
-    datefmt="%Y-%m-%d %H:%M:%S",
-    level=os.environ.get("LOGLEVEL", "INFO").upper(),
-    stream=sys.stdout,
-)
-logger = logging.getLogger("fairseq_cli.train")
-
-import numpy as np
-import torch
-from fairseq import (
-    # checkpoint_utils,
-    options,
-    quantization_utils,
-    tasks,
-    utils,
-)
-from fairseq.data import iterators
-from fairseq.data.plasma_utils import PlasmaStore
-from fairseq.dataclass.configs import FairseqConfig
-from fairseq.dataclass.utils import convert_namespace_to_omegaconf
-from fairseq.distributed import fsdp_enable_wrap, fsdp_wrap, utils as distributed_utils
-from fairseq.file_io import PathManager
-from fairseq.logging import meters, metrics, progress_bar
-from fairseq.model_parallel.megatron_trainer import MegatronTrainer
-# from fairseq.trainer import Trainer
-from omegaconf import DictConfig, OmegaConf
-
-from utils import checkpoint_utils
-from trainer import Trainer
-
-from utils.utils import print_trainable_params_percentage, setup_for_distributed
-
-import psutil
-
-def main(cfg: FairseqConfig) -> None:
-    print(distributed_utils.is_master(cfg.distributed_training))
-    print(cfg.distributed_training)
-    setup_for_distributed(distributed_utils.is_master(cfg.distributed_training))
-
-    if isinstance(cfg, argparse.Namespace):
-        cfg = convert_namespace_to_omegaconf(cfg)
-
-    utils.import_user_module(cfg.common)
-
-    if distributed_utils.is_master(cfg.distributed_training) and "job_logging_cfg" in cfg:
-        # make hydra logging work with ddp (see # see https://github.com/facebookresearch/hydra/issues/1126)
-        logging.config.dictConfig(OmegaConf.to_container(cfg.job_logging_cfg))
-
-    assert (
-        cfg.dataset.max_tokens is not None or cfg.dataset.batch_size is not None
-    ), "Must specify batch size either with --max-tokens or --batch-size"
-    metrics.reset()
-
-    if cfg.common.log_file is not None:
-        handler = logging.FileHandler(filename=cfg.common.log_file)
-        logger.addHandler(handler)
-
-    np.random.seed(cfg.common.seed)
-    utils.set_torch_seed(cfg.common.seed)
-
-    if distributed_utils.is_master(cfg.distributed_training):
-        print(cfg.checkpoint.save_dir)
-        checkpoint_utils.verify_checkpoint_directory(cfg.checkpoint.save_dir)
-
-    
-    # Print args
-    logger.info(cfg)
-
-
-    if cfg.checkpoint.write_checkpoints_asynchronously:
-        try:
-            import iopath  # noqa: F401
-        except ImportError:
-            logging.exception(
-                "Asynchronous checkpoint writing is specified but iopath is "
-                "not installed: `pip install iopath`"
-            )
-            return
-
-    # Setup task, e.g., translation, language modeling, etc.
-    task = tasks.setup_task(cfg.task)
-
-    assert cfg.criterion, "Please specify criterion to train a model"
-
-    # Build model and criterion
-    if cfg.distributed_training.ddp_backend == "fully_sharded":
-        with fsdp_enable_wrap(cfg.distributed_training):
-            model = fsdp_wrap(task.build_model(cfg.model))
-    else:
-        model = task.build_model(cfg.model)
-
-    # bitfit
-    if cfg.model.bitfit:
-        for name, param in model.named_parameters():
-            if ("layer_norm" in name and "bias" in name) or ("fc" in name and "bias" in name):
-                param.requires_grad = True
-            else:
-                param.requires_grad = False
-
-    criterion = task.build_criterion(cfg.criterion)
-
-    logger.info(model)
-    logger.info("task: {}".format(task.__class__.__name__))
-    logger.info("model: {}".format(model.__class__.__name__))
-    logger.info("criterion: {}".format(criterion.__class__.__name__))
-    logger.info(
-        "num. shared model params: {:,} (num. trained: {:,})".format(
-            sum(p.numel() for p in model.parameters() if not getattr(p, "expert", False)),
-            sum(p.numel() for p in model.parameters() if not getattr(p, "expert", False) and p.requires_grad)
-        )
-    )
-
-    logger.info(
-        "num. expert model params: {} (num. trained: {})".format(
-            sum(p.numel() for p in model.parameters() if getattr(p, "expert", False)),
-            sum(p.numel() for p in model.parameters() if getattr(p, "expert", False) and p.requires_grad),
-        )
-    )
-
-    # Load valid dataset (we load training data below, based on the latest checkpoint)
-    # We load the valid dataset AFTER building the model
-    # data_utils.raise_if_valid_subsets_unintentionally_ignored(cfg)
-    if cfg.dataset.combine_valid_subsets:
-        task.load_dataset("valid", combine=True, epoch=1)
-    else:
-        for valid_sub_split in cfg.dataset.valid_subset.split(","):
-            task.load_dataset(valid_sub_split, combine=False, epoch=1)
-
-    # (optionally) Configure quantization
-    if cfg.common.quantization_config_path is not None:
-        quantizer = quantization_utils.Quantizer(
-            config_path=cfg.common.quantization_config_path,
-            max_epoch=cfg.optimization.max_epoch,
-            max_update=cfg.optimization.max_update,
-        )
-    else:
-        quantizer = None
-
-
-
-
-    # for n, p in model.named_parameters():
-    #     if not p.requires_grad:
-    #         print(n)
-
-    # Build trainer
-    if cfg.common.model_parallel_size == 1:
-        trainer = Trainer(cfg, task, model, criterion, quantizer)
-    else:
-        trainer = MegatronTrainer(cfg, task, model, criterion)
-    logger.info(
-        "training on {} devices (GPUs/TPUs)".format(
-            cfg.distributed_training.distributed_world_size
-        )
-    )
-    logger.info(
-        "max tokens per device = {} and max sentences per device = {}".format(
-            cfg.dataset.max_tokens,
-            cfg.dataset.batch_size,
-        )
-    )
-
-
-    # Load the latest checkpoint if one is available and restore the
-    # corresponding train iterator
-    strict = getattr(cfg.model, 'strict', True)
-    logger.info('load checkpoint, strict:{}'.format(strict))
-    extra_state, epoch_itr = checkpoint_utils.load_checkpoint(
-        cfg.checkpoint,
-        trainer,
-        strict=strict,
-        # don't cache epoch iterators for sharded datasets
-        disable_iterator_cache=True,
-        load_on_cuda=cfg.checkpoint.load_on_cuda,
-    )
-    if cfg.common.tpu:
-        import torch_xla.core.xla_model as xm
-        xm.rendezvous("load_checkpoint")  # wait for all workers
-
-    max_epoch = cfg.optimization.max_epoch or math.inf
-    if max_epoch > 0 and max_epoch != math.inf:
-        total_num_updates = sum(
-            math.ceil(len(epoch_itr) / cfg.optimization.update_freq[i])
-            if i < len(cfg.optimization.update_freq) else
-            math.ceil(len(epoch_itr) / cfg.optimization.update_freq[-1])
-            for i in range(max_epoch)
-        )
-        trainer.lr_reinit(total_num_updates, trainer.get_num_updates())
-
-    # if getattr(cfg.model, "freeze_encoder", False):
-    #     for idx, layer in enumerate(model.encoder.layers):
-    #         layer.requires_grad_(False)
-    # if getattr(cfg.model, "freeze_decoder", False):
-    #     for idx, layer in enumerate(model.decoder.layers):
-    #         layer.requires_grad_(False)  
-
-    # if hasattr(cfg.model, 'progressive') or getattr(cfg.model, "freeze_perception", False):
-    #     custom_unfreeze(trainer, epoch_itr, cfg.model)
-
-    # if hasattr(cfg.model, 'only_linear_proj') and getattr(cfg.model, "only_linear_proj", False):
-    #     model.requires_grad_(False)
-    #     model.encoder.embed_tokens.weight.requires_grad = True
-    #     model.decoder.embed_tokens.weight.requires_grad = True
-
-    #     if getattr(cfg.model, "freeze_encoder_embedding", False) or getattr(
-    #             cfg.model, "encoder_prompt", False) or getattr(cfg.model, "decoder_prompt", False) or getattr(cfg.model, "adapter", False):    
-    #         model.encoder.embed_tokens.weight.requires_grad = False
-    #     if getattr(cfg.model, "freeze_decoder_embedding", False) or getattr(
-    #             cfg.model, "encoder_prompt", False) or getattr(cfg.model, "decoder_prompt", False) or getattr(cfg.model, "adapter", False):    
-    #         model.decoder.embed_tokens.weight.requires_grad = False
-
-        
-    #     model.encoder.image_proj.requires_grad_(True)
-    #     if getattr(cfg.model, "video_encoder_name", None):
-    #         model.encoder.video_proj.requires_grad_(True)
-    #     if getattr(cfg.model, "audio_encoder_name", None):
-    #         model.encoder.audio_proj.requires_grad_(True)
-
-
-    print_trainable_params_percentage(model)
-
-    lr = trainer.get_lr()
-
-    train_meter = meters.StopwatchMeter()
-    train_meter.start()
-    while epoch_itr.next_epoch_idx <= max_epoch:
-        if lr <= cfg.optimization.stop_min_lr:
-            logger.info(
-                f"stopping training because current learning rate ({lr}) is smaller "
-                "than or equal to minimum learning rate "
-                f"(--stop-min-lr={cfg.optimization.stop_min_lr})"
-            )
-            break
-
-        # train for one epoch
-
-        valid_losses, should_stop = train(cfg, trainer, task, epoch_itr)
-        if should_stop:
-            break
-
-        # only use first validation loss to update the learning rate
-        lr = trainer.lr_step(epoch_itr.epoch, valid_losses[0])
-
-        epoch_itr = trainer.get_train_iterator(
-            epoch_itr.next_epoch_idx,
-            # sharded data: get train iterator for next epoch
-            load_dataset=True,
-            # don't cache epoch iterators for sharded datasets
-            disable_iterator_cache=True,
-        )
-    train_meter.stop()
-    logger.info("done training in {:.1f} seconds".format(train_meter.sum))
-
-    # ioPath implementation to wait for all asynchronous file writes to complete.
-    if cfg.checkpoint.write_checkpoints_asynchronously:
-        logger.info(
-            "ioPath PathManager waiting for all asynchronous checkpoint "
-            "writes to finish."
-        )
-        PathManager.async_close()
-        logger.info("ioPath PathManager finished waiting.")
-
-
-def should_stop_early(cfg: DictConfig, valid_loss: float) -> bool:
-    # skip check if no validation was done in the current epoch
-    if valid_loss is None:
-        return False
-    if cfg.checkpoint.patience <= 0:
-        return False
-
-    def is_better(a, b):
-        return a > b if cfg.checkpoint.maximize_best_checkpoint_metric else a < b
-
-    prev_best = getattr(should_stop_early, "best", None)
-    if prev_best is None or is_better(valid_loss, prev_best):
-        should_stop_early.best = valid_loss
-        should_stop_early.num_runs = 0
-        return False
-    else:
-        should_stop_early.num_runs += 1
-        if should_stop_early.num_runs >= cfg.checkpoint.patience:
-            logger.info(
-                "early stop since valid performance hasn't improved for last {} runs".format(
-                    cfg.checkpoint.patience
-                )
-            )
-            return True
-        else:
-            return False
-
-
-@metrics.aggregate("train")
-def train(
-    cfg: DictConfig, trainer: Trainer, task: tasks.FairseqTask, epoch_itr
-) -> Tuple[List[Optional[float]], bool]:
-    """Train the model for one epoch and return validation losses."""
-    # Initialize data iterator
-    itr = epoch_itr.next_epoch_itr(
-        fix_batches_to_gpus=cfg.distributed_training.fix_batches_to_gpus,
-        shuffle=(epoch_itr.next_epoch_idx > cfg.dataset.curriculum),
-    )
-    update_freq = (
-        cfg.optimization.update_freq[epoch_itr.epoch - 1]
-        if epoch_itr.epoch <= len(cfg.optimization.update_freq)
-        else cfg.optimization.update_freq[-1]
-    )
-    itr = iterators.GroupedIterator(itr, update_freq)
-    if cfg.common.tpu:
-        itr = utils.tpu_data_loader(itr)
-    progress = progress_bar.progress_bar(
-        itr,
-        log_format=cfg.common.log_format,
-        log_file=cfg.common.log_file,
-        log_interval=cfg.common.log_interval,
-        epoch=epoch_itr.epoch,
-        tensorboard_logdir=(
-            cfg.common.tensorboard_logdir
-            if distributed_utils.is_master(cfg.distributed_training)
-            else None
-        ),
-        default_log_format=("tqdm" if not cfg.common.no_progress_bar else "simple"),
-        wandb_project=(
-            cfg.common.wandb_project
-            if distributed_utils.is_master(cfg.distributed_training)
-            else None
-        ),
-        wandb_run_name=os.environ.get(
-            "WANDB_NAME", os.path.basename(cfg.checkpoint.save_dir)
-        ),
-        azureml_logging=(
-            cfg.common.azureml_logging
-            if distributed_utils.is_master(cfg.distributed_training)
-            else False
-        ),
-    )
-    progress.update_config(_flatten_config(cfg))
-
-    trainer.begin_epoch(epoch_itr.epoch)
-
-    valid_subsets = cfg.dataset.valid_subset.split(",")
-    should_stop = False
-    num_updates = trainer.get_num_updates()
-    logger.info("Start iterating over samples")
-    for i, samples in enumerate(progress):
-        with metrics.aggregate("train_inner"), torch.autograd.profiler.record_function(
-            "train_step-%d" % i
-        ):
-            log_output = trainer.train_step(samples)
-
-        if log_output is not None:  # not OOM, overflow, ...
-            # log mid-epoch stats
-            num_updates = trainer.get_num_updates()
-            if num_updates % cfg.common.log_interval == 0:
-                stats = get_training_stats(metrics.get_smoothed_values("train_inner"))
-                progress.log(stats, tag="train_inner", step=num_updates)
-
-                # reset mid-epoch stats after each log interval
-                # the end-of-epoch stats will still be preserved
-                metrics.reset_meters("train_inner")
-
-        end_of_epoch = not itr.has_next()
-        valid_losses, should_stop = validate_and_save(
-            cfg, trainer, task, epoch_itr, valid_subsets, end_of_epoch
-        )
-
-        # if (hasattr(cfg.model, 'progressive') or hasattr(cfg.model, 'only_linear_proj') or hasattr(cfg.model, 'freeze_perception')) and end_of_epoch:
-        #     custom_unfreeze(trainer, epoch_itr, cfg.model)
-        #     print_trainable_params_percentage(trainer.model)
-
-        if should_stop:
-            break
-
-
-        # print(i, len(progress))
-        # if i > 5:
-        #     break
-    # log end-of-epoch stats
-    logger.info("end of epoch {} (average epoch stats below)".format(epoch_itr.epoch))
-    stats = get_training_stats(metrics.get_smoothed_values("train"))
-    progress.print(stats, tag="train", step=num_updates)
-    print_trainable_params_percentage(trainer.model)
-    # reset epoch-level meters
-    metrics.reset_meters("train")
-    return valid_losses, should_stop
-
-# progressive training
-def custom_unfreeze(trainer, epoch_itr, cfg):
-    model = trainer.model
-    epoch = epoch_itr.epoch
-    print("Epoch, ", epoch)
-    ## unfreeze epochs
-    unfreeze_epoch_encoder = cfg.unfreeze_epoch_encoder
-    unfreeze_epoch_decoder = cfg.unfreeze_epoch_decoder
-
-    unfreeze_epoch_image = cfg.unfreeze_epoch_image
-    unfreeze_epoch_video = cfg.unfreeze_epoch_video
-    unfreeze_epoch_audio = cfg.unfreeze_epoch_audio
-
-
-    if getattr(cfg, "only_linear_proj", False):
-        unfreeze_epoch = cfg.unfreeze_epoch
-        if epoch >= unfreeze_epoch:
-            model.requires_grad_(True)
-        
-            if getattr(cfg, "freeze_encoder_embedding", False) or getattr(
-                    cfg, "encoder_prompt", False) or getattr(cfg, "decoder_prompt", False) or getattr(cfg, "adapter", False):    
-                model.encoder.embed_tokens.weight.requires_grad = False
-            if getattr(cfg, "freeze_decoder_embedding", False) or getattr(
-                    cfg, "encoder_prompt", False) or getattr(cfg, "decoder_prompt", False) or getattr(cfg, "adapter", False):    
-                model.decoder.embed_tokens.weight.requires_grad = False
-
-        if trainer._ema is not None:
-            trainer._ema.requires_grad_(False) 
-        print_trainable_params_percentage(model)
-        return 
-
-    if getattr(cfg, "freeze_perception", False):
-
-        if hasattr(model.encoder, 'embed_images'):
-            if epoch >= unfreeze_epoch_image:
-                grad = True
-            else:
-                grad = False
-            model.encoder.embed_images.requires_grad_(grad) 
-            print('model.encoder.embed_images.requires_grad', grad)
-        if hasattr(model.encoder, 'embed_videos'):
-            if epoch >= unfreeze_epoch_video:
-                grad = True
-            else:
-                grad = False
-            model.encoder.embed_videos.requires_grad_(grad) 
-            print('model.encoder.embed_videos.requires_grad', grad)
-
-        if hasattr(model.encoder, 'embed_audios'):
-            if epoch >= unfreeze_epoch_audio:
-                grad = True
-            else:
-                grad = False
-            model.encoder.embed_audios.requires_grad_(grad) 
-            print('model.encoder.embed_audios.requires_grad', grad)
-
-        if trainer._ema is not None:
-            trainer._ema.requires_grad_(False) 
-        return 
-
-    if epoch >= unfreeze_epoch_encoder:
-        grad=True
-    else:
-        grad=False
-    for l in model.encoder.layers:
-        l.requires_grad_(grad) 
-    print('model.encoder.layers.requires_grad', grad)
-
-    if epoch >= unfreeze_epoch_decoder:
-        grad=True
-    else:
-        grad=False
-    for l in model.decoder.layers:
-        l.requires_grad_(grad) 
-    print('model.decoder.layers.requires_grad', grad)
-
-    if getattr(cfg, "freeze_encoder_embedding", False) or getattr(
-            cfg, "encoder_prompt", False) or getattr(cfg, "decoder_prompt", False) or getattr(cfg, "adapter", False):    
-        model.encoder.embed_tokens.weight.requires_grad = False
-    if getattr(cfg, "freeze_decoder_embedding", False) or getattr(
-            cfg, "encoder_prompt", False) or getattr(cfg, "decoder_prompt", False) or getattr(cfg, "adapter", False):    
-        model.decoder.embed_tokens.weight.requires_grad = False
-
-        if getattr(cfg, "encoder_prompt", False):
-            model.encoder.encoder_prompt_encoder.requires_grad_(True)
-        if getattr(cfg, "decoder_prompt", False):
-            model.decoder.decoder_prompt_encoder.requires_grad_(True)
-        if getattr(cfg, "adapter", False):
-            for idx, layer in enumerate(model.encoder.layers):
-                layer.adapter.requires_grad_(True)
-            for idx, layer in enumerate(model.decoder.layers):
-                layer.adapter.requires_grad_(True)  
-
-    if hasattr(model.encoder, 'embed_images'):
-        if epoch >= unfreeze_epoch_image:
-            grad = True
-        else:
-            grad = False
-        model.encoder.embed_images.requires_grad_(grad) 
-        print('model.encoder.embed_images.requires_grad', grad)
-    if hasattr(model.encoder, 'embed_videos'):
-        if epoch >= unfreeze_epoch_video:
-            grad = True
-        else:
-            grad = False
-        model.encoder.embed_videos.requires_grad_(grad) 
-        print('model.encoder.embed_videos.requires_grad', grad)
-
-    if hasattr(model.encoder, 'embed_audios'):
-        if epoch >= unfreeze_epoch_audio:
-            grad = True
-        else:
-            grad = False
-        model.encoder.embed_audios.requires_grad_(grad) 
-        print('model.encoder.embed_audios.requires_grad', grad)
-
-    if trainer._ema is not None:
-        trainer._ema.requires_grad_(False) 
-
-def _flatten_config(cfg: DictConfig):
-    config = OmegaConf.to_container(cfg)
-    # remove any legacy Namespaces and replace with a single "args"
-    namespace = None
-    for k, v in list(config.items()):
-        if isinstance(v, argparse.Namespace):
-            namespace = v
-            del config[k]
-    if namespace is not None:
-        config["args"] = vars(namespace)
-    return config
-
-
-def validate_and_save(
-    cfg: DictConfig,
-    trainer: Trainer,
-    task: tasks.FairseqTask,
-    epoch_itr,
-    valid_subsets: List[str],
-    end_of_epoch: bool,
-) -> Tuple[List[Optional[float]], bool]:
-    num_updates = trainer.get_num_updates()
-    max_update = cfg.optimization.max_update or math.inf
-
-    # Stopping conditions (and an additional one based on validation loss later
-    # on)
-    should_stop = False
-    if num_updates >= max_update:
-        should_stop = True
-        logger.info(
-            f"Stopping training due to "
-            f"num_updates: {num_updates} >= max_update: {max_update}"
-        )
-
-    training_time_hours = trainer.cumulative_training_time() / (60 * 60)
-    if (
-        cfg.optimization.stop_time_hours > 0
-        and training_time_hours > cfg.optimization.stop_time_hours
-    ):
-        should_stop = True
-        logger.info(
-            f"Stopping training due to "
-            f"cumulative_training_time: {training_time_hours} > "
-            f"stop_time_hours: {cfg.optimization.stop_time_hours} hour(s)"
-        )
-
-    do_save = (
-        (end_of_epoch and epoch_itr.epoch % cfg.checkpoint.save_interval == 0)
-        or should_stop
-        or (
-            cfg.checkpoint.save_interval_updates > 0
-            and num_updates > 0
-            and num_updates % cfg.checkpoint.save_interval_updates == 0
-            and num_updates >= cfg.dataset.validate_after_updates
-        )
-    )
-    do_validate = (
-        (not end_of_epoch and do_save)  # validate during mid-epoch saves
-        or (end_of_epoch and epoch_itr.epoch % cfg.dataset.validate_interval == 0)
-        or should_stop
-        or (
-            cfg.dataset.validate_interval_updates > 0
-            and num_updates > 0
-            and num_updates % cfg.dataset.validate_interval_updates == 0
-        )
-    ) and not cfg.dataset.disable_validation and num_updates >= cfg.dataset.validate_after_updates
-
-    # Validate
-    valid_losses = [None]
-    if do_validate:
-        valid_losses = validate(cfg, trainer, task, epoch_itr, valid_subsets)
-
-    should_stop |= should_stop_early(cfg, valid_losses[0])
-
-    # Save checkpoint
-    if do_save or should_stop:
-        checkpoint_utils.save_checkpoint(
-            cfg.checkpoint, trainer, epoch_itr, valid_losses[0], save_on_cuda=cfg.checkpoint.save_on_cuda,
-        )
-
-    return valid_losses, should_stop
-
-
-def get_training_stats(stats: Dict[str, Any]) -> Dict[str, Any]:
-    stats["wall"] = round(metrics.get_meter("default", "wall").elapsed_time, 0)
-    return stats
-
-
-def validate(
-    cfg: DictConfig,
-    trainer: Trainer,
-    task: tasks.FairseqTask,
-    epoch_itr,
-    subsets: List[str],
-) -> List[Optional[float]]:
-    """Evaluate the model on the validation set(s) and return the losses."""
-
-    if cfg.dataset.fixed_validation_seed is not None:
-        # set fixed seed for every validation
-        utils.set_torch_seed(cfg.dataset.fixed_validation_seed)
-
-    trainer.begin_valid_epoch(epoch_itr.epoch)
-    valid_losses = []
-    for subset in subsets:
-        logger.info('begin validation on "{}" subset'.format(subset))
-
-        # Initialize data iterator
-        itr = trainer.get_valid_iterator(subset).next_epoch_itr(
-            shuffle=False, set_dataset_epoch=False  # use a fixed valid set
-        )
-        if cfg.common.tpu:
-            itr = utils.tpu_data_loader(itr)
-        progress = progress_bar.progress_bar(
-            itr,
-            log_format=cfg.common.log_format,
-            log_interval=cfg.common.log_interval,
-            epoch=epoch_itr.epoch,
-            prefix=f"valid on '{subset}' subset",
-            tensorboard_logdir=(
-                cfg.common.tensorboard_logdir
-                if distributed_utils.is_master(cfg.distributed_training)
-                else None
-            ),
-            default_log_format=("tqdm" if not cfg.common.no_progress_bar else "simple"),
-            wandb_project=(
-                cfg.common.wandb_project
-                if distributed_utils.is_master(cfg.distributed_training)
-                else None
-            ),
-            wandb_run_name=os.environ.get(
-                "WANDB_NAME", os.path.basename(cfg.checkpoint.save_dir)
-            ),
-        )
-
-        # create a new root metrics aggregator so validation metrics
-        # don't pollute other aggregators (e.g., train meters)
-        with metrics.aggregate(new_root=True) as agg:
-            for i, sample in enumerate(progress):
-                if cfg.dataset.max_valid_steps is not None and i > cfg.dataset.max_valid_steps:
-                    break
-                try:
-                    # print(sample)
-                    trainer.valid_step(sample)
-                except IndexError:
-                    # print(sample)
-                    print('didnt pass')
-                    trainer.valid_step(sample)
-                    continue
-
-        # log validation stats
-        if hasattr(task, 'get_valid_stats'):
-            stats = task.get_valid_stats(cfg, trainer, agg.get_smoothed_values())
-        else:
-            stats = agg.get_smoothed_values()
-        stats = get_valid_stats(cfg, trainer, stats)
-
-        if hasattr(task, "post_validate"):
-            task.post_validate(trainer.get_model(), stats, agg)
-
-
-        progress.print(stats, tag=subset, step=trainer.get_num_updates())
-        
-        valid_losses.append(stats[cfg.checkpoint.best_checkpoint_metric])
-    return valid_losses
-
-
-def get_valid_stats(
-    cfg: DictConfig, trainer: Trainer, stats: Dict[str, Any]
-) -> Dict[str, Any]:
-    stats["num_updates"] = trainer.get_num_updates()
-    if hasattr(checkpoint_utils.save_checkpoint, "best"):
-        key = "best_{0}".format(cfg.checkpoint.best_checkpoint_metric)
-        best_function = max if cfg.checkpoint.maximize_best_checkpoint_metric else min
-        stats[key] = best_function(
-            checkpoint_utils.save_checkpoint.best,
-            stats[cfg.checkpoint.best_checkpoint_metric],
-        )
-    return stats
-
-
-def cli_main(
-    modify_parser: Optional[Callable[[argparse.ArgumentParser], None]] = None
-) -> None:
-    parser = options.get_training_parser()
-    args = options.parse_args_and_arch(parser, modify_parser=modify_parser)
-    print(args)
-    cfg = convert_namespace_to_omegaconf(args)
-
-    if cfg.common.use_plasma_view:
-        server = PlasmaStore(path=cfg.common.plasma_path)
-        logger.info(f"Started plasma server pid {server.server.pid} {cfg.common.plasma_path}")
-
-    if args.profile:
-        with torch.cuda.profiler.profile():
-            with torch.autograd.profiler.emit_nvtx():
-                distributed_utils.call_main(cfg, main)
-    else:
-        distributed_utils.call_main(cfg, main)
-
-    # if cfg.common.use_plasma_view:
-    #     server.server.kill()
-
-
-if __name__ == "__main__":
-    cli_main()
- 

trainer.py

@@ -1,1569 +0,0 @@
-# Copyright 2022 The OFA-Sys Team. 
-# All rights reserved.
-# This source code is licensed under the Apache 2.0 license 
-# found in the LICENSE file in the root directory.
-
-"""
-Train a network across multiple GPUs.
-"""
-
-import contextlib
-import logging
-import sys
-import time
-from argparse import Namespace
-from itertools import chain
-from typing import Any, Dict, List
-
-import torch
-from fairseq import models, optim, utils
-from fairseq.dataclass.configs import FairseqConfig
-from fairseq.dataclass.utils import convert_namespace_to_omegaconf
-from fairseq.distributed import utils as distributed_utils
-from fairseq.file_io import PathManager
-from fairseq.logging import meters, metrics
-from fairseq.models.ema import build_ema
-from fairseq.nan_detector import NanDetector
-from fairseq.optim import lr_scheduler
-from omegaconf import OmegaConf
-
-from utils import checkpoint_utils
-import torch.nn as nn 
-
-logger = logging.getLogger(__name__)
-
-
-class Trainer(object):
-    """Main class for data parallel training.
-
-    This class supports synchronous distributed data parallel training,
-    where multiple workers each have a full model replica and gradients
-    are accumulated across workers before each update. We use
-    :class:`~torch.nn.parallel.DistributedDataParallel` to handle
-    communication of the gradients across workers.
-    """
-
-    def __init__(self, cfg: FairseqConfig, task, model, criterion, quantizer=None):
-
-        if isinstance(cfg, Namespace):
-            logger.warning(
-                "argparse.Namespace configuration is deprecated! Automatically converting to OmegaConf"
-            )
-            cfg = convert_namespace_to_omegaconf(cfg)
-
-        self.cfg = cfg
-        self.task = task
-
-        # catalog shared parameters
-        shared_params = _catalog_shared_params(model)
-        self.tpu = cfg.common.tpu
-        self.cuda = torch.cuda.is_available() and not cfg.common.cpu and not self.tpu
-        if self.cuda:
-            self.device = torch.device("cuda")
-        elif self.tpu:
-            self.device = utils.get_tpu_device()
-        else:
-            self.device = torch.device("cpu")
-
-        if self.is_fsdp:
-            import fairscale
-            if self.cfg.common.bf16:
-                raise ValueError(
-                    "FullyShardedDataParallel is not compatible with --bf16 or "
-                    "--memory-efficient-bf16"
-                )
-            if self.cfg.distributed_training.zero_sharding != "none":
-                raise ValueError(
-                    "FullyShardedDataParallel is not compatible with --zero-sharding "
-                    "option (it's already built in)"
-                )
-            if max(self.cfg.optimization.update_freq) > 1 and fairscale.__version__ < "0.4.0":
-                raise RuntimeError(
-                    "Please update to fairscale 0.4.0 or newer when combining "
-                    "--update-freq with FullyShardedDataParallel"
-                )
-        else:
-            if (
-                hasattr(self.cfg.distributed_training, "cpu_offload")
-                and self.cfg.distributed_training.cpu_offload
-            ):
-                raise ValueError("--cpu-offload requires --ddp-backend=fully_sharded")
-
-        # copy model and criterion to current device/dtype
-        self._criterion = criterion
-        self._model = model
-        if not self.is_fsdp:
-            if cfg.common.fp16:
-                assert not cfg.common.amp, "Cannot use fp16 and AMP together"
-                self._criterion = self._criterion.half()
-                self._model = self._model.half()
-                
-                if hasattr(self._model.encoder, 'embed_audios'):
-                    from torchlibrosa.stft import Spectrogram, LogmelFilterBank
-                    for layer in self._model.modules():
-                        if isinstance(layer, LogmelFilterBank) or isinstance(layer, Spectrogram):
-                            layer.float()
-                            print(layer)
-                # for layer in self._model.modules():
-                #     if isinstance(layer, nn.BatchNorm2d):
-                #         layer.float()
-                #         print(layer)
-
-                
-            elif cfg.common.bf16:
-                self._criterion = self._criterion.to(dtype=torch.bfloat16)
-                self._model = self._model.to(dtype=torch.bfloat16)
-            elif cfg.common.amp:
-                self._amp_retries = 0
-        if (
-            not cfg.distributed_training.pipeline_model_parallel
-            # the DistributedFairseqModel wrapper will handle moving to device,
-            # so only handle cases which don't use the wrapper
-            and not self.use_distributed_wrapper
-        ):
-            self._criterion = self._criterion.to(device=self.device)
-            self._model = self._model.to(device=self.device)
-        self.pipeline_model_parallel = cfg.distributed_training.pipeline_model_parallel
-        self.last_device = None
-        if self.cuda and self.pipeline_model_parallel:
-            self.last_device = torch.device(
-                cfg.distributed_training.pipeline_devices[-1]
-            )
-
-        # check that shared parameters are preserved after device transfer
-        for shared_param in shared_params:
-            ref = _get_module_by_path(self._model, shared_param[0])
-            for path in shared_param[1:]:
-                logger.info(
-                    "detected shared parameter: {} <- {}".format(shared_param[0], path)
-                )
-                _set_module_by_path(self._model, path, ref)
-
-        self._dummy_batch = None  # indicates we don't have a dummy batch at first
-        self._lr_scheduler = None
-        self._num_updates = 0
-        self._num_xla_compiles = 0  # for TPUs
-        self._optim_history = None
-        self._optimizer = None
-        self._warn_once = set()
-        self._wrapped_criterion = None
-        self._wrapped_model = None
-        self._ema = None
-
-        # TODO(myleott): support tpu
-        if self.cuda and self.data_parallel_world_size > 1:
-            self._grad_norm_buf = torch.cuda.DoubleTensor(self.data_parallel_world_size)
-        else:
-            self._grad_norm_buf = None
-
-        self.quantizer = quantizer
-        if self.quantizer is not None:
-            self.quantizer.set_trainer(self)
-
-        # get detailed cuda environment
-        if self.cuda:
-            self.cuda_env = utils.CudaEnvironment()
-            if self.data_parallel_world_size > 1:
-                self.cuda_env_arr = distributed_utils.all_gather_list(
-                    self.cuda_env, group=distributed_utils.get_global_group()
-                )
-            else:
-                self.cuda_env_arr = [self.cuda_env]
-            if self.data_parallel_rank == 0:
-                utils.CudaEnvironment.pretty_print_cuda_env_list(self.cuda_env_arr)
-        else:
-            self.cuda_env = None
-            self.cuda_env_arr = None
-
-        metrics.log_start_time("wall", priority=790, round=0)
-
-        self._start_time = time.time()
-        self._previous_training_time = 0
-        self._cumulative_training_time = None
-
-    def reinitialize(self):
-        """Reinitialize the Trainer, typically after model params change."""
-        self._lr_scheduler = None
-        self._optimizer = None
-        self._wrapped_criterion = None
-        self._wrapped_model = None
-
-    @property
-    def data_parallel_world_size(self):
-        if self.cfg.distributed_training.distributed_world_size == 1:
-            return 1
-        return distributed_utils.get_data_parallel_world_size()
-
-    @property
-    def data_parallel_process_group(self):
-        return distributed_utils.get_data_parallel_group()
-
-    @property
-    def data_parallel_rank(self):
-        if self.cfg.distributed_training.distributed_world_size == 1:
-            return 0
-        return distributed_utils.get_data_parallel_rank()
-
-    @property
-    def is_data_parallel_master(self):
-        # NOTE: this returns true for all model parallel replicas with data
-        # parallel rank 0
-        return self.data_parallel_rank == 0
-
-    @property
-    def use_distributed_wrapper(self) -> bool:
-        return (
-            self.data_parallel_world_size > 1 and not self.cfg.optimization.use_bmuf
-        ) or (
-            self.is_fsdp and self.cfg.distributed_training.cpu_offload
-        )
-
-    @property
-    def should_save_checkpoint_on_current_rank(self) -> bool:
-        """Indicates whether to save checkpoints on the current DDP rank."""
-        if (
-            self.is_fsdp and self.cfg.distributed_training.use_sharded_state
-        ) or getattr(self.cfg.model, "base_layers", 0) > 0:
-            return True
-        else:
-            return self.is_data_parallel_master
-
-    @property
-    def always_call_state_dict_during_save_checkpoint(self) -> bool:
-        if self.is_fsdp and not self.cfg.distributed_training.use_sharded_state:
-            # FSDP calls communication collective when consolidating checkpoints
-            return True
-        else:
-            return False
-
-    @property
-    def checkpoint_suffix(self) -> str:
-        """Suffix to add to the checkpoint file name."""
-        if self.is_fsdp and self.cfg.distributed_training.use_sharded_state:
-            return self.cfg.checkpoint.checkpoint_suffix + "-shard{0}".format(
-                self.data_parallel_rank
-            )
-        else:
-            return self.cfg.checkpoint.checkpoint_suffix or ""
-
-    @property
-    def criterion(self):
-        if self._wrapped_criterion is None:
-            if utils.has_parameters(self._criterion) and self.use_distributed_wrapper:
-                self._wrapped_criterion = models.DistributedFairseqModel(
-                    self.cfg.distributed_training,
-                    self._criterion,
-                    process_group=self.data_parallel_process_group,
-                    device=self.device,
-                )
-            else:
-                self._wrapped_criterion = self._criterion
-        return self._wrapped_criterion
-
-    @property
-    def model(self):
-        if self._wrapped_model is None:
-            if self.use_distributed_wrapper:
-                self._wrapped_model = models.DistributedFairseqModel(
-                    self.cfg.distributed_training,
-                    self._model,
-                    process_group=self.data_parallel_process_group,
-                    device=self.device,
-                )
-            else:
-                self._wrapped_model = self._model
-        return self._wrapped_model
-
-    @property
-    def ema(self):
-        if self._ema is None:
-            self._build_ema()
-        return self._ema
-
-    def _build_ema(self):
-        if self.cfg.ema.store_ema:
-            self._ema = build_ema(self._model, self.cfg.ema, self.device)
-            logger.info(
-                "Exponential Moving Average Shadow Model is initialized."
-            )
-
-    @property
-    def optimizer(self):
-        if self._optimizer is None:
-            self._build_optimizer()
-        return self._optimizer
-
-    @property
-    def lr_scheduler(self):
-        if self._lr_scheduler is None:
-            self._build_optimizer()  # this will initialize self._lr_scheduler
-        return self._lr_scheduler
-
-    def _build_optimizer(self):
-        # params = list(self.model.parameters())
-        # print("len of model param:", len(params))
-        # params += list(
-        #     filter(
-        #         lambda p: p.requires_grad,
-        #         chain(self.criterion.parameters()),
-        #     )
-        # )
-
-        
-        params = list(
-            filter(
-                lambda p: p.requires_grad,
-                chain(self.model.parameters(), self.criterion.parameters()),
-            )
-        )
-        print("len of optim param:", len(params))
-
-        if self.is_fsdp and self.cfg.common.fp16:
-            # FullyShardedDataParallel always uses MemoryEfficientFP16 wrapper,
-            # mostly for the grad scaling. But if we don't have the
-            # --memory-efficient-fp16 flag set, then we're effectively doing
-            # regular --fp16 and can allow the use of optimizers that would
-            # otherwise be unsupported by MemoryEfficientFP16Optimizer.
-            allow_unsupported = not self.cfg.common.memory_efficient_fp16
-            self._optimizer = optim.MemoryEfficientFP16Optimizer.build_optimizer(
-                self.cfg, params, allow_unsupported=allow_unsupported
-            )
-        elif self.cfg.common.fp16 or self.cfg.common.bf16 or self.cfg.common.amp:
-            if self.cuda and torch.cuda.get_device_capability(0)[0] < 7:
-                logger.info(
-                    "NOTE: your device does NOT support faster training with --fp16 or --amp, "
-                    "please switch to FP32 which is likely to be faster"
-                )
-            if (
-                self.cfg.common.memory_efficient_fp16
-                or self.cfg.common.memory_efficient_bf16
-            ):
-                self._optimizer = optim.MemoryEfficientFP16Optimizer.build_optimizer(
-                    self.cfg, params
-                )
-            elif self.cfg.common.amp:
-                self._optimizer = optim.AMPOptimizer.build_optimizer(self.cfg, params)
-            else:
-                self._optimizer = optim.FP16Optimizer.build_optimizer(self.cfg, params)
-        else:
-            if self.cuda and torch.cuda.get_device_capability(0)[0] >= 7:
-                logger.info("NOTE: your device may support faster training with --fp16 or --amp")
-            self._optimizer = optim.build_optimizer(self.cfg.optimizer, params)
-
-        if self.is_fsdp:
-            assert (
-                not self.cfg.optimization.use_bmuf
-            ), "--ddp-backend=fully_sharded is not compatible with BMUF"
-            assert self._optimizer.supports_flat_params, (
-                "--ddp-backend=fully_sharded is only compatible with pointwise "
-                "optimizers (e.g., Adam, AdamW, Adadelta, Adamax, SGD, etc.). "
-                "However, the sharding will result in slightly different results when "
-                "using non-pointwise optimizers (e.g., Adagrad, Adafactor, LAMB)"
-            )
-
-        if self.cfg.optimization.use_bmuf:
-            self._optimizer = optim.FairseqBMUF(
-                self.cfg.bmuf,
-                self._optimizer,
-            )
-
-        if self.cfg.distributed_training.zero_sharding == "os":
-            if (
-                self.cfg.common.fp16
-                and not self.cfg.common.memory_efficient_fp16
-                and not self.cfg.common.memory_efficient_bf16
-            ) and not self.cfg.common.fp16_no_flatten_grads:
-                raise ValueError(
-                    "ZeRO is incomptabile with fp16 and flattened grads. "
-                    "Please use --fp16-no-flatten-grads"
-                )
-            else:
-                optim.shard_(self._optimizer, self.data_parallel_process_group)
-
-        # We should initialize the learning rate scheduler immediately after
-        # building the optimizer, so that the initial learning rate is set.
-        self._lr_scheduler = lr_scheduler.build_lr_scheduler(
-            self.cfg.lr_scheduler,
-            self.optimizer,
-        )
-        self._lr_scheduler.step_update(0)
-
-    @property
-    def is_fsdp(self):
-        return self.cfg.distributed_training.ddp_backend == "fully_sharded"
-
-    def consolidate_optimizer(self):
-        """For OSS, we need to consolidate the state dict."""
-        if self.cfg.checkpoint.no_save_optimizer_state:
-            return
-        self._gathered_optim_state = None
-        if hasattr(self.optimizer.optimizer, "consolidate_state_dict"):
-            self.optimizer.optimizer.consolidate_state_dict()
-        elif self.is_fsdp and not self.model.use_sharded_state:
-            st = self.model.gather_full_optim_state_dict(
-                self.optimizer
-            )  # only returns on rank 0
-            self._gathered_optim_state = st
-
-    def state_dict(self):
-        state_dict = {
-            "args": None,  # legacy
-            "cfg": (
-                OmegaConf.to_container(self.cfg, resolve=True, enum_to_str=True)
-                if OmegaConf.is_config(self.cfg)
-                else self.cfg
-            ),
-            "model": self.model.state_dict(),
-            "criterion": (
-                self.criterion.state_dict()
-                if utils.has_parameters(self.criterion)
-                else None
-            ),
-            "optimizer_history": (self._optim_history or [])
-            + [
-                {
-                    "criterion_name": self.get_criterion().__class__.__name__,
-                    "optimizer_name": self.optimizer.__class__.__name__,
-                    "lr_scheduler_state": self.lr_scheduler.state_dict(),
-                    "num_updates": self.get_num_updates(),
-                }
-            ],
-            "task_state": self.task.state_dict() if self.task is not None else {},
-            "extra_state": {
-                "metrics": metrics.state_dict(),
-                "previous_training_time": self.cumulative_training_time(),
-            },
-        }
-        if self.cfg.ema.store_ema:
-            # Save EMA model state as extra state
-            state_dict["extra_state"]["ema"] = self.ema.get_model().state_dict()
-            if self.cfg.ema.ema_fp32:
-                # Save EMA params in fp32
-                state_dict["extra_state"]["ema_fp32_params"] = self.ema.fp32_params
-        if not self.cfg.checkpoint.no_save_optimizer_state:
-            if self._gathered_optim_state is not None:
-                state_dict["last_optimizer_state"] = self._gathered_optim_state
-                self._gathered_optim_state = None
-            else:
-                state_dict["last_optimizer_state"] = self.optimizer.state_dict()
-        if self.is_fsdp:
-            # save meta data for recombining checkpoint upon loading
-            state_dict["fsdp_metadata"] = self.model.local_metadata_dict()
-        return state_dict
-
-    def save_checkpoint(self, filename, extra_state, save_on_cuda=False):
-        """Save all training state in a checkpoint file."""
-        logger.info(f"Saving checkpoint to {filename}")
-        # call state_dict on all ranks in case it needs internal communication
-        if not save_on_cuda:
-            state_dict = utils.move_to_cpu(self.state_dict())
-        else:
-            print("Save on cuda")
-            state_dict = self.state_dict()
-        state_dict["extra_state"].update(extra_state)
-        if self.should_save_checkpoint_on_current_rank:
-            checkpoint_utils.torch_persistent_save(
-                state_dict,
-                filename,
-                async_write=self.cfg.checkpoint.write_checkpoints_asynchronously,
-            )
-        logger.info(f"Finished saving checkpoint to {filename}")
-
-    def load_checkpoint(
-        self,
-        filename,
-        reset_optimizer=False,
-        reset_lr_scheduler=False,
-        optimizer_overrides=None,
-        reset_meters=False,
-        strict=True,
-        load_on_cuda=False,
-    ):
-        """
-        Load all training state from a checkpoint file.
-        rank = 0 will load the checkpoint, and then broadcast it to all
-        other ranks.
-        """
-        extra_state, self._optim_history, last_optim_state = None, [], None
-
-        logger.info(f"Preparing to load checkpoint {filename}")
-        is_distributed = self.data_parallel_world_size > 1
-        bexists = PathManager.isfile(filename)
-        if bexists:
-            load_on_all_ranks = (
-                self.cfg.checkpoint.load_checkpoint_on_all_dp_ranks
-                # TPUs don't support broadcast yet, so load checkpoints
-                # on every worker for now
-                or self.tpu
-                # FSDP requires loading checkpoint shards on all ranks
-                or (self.is_fsdp and self.cfg.distributed_training.use_sharded_state)
-                or getattr(self.cfg.model, "base_layers", 0) > 0
-            )
-
-            if load_on_all_ranks or self.data_parallel_rank == 0:
-                state = checkpoint_utils.load_checkpoint_to_cpu(
-                    filename, load_on_all_ranks=load_on_all_ranks, strict=strict,load_on_cuda=load_on_cuda,
-                )
-                last_optim_state = state.get("last_optimizer_state", None)
-
-                # If doing zero_sharding, do not broadcast global optimizer
-                # state. Later we will broadcast sharded states to each rank
-                # to avoid memory from exploding.
-                if (
-                    not load_on_all_ranks
-                    and self.cfg.distributed_training.zero_sharding == "os"
-                    and "last_optimizer_state" in state
-                    and is_distributed
-                ):
-                    state["last_optimizer_state"] = "SHARDED"
-            else:
-                last_optim_state = None
-                state = None
-
-            if is_distributed and not load_on_all_ranks:    # .contiguous()
-                state = distributed_utils.broadcast_object(
-                    state,
-                    src_rank=0,
-                    group=self.data_parallel_process_group,
-                    dist_device=self.device,
-                )
-                if self.data_parallel_rank > 0:
-                    last_optim_state = state.get("last_optimizer_state", None)
-
-            # load model parameters
-            try:
-                if self.cfg.checkpoint.use_ema_weights_to_init_param and "extra_state" in state and "ema" in state["extra_state"]:
-                    logger.info("use_ema_weights_to_init_param = True, will use EMA weights in the ckpt to init the model param...")
-                    ema_state_dict = state["extra_state"]["ema_fp32_params"] if "ema_fp32_params" in state["extra_state"] else state["extra_state"]["ema"]
-                    msg = self.model.load_state_dict(
-                        ema_state_dict, strict=strict, model_cfg=self.cfg.model
-                    )
-                else:
-                    msg = self.model.load_state_dict(
-                        state["model"], strict=strict, model_cfg=self.cfg.model
-                    )
-                logger.info(msg)
-
-                # save memory for later steps
-                if not (self.cfg.ema.store_ema and (self.cfg.checkpoint.use_latest_weights_to_init_ema or not ("extra_state" in state and "ema" in state["extra_state"]))):
-                    del state["model"]
-                if utils.has_parameters(self.get_criterion()) and 'criterion' in  state:
-                    self.get_criterion().load_state_dict(
-                        state["criterion"], strict=strict
-                    )
-                    del state["criterion"]
-
-            except Exception:
-                raise Exception(
-                    "Cannot load model parameters from checkpoint {}; "
-                    "please ensure that the architectures match.".format(filename)
-                )
-            extra_state = state.get("extra_state", None)
-            self._optim_history = state.get("optimizer_history", None)
-
-        if last_optim_state is not None and not reset_optimizer:
-            # rebuild optimizer after loading model, since params may have changed
-            self._build_optimizer()
-
-            # only reload optimizer and lr_scheduler if they match
-            last_optim = self._optim_history[-1]
-            assert (
-                last_optim["criterion_name"] == self.get_criterion().__class__.__name__
-            ), f"Criterion does not match; please reset the optimizer (--reset-optimizer). {last_optim['criterion_name']} vs {self.get_criterion().__class__.__name__}"
-            assert (
-                last_optim["optimizer_name"] == self.optimizer.__class__.__name__
-            ), f"Optimizer does not match; please reset the optimizer (--reset-optimizer). {last_optim['optimizer_name']} vs {self.optimizer.__class__.__name__}"
-
-            if not reset_lr_scheduler:
-                self.lr_scheduler.load_state_dict(last_optim["lr_scheduler_state"])
-
-            if self.is_fsdp and not self.model.use_sharded_state:
-                # if use_sharded_state, the last_optim_state is already sharded, skip this
-                last_optim_state = self.model.get_shard_from_optim_state_dict(
-                    last_optim_state
-                )
-            elif not load_on_all_ranks and is_distributed:
-                last_optim_state = self.optimizer.broadcast_global_state_dict(
-                    last_optim_state
-                )
-
-            self.optimizer.load_state_dict(last_optim_state, optimizer_overrides)
-
-            self.set_num_updates(last_optim["num_updates"])
-
-        if extra_state is not None:
-            itr_state = extra_state["train_iterator"]
-            epoch = itr_state["epoch"]
-
-            if "previous_training_time" in extra_state:
-                self._previous_training_time = extra_state["previous_training_time"]
-                self._start_time = time.time()
-
-            self.lr_step(epoch)
-
-            if (
-                itr_state.get("version", 1) >= 2
-                and itr_state["iterations_in_epoch"] == 0
-            ):
-                # reset meters at start of epoch
-                reset_meters = True
-
-            if "metrics" in extra_state and not reset_meters:
-                metrics.load_state_dict(extra_state["metrics"])
-
-                # reset TimeMeters, since their start times don't make sense anymore
-                for meter in metrics.get_meters("default"):
-                    if isinstance(meter, meters.TimeMeter):
-                        meter.reset()
-
-            if self.cfg.ema.store_ema:
-                if self.cfg.checkpoint.use_latest_weights_to_init_ema or "ema" not in extra_state:
-                    if "ema" not in extra_state:
-                        logger.warn(
-                            "EMA not found in checkpoint. But store_ema is True. "
-                            "EMA is re-initialized from checkpoint."
-                        )
-                    elif self.cfg.checkpoint.use_latest_weights_to_init_ema:
-                        logger.info(
-                            "use_latest_weights_to_init_ema = True. EMA is re-initialized from checkpoint."
-                        )
-                    self.ema.restore(state["model"], build_fp32_params=self.cfg.ema.ema_fp32)
-                    del state["model"]
-                else:
-                    logger.info(
-                        "Loading EMA from checkpoint"
-                    )
-                    self.ema.restore(extra_state["ema"], build_fp32_params=False)
-
-                    if self.cfg.ema.ema_fp32:
-                        if "ema_fp32_params" in extra_state:
-                            logger.info(
-                                "Loading EMA fp32 params from checkpoint"
-                            )
-                            self.ema.build_fp32_params(extra_state["ema_fp32_params"])
-                        else:
-                            logger.info(
-                                "Building EMA fp32 params from EMA model in checkpoint"
-                            )
-                            self.ema.build_fp32_params()
-
-            logger.info(
-                "Loaded checkpoint {} (epoch {} @ {} updates)".format(
-                    filename, epoch, self.get_num_updates()
-                )
-            )
-
-        else:
-            logger.info("No existing checkpoint found {}".format(filename))
-
-        # print("delete state ...")
-        # del state  # dereference seems crucial
-        # torch.cuda.empty_cache()
-
-        return extra_state
-
-    def get_train_iterator(
-        self,
-        epoch,
-        combine=True,
-        load_dataset=True,
-        data_selector=None,
-        shard_batch_itr=True,
-        disable_iterator_cache=False,
-    ):
-        """Return an EpochBatchIterator over the training set for a given epoch."""
-        if load_dataset:
-            logger.info("loading train data for epoch {}".format(epoch))
-            self.task.load_dataset(
-                self.cfg.dataset.train_subset,
-                epoch=epoch,
-                combine=combine,
-                data_selector=data_selector,
-                tpu=self.tpu,
-            )
-        batch_iterator = self.task.get_batch_iterator(
-            dataset=self.task.dataset(self.cfg.dataset.train_subset),
-            max_tokens=self.cfg.dataset.max_tokens,
-            max_sentences=self.cfg.dataset.batch_size,
-            max_positions=utils.resolve_max_positions(
-                self.task.max_positions(),
-                self.model.max_positions(),
-                self.cfg.dataset.max_tokens,
-            ),
-            ignore_invalid_inputs=True,
-            required_batch_size_multiple=self.cfg.dataset.required_batch_size_multiple,
-            seed=self.cfg.common.seed,
-            num_shards=self.data_parallel_world_size if shard_batch_itr else 1,
-            shard_id=self.data_parallel_rank if shard_batch_itr else 0,
-            num_workers=self.cfg.dataset.num_workers,
-            epoch=epoch,
-            data_buffer_size=self.cfg.dataset.data_buffer_size,
-            disable_iterator_cache=disable_iterator_cache,
-        )
-        self.reset_dummy_batch(batch_iterator.first_batch)
-        batch_iterator.dataset.dataset._seek()
-        return batch_iterator
-
-    def get_valid_iterator(
-        self,
-        subset,
-        disable_iterator_cache=False,
-    ):
-        """Return an EpochBatchIterator over given validation subset for a given epoch."""
-        self.task.dataset(subset).dataset._seek()
-        batch_iterator = self.task.get_batch_iterator(
-            dataset=self.task.dataset(subset),
-            max_tokens=self.cfg.dataset.max_tokens_valid,
-            max_sentences=self.cfg.dataset.batch_size_valid,
-            max_positions=utils.resolve_max_positions(
-                self.task.max_positions(),
-                self.model.max_positions(),
-            ),
-            ignore_invalid_inputs=self.cfg.dataset.skip_invalid_size_inputs_valid_test,
-            required_batch_size_multiple=self.cfg.dataset.required_batch_size_multiple,
-            seed=self.cfg.common.seed,
-            num_shards=self.data_parallel_world_size,
-            shard_id=self.data_parallel_rank,
-            num_workers=self.cfg.dataset.num_workers,
-            # always pass a fixed "epoch" to keep validation data consistent
-            # across training epochs
-            epoch=1,
-            data_buffer_size=self.cfg.dataset.data_buffer_size,
-            disable_iterator_cache=disable_iterator_cache,
-        )
-        self.reset_dummy_batch(batch_iterator.first_batch)
-        batch_iterator.dataset.dataset._seek()
-        return batch_iterator
-
-    def begin_epoch(self, epoch):
-        """Called at the beginning of each epoch."""
-        logger.info("begin training epoch {}".format(epoch))
-
-        self.lr_step_begin_epoch(epoch)
-
-        if self.quantizer is not None:
-            self.quantizer.begin_epoch(epoch)
-
-        # task specific setup per epoch
-        self.task.begin_epoch(epoch, self.get_model())
-
-        if self.tpu:
-            import torch_xla.core.xla_model as xm
-
-            xm.rendezvous("begin_epoch")  # wait for all workers
-            xm.mark_step()
-
-    def begin_valid_epoch(self, epoch):
-        """Called at the beginning of each validation epoch."""
-
-        # task specific setup per validation epoch
-        self.task.begin_valid_epoch(epoch, self.get_model())
-
-    def reset_dummy_batch(self, batch):
-        self._dummy_batch = batch
-
-    @metrics.aggregate("train")
-    def train_step(self, samples, raise_oom=False):
-        """Do forward, backward and parameter update."""
-        self._set_seed()
-        self.model.train()
-        self.criterion.train()
-        self.zero_grad()
-
-        metrics.log_start_time("train_wall", priority=800, round=0)
-
-        # If EMA is enabled through store_ema=True
-        # and task.uses_ema is True, pass the EMA model as a keyword
-        # argument to the task.
-        extra_kwargs = {}
-        if self.cfg.ema.store_ema and getattr(self.task, "uses_ema", False):
-            extra_kwargs["ema_model"] = self.ema.get_model()
-
-        # forward and backward pass
-        logging_outputs, sample_size, ooms = [], 0, 0
-        for i, sample in enumerate(samples):  # delayed update loop
-            sample, is_dummy_batch = self._prepare_sample(sample)
-
-            def maybe_no_sync():
-                """
-                Whenever *samples* contains more than one mini-batch, we
-                want to accumulate gradients locally and only call
-                all-reduce in the last backwards pass.
-                """
-                if (
-                    self.data_parallel_world_size > 1
-                    and hasattr(self.model, "no_sync")
-                    and i < len(samples) - 1
-                    # The no_sync context manager results in increased memory
-                    # usage with FSDP, since full-size gradients will be
-                    # accumulated on each GPU. It's typically a better tradeoff
-                    # to do the extra communication with FSDP.
-                    and not self.is_fsdp
-                ):
-                    return self.model.no_sync()
-                else:
-                    return contextlib.ExitStack()  # dummy contextmanager
-
-            try:
-                with maybe_no_sync():
-                    # forward and backward
-                    loss, sample_size_i, logging_output = self.task.train_step(
-                        sample=sample,
-                        model=self.model,
-                        criterion=self.criterion,
-                        optimizer=self.optimizer,
-                        update_num=self.get_num_updates(),
-                        ignore_grad=is_dummy_batch,
-                        **extra_kwargs,
-                    )
-                    del loss
-
-                logging_outputs.append(logging_output)
-                sample_size += sample_size_i
-
-                # emptying the CUDA cache after the first step can
-                # reduce the chance of OOM
-                if self.cuda and self.get_num_updates() == 0:
-                    torch.cuda.empty_cache()
-            except RuntimeError as e:
-                if "out of memory" in str(e):
-                    self._log_oom(e)
-                    if raise_oom:
-                        raise e
-                    logger.warning(
-                        "attempting to recover from OOM in forward/backward pass"
-                    )
-                    ooms += 1
-                    self.zero_grad()
-                    if self.cuda:
-                        torch.cuda.empty_cache()
-                    if self.cfg.distributed_training.distributed_world_size == 1:
-                        return None
-                else:
-                    raise e
-
-            if self.tpu and i < len(samples) - 1:
-                # tpu-comment: every XLA operation before marking step is
-                # appended to the IR graph, and processing too many batches
-                # before marking step can lead to OOM errors.
-                # To handle gradient accumulation use case, we explicitly
-                # mark step here for every forward pass without a backward pass
-                self._xla_markstep_and_send_to_cpu()
-
-        if is_dummy_batch:
-            if torch.is_tensor(sample_size):
-                sample_size.zero_()
-            else:
-                sample_size *= 0.0
-
-        if torch.is_tensor(sample_size):
-            sample_size = sample_size.float()
-        else:
-            sample_size = float(sample_size)
-
-        # gather logging outputs from all replicas
-        if self._sync_stats():
-            train_time = self._local_cumulative_training_time()
-            logging_outputs, (
-                sample_size,
-                ooms,
-                total_train_time,
-            ) = self._aggregate_logging_outputs(
-                logging_outputs, sample_size, ooms, train_time, ignore=is_dummy_batch
-            )
-            self._cumulative_training_time = (
-                total_train_time / self.data_parallel_world_size
-            )
-
-        overflow = False
-        try:
-            with torch.autograd.profiler.record_function("reduce-grads"):
-                # reduce gradients across workers
-                self.optimizer.all_reduce_grads(self.model)
-                if utils.has_parameters(self.criterion):
-                    self.optimizer.all_reduce_grads(self.criterion)
-
-            with torch.autograd.profiler.record_function("multiply-grads"):
-                # multiply gradients by (data_parallel_size / sample_size) since
-                # DDP normalizes by the number of data parallel workers for
-                # improved fp16 precision.
-                # Thus we get (sum_of_gradients / sample_size) at the end.
-                # In case of fp16, this step also undoes loss scaling.
-                # (Debugging note: Some optimizers perform this scaling on the
-                # fly, so inspecting model.parameters() or optimizer.params may
-                # still show the original, unscaled gradients.)
-                numer = (
-                    self.data_parallel_world_size
-                    if not self.cfg.optimization.use_bmuf or self._sync_stats()
-                    else 1
-                )
-                self.optimizer.multiply_grads(numer / (sample_size or 1.0))
-                # Note: (sample_size or 1.0) handles the case of a zero gradient, in a
-                # way that avoids CPU/device transfers in case sample_size is a GPU or
-                # TPU object. The assumption is that the gradient itself is also 0.
-
-            with torch.autograd.profiler.record_function("clip-grads"):
-                # clip grads
-                grad_norm = self.clip_grad_norm(self.cfg.optimization.clip_norm)
-
-            # check that grad norms are consistent across workers
-            # on tpu check tensor is slow
-            if not self.tpu:
-                if (
-                    not self.cfg.optimization.use_bmuf
-                    and self.cfg.distributed_training.ddp_backend != "slow_mo"
-                ):
-                    self._check_grad_norms(grad_norm)
-                if not torch.isfinite(grad_norm).all():
-                    # in case of AMP, if gradients are Nan/Inf then
-                    # optimizer step is still required
-                    if self.cfg.common.amp:
-                        overflow = True
-                    else:
-                        # check local gradnorm single GPU case, trigger NanDetector
-                        raise FloatingPointError("gradients are Nan/Inf")
-
-            with torch.autograd.profiler.record_function("optimizer"):
-                # take an optimization step
-                self.task.optimizer_step(
-                    self.optimizer, model=self.model, update_num=self.get_num_updates()
-                )
-                if self.cfg.common.amp and overflow:
-                    if self._amp_retries == self.cfg.common.amp_batch_retries:
-                        logger.info("AMP: skipping this batch.")
-                        self._amp_retries = 0
-                    else:
-                        self._amp_retries += 1
-                        return self.train_step(samples, raise_oom)  # recursion to feed in same batch
-
-        except FloatingPointError:
-            # re-run the forward and backward pass with hooks attached to print
-            # out where it fails
-            self.zero_grad()
-            with NanDetector(self.get_model()):
-                for _, sample in enumerate(samples):
-                    sample, _ = self._prepare_sample(sample)
-                    self.task.train_step(
-                        sample,
-                        self.model,
-                        self.criterion,
-                        self.optimizer,
-                        self.get_num_updates(),
-                        ignore_grad=False,
-                        **extra_kwargs,
-                    )
-            raise
-        except OverflowError as e:
-            overflow = True
-            logger.info(
-                f"NOTE: gradient overflow detected, ignoring gradient, {str(e)}"
-            )
-            grad_norm = torch.tensor(0.0).cuda()
-            self.zero_grad()
-        except RuntimeError as e:
-            if "out of memory" in str(e):
-                self._log_oom(e)
-                logger.error("OOM during optimization, irrecoverable")
-            raise e
-
-        # Some distributed wrappers (e.g., SlowMo) need access to the optimizer
-        # after the step
-        if hasattr(self.model, "perform_additional_optimizer_actions"):
-            if hasattr(self.optimizer, "fp32_params"):
-                self.model.perform_additional_optimizer_actions(
-                    self.optimizer.optimizer, self.optimizer.fp32_params
-                )
-            else:
-                self.model.perform_additional_optimizer_actions(
-                    self.optimizer.optimizer
-                )
-
-        logging_output = None
-        if not overflow or self.cfg.distributed_training.ddp_backend == "slow_mo":
-            self.set_num_updates(self.get_num_updates() + 1)
-
-            if self.cfg.ema.store_ema:
-                # Step EMA forward with new model.
-                self.ema.step(
-                    self.get_model(),
-                    self.get_num_updates(),
-                )
-                metrics.log_scalar(
-                    "ema_decay",
-                    self.ema.get_decay(),
-                    priority=10000,
-                    round=5,
-                    weight=0,
-                )
-
-            if self.tpu:
-                import torch_xla.core.xla_model as xm
-
-                # mark step on TPUs
-                self._xla_markstep_and_send_to_cpu()
-
-                # only log stats every log_interval steps
-                # this causes wps to be misreported when log_interval > 1
-                logging_output = {}
-                if self.get_num_updates() % self.cfg.common.log_interval == 0:
-                    # log memory usage
-                    mem_info = xm.get_memory_info(self.device)
-                    gb_free = mem_info["kb_free"] / 1024 / 1024
-                    gb_total = mem_info["kb_total"] / 1024 / 1024
-                    metrics.log_scalar(
-                        "gb_free", gb_free, priority=1500, round=1, weight=0
-                    )
-                    metrics.log_scalar(
-                        "gb_total", gb_total, priority=1600, round=1, weight=0
-                    )
-                    logging_outputs = self._xla_markstep_and_send_to_cpu(
-                        logging_outputs
-                    )
-                    logging_output = self._reduce_and_log_stats(
-                        logging_outputs, sample_size, grad_norm
-                    )
-
-                # log whenever there's an XLA compilation, since these
-                # slow down training and may indicate opportunities for
-                # optimization
-                self._check_xla_compilation()
-            else:
-                if self.cuda and self.cuda_env is not None:
-                    # log minimum free memory over the iteration
-                    gb_used = torch.cuda.max_memory_allocated() / 1024 / 1024 / 1024
-                    torch.cuda.reset_peak_memory_stats()
-                    gb_free = self.cuda_env.total_memory_in_GB - gb_used
-                    metrics.log_scalar(
-                        "gb_free", gb_free, priority=1500, round=1, weight=0
-                    )
-
-                # log stats
-                logging_output = self._reduce_and_log_stats(
-                    logging_outputs, sample_size, grad_norm
-                )
-
-                # clear CUDA cache to reduce memory fragmentation
-                if (
-                    self.cuda
-                    and self.cfg.common.empty_cache_freq > 0
-                    and (
-                        (self.get_num_updates() + self.cfg.common.empty_cache_freq - 1)
-                        % self.cfg.common.empty_cache_freq
-                    )
-                    == 0
-                ):
-                    torch.cuda.empty_cache()
-
-        if self.cfg.common.fp16 or self.cfg.common.amp:
-            metrics.log_scalar(
-                "loss_scale",
-                (
-                    self.optimizer.scaler.loss_scale
-                    if self.cfg.common.fp16
-                    else self.optimizer.scaler.get_scale()
-                ),
-                priority=700,
-                round=4,
-                weight=0,
-            )
-
-        metrics.log_stop_time("train_wall")
-        return logging_output
-
-    @metrics.aggregate("valid")
-    def valid_step(self, sample, raise_oom=False):
-        """Do forward pass in evaluation mode."""
-        if self.tpu:
-            import torch_xla.core.xla_model as xm
-
-            xm.rendezvous("valid_step")  # wait for all workers
-
-        # If EMA is enabled through store_ema=True
-        # and task.uses_ema is True, pass the EMA model as a keyword
-        # argument to the task.
-        extra_kwargs = {}
-        if self.cfg.ema.store_ema and getattr(self.task, "uses_ema", False):
-            extra_kwargs["ema_model"] = self.ema.get_model()
-
-        with torch.no_grad():
-            self.model.eval()
-            self.criterion.eval()
-
-            sample, is_dummy_batch = self._prepare_sample(sample)
-
-            try:
-                _loss, sample_size, logging_output = self.task.valid_step(
-                    sample, self.model, self.criterion, **extra_kwargs
-                )
-            except RuntimeError as e:
-                if "out of memory" in str(e):
-                    self._log_oom(e)
-                    if not raise_oom:
-                        logger.warning(
-                            "ran out of memory in validation step, retrying batch"
-                        )
-                        for p in self.model.parameters():
-                            if p.grad is not None:
-                                p.grad = None  # free some memory
-                        if self.cuda:
-                            torch.cuda.empty_cache()
-                        return self.valid_step(sample, raise_oom=True)
-                raise e
-            
-
-            logging_outputs = [logging_output]
-            if is_dummy_batch:
-                if torch.is_tensor(sample_size):
-                    sample_size.zero_()
-                else:
-                    sample_size *= 0.0
-
-        # gather logging outputs from all replicas
-        if self.data_parallel_world_size > 1:
-            logging_outputs, (sample_size,) = self._aggregate_logging_outputs(
-                logging_outputs,
-                sample_size,
-                ignore=is_dummy_batch,
-            )
-
-        # log validation stats
-        if self.tpu:
-            logging_outputs = self._xla_markstep_and_send_to_cpu(logging_outputs)
-        logging_output = self._reduce_and_log_stats(logging_outputs, sample_size)
-
-        return logging_output
-
-    def zero_grad(self):
-        self.optimizer.zero_grad()
-
-    def lr_step_begin_epoch(self, epoch):
-        """Adjust the learning rate at the beginning of the epoch."""
-        self.lr_scheduler.step_begin_epoch(epoch)
-        # prefer updating the LR based on the number of steps
-        return self.lr_step_update()
-
-    def lr_reinit(self, total_updates, num_updates):
-        self.lr_scheduler.reinit(total_updates, num_updates)
-
-    def lr_step(self, epoch, val_loss=None):
-        """Adjust the learning rate at the end of the epoch."""
-        self.lr_scheduler.step(epoch, val_loss)
-        # prefer updating the LR based on the number of steps
-        return self.lr_step_update()
-
-    def lr_step_update(self):
-        """Update the learning rate after each update."""
-        new_lr = self.lr_scheduler.step_update(self.get_num_updates())
-        if isinstance(new_lr, dict):
-            for k, v in new_lr.items():
-                metrics.log_scalar(f"lr_{k}", v, weight=0, priority=300)
-            new_lr = new_lr.get("default", next(iter(new_lr.values())))
-        else:
-            metrics.log_scalar("lr", new_lr, weight=0, priority=300)
-        return new_lr
-
-    def get_lr(self):
-        """Get the current learning rate."""
-        return self.optimizer.get_lr()
-
-    def get_model(self):
-        """Get the (non-wrapped) model instance."""
-        return self._model
-
-    def get_criterion(self):
-        """Get the (non-wrapped) criterion instance."""
-        return self._criterion
-
-    def get_meter(self, name):
-        """[deprecated] Get a specific meter by name."""
-        from fairseq import meters
-
-        if "get_meter" not in self._warn_once:
-            self._warn_once.add("get_meter")
-            utils.deprecation_warning(
-                "Trainer.get_meter is deprecated. Please use fairseq.metrics instead."
-            )
-
-        train_meters = metrics.get_meters("train")
-        if train_meters is None:
-            train_meters = {}
-
-        if name == "train_loss" and "loss" in train_meters:
-            return train_meters["loss"]
-        elif name == "train_nll_loss":
-            # support for legacy train.py, which assumed this meter is
-            # always initialized
-            m = train_meters.get("nll_loss", None)
-            return m or meters.AverageMeter()
-        elif name == "wall":
-            # support for legacy train.py, which assumed this meter is
-            # always initialized
-            m = metrics.get_meter("default", "wall")
-            return m or meters.TimeMeter()
-        elif name == "wps":
-            m = metrics.get_meter("train", "wps")
-            return m or meters.TimeMeter()
-        elif name in {"valid_loss", "valid_nll_loss"}:
-            # support for legacy train.py, which assumed these meters
-            # are always initialized
-            k = name[len("valid_") :]
-            m = metrics.get_meter("valid", k)
-            return m or meters.AverageMeter()
-        elif name == "oom":
-            return meters.AverageMeter()
-        elif name in train_meters:
-            return train_meters[name]
-        return None
-
-    def get_num_updates(self):
-        """Get the number of parameters updates."""
-        return self._num_updates
-
-    def set_num_updates(self, num_updates):
-        """Set the number of parameters updates."""
-        self._num_updates = num_updates
-        self.lr_step_update()
-        if self.quantizer:
-            self.quantizer.step_update(self._num_updates)
-        metrics.log_scalar("num_updates", self._num_updates, weight=0, priority=200)
-
-    def clip_grad_norm(self, clip_norm):
-        def agg_norm_fn(total_norm):
-            total_norm = total_norm.cuda().float() ** 2
-            total_norm = distributed_utils.all_reduce(
-                total_norm, group=self.data_parallel_process_group
-            )
-            return total_norm ** 0.5
-
-        should_agg_norm = (
-            self.is_fsdp
-            and (
-                self.data_parallel_process_group is not None
-                or torch.distributed.is_initialized()
-            )
-        )
-        return self.optimizer.clip_grad_norm(
-            clip_norm, aggregate_norm_fn=agg_norm_fn if should_agg_norm else None
-        )
-
-    def cumulative_training_time(self):
-        if self._cumulative_training_time is None:
-            # single GPU
-            return self._local_cumulative_training_time()
-        else:
-            return self._cumulative_training_time
-
-    def _local_cumulative_training_time(self):
-        """Aggregate training time in seconds."""
-        return time.time() - self._start_time + self._previous_training_time
-
-    def _fp_convert_sample(self, sample):
-        def apply_half(t):
-            if t.dtype is torch.float32:
-                return t.to(dtype=torch.half)
-            return t
-
-        def apply_bfloat16(t):
-            if t.dtype is torch.float32:
-                return t.to(dtype=torch.bfloat16)
-            return t
-
-        if self.cfg.common.fp16:
-            sample = utils.apply_to_sample(apply_half, sample)
-
-        if self.cfg.common.bf16:
-            sample = utils.apply_to_sample(apply_bfloat16, sample)
-
-        return sample
-
-    def _prepare_sample(self, sample, is_dummy=False):
-        if sample == "DUMMY":
-            raise Exception(
-                "Trying to use an uninitialized 'dummy' batch. This usually indicates "
-                "that the total number of batches is smaller than the number of "
-                "participating GPUs. Try reducing the batch size or using fewer GPUs."
-            )
-
-        if sample is None or len(sample) == 0:
-            assert (
-                self._dummy_batch is not None and len(self._dummy_batch) > 0
-            ), "Invalid dummy batch: {}".format(self._dummy_batch)
-            sample, _ = self._prepare_sample(self._dummy_batch, is_dummy=True)
-            return sample, True
-
-        # Given that PCIe/NVLink bandwidth is significantly smaller than DRAM bandwidth
-        # it makes sense to do the format conversion on the CPU and then transfer
-        # a smaller buffer to the device. This also saves GPU memory capacity.
-
-        if self.cfg.common.on_cpu_convert_precision:
-            sample = self._fp_convert_sample(sample)
-
-        if self.cuda:
-            if self.pipeline_model_parallel:
-                if 'target' in sample:
-                    sample['target'] = utils.move_to_cuda(sample['target'], device=self.last_device)
-            else:
-                sample = utils.move_to_cuda(sample)
-        elif self.tpu and is_dummy:
-            # the dummy batch may not be on the appropriate device
-            sample = utils.move_to_cuda(sample, device=self.device)
-
-        if not self.cfg.common.on_cpu_convert_precision:
-            sample = self._fp_convert_sample(sample)
-
-        if self._dummy_batch == "DUMMY":
-            self._dummy_batch = sample
-
-        return sample, False
-
-    def _set_seed(self):
-        # Set seed based on args.seed and the update number so that we get
-        # reproducible results when resuming from checkpoints
-        seed = self.cfg.common.seed + self.get_num_updates()
-        utils.set_torch_seed(seed)
-
-    def _sync_stats(self):
-        # Return True if it's using multiple GPUs and DDP or multiple GPUs with
-        # BMUF and it's a bmuf sync with warmup iterations completed before.
-        if self.data_parallel_world_size == 1:
-            return False
-        elif self.cfg.optimization.use_bmuf:
-            return (
-                self.get_num_updates() + 1
-            ) % self.cfg.bmuf.global_sync_iter == 0 and (
-                self.get_num_updates() + 1
-            ) > self.cfg.bmuf.warmup_iterations
-        else:
-            return True
-
-    def _log_oom(self, exc):
-        msg = "OOM: Ran out of memory with exception: {}".format(exc)
-        logger.warning(msg)
-        if torch.cuda.is_available() and hasattr(torch.cuda, "memory_summary"):
-            for device_idx in range(torch.cuda.device_count()):
-                logger.warning(torch.cuda.memory_summary(device=device_idx))
-        sys.stderr.flush()
-
-    def _aggregate_logging_outputs(
-        self,
-        logging_outputs: List[Dict[str, Any]],
-        *extra_stats_to_sum,
-        ignore=False,
-    ):
-        if self.task.__class__.logging_outputs_can_be_summed(self.get_criterion()):
-            return self._fast_stat_sync_sum(
-                logging_outputs, *extra_stats_to_sum, ignore=ignore
-            )
-        else:
-            return self._all_gather_list_sync(
-                logging_outputs, *extra_stats_to_sum, ignore=ignore
-            )
-
-    def _all_gather_list_sync(
-        self,
-        logging_outputs: List[Dict[str, Any]],
-        *extra_stats_to_sum,
-        ignore=False,
-    ):
-        """
-        Sync logging outputs across workers. all_gather_list_sync is
-        suitable when logging outputs are complex types.
-        """
-        if self.tpu:
-            raise NotImplementedError
-        if ignore:
-            logging_outputs = []
-        results = list(
-            zip(
-                *distributed_utils.all_gather_list(
-                    [logging_outputs] + list(extra_stats_to_sum),
-                    max_size=getattr(self.cfg.common, "all_gather_list_size", 16384),
-                    group=self.data_parallel_process_group,
-                )
-            )
-        )
-        logging_outputs, extra_stats_to_sum = results[0], results[1:]
-        logging_outputs = list(chain.from_iterable(logging_outputs))
-        extra_stats_to_sum = [sum(s) for s in extra_stats_to_sum]
-        return logging_outputs, extra_stats_to_sum
-
-    def _fast_stat_sync_sum(
-        self,
-        logging_outputs: List[Dict[str, Any]],
-        *extra_stats_to_sum,
-        ignore=False,
-    ):
-        """
-        Sync logging outputs across workers. fast_stat_sync_sum is
-        faster than all_gather_list_sync, but is only suitable when
-        logging outputs are scalars and can be summed. Note that
-        *logging_outputs* cannot contain any nested dicts/lists.
-        """
-        data = {}
-        for i, stat in enumerate(extra_stats_to_sum):
-            data["extra_stats_" + str(i)] = stat
-        if len(logging_outputs) > 0:
-            log_keys = list(logging_outputs[0].keys())
-            for k in log_keys:
-                if not ignore:
-                    v = sum(log[k] for log in logging_outputs if k in log)
-                else:
-                    v = logging_outputs[0][k]
-                    v = torch.zeros_like(v) if torch.is_tensor(v) else 0
-                data["logging_outputs_" + k] = v
-        else:
-            log_keys = None
-
-        data = distributed_utils.all_reduce_dict(
-            data, device=self.device, group=self.data_parallel_process_group
-        )
-
-        extra_stats_to_sum = [
-            data["extra_stats_" + str(i)] for i in range(len(extra_stats_to_sum))
-        ]
-        if log_keys is not None:
-            logging_outputs = [{k: data["logging_outputs_" + k] for k in log_keys}]
-        else:
-            logging_outputs = []
-        return logging_outputs, extra_stats_to_sum
-
-    def _check_grad_norms(self, grad_norm):
-        """Check that grad norms are consistent across workers."""
-        if self._grad_norm_buf is not None:
-            self._grad_norm_buf.zero_()
-            self._grad_norm_buf[self.data_parallel_rank] = grad_norm
-            distributed_utils.all_reduce(
-                self._grad_norm_buf, group=self.data_parallel_process_group
-            )
-
-            def is_consistent(tensor):
-                max_abs_diff = torch.max(torch.abs(tensor - tensor[0]))
-                return (
-                    (torch.isfinite(tensor).all()
-                     and (max_abs_diff / (tensor[0] + 1e-6) < 1e-6).all())
-                    or
-                    (self.cfg.common.amp and not torch.isfinite(tensor).all())
-                    # in case of amp non-finite grads are fine
-                )
-
-            if not is_consistent(self._grad_norm_buf):
-                pretty_detail = "\n".join(
-                    "rank {:3d} = {:.8f}".format(r, n)
-                    for r, n in enumerate(self._grad_norm_buf.tolist())
-                )
-                error_detail = "grad_norm across the workers:\n{}\n".format(
-                    pretty_detail
-                )
-                # use FloatingPointError to trigger NanDetector
-                raise FloatingPointError(
-                    "Fatal error: gradients are inconsistent between workers. "
-                    "Try --ddp-backend=legacy_ddp. "
-                    "Or are you mixing up different generation of GPUs in training?"
-                    + "\n"
-                    + "-" * 80
-                    + "\n{}\n".format(error_detail)
-                    + "-" * 80
-                )
-
-    def _reduce_and_log_stats(self, logging_outputs, sample_size, grad_norm=None):
-        if grad_norm is not None and (
-            not torch.is_tensor(grad_norm) or torch.isfinite(grad_norm)
-        ):
-            metrics.log_speed("ups", 1.0, priority=100, round=2)
-            metrics.log_scalar("gnorm", grad_norm, priority=400, round=3)
-            if self.cfg.optimization.clip_norm > 0:
-                metrics.log_scalar(
-                    "clip",
-                    torch.where(
-                        grad_norm > self.cfg.optimization.clip_norm,
-                        grad_norm.new_tensor(100),
-                        grad_norm.new_tensor(0),
-                    ),
-                    priority=500,
-                    round=1,
-                )
-
-        with metrics.aggregate() as agg:
-            if logging_outputs is not None:
-                self.task.reduce_metrics(logging_outputs, self.get_criterion())
-                del logging_outputs
-
-            # extra warning for criterions that don't properly log a loss value
-            if "loss" not in agg:
-                if "loss" not in self._warn_once:
-                    self._warn_once.add("loss")
-                    logger.warning(
-                        "Criterion.reduce_metrics did not log a 'loss' value, "
-                        "which may break some functionality"
-                    )
-                metrics.log_scalar("loss", -1)
-
-            # support legacy interface
-            if self.tpu:
-                logging_output = {}
-            else:
-                logging_output = agg.get_smoothed_values()
-                logging_output["sample_size"] = sample_size
-                for key_to_delete in ["ppl", "wps", "wpb", "bsz"]:
-                    if key_to_delete in logging_output:
-                        del logging_output[key_to_delete]
-            return logging_output
-
-    def _check_xla_compilation(self):
-        import torch_xla.debug.metrics as met
-
-        compile_stats = met.metric_data("CompileTime")
-        if compile_stats is None:
-            return
-        num_xla_compiles = compile_stats[0]
-        if num_xla_compiles > self._num_xla_compiles:
-            logger.warning(
-                "XLA compilation detected on device #{}; too many of these can lead "
-                "to slow training, but we expect a few in the beginning".format(
-                    self.cfg.distributed_training.distributed_rank
-                )
-            )
-        self._num_xla_compiles = num_xla_compiles
-
-    def _xla_markstep_and_send_to_cpu(self, data=None):
-        import torch_xla.core.xla_model as xm
-
-        xm.mark_step()
-        if data is not None:
-            from fairseq.utils import xla_device_to_cpu
-
-            return xla_device_to_cpu(data)
-
-
-def _catalog_shared_params(module, memo=None, prefix=""):
-    if memo is None:
-        first_call = True
-        memo = {}
-    else:
-        first_call = False
-    for name, param in module._parameters.items():
-        param_prefix = prefix + ("." if prefix else "") + name
-        if param not in memo:
-            memo[param] = []
-        memo[param].append(param_prefix)
-    for name, m in module._modules.items():
-        if m is None:
-            continue
-        submodule_prefix = prefix + ("." if prefix else "") + name
-        _catalog_shared_params(m, memo, submodule_prefix)
-    if first_call:
-        return [x for x in memo.values() if len(x) > 1]
-
-
-def _get_module_by_path(module, path):
-    path = path.split(".")
-    for name in path:
-        module = getattr(module, name)
-    return module
-
-
-def _set_module_by_path(module, path, value):
-    path = path.split(".")
-    for name in path[:-1]:
-        module = getattr(module, name)
-    setattr(module, path[-1], value)

transformers.md

@@ -1,69 +0,0 @@
-# Use in huggingface transformers (Beta)
-
-[**Colab Notebook**](https://colab.research.google.com/drive/1Ho81RBV8jysZ7e0FhsSCk_v938QeDuy3?usp=sharing)
-![image](https://user-images.githubusercontent.com/27664428/190052470-56679999-571b-4d46-a9a8-e567b78e20d1.png)
-
-
-We now support inference of OFA on the huggingface transformers. In the near future, we will provide the codes for training. 
-
-Model checkpoints are stored in our [huggingface models](https://huggingface.co/OFA-Sys). Specifically, 5 versions of the pretrained OFA models, namely OFA-tiny, OFA-medium, OFA-base, OFA-large, and OFA-huge have been already uploaded. For more information about the models, please refer to the Model Card on our [README](https://github.com/OFA-Sys/OFA). 
-Note that each directory includes 4 files, namely `config.json` which consists of model configuration, `vocab.json` and `merge.txt` for our OFA tokenizer, and lastly `pytorch_model.bin` which consists of model weights. There is no need to worry about the mismatch between Fairseq and transformers, since we have addressed the issue yet. 
-
-To use it in transformers, you can first refer to this notebook ([link](https://colab.research.google.com/drive/1Ho81RBV8jysZ7e0FhsSCk_v938QeDuy3?usp=sharing)). For more information, you can find codes in this branch https://github.com/OFA-Sys/OFA/tree/feature/add_transformers. 
-
-In the following, we introduce the details in our provided notebook and illustrate how to use OFA in Transformers. 
-
-First, install the transformers and download the models (take OFA-tiny as an example) as shown below.
-
-```
-git clone --single-branch --branch feature/add_transformers https://github.com/OFA-Sys/OFA.git
-pip install OFA/transformers/
-git clone https://huggingface.co/OFA-Sys/OFA-tiny 
-```
-
-Next, refer the path to OFA-tiny to `ckpt_dir`, and prepare an image for the testing example below. Also, ensure that you have pillow and torchvision in your environment. Check if there is the directory `generate` in your model directory `transformers/src/transformers/models/ofa` to ensure that you can use the sequence generator that we provide. 
-
-```
->>> from PIL import Image
->>> from torchvision import transforms
->>> from transformers import OFATokenizer, OFAModel
->>> from transformers.models.ofa.generate import sequence_generator
-
->>> mean, std = [0.5, 0.5, 0.5], [0.5, 0.5, 0.5]
->>> resolution = 256
->>> patch_resize_transform = transforms.Compose([
-        lambda image: image.convert("RGB"),
-        transforms.Resize((resolution, resolution), interpolation=Image.BICUBIC),
-        transforms.ToTensor(), 
-        transforms.Normalize(mean=mean, std=std)
-    ])
-
-
->>> tokenizer = OFATokenizer.from_pretrained(ckpt_dir)
-
->>> txt = " what does the image describe?"
->>> inputs = tokenizer([txt], return_tensors="pt").input_ids
->>> img = Image.open(path_to_image)
->>> patch_img = patch_resize_transform(img).unsqueeze(0)
-
-
->>> # using the generator of fairseq version
->>> model = OFAModel.from_pretrained(ckpt_dir, use_cache=True)
->>> generator = sequence_generator.SequenceGenerator(
-                    tokenizer=tokenizer,
-                    beam_size=5,
-                    max_len_b=16,
-                    min_len=0,
-                    no_repeat_ngram_size=3,
-                )
->>> data = {}
->>> data["net_input"] = {"input_ids": inputs, 'patch_images': patch_img, 'patch_masks':torch.tensor([True])}
->>> gen_output = generator.generate([model], data)
->>> gen = [gen_output[i][0]["tokens"] for i in range(len(gen_output))]
-
->>> # using the generator of huggingface version
->>> model = OFAModel.from_pretrained(ckpt_dir, use_cache=False)
->>> gen = model.generate(inputs, patch_images=patch_img, num_beams=5, no_repeat_ngram_size=3) 
-
->>> print(tokenizer.batch_decode(gen, skip_special_tokens=True))
-```