init

Image_gen.ipynb

@@ -0,0 +1,301 @@
+{
+ "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
+}

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright 1999-2022 Alibaba Group Holding Ltd.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README.md

@@ -1,12 +1,618 @@
----
-title: UnIVAL
-emoji: 🌖
-colorFrom: green
-colorTo: green
-sdk: gradio
-sdk_version: 3.35.2
-app_file: app.py
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+<!---
+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.
+-->
+
+todo:
+models
+data
+all readme
+animation
+
+readme of:
+rewarded soups
+and others  
+
+<p align="center">
+    <br>
+    <img src="examples/OFA_logo_tp_path.svg" width="150" />
+    <br>
+<p>
+<br>
+
+<p align="center">
+        <a href="modelscope.md">ModelScope</a>&nbsp ｜ &nbsp<a href="checkpoints.md">Checkpoints</a>&nbsp ｜ &nbsp<a href="colab.md">Colab</a>&nbsp ｜ &nbsp<a href="https://huggingface.co/ofa-sys">Demo</a>&nbsp ｜ &nbsp<a href="http://arxiv.org/abs/2202.03052">Paper </a>&nbsp ｜ &nbspBlog
+</p>
+
+<p align="center">
+    <br>
+    <img src="examples/demo.gif" width="800" />
+    <br>
+<p>
+
+[colab]: <https://colab.research.google.com/assets/colab-badge.svg>
+
+OFA is a unified sequence-to-sequence pretrained model (support **English** and **Chinese**) that unifies modalities (i.e., cross-modality, vision, language) and tasks (**finetuning** and **prompt tuning** are supported): image captioning (1st at the [MSCOCO Leaderboard](https://competitions.codalab.org/competitions/3221#results)), VQA ([link](https://eval.ai/web/challenges/challenge-page/830/leaderboard/2278)), visual grounding, text-to-image generation, text classification, text generation, image classification, etc. We provide **step-by-step** instructions for pretraining and finetuning and corresponding checkpoints (check official ckpt \[[EN](checkpoints.md)|[CN](checkpoints_cn.md)\] or [huggingface ckpt](https://huggingface.co/OFA-Sys)).
+
+We sincerely welcome contributions to our project. Feel free to contact us or send us issues / PRs!
+<br></br>
+
+# Our installation
+
+after installling pycocoevalcap, donwload needed models:
+```
+python -c "from pycocoevalcap.spice.spice import Spice; tmp = Spice()"
+
+```
+
+# Online Demos
+We provide online demo via Hugging Face Spaces for you to interact with our pretrained and finetuned models. Below are the links to the demos:
+* Image Captioning \[[ModelScope](https://modelscope.cn/#/models/damo/ofa_image-caption_coco_large_en/summary)  |  [Spaces](https://huggingface.co/spaces/OFA-Sys/OFA-Image_Caption)\]
+* Visual Grounding \[[ModelScope](https://modelscope.cn/#/models/damo/ofa_visual-grounding_refcoco_large_en/summary) | [Spaces](https://huggingface.co/spaces/OFA-Sys/OFA-Visual_Grounding)\]
+* Visual Question Answering \[[ModelScope](https://modelscope.cn/#/models/damo/ofa_visual-question-answering_pretrain_large_en/summary) | [Spaces](https://huggingface.co/spaces/OFA-Sys/OFA-Visual_Question_Answering)\]
+* Text-to-Image Generation \[[ModelScope](https://modelscope.cn/#/models/damo/ofa_text-to-image-synthesis_coco_large_en/summary) | [Spaces](https://huggingface.co/spaces/OFA-Sys/OFA-Text2Image_Generation)\]
+* Generic Interface \[[Spaces](https://huggingface.co/spaces/OFA-Sys/OFA-Generic_Interface)\]
+
+Also we provide Colab notebooks for you to better perceive the procedures. Click [here](colab.md) to check them out!
+<br></br>
+
+# Use in Huggingface Transformers
+We support the inference of OFA in Huggingface Transformers. Check the [README](transformers.md) and [Colab Notebook](https://colab.research.google.com/drive/1Ho81RBV8jysZ7e0FhsSCk_v938QeDuy3?usp=sharing) for more information. Codes are released in this branch https://github.com/OFA-Sys/OFA/tree/feature/add_transformers
+<br><br>
+
+
+# News
+* 2022.8.22: Released checkpoints and demos of **OFA** and **Chinese CLIP** on [ModelScope](https://modelscope.cn/). Check the [README](modelscope.md) for more details!
+* 2022.8.16: Released the **Chinese** version of OFA. **OFA-CN** needs only switching to `bpe_dir=../../utils/BERT_CN_dict` and `bpe=bert` and using our provided Chinese checkpoints in [checkpoints_cn.md](checkpoints_cn.md). Temporarily, we only provide base-size and large-size pretrained checkpoints and finetuned checkpoints on [MUGE Caption](https://tianchi.aliyun.com/muge) and the Chinese version of RefCOCO(-/+/g) (to release soon). 
+* 2022.8.5: Released support of **prompt tuning** for OFA. Check our paper [here](https://arxiv.org/abs/2208.02532)! Please see the [prompt_tuning.md](prompt_tuning.md) for further details.
+* 2022.7.7: Updated support of OFA on **huggingface transformers** (fixed bugs in forward, add sequence generator from Fairseq to ensure performance, etc.). Refer to the doc [transformers.md](transformers.md) and the branch `feature/add_transformers`. 
+* 2022.6.17: Released the pretrained checkpoint of **OFA-Huge**. To use it, set `--arch=ofa_huge` in the script.
+* 2022.5.15: OFA was accepted by **ICML 2022**
+* 2022.4.28: Add support of inference on **huggingface transformers**. For how to use it, please refer to the doc [transformers.md](transformers.md) and our [huggingface models](https://huggingface.co/OFA-Sys). 
+* 2022.4.16: Released lightweight pretrained models **OFA-Medium** (~93M params) and **OFA-Tiny** (~33M params) in [checkpoints.md](checkpoints.md). To use them, you just need to load the corresponding checkpoint and set `--arch=ofa_medium` or `--arch=ofa_tiny` in the scripts.
+
+<details>
+    <summary><b>More News</b></summary>
+    <p>
+        <ul>
+            <li>2022.3.23: Added [Encouraging Loss](https://arxiv.org/pdf/2110.06537.pdf) as a feature. See [README_EncouragingLoss.md](README_EncouragingLoss.md). Leveraging this feature, OFA-Large has achieved improved results in both VQA (**test-std acc: 80.67**) and Image Classification (**test acc: 85.6**) recently.</li>
+            <li>2022.3.21: Released codes for pretraining OFA.</li>
+            <li>2022.3.18: Released the finetuned <b>OFA-Base</b> (~180M parameters) checkpoints and running scripts for vision & language tasks, including: <b>Caption (146.4 CIDEr), VQA (78.07 on test-std), SNLI-VE (89.3 on dev), RefCOCO (90.67 on testA), RefCOCO+ (87.15 on testA) and RefCOCOg (82.31 on test-u)</b>.</li>
+            <li>2022.3.11: Released the finetuning & inference code/checkpoints for <b>Gigaword</b>.</li>
+            <li>2022.3.08: Released the pretrained checkpoint of <b>OFA-Base</b> in <a href="https://github.com/OFA-Sys/OFA/blob/main/checkpoints.md">checkpoints.md</a>. To use OFA-Base, you just need to load <code>ofa_base.pt</code> and change <code>--arch=ofa_large</code> to <code>--arch=ofa_base</code> in the training scripts.</li>
+            <li>2022.3.07: Released the finetuning & inference code/checkpoints for <b>Image Classification</b>, which achieves <b>85.0</b> accuracy on ImageNet-1K, slightly better than reported in OFA paper.</li>
+            <li>2022.3.04: Released the finetuning & inference code/checkpoints for <b>Text-to-Image Generation</b>.</li>
+            <li>2022.3.03: Released the finetuning & inference code/checkpoints for <b>SNLI-VE</b> and <b>GLUE</b>.</li>
+            <li>2022.2.22: Released the finetuning & inference code/checkpoints for <b>Visual Question Answering</b>, which can reproduce <b>the reported VQA accuracy in OFA paper (80.02 on test-std)</b>. Check our results on the <a href="https://eval.ai/web/challenges/challenge-page/830/leaderboard/2278">VQA Challenge</a>.</li>
+            <li>2022.2.15: Released finetuning & inference code/checkpoints for <b>Referring Expression Comprehension</b></li>
+            <li>2022.2.10: Released the inference code & finetuned checkpoint for <b>Image captioning</b>, which can reproduce <b>the results on COCO Karparthy test split (149.6 CIDEr)</b>. OFA also achieves No.1 on the COCO image captioning online leaderboard <a href='https://competitions.codalab.org/competitions/3221#results'>Link</a> (marked as M6-Team).</li>
+        </ul>
+    </p>
+</details>
+<br></br>
+
+
+# Model Card
+We list the parameters and pretrained checkpoints of OFAs below. For finetuned checkpoints, please refer to [checkpoints.md](checkpoints.md). 
+
+<table border="1" width="100%">
+    <tr align="center">
+        <th>Model</th><th>Ckpt</th><th>Params</th><th>Backbone</th><th>Hidden size</th><th>Intermediate size</th><th>Num. of heads</th><th>Enc layers</th><th>Dec layers</th>
+    </tr>
+    <tr align="center">
+        <td>OFA<sub>Tiny</sub></td><td><a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/checkpoints/ofa_tiny.pt">Download</a></td><td>33M</td><td>ResNet50</td><td>256</td><td>1024</td><td>4</td><td>4</td><td>4</td>
+    </tr>
+    <tr align="center">
+        <td>OFA<sub>Medium</sub></td><td><a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/checkpoints/ofa_medium.pt">Download</a></td><td>93M</td><td>ResNet101</td><td>512</td></td><td>2048</td><td>8</td><td>4</td><td>4</td>
+    </tr>
+    <tr align="center">
+        <td>OFA<sub>Base</sub></td><td><a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/checkpoints/ofa_base.pt">Download</a></td><td>180M</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><a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/checkpoints/ofa_large.pt">Download</a></td><td>470M</td><td>ResNet152</td><td>1024</td></td><td>4096</td><td>16</td><td>12</td><td>12</td>
+    </tr>
+    <tr align="center">
+        <td>OFA<sub>Huge</sub></td><td><a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/checkpoints/ofa_huge.pt">Download</a></td><td>930M</td><td>ResNet152</td><td>1280</td></td><td>5120</td><td>16</td><td>24</td><td>12</td>
+    </tr>
+</table>
+<br></br>
+
+# Results
+Below we demonstrate the results of OFAs on cross-modal understanding and generation. 
+
+<table border="1" width="100%">
+    <tr align="center">
+        <th>Task</th><th>Image Captioning</th><th>VQA</th><th>Visual Entailment</th><th colspan="3">Referring Expression Comprehension</th>
+    </tr>
+    <tr align="center">
+        <td>Dataset</td><td>COCO</td><td>VQA v2</td><td>SNLI-VE</td><td>RefCOCO</td><td>RefCOCO+</td><td>RefCOCOg</td>
+    </tr>
+    <tr align="center">
+        <td>Split</td><td>Karpathy test (CE/CIDEr)</td><td>test-dev/test-std</td><td>val/test</td><td>val/test-a/test-b</td><td>val/test-a/test-b</td><td>val-u/test-u</td>
+    </tr>
+    <tr align="center">
+        <td>Metric</td><td>CIDEr</td><td>Acc.</td><td>Acc.</td><td colspan="3">Acc.</td>
+    </tr>
+    <tr align="center">
+        <td>OFA<sub>Tiny</sub></td><td>119.0 / 128.7</td><td>70.3 / 70.4</td><td>85.3 / 85.2</td><td>80.20 / 84.07 / 75.00</td><td>68.22 / 75.13 / 57.66</td><td>72.02 / 69.74</td>
+    </tr>
+    <tr align="center">
+        <td>OFA<sub>Medium</sub></td><td>130.4 / 140.3</td><td>75.4 / 75.5</td><td>86.6 / 87.0</td><td>85.34 / 87.68 / 77.92</td><td>76.09 / 83.04 / 66.25</td><td>78.76 / 78.58</td>
+    </tr>
+    <tr align="center">
+        <td>OFA<sub>Base</sub></td><td>138.2 / 146.7</td><td>78.0 / 78.1</td><td>89.3 / 89.2</td><td>88.48 / 90.67 / 83.30</td><td>81.39 / 87.15 / 74.29</td><td>82.29 / 82.31</td>
+    </tr>
+    <tr align="center">
+        <td>OFA<sub>Large</sub></td><td>142.2 / 150.7</td><td>80.4 / 80.7</td><td>90.3 / 90.2</td><td>90.05 / 92.93 / 85.26</td><td>85.80 / 89.87 / 79.22</td><td>85.89 / 86.55</td>
+    </tr>
+    <tr align="center">
+        <td>OFA<sub>Huge</sub></td><td>145.3 / 154.9</td><td>82.0 / 82.0</td><td>91.0 / 91.2</td><td>92.04 / 94.03 / 88.44</td><td>87.86 / 91.70 / 80.71</td><td>88.07 / 88.78</td>
+    </tr>
+</table>
+<br></br>
+
+# Requirements
+* python 3.7.4
+* pytorch 1.8.1
+* torchvision 0.9.1
+* JAVA 1.8 (for COCO evaluation)
+<br></br>
+
+# Installation
+```bash
+git clone https://github.com/OFA-Sys/OFA
+pip install -r requirements.txt
+```
+<br></br>
+
+# Datasets and Checkpoints
+See [datasets.md](datasets.md) and [checkpoints.md](checkpoints.md).
+<br></br>
+
+# Training & Inference
+Below we provide methods for training and inference on different tasks. We provide both pretrained OFA-Large and OFA-Base in [checkpoints.md](checkpoints.md). The scripts mentioned in this section are prepared for OFA-Large. For reproducing the downstreaming results of OFA-Base, we have also provided the corresponding finetuning and inference scripts for OFA-Base in the `run_scripts/` folder.
+
+We recommend that your workspace directory should be organized like this: 
+```
+OFA/
+├── checkpoints/
+│   ├── ofa_base.pt
+│   ├── ofa_large.pt
+│   ├── caption_large_best_clean.pt
+│   └── ...
+├── criterions/
+├── data/
+├── dataset/
+│   ├── caption_data/
+│   ├── gigaword_data/
+│   └── ...
+├── fairseq/
+├── models/
+├── run_scripts/
+├── tasks/
+├── train.py
+├── trainer.py
+└── utils/
+```
+
+
+## Image Processing
+To ensure the efficiency of processing data, we did not store images with small files, but instead we encode them to base64 strings.
+Transforming image files to base64 strings is simple. Run the following code:
+```python
+from PIL import Image
+from io import BytesIO
+import base64
+
+img = Image.open(file_name) # path to file
+img_buffer = BytesIO()
+img.save(img_buffer, format=img.format)
+byte_data = img_buffer.getvalue()
+base64_str = base64.b64encode(byte_data) # bytes
+base64_str = base64_str.decode("utf-8") # str
+```
+
+## Pretraining
+Below we provide methods for pretraining OFA.
+
+<details>
+    <summary><b>1. Prepare the Dataset</b></summary>
+    <p>
+        To pretrain OFA, you should first download the dataset we provide (<a href="https://ofa-beijing.oss-cn-beijing.aliyuncs.com/datasets/pretrain_data/pretrain_data_examples.zip">pretrain_data_examples.zip</a>, a small subset of the original pretraining data). For your customed pretraining datasets, please prepare your training samples into the same format. <code>pretrain_data_examples.zip</code> contains 4 TSV files: <code>vision_language_examples.tsv</code>, <code>text_examples.tsv</code>, <code>image_examples.tsv</code> and <code>detection_examples.tsv</code>. Details of these files are as follows: 
+        <br />
+        <ul type="circle">
+            <li><b>vision_language_examples.tsv</b>:
+    Each line contains uniq-id, image (base64 string), caption, question, answer, ground-truth objects (objects appearing in the caption or question), dataset name (source of the data) and task type (caption, qa or visual gronunding). Prepared for the pretraining tasks of visual grounding, grounded captioning, image-text matching, image captioning and visual question answering. </li>
+            <li><b>text_examples.tsv</b>: Each line contains uniq-id and text. Prepared for the pretraining task of text infilling. </li> 
+            <li><b>image_examples.tsv</b>: Each line contains uniq-id, image (base64 string, should be resized to 256*256 resolution) and image-code (generate the sparse codes for the central part of image through VQ-GAN). Prepared for the pretraining task of image infilling. </li>
+            <li><b>detection_examples.tsv</b>: Each line contains uniq-id, image (base64 string) and bounding box annotations (contains the top-left and bottom-right coordinates of the bounding box, object_id and object_name, seperated by commas). Prepared for the pretraining task of detection. </li>
+        </ul>
+        In addition, the folder negative_sample in pretrain_data_examples.zip contains three files <code>all_captions.txt</code>, <code>object.txt</code> and <code>type2ans.json</code>. The data in these files are used as negative samples for the image-text matching (ITM) task.
+    </p>
+</details>
+<details>
+    <summary><b>2. Pretraining</b></summary>
+    <p>
+        By default, the pretraining script will attempt to restore the released pretrained checkpoints of OFA-Base or OFA-Large and perform continuous pretraining. Continuous pretraining is more recommended, which achieves much better results compared with pretraining from scratch. For continuous pretraining, please download the pretrained weights in advance (see <a href='checkpoints.md'>checkpoints.md</a>) and put them in the correct directory <code>OFA/checkpoints/</code>. If not, the pretraining will begin from scratch.
+    </p>
+<pre>
+cd run_scripts/pretraining
+bash pretrain_ofa_large.sh # Pretrain OFA-Large. For OFA-Base, use pretrain_ofa_base.sh
+</pre>
+    <p>
+        If the pretrained OFA checkpoint is restored successfully, you will see the following information in the log:
+    </p>
+<pre>
+INFO: Loaded checkpoint ../../checkpoints/ofa_large.pt
+</pre>
+</details>
+
+## Image Captioning
+We provide procedures to reproduce our results of image captioning on our paper below.
+<details>
+    <summary><b>1. Prepare the Dataset & Checkpoints</b></summary>
+    <p>
+        Download data (see <a href='datasets.md'>datasets.md</a>) and models (see <a href='checkpoints.md'>checkpoints.md</a>) and put them in the correct directory. The dataset zipfile <code>caption_data.zip</code> contains caption_stage1_train.tsv, caption_stage2_train.tsv, caption_val.tsv and caption_test.tsv. Each image corresponds to only 1 caption in <code>caption_stage1_train.tsv</code> and corresponds to multiple captions in other TSV files (about 5 captions per image). Each line of the dataset represents a caption sample with the following format. The information of uniq-id, image-id, caption, predicted object labels (taken from <a href='https://github.com/pzzhang/VinVL'>VinVL</a>, not used), image base64 string are separated by tabs.
+    </p>
+<pre>
+162365  12455   the sun sets over the trees beyond some docks.  sky&&water&&dock&&pole  /9j/4AAQSkZJ....UCP/2Q==
+</pre>
+</details>
+<details>
+    <summary><b>2. Finetuning</b></summary>
+    <p>
+        Following previous standard practice, we divide the finetuning process of image captioning into two stages. In stage 1, we finetune OFA with cross-entropy loss on 4 NVIDIA-V100 GPUs with 32GB memory (expected to obtain ~139.5 CIDEr on the validation set at this stage). In stage 2, we select the best checkpoint of stage 1 and train with CIDEr optimization on 8 NVIDIA-V100 GPUs. <b>Note that CIDEr optimization is very unstable and requires careful hyperparameter tuning. If you encounter training errors in the stage2 finetuning, you can increase the batch size or reduce the learning rate. If neither of these works, you can directly set </b><code>--freeze-resnet</code><b> to freeze the inner states of batch normalization.</b>
+    </p>
+<pre>
+cd run_scripts/caption
+nohup sh train_caption_stage1.sh > train_stage1.out &  # stage 1, train with cross-entropy loss
+nohup sh train_caption_stage2.sh > train_stage2.out &  # stage 2, load the best ckpt of stage1 and train with CIDEr optimization 
+</pre>
+</details>
+<details>
+    <summary><b>3. Inference</b></summary>
+    <p>
+        Run the following commands to get your results and evaluate your model.
+    </p>
+<pre>
+cd run_scripts/caption ; sh evaluate_caption.sh  # inference & evaluate
+</pre>
+</details>
+
+## Text-to-Image Generation 
+This part provides procedures for the finetuning and inference of text-to-image generation. See below.
+
+<details>
+    <summary><b>1. Prepare the Dataset & Checkpoints</b></summary>
+    <p>
+        Download data (see <a href="datasets.md">datasets.md</a>) and models (see <a href="checkpoints.md">checkpoints.md</a>) and put them in the correct directory. The dataset zipfile <code>coco_image_gen.zip</code> contains <code>coco_vqgan_train.tsv</code>, <code>coco_vqgan_dev.tsv</code> and <code>coco_vqgan_full_test.tsv</code>. Each line of the dataset represents a sample with the following format. The information of uniq-id, image-code (produced by <a href="https://github.com/CompVis/taming-transformers">vqgan</a>, a list of integers separated by single-whitespaces), lowercased caption are separated by tabs.
+    </p>
+<pre>
+1	6674 4336 4532 5334 3251 5461 3615 2469 ...4965 4190 1846	the people are posing for a group photo.
+</pre>
+    <p>
+        The checkpoint zipfile <code>image_gen_large_best.zip</code> contains <code>image_gen_large_best.pt</code>, <code>vqgan/last.ckpt</code>, <code>vqgan/model.yaml</code> and <code>clip/Vit-B-16.pt</code>. 
+    </p>
+</details>
+<details>
+    <summary><b>2. Shuffle the Training Data</b></summary>
+    <p>
+        (Optional, but achieves better result): If the disk storage is sufficient, we recommend to prepare the shuffled training data for each epoch in advance. 
+    </p>
+<pre>
+cd dataset/image_gen
+ln coco_vqgan_train.tsv coco_vqgan_train_1.tsv
+for idx in `seq 1 9`;do shuf coco_vqgan_train_${idx}.tsv > coco_vqgan_train_$[${idx}+1].tsv;done # each file is used for an epoch
+</pre>
+</details>
+<details>
+    <summary><b>3. Finetuning</b></summary>
+    <p>
+        Following previous practice, we divide the finetuning process of image generating into two stages. In stage 1, we finetune OFA with cross-entropy loss on 4 8-V100-32G-GPU servers (expected to obtain ~32.5+ CLIP Score on the validation set at this stage). In stage 2, we select the last checkpoint of stage 1 and train with CLIP Score optimization on 4 8-V100-32G-GPU servers (expected to obtain ~34.0+ CLIP Score on the validation set at this stage). During the validation, the generated image will be dumped into <code>_GEN_IMAGE_PATH_</code>. 
+    </p>
+<pre>
+# run on each worker after the distributed and data configs have been correctly set following the guide in train_image_gen_stage1_distributed.sh 
+cd run_scripts/image_gen
+nohup sh train_image_gen_stage1_distributed.sh # stage 1, train with cross-entropy loss
+nohup sh train_image_gen_stage2_distributed.sh # stage 2, load the last ckpt of stage1 and train with CLIP Score optimization 
+</pre>
+</details>
+<details>
+    <summary><b>4. Inference</b></summary>
+    <p>
+        Run the command below to generate your images. 
+    </p>
+<pre>
+cd run_scripts/image_gen ; sh evaluate_image_gen.sh  # inference & evaluate (FID, IS and CLIP Score)
+</pre>
+</details>
+
+## Visual Question Answering
+Here we provide the finetuning and inference codes to reproduce the VQAv2 result reported in our paper (**test-std 80.02**). We believe much improvement on accuracy can still be achieved based on this codebase :)
+<details>
+    <summary><b>1. Prepare the Dataset & Checkpoints</b></summary>
+    <p>
+        Download data (see <a href="datasets.md">datasets.md</a>) and models (see <a href="checkpoints.md">checkpoints.md</a>) and put them in the correct directory. The dataset zipfile <code>vqa_data.zip</code> is around 100G and the decompressed data costs around 135G disk storage, which contains the training, validation and testing samples together with other necessary data resources. (Since <code>vqa_data.zip</code> is large in size, we have also provided chunked parts of the dataset files for more convenient and stable downloading. Please refer to <a href="https://github.com/OFA-Sys/OFA/issues/68#issuecomment-1096837349">issue #68</a>.) Following common practice, VG-QA samples are also included in the training data. To adapt to the seq2seq paradigm of OFA, we transform original VQA training questions with multiple golden answers into multiple training samples. For the original VQA validation set, we keep around 10k samples for our validation and utilize the other samples for training. Each line of the dataset represents a VQA sample with the following format. The information of question-id, image-id, question, answer (with confidence), predicted object labels (taken from <a href="https://github.com/pzzhang/VinVL">VinVL</a>, slightly brings around +0.1 accuracy improvement), image base64 string are separated by tabs. 
+    </p>
+<pre>
+79459   79459   is this person wearing shorts?  0.6|!+no    house&&short&&...&&sky  /9j/4AAQS...tigZ/9k=
+</pre>
+    <p>
+        For fine-tuning on customed VQA-formulated tasks, please refer to issue <a href="https://github.com/OFA-Sys/OFA/issues/76">#76</a>, <a href="https://github.com/OFA-Sys/OFA/issues/105">#105</a> and <a href="https://github.com/OFA-Sys/OFA/issues/73">#73</a> for more information.
+    </p>
+</details>
+<details>
+    <summary><b>2. Shuffle the Training Data</b></summary>
+    <p>
+        (Optional, but achieves better finetuning accuracy): If the disk storage is sufficient, we recommend to prepare the shuffled training data for each epoch in advance. In our experiments, we use shuffling which brings around <b>+0.3</b> improvement on VQA accuracy.
+    </p>
+<pre>
+cd dataset/vqa_data
+ln vqa_train.tsv vqa_train_1.tsv
+for idx in `seq 1 9`;do shuf vqa_train_${idx}.tsv > vqa_train_$[${idx}+1].tsv;done # each file is used for an epoch
+</pre>
+</details>
+<details>
+    <summary><b>3. Finetuning</b></summary>
+    <p>
+        In our experiments, the VQA finetuning is performed on 4 8-A100-GPU servers (<i>with RDMA</i>). Here provides the finetuning script <code>train_vqa_distributed.sh</code>, which supports multi-server distributed training (as well as single-server training). Please refer to the comments in the beginning of the script and set the configs correctly according to your distribution environment. If you have shuffled the training data in the previous step, please correctly specify the training data path following the guide in the script comments. <b>The command should be run on each worker.</b> 
+    </p>
+<pre>
+# run on each worker after the distributed and data configs have been correctly set following the guide in train_vqa_distributed.sh 
+cd run_scripts/vqa
+bash train_vqa_distributed.sh 
+</pre>
+    <p>
+        In our experiments, the finetuning costs around 36 hours (for 12 epochs). After each epoch, an evaluation on validation set is performed. The best validation accuracy during finetuning will be around 80.8. The log is saved in <code>${log_dir}</code>.
+    </p>
+    <p>
+        <i>(Update on validation time-cost)</i> As will be mentioned in the <i>4. Inference</i> section, we prepare 2 types of inference: beam-search and all-candidate inference. By default, all-candidate inference is used for validation during fine-tuning, which achieves better accuracy but costs much time. Now we have added a new option in the training scripts called <code>--val-inference-type</code> to switch the validation inference type during fine-tuning. If you feel the validation takes too long, you can refer to <a href="https://github.com/OFA-Sys/OFA/pull/79">PR #79</a> to activate beam-search validation, which significantly takes much less time, with around 0.5-0.6 validation score degradation compared with all-candidate validation.
+    </p>
+</details>
+<details>
+    <summary><b>4. Inference</b></summary>
+    <p>
+        We provide 2 types of inference, <b>beam-search</b> (much faster but gets sub-optimal accuracy) and <b>all-candidate evaluation</b> (slower but best accuracy). <br></br>
+        For beam-search inference, use the script <code>evaluate_vqa_beam.sh</code>. Refer to the command below. The inference on test set costs around 16 GPU hours. After inference on test set, the result JSON file will be dumped in the <code>${result_path}</code> defined in the shell script. You can submit the result <code>test_predict.json</code> to <a href="https://eval.ai/web/challenges/challenge-page/830/overview">EvalAI</a>. Using our released finetuned checkpoint, beam-search inference will get 80.15 validation accuracy, 79.36 test-dev accuracy and 79.48 test-std accuracy (around 0.6 lower than all-candidate evaluation).
+    </p>
+<pre>
+cd run_scripts/vqa
+bash evaluate_vqa_beam.sh val # specify 'val' or 'test'
+</pre>
+    <p>
+        For all-candidate evaluation, we recommend to use the distributed script <code>evaluate_vqa_allcand_distributed.sh</code>. Please refer to the guide in the script to set the distributed configs before running. The result JSON file will be dumped in the <code>${result_path}</code> defined in the shell script of rank-0 server. All-candidate evaluation computes scores on all the candidate answers in the VQA dataset, which achieves <b>80.82</b> validation accuracy, <b>79.87</b> test-dev accuracy and <b>80.02</b> test-std accuracy, reproducing our reported results in the paper. However, the inference on test set costs around 1k GPU hours, which is much slower.
+    </p>
+<pre>
+# run on each worker after the distributed configs have been correctly set following the guide in evaluate_vqa_allcand_distributed.sh
+cd run_scripts/vqa
+bash evaluate_vqa_allcand_distributed.sh val # specify 'val' or 'test'
+</pre>
+</details>
+
+## Referring Expression Comprehension 
+Here provides procedures for you to prepare data, train, and evaluate your model on visual grounding. 
+<details>
+    <summary><b>1. Prepare the Dataset & Checkpoints</b></summary>
+    <p>
+        Download data (see <a href='datasets.md'>datasets.md</a>) and models (see <a href='checkpoints.md'>checkpoints.md</a>) and put them in the correct directory. We provide RefCOCO (split by UNC), RefCOCO+ (split by UNC) and RefCOCOg (split by UMD) datasets. See <a href='https://www.tensorflow.org/datasets/catalog/ref_coco'>RefCOCO</a> and <a href="https://github.com/lichengunc/refer">Refer</a> for more details. Note that in the original dataset, each region-coord (or bounding box) may corresponds to multiple descriptive texts. We split these texts into multiple samples so that the region-coord in each sample corresponds to only one text. Each line of the processed dataset represents a sample with the following format. The information of uniq-id, image-id, text, region-coord (separated by commas), image base64 string are separated by tabs.
+    </p>
+<pre>
+79_1    237367  A woman in a white blouse holding a glass of wine.  230.79,121.75,423.66,463.06 9j/4AAQ...1pAz/9k=
+</pre>
+</details>
+<details>
+    <summary><b>2. Finetuning</b></summary>
+    <p>
+        Unlike the original paper, we finetune OFA with a drop-path rate of 0.2, and found that training with this hyper-parameter achieves better results. We will update the reported results of the paper later.
+    </p>
+<pre>
+cd run_scripts/refcoco
+nohup sh train_refcoco.sh > train_refcoco.out &  # finetune for refcoco
+nohup sh train_refcocoplus.sh > train_refcocoplus.out &  # finetune for refcoco+
+nohup sh train_refcocog.sh > train_refcocog.out &  # finetune for refcocog
+</pre>
+</details>
+<details>
+    <summary><b>3. Inference</b></summary>
+    <p>
+        Run the following commands for the evaluation. 
+    </p>
+<pre>
+cd run_scripts/refcoco ; sh evaluate_refcoco.sh  # inference & evaluate for refcoco/refcoco+/refcocog
+</pre>
+</details>
+
+## Visual Entailment
+We provide steps for you to reproduce our results in visual entailment. See the details below.
+
+<details>
+    <summary><b>1. Prepare the Dataset & Checkpoints</b></summary>
+    <p>
+        Download data (see <a href="datasets.md">datasets.md</a>) and models (see <a href="checkpoints.md">checkpoints.md</a>) and put them in the correct directory. Each line of the processed dataset represents a sample with the following format. The information of uniq-id, image-id, image base64 string, hypothesis, caption (or text premise), label are separated by tabs.
+    </p>
+<pre>
+252244149.jpg#1r1n  252244149   /9j/4AAQ...MD/2Q==   a man in pink and gold is chewing on a wooden toothpick.   a man in pink is chewing a toothpick on the subway.   neutral 
+</pre>
+</details>
+<details>
+    <summary><b>2. Finetuning</b></summary>
+    <p>
+        In our experiments, the SNLI-VE finetuning is performed on 8 NVIDIA-V100 GPUs with 32GB memory. In this task, we experimented with only a few sets of hyperparameters. We believe that proper hyperparameter tuning can lead to further accuracy improvement.
+    </p>
+<pre>
+cd run_scripts/snli_ve
+nohup sh train_snli_ve.sh > train_snli_ve.out &  # finetune for snli_ve
+</pre>
+</details>
+<details>
+    <summary><b>3. Inference</b></summary>
+    <p>
+        Run the following command to obtain the results.
+    </p>
+<pre>
+cd run_scripts/snli_ve ; sh evaluate_snli_ve.sh dev  # specify 'dev' or 'test'
+</pre>
+</details>
+   
+## GLUE
+Here we provide steps for you to finetune and evaluate our model on language understanding tasks. We demonstrate our practice for the GLUE benchmark. 
+
+<details>
+    <summary><b>1. Prepare the Dataset & Checkpoints</b></summary>
+    <p>
+        Download data (see <a href="datasets.md">datasets.md</a>) and models (see <a href="checkpoints.md">checkpoints.md</a>) and put them in the correct directory. we provide 7 language understanding datasets from GLUE benchmark, including COLA, MNLI, MRPC, QNLI, QQP, RTE and SST2. More details about these datasets can be found in this <a href="https://openreview.net/pdf?id=rJ4km2R5t7">link</a>.
+    </p>
+</details>
+<details>
+    <summary><b>2. Finetuning</b></summary>
+    <p>
+        For each task, we have tried multiple sets of hyperparameters (including learning rate, batch size, training epochs). The results under different sets of hyperparameters can be found in <code>${log_dir}</code>.
+    </p>
+<pre>
+cd run_scripts/glue
+nohup sh train_cola.sh > train_cola.out &  # finetune for cola
+nohup sh train_mnli.sh > train_mnli.out &  # finetune for mnli
+nohup sh train_mrpc.sh > train_mrpc.out &  # finetune for mrpc
+nohup sh train_qnli.sh > train_qnli.out &  # finetune for qnli
+nohup sh train_qqp.sh > train_qqp.out &  # finetune for qqp
+nohup sh train_rte.sh > train_rte.out &  # finetune for rte
+nohup sh train_sst2.sh > train_sst2.out &  # finetune for sst2
+</pre>
+</details>
+
+## Image Classification on ImageNet-1K
+We provide the finetuning and inference codes which reproduce **85.0 ImageNet-1K accuracy**, slightly better than reported in our paper. 
+
+<details>
+    <summary><b>1. Prepare the Dataset & Checkpoints</b></summary>
+    <p>
+        Download data (see <a href="datasets.md">datasets.md</a>) and models (see <a href="checkpoints.md">checkpoints.md</a>) and put them in the correct directory. Our provided data is derived from the original <a href="http://image-net.org/">ImageNet-1K</a> (ILSVRC2012 train & validation) dataset and shares the same data split with it. To formulate the classification task into seq2seq paradigm, we use the <a href="https://github.com/HoldenCaulfieldRye/caffe/blob/master/data/ilsvrc12/synset_words.txt">synset words</a> provided by Caffe as the generation target for each image class. Each line of the processed dataset represents a sample with the following format. The information of image base64 string, classification label (1-indexed, conform to the order in <code>synset_words.txt</code>), synset words of the label are separated by tabs.
+    </p>
+<pre>
+_9j_4AAQS...fzX__Z  769 rugby ball
+</pre>
+</details>
+<details>
+    <summary><b>2. Shuffle the Training Data</b></summary>
+    <p>
+        (Optional, but achieves better finetuning accuracy): If the disk storage is sufficient, we recommend to prepare the shuffled training data for each epoch in advance. In our experiments, we use shuffling which brings around <b>+0.2</b> improvement on ImageNet-1K accuracy.
+    </p>
+<pre>
+cd dataset/imagenet_1k_data
+ln imagenet_1k_train.tsv imagenet_1k_train_1.tsv
+for idx in `seq 1 9`;do shuf imagenet_1k_train_${idx}.tsv > imagenet_1k_train_$[${idx}+1].tsv;done # each file is used for an epoch one by one
+</pre>
+</details>
+<details>
+    <summary><b>3. Finetuning</b></summary>
+    <p>
+        In our experiments, the ImageNet-1K finetuning is performed on 2 8-A100-GPU servers (<i>with RDMA</i>). Here provides the finetuning script <code>train_imagenet_distributed.sh</code>, which supports multi-server distributed training (as well as single-server training). Please refer to the comments in the beginning of the script and set the configs correctly according to your distribution environment. If you have shuffled the training data in the previous step, please correctly specify the training data path following the guide in the script comments. <b>The command should be run on each worker.</b> For quick evaluation during finetuning, by default we sample 20% of the original validation split and report accuracy on this subset after each epoch. The accuracy on the validation subset is generally ±0.1 relative to accuracy on the whole validation split.
+    </p>
+<pre>
+# run on each worker after the distributed and data configs have been correctly set following the guide in train_imagenet_distributed.sh
+cd run_scripts/image_classify
+bash train_imagenet_distributed.sh
+</pre>
+    <p>
+        In our experiments, the finetuning costs around 80 hours (for 32 epochs). The best accuracy on validation subset during finetuning will be around 85.0. The log is saved in <code>${log_dir}</code>.
+    </p>
+</details>
+<details>
+    <summary><b>4. Inference</b></summary>
+    <p>
+        To get the validation accuracy on the whole ImageNet-1K validation set, run the following command. The evaluation costs around 10 GPU hours. The accuracy will be reported in the stdout (expected to be around <b>85.0</b>).
+    </p>
+<pre>
+cd run_scripts/image_classify ; sh evaluate_imagenet.sh  # inference & evaluate for imagenet-1k
+</pre>
+</details> 
+
+## Gigaword
+We provide steps for you to reproduce our results in Gigaword. See the details below.
+
+<details>
+    <summary><b>1. Prepare the Dataset & Checkpoints</b></summary>
+    <p>
+        Download data (see <a href="datasets.md">datasets.md</a>) and models (see <a href="checkpoints.md">checkpoints.md</a>) and put them in the correct directory. The original dataset is taken from <a href="https://github.com/microsoft/unilm/">UniLM</a> and we organized the data into the tsv format. Each line of the processed dataset represents a sample with the following format. The information of source and target texts are separated by tabs.
+    </p>
+<pre>
+factory orders for manufactured goods rose #.# percent in september...  us september factory orders up #.# percent
+</pre>
+</details>
+<details>
+    <summary><b>2. Finetuning</b></summary>
+    <p>
+        Run the following command to train the model.
+    </p>
+<pre>
+cd run_scripts/gigaword
+nohup sh train_gigaword.sh > train_gigaword.out &  # finetune for gigaword
+</pre>
+</details>
+<details>
+    <summary><b>3. Inference</b></summary>
+    <p>
+        Run the following command to obtain the results (~36.43 rougeL).
+    </p>
+<pre>
+cd run_scripts/gigaword ; sh evaluate_gigaword.sh  # inference & evaluate for gigaword
+</pre>
+</details> 
+
+<br></br>
+
+# Gallery
+Below we provide examples of OFA in text-to-image generation and open-ended VQA. Also, we demonstrate its performance in unseen task (Grounded QA) as well as unseen domain (Visual Grounding on images from unseen domains). 
+
+## Text-to-Image Generation
+
+![case1](examples/case1.png)
+
+
+## Open-Ended VQA
+![open_vqa](examples/open_vqa.png)
+
+## Grounded QA (unseen task)
+![grounded_qa](examples/grounded_qa.png)
+
+## Visual Grounding (unseen domain)
+![vg](examples/viusal_grounding.png)
+<br></br>
+
+# Related Codebase
+* [Fairseq](https://github.com/pytorch/fairseq)
+* [taming-transformers](https://github.com/CompVis/taming-transformers)
+<br></br>
+
+
+# Getting Involved
+Feel free to submit Github issues or pull requests. Welcome to contribute to our project!
+
+To contact us, never hestitate to send an email to `zheluo.wp@alibaba-inc.com` or `junyang.ljy@alibaba-inc.com`!
+<br></br>
+
+
+# Citation
+Please cite our paper if you find it helpful :)
+
+```
+@article{wang2022ofa,
+  author    = {Peng Wang and
+               An Yang and
+               Rui Men and
+               Junyang Lin and
+               Shuai Bai and
+               Zhikang Li and
+               Jianxin Ma and
+               Chang Zhou and
+               Jingren Zhou and
+               Hongxia Yang},
+  title     = {OFA: Unifying Architectures, Tasks, and Modalities Through a Simple Sequence-to-Sequence
+               Learning Framework},
+  journal   = {CoRR},
+  volume    = {abs/2202.03052},
+  year      = {2022}
+}
+```
+<br></br>

README_EncouragingLoss.md

@@ -0,0 +1,34 @@
+# 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.

app.py

@@ -0,0 +1,297 @@
+import os
+
+os.system('cd fairseq;'
+          'pip install ./; cd ..')
+os.system('ls -l')
+
+import torch
+import numpy as np
+import gradio as gr
+import cv2
+from PIL import Image
+from torchvision import transforms
+
+from fairseq import utils, tasks, options
+from fairseq import checkpoint_utils
+from fairseq.dataclass.utils import convert_namespace_to_omegaconf
+
+from tasks.mm_tasks.caption import CaptionTask
+from tasks.mm_tasks.refcoco import RefcocoTask
+from tasks.mm_tasks.vqa_gen import VqaGenTask
+
+
+def move2gpu(models, cfg):
+    for model in models:
+        model.eval()
+        if use_fp16:
+            model.half()
+        if use_cuda and not cfg.distributed_training.pipeline_model_parallel:
+            model.cuda()
+        model.prepare_for_inference_(cfg)
+
+
+def construct_transform(patch_image_size):
+    mean = [0.5, 0.5, 0.5]
+    std = [0.5, 0.5, 0.5]
+
+    patch_resize_transform = transforms.Compose([
+        lambda image: image.convert("RGB"),
+        transforms.Resize((patch_image_size, patch_image_size), interpolation=Image.BICUBIC),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=mean, std=std),
+    ])
+
+    return patch_resize_transform
+
+
+# Register tasks
+tasks.register_task('caption', CaptionTask)
+tasks.register_task('refcoco', RefcocoTask)
+tasks.register_task('vqa_gen', VqaGenTask)
+# turn on cuda if GPU is available
+use_cuda = torch.cuda.is_available()
+# use fp16 only when GPU is available
+use_fp16 = False
+
+# # download checkpoints
+# os.system('wget https://ofa-silicon.oss-us-west-1.aliyuncs.com/checkpoints/caption_demo.pt; '
+#           'mkdir -p checkpoints; mv caption_demo.pt checkpoints/caption_demo.pt')
+# os.system('wget https://ofa-silicon.oss-us-west-1.aliyuncs.com/checkpoints/refcoco_demo.pt; '
+#           'mkdir -p checkpoints; mv refcoco_demo.pt checkpoints/refcoco_demo.pt')
+# os.system('wget https://ofa-silicon.oss-us-west-1.aliyuncs.com/checkpoints/general_demo.pt; '
+#           'mkdir -p checkpoints; mv general_demo.pt checkpoints/general_demo.pt')
+
+
+checkpoint_path = 'checkpoints/unival_s2_hs/checkpoint1.pt'
+
+# Load ckpt & config for Image Captioning
+caption_overrides={"eval_cider":False, "beam":5, "max_len_b":22, "no_repeat_ngram_size":3, "seed":7, "unnormalized": False,
+           "bpe_dir":"utils/BPE", "video_model_path": None,}
+
+caption_models, caption_cfg, caption_task = checkpoint_utils.load_model_ensemble_and_task(
+    utils.split_paths(checkpoint_path),
+    arg_overrides=caption_overrides
+)
+
+# Load ckpt & config for Refcoco
+refcoco_overrides = {"bpe_dir":"utils/BPE", "video_model_path": None}
+
+refcoco_models, refcoco_cfg, refcoco_task = checkpoint_utils.load_model_ensemble_and_task(
+    utils.split_paths(checkpoint_path),
+    arg_overrides=refcoco_overrides
+)
+refcoco_cfg.common.seed = 7
+refcoco_cfg.generation.beam = 5
+refcoco_cfg.generation.min_len = 4
+refcoco_cfg.generation.max_len_a = 0
+refcoco_cfg.generation.max_len_b = 4
+refcoco_cfg.generation.no_repeat_ngram_size = 3
+
+# Load pretrained ckpt & config for VQA
+parser = options.get_generation_parser()
+input_args = ["", "--task=vqa_gen", "--beam=100", "--unnormalized", f"--path={checkpoint_path}", "--bpe-dir=utils/BPE"]
+args = options.parse_args_and_arch(parser, input_args)
+vqa_cfg = convert_namespace_to_omegaconf(args)
+vqa_task = tasks.setup_task(vqa_cfg.task)
+vqa_models, vqa_cfg = checkpoint_utils.load_model_ensemble(
+    utils.split_paths(vqa_cfg.common_eval.path),
+    task=vqa_task
+)
+
+# Load pretrained ckpt & config for Generic Interface
+parser = options.get_generation_parser()
+input_args = ["", "--task=refcoco", "--beam=10", f"--path={checkpoint_path}", "--bpe-dir=utils/BPE", "--no-repeat-ngram-size=3", "--patch-image-size=384"]
+args = options.parse_args_and_arch(parser, input_args)
+general_cfg = convert_namespace_to_omegaconf(args)
+general_task = tasks.setup_task(general_cfg.task)
+general_models, general_cfg = checkpoint_utils.load_model_ensemble(
+    utils.split_paths(general_cfg.common_eval.path),
+    task=general_task
+)
+
+# move models to gpu
+move2gpu(caption_models, caption_cfg)
+move2gpu(refcoco_models, refcoco_cfg)
+move2gpu(vqa_models, vqa_cfg)
+move2gpu(general_models, general_cfg)
+
+# Initialize generator
+caption_generator = caption_task.build_generator(caption_models, caption_cfg.generation)
+refcoco_generator = refcoco_task.build_generator(refcoco_models, refcoco_cfg.generation)
+vqa_generator = vqa_task.build_generator(vqa_models, vqa_cfg.generation)
+vqa_generator.zero_shot = True
+vqa_generator.constraint_trie = None
+general_generator = general_task.build_generator(general_models, general_cfg.generation)
+
+# Construct image transforms
+caption_transform = construct_transform(caption_cfg.task.patch_image_size)
+refcoco_transform = construct_transform(refcoco_cfg.task.patch_image_size)
+vqa_transform = construct_transform(vqa_cfg.task.patch_image_size)
+general_transform = construct_transform(general_cfg.task.patch_image_size)
+
+# Text preprocess
+bos_item = torch.LongTensor([caption_task.src_dict.bos()])
+eos_item = torch.LongTensor([caption_task.src_dict.eos()])
+pad_idx = caption_task.src_dict.pad()
+
+
+def get_symbols_to_strip_from_output(generator):
+    if hasattr(generator, "symbols_to_strip_from_output"):
+        return generator.symbols_to_strip_from_output
+    else:
+        return {generator.bos, generator.eos}
+
+
+def decode_fn(x, tgt_dict, bpe, generator, tokenizer=None):
+    x = tgt_dict.string(x.int().cpu(), extra_symbols_to_ignore=get_symbols_to_strip_from_output(generator))
+    token_result = []
+    bin_result = []
+    img_result = []
+    for token in x.strip().split():
+      if token.startswith('<bin_'):
+        bin_result.append(token)
+      elif token.startswith('<code_'):
+        img_result.append(token)
+      else:
+        if bpe is not None:
+          token = bpe.decode('{}'.format(token))
+        if tokenizer is not None:
+          token = tokenizer.decode(token)
+        if token.startswith(' ') or len(token_result) == 0:
+          token_result.append(token.strip())
+        else:
+          token_result[-1] += token
+
+    return ' '.join(token_result), ' '.join(bin_result), ' '.join(img_result)
+
+
+def bin2coord(bins, w_resize_ratio, h_resize_ratio, cfg):
+    bin_list = [int(bin[5:-1]) for bin in bins.strip().split()]
+    coord_list = []
+    coord_list += [bin_list[0] / (cfg.task.num_bins - 1) * cfg.task.max_image_size / w_resize_ratio]
+    coord_list += [bin_list[1] / (cfg.task.num_bins - 1) * cfg.task.max_image_size / h_resize_ratio]
+    coord_list += [bin_list[2] / (cfg.task.num_bins - 1) * cfg.task.max_image_size / w_resize_ratio]
+    coord_list += [bin_list[3] / (cfg.task.num_bins - 1) * cfg.task.max_image_size / h_resize_ratio]
+    return coord_list
+
+
+def encode_text(text, length=None, append_bos=False, append_eos=False):
+    line = [
+        caption_task.bpe.encode(' {}'.format(word.strip()))
+        if not word.startswith('<code_') and not word.startswith('<bin_') else word
+        for word in text.strip().split()
+    ]
+    line = ' '.join(line)
+    s = caption_task.tgt_dict.encode_line(
+        line=line,
+        add_if_not_exist=False,
+        append_eos=False
+    ).long()
+    if length is not None:
+        s = s[:length]
+    if append_bos:
+        s = torch.cat([bos_item, s])
+    if append_eos:
+        s = torch.cat([s, eos_item])
+    return s
+
+
+def construct_sample(image: Image, instruction: str, transform):
+    patch_image = transform(image).unsqueeze(0)
+    patch_mask = torch.tensor([True])
+
+    instruction = encode_text(' {}'.format(instruction.lower().strip()), append_bos=True, append_eos=True).unsqueeze(0)
+    instruction_length = torch.LongTensor([s.ne(pad_idx).long().sum() for s in instruction])
+    sample = {
+        "id": np.array(['42']),
+        "net_input": {
+            "src_tokens": instruction,
+            "src_lengths": instruction_length,
+            "patch_images": patch_image,
+            "patch_masks": patch_mask,
+        }
+    }
+    return sample
+
+
+# Function to turn FP32 to FP16
+def apply_half(t):
+    if t.dtype is torch.float32:
+        return t.to(dtype=torch.half)
+    return t
+
+
+def inference(image, task_type, instruction):
+    if task_type == 'Image Captioning':
+        task = caption_task
+        models = caption_models
+        generator = caption_generator
+        instruction = 'what does the image describe?'
+        transform = caption_transform
+        cfg = caption_cfg
+    elif task_type == 'Visual Question Answering':
+        task = vqa_task
+        models = vqa_models
+        generator = vqa_generator
+        transform = vqa_transform
+        cfg = vqa_cfg
+    elif task_type == 'Visual Grounding':
+        task = refcoco_task
+        models = refcoco_models
+        generator = refcoco_generator
+        instruction = 'which region does the text " {} " describe?'.format(instruction)
+        transform = refcoco_transform
+        cfg = refcoco_cfg
+    elif task_type == 'General':
+        task = general_task
+        models = general_models
+        generator = general_generator
+        transform = general_transform
+        cfg = general_cfg
+    else:
+        raise NotImplementedError
+
+    # Construct input sample & preprocess for GPU if cuda available
+    sample = construct_sample(image, instruction, transform)
+    sample = utils.move_to_cuda(sample) if use_cuda else sample
+    sample = utils.apply_to_sample(apply_half, sample) if use_fp16 else sample
+
+    # Generate result
+    with torch.no_grad():
+        hypos = task.inference_step(generator, models, sample)
+        tokens, bins, imgs = decode_fn(hypos[0][0]["tokens"], task.tgt_dict, task.bpe, generator)
+
+    if bins.strip() != '':
+        w, h = image.size
+        w_resize_ratio = task.cfg.patch_image_size / w
+        h_resize_ratio = task.cfg.patch_image_size / h
+        img = np.asarray(image)
+        coord_list = bin2coord(bins, w_resize_ratio, h_resize_ratio, cfg)
+        cv2.rectangle(
+            img,
+            (int(coord_list[0]), int(coord_list[1])),
+            (int(coord_list[2]), int(coord_list[3])),
+            (0, 255, 0),
+            3
+        )
+        return img, None
+    else:
+        return None, tokens
+
+inputs = [gr.inputs.Image(type='pil'), gr.inputs.Radio(choices=['Image Captioning',"Visual Question Answering", "Visual Grounding", "General"], type="value", default="Image Captioning", label="Task"), gr.inputs.Textbox(lines=1, label="Instruction")]
+outputs = [gr.outputs.Image(type='pil'), 'text']
+examples = [
+    ['examples/pokemons.jpeg', 'Image Captioning', None],
+    ['examples/cats.jpeg', 'Visual Question Answering', 'where are the cats?'],
+    ['examples/one_piece.jpeg', 'Visual Grounding', 'a man in a straw hat and a red dress'],
+    ['examples/three_houses.jpeg', 'General', 'which region does the text " a grey car " describe?'],
+    ['examples/three_houses.jpeg', 'General', 'what color is the left car?']
+]
+
+title = "OFA"
+description = "Gradio Demo for OFA: Unifying Architectures, Tasks, and Modalities Through a Simple Sequence-to-Sequence Learning Framework"
+article = "<p style='text-align: center'><a href='http://arxiv.org/abs/2202.03052' target='_blank'>Paper</a> | <a href='https://github.com/OFA-Sys/OFA' target='_blank'>Github Repo</a></p>"
+
+io = gr.Interface(fn=inference, inputs=inputs, outputs=outputs,
+                  title=title, description=description, article=article, examples=examples, cache_examples=False)
+io.launch()

checkpoints.md

@@ -0,0 +1,36 @@
+# 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/unival_s2_hs/checkpoint1.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0b062bb0fa7c45266ee36326391e355724cccaee3119a9d3ee55d93488838a33
+size 2570641445

checkpoints_cn.md

@@ -0,0 +1,82 @@
+# 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

@@ -0,0 +1,9 @@
+# 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)

criterions/__init__.py

@@ -0,0 +1,5 @@
+from .label_smoothed_cross_entropy import AdjustLabelSmoothedCrossEntropyCriterion
+from .clip_scst_loss import ClipScstRewardCriterion
+from .label_smoothed_encouraging_loss import AdjustLabelSmoothedEncouragingLossCriterion
+from .label_smoothed_cross_entropy_scst import AdjustLabelSmoothedCrossEntropySCSTCriterion
+from .refcoco_scst_loss import RefCOCOScstRewardCriterion

criterions/clip_scst_loss.py

@@ -0,0 +1,277 @@
+# 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 math
+from dataclasses import dataclass, field
+from typing import Optional
+from PIL import Image
+from torchvision import transforms
+
+import torch
+import numpy as np
+from fairseq import metrics
+from fairseq.data import data_utils
+from fairseq.criterions import FairseqCriterion, register_criterion
+from fairseq.dataclass import FairseqDataclass
+from fairseq import utils
+from omegaconf import II
+
+from models import clip
+
+
+def custom_to_pil(x):
+    x = x.detach().cpu()
+    x = torch.clamp(x, -1., 1.)
+    x = (x + 1.) / 2.
+    x = x.permute(1, 2, 0).numpy()
+    x = (255 * x).astype(np.uint8)
+    x = Image.fromarray(x)
+    if not x.mode == "RGB":
+        x = x.convert("RGB")
+    return x
+
+
+def scst_loss(lprobs, target, reward, ignore_index=None, reduce=True):
+    loss = -lprobs.gather(dim=-1, index=target.unsqueeze(-1)).squeeze() * reward.unsqueeze(-1)
+    if ignore_index is not None:
+        pad_mask = target.eq(ignore_index)
+        loss.masked_fill_(pad_mask, 0.0)
+        ntokens = (~pad_mask).sum()
+    else:
+        loss = loss.squeeze(-1)
+        ntokens = target.numel()
+    if reduce:
+        loss = loss.sum()
+    return loss, ntokens
+
+
+@dataclass
+class ClipScstRewardCriterionConfig(FairseqDataclass):
+    ignore_prefix_size: int = field(
+        default=0,
+        metadata={"help": "Ignore first N tokens"},
+    )
+    sentence_avg: bool = II("optimization.sentence_avg")
+    constraint_range: Optional[str] = field(
+        default=None,
+        metadata={"help": "constraint range"}
+    )
+
+
+@register_criterion(
+    "clip_scst_reward_criterion", dataclass=ClipScstRewardCriterionConfig
+)
+class ClipScstRewardCriterion(FairseqCriterion):
+    CLIP_REWARD_WEIGHT = 2.5
+
+    def __init__(
+        self,
+        task,
+        sentence_avg,
+        ignore_prefix_size=0,
+        constraint_range=None
+    ):
+        super().__init__(task)
+        self.sentence_avg = sentence_avg
+        self.ignore_prefix_size = ignore_prefix_size
+
+        self.constraint_start = None
+        self.constraint_end = None
+        if constraint_range is not None:
+            constraint_start, constraint_end = constraint_range.split(',')
+            self.constraint_start = int(constraint_start)
+            self.constraint_end = int(constraint_end)
+
+    def forward(self, model, sample, update_num=0, reduce=True):
+        """Compute the loss for the given sample.
+
+        Returns a tuple with three elements:
+        1) the loss
+        2) the sample size, which is used as the denominator for the gradient
+        3) logging outputs to display while training
+        """
+        loss, score, ntokens, nsentences = self.compute_loss(model, sample, reduce=reduce)
+
+        sample_size = (
+            nsentences if self.sentence_avg else ntokens
+        )
+        logging_output = {
+            "loss": loss.data,
+            "score": score,
+            "ntokens": ntokens,
+            "nsentences": nsentences,
+            "sample_size": sample_size,
+        }
+        return loss, sample_size, logging_output
+
+    def _calculate_clip_scores(self, gen_res, gt_text, device):
+        '''
+        gen_res: generated images, list of Image
+        gt_text: input captions.
+        device: device for clip model
+        '''
+        batch_size = len(gt_text)
+        gen_res_size = len(gen_res)
+        img_per_seq = gen_res_size // batch_size
+
+        hyp_images = torch.stack(
+            [self.task.clip_preprocess(gen_image) for gen_image in gen_res], dim=0
+        ).to(device)
+
+        clip_input = clip.tokenize([text for text in gt_text]).to(device)
+        with torch.no_grad():
+            image_features = self.task.clip_model.encode_image(hyp_images)
+            text_features = self.task.clip_model.encode_text(clip_input)
+            image_features /= image_features.norm(dim=-1, keepdim=True)
+            text_features /= text_features.norm(dim=-1, keepdim=True)
+            image_features = image_features.view(batch_size, img_per_seq, -1)
+            text_features = text_features.view(batch_size, 1, -1)
+            ti_similarity = image_features @ text_features.transpose(1, 2)
+            ti_similarity = ti_similarity.view(-1)
+
+        scores = self.CLIP_REWARD_WEIGHT * ti_similarity
+        return scores
+
+    def get_generator_out(self, model, sample):
+        model.eval()
+        with torch.no_grad():
+            self.task.scst_generator.model.eval()
+            gen_out = self.task.scst_generator.generate([model], sample)
+
+        gen_target = []
+        gen_res = []
+        gt_text = []
+        for i in range(len(gen_out)):
+            with torch.no_grad():
+                tokens = torch.stack([item['tokens'][:-1] for item in gen_out[i]], dim=0)
+                tokens += -len(self.task.src_dict) + self.task.cfg.code_dict_size + self.task.cfg.num_bins
+                images = self.task.image_tokenizer.decode_code(
+                    tokens.view(-1, self.task.cfg.code_image_size // 8, self.task.cfg.code_image_size // 8)
+                )
+                images = [custom_to_pil(image) for image in images]
+
+            gen_target += [item['tokens'] for item in gen_out[i]]
+            gen_res += images
+            gt_text.append(
+                self.task.bpe.decode(
+                    self.task.tgt_dict.string(
+                        utils.strip_pad(sample['net_input']['src_tokens'][i], self.padding_idx).cpu().int()
+                    )
+                )[38:] # remove task instruction.
+            )
+
+        return gen_target, gen_res, gt_text
+
+    def get_reward_and_scores(self, gen_res, gt_text, device):
+        batch_size = len(gt_text)
+        gen_res_size = len(gen_res)
+        img_per_sample = gen_res_size // batch_size
+
+        scores = self._calculate_clip_scores(gen_res, gt_text, device)
+        sc_ = scores.reshape(batch_size, img_per_sample)
+        baseline = (sc_.sum(1, keepdim=True) - sc_) / (sc_.shape[1] - 1)
+        # sample - baseline
+        reward = scores.reshape(batch_size, img_per_sample)
+        reward = reward - baseline
+        reward = reward.view(-1)
+
+        return reward, scores
+
+    def get_net_output(self, model, sample, gen_target):
+        def merge(sample_list, eos=self.task.tgt_dict.eos(), move_eos_to_beginning=False):
+            return data_utils.collate_tokens(
+                sample_list,
+                pad_idx=self.padding_idx,
+                eos_idx=eos,
+                left_pad=False,
+                move_eos_to_beginning=move_eos_to_beginning,
+            )
+
+        batch_size = len(sample["target"])
+        gen_target_size = len(gen_target)
+        img_per_sample = gen_target_size // batch_size
+
+        model.train()
+        sample_src_tokens = torch.repeat_interleave(
+            sample['net_input']['src_tokens'], img_per_sample, dim=0
+        )
+        sample_src_lengths = torch.repeat_interleave(
+            sample['net_input']['src_lengths'], img_per_sample, dim=0
+        )
+        sample_code_masks = torch.repeat_interleave(
+            sample['net_input']['code_masks'], img_per_sample, dim=0
+        )
+        gen_prev_output_tokens = torch.as_tensor(
+            merge(gen_target, eos=self.task.tgt_dict.bos(), move_eos_to_beginning=True),
+            device=sample["target"].device, dtype=torch.int64
+        )
+        gen_target_tokens = torch.as_tensor(
+            merge(gen_target), device=sample["target"].device, dtype=torch.int64
+        )
+        net_output = model(
+            src_tokens=sample_src_tokens, src_lengths=sample_src_lengths,
+            code_masks=sample_code_masks, prev_output_tokens=gen_prev_output_tokens
+        )
+
+        return net_output, gen_target_tokens
+
+    def get_lprobs_and_target(self, model, net_output, gen_target):
+        if self.constraint_start is not None and self.constraint_end is not None:
+            net_output[0][:, :, 4:self.constraint_start] = -math.inf
+            net_output[0][:, :, self.constraint_end:] = -math.inf
+        lprobs = model.get_normalized_probs(net_output, log_probs=True)
+        if self.ignore_prefix_size > 0:
+            if getattr(lprobs, "batch_first", False):
+                lprobs = lprobs[:, self.ignore_prefix_size :, :].contiguous()
+                gen_target = gen_target[:, self.ignore_prefix_size :].contiguous()
+            else:
+                lprobs = lprobs[self.ignore_prefix_size :, :, :].contiguous()
+                gen_target = gen_target[self.ignore_prefix_size :, :].contiguous()
+        return lprobs, gen_target
+
+    def compute_loss(self, model, sample, reduce=True):
+        gen_target, gen_res, gt_text = self.get_generator_out(model, sample)
+        reward, scores = self.get_reward_and_scores(gen_res, gt_text, device=sample["target"].device)
+        net_output, gen_target_tokens = self.get_net_output(model, sample, gen_target)
+        gen_lprobs, gen_target_tokens = self.get_lprobs_and_target(model, net_output, gen_target_tokens)
+        loss, ntokens = scst_loss(gen_lprobs, gen_target_tokens, reward, ignore_index=self.padding_idx, reduce=reduce)
+        nsentences = gen_target_tokens.size(0)
+
+        return loss, scores.sum(), ntokens, nsentences
+
+    @classmethod
+    def reduce_metrics(cls, logging_outputs) -> None:
+        """Aggregate logging outputs from data parallel training."""
+        loss_sum = sum(log.get("loss", 0) for log in logging_outputs)
+        score_sum = sum(log.get("score", 0) for log in logging_outputs)
+        ntokens = sum(log.get("ntokens", 0) for log in logging_outputs)
+        nsentences = sum(log.get("nsentences", 0) for log in logging_outputs)
+        sample_size = sum(log.get("sample_size", 0) for log in logging_outputs)
+
+        metrics.log_scalar(
+            "loss", loss_sum / sample_size, sample_size, round=3
+        )
+        metrics.log_scalar(
+            "score", score_sum / nsentences, nsentences, round=3
+        )
+
+        metrics.log_scalar(
+            "ntokens", ntokens, 1, round=3
+        )
+        metrics.log_scalar(
+            "nsentences", nsentences, 1, round=3
+        )
+        metrics.log_scalar(
+            "sample_size", sample_size, 1, round=3
+        )
+
+    @staticmethod
+    def logging_outputs_can_be_summed() -> bool:
+        """
+        Whether the logging outputs returned by `forward` can be summed
+        across workers prior to calling `reduce_metrics`. Setting this
+        to True will improves distributed training speed.
+        """
+        return True

criterions/label_smoothed_cross_entropy.py

@@ -0,0 +1,346 @@
+# 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 math
+from dataclasses import dataclass, field
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+import numpy as np
+from fairseq import metrics, utils
+from fairseq.criterions import FairseqCriterion, register_criterion
+from fairseq.dataclass import FairseqDataclass
+from omegaconf import II
+
+
+@dataclass
+class AdjustLabelSmoothedCrossEntropyCriterionConfig(FairseqDataclass):
+    label_smoothing: float = field(
+        default=0.0,
+        metadata={"help": "epsilon for label smoothing, 0 means no label smoothing"},
+    )
+    report_accuracy: bool = field(
+        default=False,
+        metadata={"help": "report accuracy metric"},
+    )
+    ignore_prefix_size: int = field(
+        default=0,
+        metadata={"help": "Ignore first N tokens"},
+    )
+    ignore_eos: bool = field(
+        default=False,
+        metadata={"help": "Ignore eos token"},
+    )
+    sentence_avg: bool = II("optimization.sentence_avg")
+    drop_worst_ratio: float = field(
+        default=0.0,
+        metadata={"help": "ratio for discarding bad samples"},
+    )
+    drop_worst_after: int = field(
+        default=0,
+        metadata={"help": "steps for discarding bad samples"},
+    )
+    use_rdrop: bool = field(
+        default=False, metadata={"help": "use R-Drop"}
+    )
+    reg_alpha: float = field(
+        default=1.0, metadata={"help": "weight for R-Drop"}
+    )
+    sample_patch_num: int = field(
+        default=196, metadata={"help": "sample patches for v1"}
+    )
+    constraint_range: Optional[str] = field(
+        default=None,
+        metadata={"help": "constraint range"}
+    )
+
+
+def construct_rdrop_sample(x):
+    if isinstance(x, dict):
+        for key in x:
+            x[key] = construct_rdrop_sample(x[key])
+        return x
+    elif isinstance(x, torch.Tensor):
+        return x.repeat(2, *([1] * (x.dim()-1)))
+    elif isinstance(x, int):
+        return x * 2
+    elif isinstance(x, np.ndarray):
+        return x.repeat(2)
+    else:
+        raise NotImplementedError
+
+
+def kl_loss(p, q):
+    p_loss = F.kl_div(p, torch.exp(q), reduction='sum')
+    q_loss = F.kl_div(q, torch.exp(p), reduction='sum')
+    loss = (p_loss + q_loss) / 2
+    return loss
+
+
+def label_smoothed_nll_loss(
+        lprobs, target, epsilon, update_num, reduce=True,
+        drop_worst_ratio=0.0, drop_worst_after=0, use_rdrop=False, reg_alpha=1.0,
+        constraint_masks=None, constraint_start=None, constraint_end=None
+):
+    if target.dim() == lprobs.dim() - 1:
+        target = target.unsqueeze(-1)
+    nll_loss = -lprobs.gather(dim=-1, index=target).squeeze(-1)
+    if constraint_masks is not None:
+        smooth_loss = -lprobs.masked_fill(~constraint_masks, 0).sum(dim=-1, keepdim=True).squeeze(-1)
+        eps_i = epsilon / (constraint_masks.sum(1) - 1 + 1e-6)
+    elif constraint_start is not None and constraint_end is not None:
+        constraint_range = [0, 1, 2, 3] + list(range(constraint_start, constraint_end))
+        smooth_loss = -lprobs[:, constraint_range].sum(dim=-1, keepdim=True).squeeze(-1)
+        eps_i = epsilon / (len(constraint_range) - 1 + 1e-6)
+    else:
+        smooth_loss = -lprobs.sum(dim=-1, keepdim=True).squeeze(-1)
+        eps_i = epsilon / (lprobs.size(-1) - 1)
+    loss = (1.0 - epsilon - eps_i) * nll_loss + eps_i * smooth_loss
+    if drop_worst_ratio > 0 and update_num > drop_worst_after:
+        if use_rdrop:
+            true_batch_size = loss.size(0) // 2
+            _, indices = torch.topk(loss[:true_batch_size], k=int(true_batch_size * (1 - drop_worst_ratio)), largest=False)
+            loss = torch.cat([loss[indices], loss[indices+true_batch_size]])
+            nll_loss = torch.cat([nll_loss[indices], nll_loss[indices+true_batch_size]])
+            lprobs = torch.cat([lprobs[indices], lprobs[indices+true_batch_size]])
+        else:
+            loss, indices = torch.topk(loss, k=int(loss.shape[0] * (1 - drop_worst_ratio)), largest=False)
+            nll_loss = nll_loss[indices]
+            lprobs = lprobs[indices]
+
+    ntokens = loss.numel()
+    nll_loss = nll_loss.sum()
+    loss = loss.sum()
+    if use_rdrop:
+        true_batch_size = lprobs.size(0) // 2
+        p = lprobs[:true_batch_size]
+        q = lprobs[true_batch_size:]
+        if constraint_start is not None and constraint_end is not None:
+            constraint_range = [0, 1, 2, 3] + list(range(constraint_start, constraint_end))
+            p = p[:, constraint_range]
+            q = q[:, constraint_range]
+        loss += kl_loss(p, q) * reg_alpha
+
+    return loss, nll_loss, ntokens
+
+
+@register_criterion(
+    "adjust_label_smoothed_cross_entropy", dataclass=AdjustLabelSmoothedCrossEntropyCriterionConfig
+)
+class AdjustLabelSmoothedCrossEntropyCriterion(FairseqCriterion):
+    def __init__(
+        self,
+        task,
+        sentence_avg,
+        label_smoothing,
+        ignore_prefix_size=0,
+        ignore_eos=False,
+        report_accuracy=False,
+        drop_worst_ratio=0,
+        drop_worst_after=0,
+        use_rdrop=False,
+        reg_alpha=1.0,
+        sample_patch_num=196,
+        constraint_range=None
+    ):
+        super().__init__(task)
+        self.sentence_avg = sentence_avg
+        self.eps = label_smoothing
+        self.ignore_prefix_size = ignore_prefix_size
+        self.ignore_eos = ignore_eos
+        self.report_accuracy = report_accuracy
+        self.drop_worst_ratio = drop_worst_ratio
+        self.drop_worst_after = drop_worst_after
+        self.use_rdrop = use_rdrop
+        self.reg_alpha = reg_alpha
+        self.sample_patch_num = sample_patch_num
+
+        self.constraint_start = None
+        self.constraint_end = None
+        if constraint_range is not None:
+            constraint_start, constraint_end = constraint_range.split(',')
+            self.constraint_start = int(constraint_start)
+            self.constraint_end = int(constraint_end)
+
+    def forward(self, model, sample, update_num=0, reduce=True):
+        """Compute the loss for the given sample.
+
+        Returns a tuple with three elements:
+        1) the loss
+        2) the sample size, which is used as the denominator for the gradient
+        3) logging outputs to display while training
+        """
+        if isinstance(sample, list):
+            if self.sample_patch_num > 0:
+                sample[0]['net_input']['sample_patch_num'] = self.sample_patch_num
+            # change to support len(samples) > 2
+            loss_v1, sample_size_v1, logging_output_v1 = self.forward(model, sample[0], update_num, reduce)
+            loss_v2, sample_size_v2, logging_output_v2 = self.forward(model, sample[1], update_num, reduce)
+            loss = loss_v1 / sample_size_v1 + loss_v2 / sample_size_v2
+            sample_size = 1
+            logging_output = {
+                "loss": loss.data,
+                "loss_v1": loss_v1.data,
+                "loss_v2": loss_v2.data,
+                "nll_loss": logging_output_v1["nll_loss"].data / sample_size_v1 + logging_output_v2["nll_loss"].data / sample_size_v2,
+                "ntokens": logging_output_v1["ntokens"] + logging_output_v2["ntokens"],
+                "nsentences": logging_output_v1["nsentences"] + logging_output_v2["nsentences"],
+                "sample_size": 1,
+                "sample_size_v1": sample_size_v1,
+                "sample_size_v2": sample_size_v2,
+            }
+            return loss, sample_size, logging_output
+
+        if self.use_rdrop:
+            construct_rdrop_sample(sample)
+
+        net_output = model(**sample["net_input"])
+        loss, nll_loss, ntokens = self.compute_loss(model, net_output, sample, update_num, reduce=reduce)
+        sample_size = (
+            sample["target"].size(0) if self.sentence_avg else ntokens
+        )
+        logging_output = {
+            "loss": loss.data,
+            "nll_loss": nll_loss.data,
+            "ntokens": sample["ntokens"],
+            "nsentences": sample["nsentences"],
+            "sample_size": sample_size,
+        }
+        if self.report_accuracy:
+            n_correct, total = self.compute_accuracy(model, net_output, sample)
+            logging_output["n_correct"] = utils.item(n_correct.data)
+            logging_output["total"] = utils.item(total.data)
+
+        return loss, sample_size, logging_output
+
+    def get_lprobs_and_target(self, model, net_output, sample):
+        conf = sample['conf'][:, None, None] if 'conf' in sample and sample['conf'] is not None else 1
+        constraint_masks = None
+        if "constraint_masks" in sample and sample["constraint_masks"] is not None:
+            constraint_masks = sample["constraint_masks"]
+            net_output[0].masked_fill_(~constraint_masks, -math.inf)
+        if self.constraint_start is not None and self.constraint_end is not None:
+            net_output[0][:, :, 4:self.constraint_start] = -math.inf
+            net_output[0][:, :, self.constraint_end:] = -math.inf
+        lprobs = model.get_normalized_probs(net_output, log_probs=True) * conf
+        target = model.get_targets(sample, net_output)
+        if self.ignore_prefix_size > 0:
+            lprobs = lprobs[:, self.ignore_prefix_size :, :].contiguous()
+            target = target[:, self.ignore_prefix_size :].contiguous()
+            if constraint_masks is not None:
+                constraint_masks = constraint_masks[:, self.ignore_prefix_size :, :].contiguous()
+        if self.ignore_eos:
+            bsz, seq_len, embed_dim = lprobs.size()
+            eos_indices = target.eq(self.task.tgt_dict.eos())
+            lprobs = lprobs[~eos_indices].reshape(bsz, seq_len-1, embed_dim)
+            target = target[~eos_indices].reshape(bsz, seq_len-1)
+            if constraint_masks is not None:
+                constraint_masks = constraint_masks[~eos_indices].reshape(bsz, seq_len-1, embed_dim)
+        if constraint_masks is not None:
+            constraint_masks = constraint_masks.view(-1, constraint_masks.size(-1))
+        return lprobs.view(-1, lprobs.size(-1)), target.view(-1), constraint_masks
+
+    def compute_loss(self, model, net_output, sample, update_num, reduce=True):
+        lprobs, target, constraint_masks = self.get_lprobs_and_target(model, net_output, sample)
+        if constraint_masks is not None:
+            constraint_masks = constraint_masks[target != self.padding_idx]
+        # print(target.shape, self.padding_idx, lprobs.shape, target, lprobs)
+        lprobs = lprobs[target != self.padding_idx]
+        target = target[target != self.padding_idx]
+        loss, nll_loss, ntokens = label_smoothed_nll_loss(
+            lprobs,
+            target,
+            self.eps,
+            update_num,
+            reduce=reduce,
+            drop_worst_ratio=self.drop_worst_ratio,
+            drop_worst_after=self.drop_worst_after,
+            use_rdrop=self.use_rdrop,
+            reg_alpha=self.reg_alpha,
+            constraint_masks=constraint_masks,
+            constraint_start=self.constraint_start,
+            constraint_end=self.constraint_end
+        )
+        return loss, nll_loss, ntokens
+
+    def compute_accuracy(self, model, net_output, sample):
+        lprobs, target = self.get_lprobs_and_target(model, net_output, sample)
+        mask = target.ne(self.padding_idx)
+        n_correct = torch.sum(
+            lprobs.argmax(1).masked_select(mask).eq(target.masked_select(mask))
+        )
+        total = torch.sum(mask)
+        return n_correct, total
+
+    @classmethod
+    def reduce_metrics(cls, logging_outputs) -> None:
+        """Aggregate logging outputs from data parallel training."""
+        loss_sum = sum(log.get("loss", 0) for log in logging_outputs)
+        loss_sum_v1 = sum(log.get("loss_v1", 0) for log in logging_outputs)
+        loss_sum_v2 = sum(log.get("loss_v2", 0) for log in logging_outputs)
+        nll_loss_sum = sum(log.get("nll_loss", 0) for log in logging_outputs)
+        ntokens = sum(log.get("ntokens", 0) for log in logging_outputs)
+        nsentences = sum(log.get("nsentences", 0) for log in logging_outputs)
+        sample_size = sum(log.get("sample_size", 0) for log in logging_outputs)
+        sample_size_v1 = sum(log.get("sample_size_v1", 0) for log in logging_outputs)
+        sample_size_v2 = sum(log.get("sample_size_v2", 0) for log in logging_outputs)
+
+        metrics.log_scalar(
+            "loss", loss_sum / sample_size, sample_size, round=3
+        )
+        metrics.log_scalar(
+            "loss_v1", loss_sum_v1 / max(sample_size_v1, 1), max(sample_size_v1, 1), round=3
+        )
+        metrics.log_scalar(
+            "loss_v2", loss_sum_v2 / max(sample_size_v2, 1), max(sample_size_v2, 1), round=3
+        )
+        metrics.log_scalar(
+            "nll_loss", nll_loss_sum / sample_size, ntokens, round=3
+        )
+        metrics.log_derived(
+            "ppl", lambda meters: utils.get_perplexity(meters["nll_loss"].avg)
+        )
+
+        metrics.log_scalar(
+            "ntokens", ntokens, 1, round=3
+        )
+        metrics.log_scalar(
+            "nsentences", nsentences, 1, round=3
+        )
+        metrics.log_scalar(
+            "sample_size", sample_size, 1, round=3
+        )
+        metrics.log_scalar(
+            "sample_size_v1", sample_size_v1, 1, round=3
+        )
+        metrics.log_scalar(
+            "sample_size_v2", sample_size_v2, 1, round=3
+        )
+
+        total = utils.item(sum(log.get("total", 0) for log in logging_outputs))
+        if total > 0:
+            metrics.log_scalar("total", total)
+            n_correct = utils.item(
+                sum(log.get("n_correct", 0) for log in logging_outputs)
+            )
+            metrics.log_scalar("n_correct", n_correct)
+            metrics.log_derived(
+                "accuracy",
+                lambda meters: round(
+                    meters["n_correct"].sum * 100.0 / meters["total"].sum, 3
+                )
+                if meters["total"].sum > 0
+                else float("nan"),
+            )
+
+    @staticmethod
+    def logging_outputs_can_be_summed() -> bool:
+        """
+        Whether the logging outputs returned by `forward` can be summed
+        across workers prior to calling `reduce_metrics`. Setting this
+        to True will improves distributed training speed.
+        """
+        return True

criterions/label_smoothed_cross_entropy_scst.py

@@ -0,0 +1,555 @@
+# 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 math
+from dataclasses import dataclass, field
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+import numpy as np
+from fairseq import metrics, utils
+from fairseq.criterions import FairseqCriterion, register_criterion
+from fairseq.dataclass import FairseqDataclass
+from omegaconf import II
+
+
+from mapcalc import calculate_map, calculate_map_range
+
+@dataclass
+class AdjustLabelSmoothedCrossEntropySCSTCriterionConfig(FairseqDataclass):
+    label_smoothing: float = field(
+        default=0.0,
+        metadata={"help": "epsilon for label smoothing, 0 means no label smoothing"},
+    )
+    report_accuracy: bool = field(
+        default=False,
+        metadata={"help": "report accuracy metric"},
+    )
+    ignore_prefix_size: int = field(
+        default=0,
+        metadata={"help": "Ignore first N tokens"},
+    )
+    ignore_eos: bool = field(
+        default=False,
+        metadata={"help": "Ignore eos token"},
+    )
+    sentence_avg: bool = II("optimization.sentence_avg")
+    drop_worst_ratio: float = field(
+        default=0.0,
+        metadata={"help": "ratio for discarding bad samples"},
+    )
+    drop_worst_after: int = field(
+        default=0,
+        metadata={"help": "steps for discarding bad samples"},
+    )
+    use_rdrop: bool = field(
+        default=False, metadata={"help": "use R-Drop"}
+    )
+    reg_alpha: float = field(
+        default=1.0, metadata={"help": "weight for R-Drop"}
+    )
+    sample_patch_num: int = field(
+        default=196, metadata={"help": "sample patches for v1"}
+    )
+    constraint_range: Optional[str] = field(
+        default=None,
+        metadata={"help": "constraint range"}
+    )
+    acc_thresh: Optional[float] = field(
+        default=None, metadata={"help": "acc thresh for refcoco"}
+    )
+    metric: Optional[str] = field(
+        default='acc',
+        metadata={"help": "metric"}
+    )
+
+    max_area_size: Optional[float] = field(
+        default=None, metadata={"help": "max_area_size"}
+    )
+
+    min_area_size: Optional[float] = field(
+        default=None, metadata={"help": "min_area_size"}
+    )
+    logprob: Optional[bool] = field(
+        default=False, metadata={"help": "maximise log prob"}
+    )
+
+    pos_reward: Optional[float] = field(
+        default=None, metadata={"help": "pos_reward"}
+    )
+
+    neg_reward: Optional[float] = field(
+        default=None, metadata={"help": "neg_reward"}
+    )
+
+    reinforce: Optional[bool] = field(
+        default=False, metadata={"help": "reinforce"}
+    )
+
+    lambda_reinforce: Optional[float] = field(
+        default=0, metadata={"help": "lambda_reinforce"}
+    )
+
+
+
+def construct_rdrop_sample(x):
+    if isinstance(x, dict):
+        for key in x:
+            x[key] = construct_rdrop_sample(x[key])
+        return x
+    elif isinstance(x, torch.Tensor):
+        return x.repeat(2, *([1] * (x.dim()-1)))
+    elif isinstance(x, int):
+        return x * 2
+    elif isinstance(x, np.ndarray):
+        return x.repeat(2)
+    else:
+        raise NotImplementedError
+
+
+def kl_loss(p, q):
+    p_loss = F.kl_div(p, torch.exp(q), reduction='sum')
+    q_loss = F.kl_div(q, torch.exp(p), reduction='sum')
+    loss = (p_loss + q_loss) / 2
+    return loss
+
+
+def label_smoothed_nll_loss(
+        lprobs, target, epsilon, update_num, reduce=True,
+        drop_worst_ratio=0.0, drop_worst_after=0, use_rdrop=False, reg_alpha=1.0,
+        constraint_masks=None, constraint_start=None, constraint_end=None
+):
+    if target.dim() == lprobs.dim() - 1:
+        target = target.unsqueeze(-1)
+    nll_loss = -lprobs.gather(dim=-1, index=target).squeeze(-1)
+    if constraint_masks is not None:
+        smooth_loss = -lprobs.masked_fill(~constraint_masks, 0).sum(dim=-1, keepdim=True).squeeze(-1)
+        eps_i = epsilon / (constraint_masks.sum(1) - 1 + 1e-6)
+    elif constraint_start is not None and constraint_end is not None:
+        constraint_range = [0, 1, 2, 3] + list(range(constraint_start, constraint_end))
+        smooth_loss = -lprobs[:, constraint_range].sum(dim=-1, keepdim=True).squeeze(-1)
+        eps_i = epsilon / (len(constraint_range) - 1 + 1e-6)
+    else:
+        smooth_loss = -lprobs.sum(dim=-1, keepdim=True).squeeze(-1)
+        eps_i = epsilon / (lprobs.size(-1) - 1)
+    loss = (1.0 - epsilon - eps_i) * nll_loss + eps_i * smooth_loss
+    if drop_worst_ratio > 0 and update_num > drop_worst_after:
+        if use_rdrop:
+            true_batch_size = loss.size(0) // 2
+            _, indices = torch.topk(loss[:true_batch_size], k=int(true_batch_size * (1 - drop_worst_ratio)), largest=False)
+            loss = torch.cat([loss[indices], loss[indices+true_batch_size]])
+            nll_loss = torch.cat([nll_loss[indices], nll_loss[indices+true_batch_size]])
+            lprobs = torch.cat([lprobs[indices], lprobs[indices+true_batch_size]])
+        else:
+            loss, indices = torch.topk(loss, k=int(loss.shape[0] * (1 - drop_worst_ratio)), largest=False)
+            nll_loss = nll_loss[indices]
+            lprobs = lprobs[indices]
+
+    ntokens = loss.numel()
+    nll_loss = nll_loss.sum()
+    # loss = loss.sum()
+    if use_rdrop:
+        true_batch_size = lprobs.size(0) // 2
+        p = lprobs[:true_batch_size]
+        q = lprobs[true_batch_size:]
+        if constraint_start is not None and constraint_end is not None:
+            constraint_range = [0, 1, 2, 3] + list(range(constraint_start, constraint_end))
+            p = p[:, constraint_range]
+            q = q[:, constraint_range]
+        loss = loss + ((kl_loss(p, q) * reg_alpha)/loss.shape[0])
+
+    return loss, nll_loss, ntokens
+
+
+@register_criterion(
+    "adjust_label_smoothed_cross_entropy_scst", dataclass=AdjustLabelSmoothedCrossEntropySCSTCriterionConfig
+)
+class AdjustLabelSmoothedCrossEntropySCSTCriterion(FairseqCriterion):
+    def __init__(
+        self,
+        task,
+        sentence_avg,
+        label_smoothing,
+        ignore_prefix_size=0,
+        ignore_eos=False,
+        report_accuracy=False,
+        drop_worst_ratio=0,
+        drop_worst_after=0,
+        use_rdrop=False,
+        reg_alpha=1.0,
+        sample_patch_num=196,
+        constraint_range=None,
+        acc_thresh=None,
+        metric='acc',
+        max_area_size=None,
+        min_area_size=None,
+        logprob=False,
+        pos_reward=None,
+        neg_reward=None,
+        reinforce=False,
+        lambda_reinforce=0,
+    ):
+        super().__init__(task)
+        self.sentence_avg = sentence_avg
+        self.eps = label_smoothing
+        self.ignore_prefix_size = ignore_prefix_size
+        self.ignore_eos = ignore_eos
+        self.report_accuracy = report_accuracy
+        self.drop_worst_ratio = drop_worst_ratio
+        self.drop_worst_after = drop_worst_after
+        self.use_rdrop = use_rdrop
+        self.reg_alpha = reg_alpha
+        self.sample_patch_num = sample_patch_num
+
+        
+
+        self.constraint_start = None
+        self.constraint_end = None
+        if constraint_range is not None:
+            constraint_start, constraint_end = constraint_range.split(',')
+            self.constraint_start = int(constraint_start)
+            self.constraint_end = int(constraint_end)
+
+        self.acc_thresh = acc_thresh
+        self.metric = metric
+        self.min_area_size = min_area_size
+        self.max_area_size = max_area_size
+        self.logprob = logprob
+
+        self.pos_reward = pos_reward
+        self.neg_reward = neg_reward
+
+        self.reinforce = reinforce
+        self.lambda_reinforce = lambda_reinforce
+
+    def get_generator_out(self, model, sample):
+
+        model.eval()
+        with torch.no_grad():
+            self.task.scst_generator.model.eval()
+            gen_out = self.task.scst_generator.generate([model], sample)
+
+        hyps, refs = [], []
+        for i in range(len(gen_out)):
+            hyps.append(gen_out[i][0]["tokens"][:-1] - len(self.task.src_dict) + self.task.cfg.num_bins)
+            refs.append(sample["target"][i][:-1] - len(self.task.src_dict) + self.task.cfg.num_bins)
+
+        return torch.stack(hyps, dim=0), torch.stack(refs, dim=0)
+    
+    def _calculate_map_score(self, hyps, refs, thresh=0.5):
+        
+        
+        ground_truth = {
+            'boxes': refs.cpu().numpy().tolist(),
+
+            'labels': [1 for i in range(refs.shape[0])]
+            }
+
+        result_dict = {
+            'boxes': hyps.cpu().numpy().tolist(),
+
+            'labels': [1 for i in range(hyps.shape[0])], 
+            }
+
+        score = calculate_map(ground_truth, result_dict, thresh)
+
+        score = torch.tensor(score).unsqueeze(0).repeat(refs.shape[0]).to(hyps.device)
+        return score
+        
+    def _calculate_ap_score(self, hyps, refs, thresh=0.5, min_area_size=None, max_area_size=None):
+        interacts = torch.cat(
+            [torch.where(hyps[:, :2] < refs[:, :2], refs[:, :2], hyps[:, :2]),
+             torch.where(hyps[:, 2:] < refs[:, 2:], hyps[:, 2:], refs[:, 2:])],
+            dim=1
+        )
+        area_predictions = (hyps[:, 2] - hyps[:, 0]) * (hyps[:, 3] - hyps[:, 1]) ## x1, y1, x2, y2, x1 < x2
+        area_targets = (refs[:, 2] - refs[:, 0]) * (refs[:, 3] - refs[:, 1])
+        interacts_w = interacts[:, 2] - interacts[:, 0]
+        interacts_h = interacts[:, 3] - interacts[:, 1]
+        area_interacts = interacts_w * interacts_h
+        ious = area_interacts / (area_predictions + area_targets - area_interacts + 1e-6)
+        
+
+        if max_area_size is not None and min_area_size is not None:
+            ious =  ious * (torch.logical_or(area_targets < max_area_size, area_targets > min_area_size).float())
+
+        elif min_area_size is not None:
+            ious =  ious * (area_targets > min_area_size).float()
+
+        elif max_area_size is not None:
+            ious =  ious * (area_targets < max_area_size).float()
+
+        if thresh is None:
+            return ious
+        else:
+            return ((ious >= thresh) & (interacts_w > 0) & (interacts_h > 0)).float()
+
+    def reward_step(self, sample, model):
+        
+        hyps, refs = self.get_generator_out(model, sample)
+        hyps = hyps / (self.task.cfg.num_bins - 1) * self.task.cfg.max_image_size
+        refs = refs / (self.task.cfg.num_bins - 1) * self.task.cfg.max_image_size
+        hyps[:, ::2] /= sample['w_resize_ratios'].unsqueeze(1)
+        hyps[:, 1::2] /= sample['h_resize_ratios'].unsqueeze(1)
+        refs[:, ::2] /= sample['w_resize_ratios'].unsqueeze(1)
+        refs[:, 1::2] /= sample['h_resize_ratios'].unsqueeze(1)
+
+        # scores = self._calculate_ap_score(hyps, refs)
+        if self.metric == 'acc':
+            scores = self._calculate_ap_score(hyps, sample['region_coords'].float(), thresh=self.acc_thresh, 
+                                              min_area_size=self.min_area_size, max_area_size=self.max_area_size)
+        elif self.metric == 'map':
+            scores = self._calculate_map_score(hyps, sample['region_coords'].float(), thresh=self.acc_thresh)
+        else:
+            raise NotImplemented
+        
+        # logging_output["_score_sum"] = scores.sum().item()
+        # logging_output["_score_cnt"] = scores.size(0)
+
+        if self.pos_reward:
+            scores = torch.where(scores > 0, self.pos_reward, scores)
+        if self.neg_reward:
+            scores = torch.where(scores == 0, self.neg_reward, scores)
+
+
+        return scores
+    
+    def forward(self, model, sample, update_num=0, reduce=True):
+        """Compute the loss for the given sample.
+
+        Returns a tuple with three elements:
+        1) the loss
+        2) the sample size, which is used as the denominator for the gradient
+        3) logging outputs to display while training
+        """
+        if isinstance(sample, list):
+            if self.sample_patch_num > 0:
+                sample[0]['net_input']['sample_patch_num'] = self.sample_patch_num
+            # change to support len(samples) > 2
+            loss_v1, sample_size_v1, logging_output_v1 = self.forward(model, sample[0], update_num, reduce)
+            loss_v2, sample_size_v2, logging_output_v2 = self.forward(model, sample[1], update_num, reduce)
+            loss = loss_v1 / sample_size_v1 + loss_v2 / sample_size_v2
+            sample_size = 1
+            logging_output = {
+                "loss": loss.data,
+                "loss_v1": loss_v1.data,
+                "loss_v2": loss_v2.data,
+                "nll_loss": logging_output_v1["nll_loss"].data / sample_size_v1 + logging_output_v2["nll_loss"].data / sample_size_v2,
+                "ntokens": logging_output_v1["ntokens"] + logging_output_v2["ntokens"],
+                "nsentences": logging_output_v1["nsentences"] + logging_output_v2["nsentences"],
+                "sample_size": 1,
+                "sample_size_v1": sample_size_v1,
+                "sample_size_v2": sample_size_v2,
+                "reward": (logging_output_v1["reward"] + logging_output_v2["reward"])/2,
+            }
+            return loss, sample_size, logging_output
+
+        if self.use_rdrop:
+            construct_rdrop_sample(sample)
+
+        ### scst
+        reward = self.reward_step(sample, model) # shape = bs
+        model.train()
+        net_output = model(**sample["net_input"])
+        loss, nll_loss, ntokens = self.compute_loss(model, net_output, sample, update_num, reduce=reduce, reward=reward)
+
+
+        
+        
+        # loss = loss*reward
+        
+        loss = loss.sum()
+        sample_size = (
+            sample["target"].size(0) if self.sentence_avg else ntokens
+        )
+        logging_output = {
+            "loss": loss.data,
+            "nll_loss": nll_loss.data,
+            "ntokens": sample["ntokens"],
+            "nsentences": sample["nsentences"],
+            "sample_size": sample_size,
+            "reward": reward.mean(),
+        }
+        if self.report_accuracy:
+            n_correct, total = self.compute_accuracy(model, net_output, sample)
+            logging_output["n_correct"] = utils.item(n_correct.data)
+            logging_output["total"] = utils.item(total.data)
+
+        return loss, sample_size, logging_output
+
+    def get_lprobs_and_target(self, model, net_output, sample, reward=None):
+        conf = sample['conf'][:, None, None] if 'conf' in sample and sample['conf'] is not None else 1
+        constraint_masks = None
+        if "constraint_masks" in sample and sample["constraint_masks"] is not None:
+            constraint_masks = sample["constraint_masks"]
+            net_output[0].masked_fill_(~constraint_masks, -math.inf)
+        if self.constraint_start is not None and self.constraint_end is not None:
+            net_output[0][:, :, 4:self.constraint_start] = -math.inf
+            net_output[0][:, :, self.constraint_end:] = -math.inf
+        lprobs = model.get_normalized_probs(net_output, log_probs=True) * conf
+        target = model.get_targets(sample, net_output)
+        if self.ignore_prefix_size > 0:
+            lprobs = lprobs[:, self.ignore_prefix_size :, :].contiguous()
+            target = target[:, self.ignore_prefix_size :].contiguous()
+            if constraint_masks is not None:
+                constraint_masks = constraint_masks[:, self.ignore_prefix_size :, :].contiguous()
+        if self.ignore_eos:
+            bsz, seq_len, embed_dim = lprobs.size()
+            eos_indices = target.eq(self.task.tgt_dict.eos())
+            lprobs = lprobs[~eos_indices].reshape(bsz, seq_len-1, embed_dim)
+            target = target[~eos_indices].reshape(bsz, seq_len-1)
+            if constraint_masks is not None:
+                constraint_masks = constraint_masks[~eos_indices].reshape(bsz, seq_len-1, embed_dim)
+        if constraint_masks is not None:
+            constraint_masks = constraint_masks.view(-1, constraint_masks.size(-1))
+        
+        if reward is not None:
+            reward = reward.unsqueeze(1).unsqueeze(1)
+            lprobs = lprobs*reward
+        return lprobs.view(-1, lprobs.size(-1)), target.view(-1), constraint_masks
+
+    def compute_loss(self, model, net_output, sample, update_num, reduce=True, reward=None):
+        lprobs, target, constraint_masks = self.get_lprobs_and_target(model, net_output, sample, reward=reward)
+
+        if constraint_masks is not None:
+            constraint_masks = constraint_masks[target != self.padding_idx]
+        # print(target.shape, self.padding_idx, lprobs.shape, target, lprobs)
+        lprobs = lprobs[target != self.padding_idx]
+        target = target[target != self.padding_idx]
+
+        
+        loss, nll_loss, ntokens = label_smoothed_nll_loss(
+            lprobs,
+            target,
+            self.eps,
+            update_num,
+            reduce=reduce,
+            drop_worst_ratio=self.drop_worst_ratio,
+            drop_worst_after=self.drop_worst_after,
+            use_rdrop=self.use_rdrop,
+            reg_alpha=self.reg_alpha,
+            constraint_masks=constraint_masks,
+            constraint_start=self.constraint_start,
+            constraint_end=self.constraint_end
+        )
+        
+        if self.logprob and self.reinforce:
+            # print(-lprobs.max(dim=-1)[0].squeeze(-1).sum(), loss)
+            if self.lambda_reinforce > 0:
+                lprobs_, target_, constraint_masks_ = self.get_lprobs_and_target(model, net_output, sample, reward=None)
+                
+                loss_, _, ntokens = label_smoothed_nll_loss(
+                    lprobs_,
+                    target_,
+                    self.eps,
+                    update_num,
+                    reduce=reduce,
+                    drop_worst_ratio=self.drop_worst_ratio,
+                    drop_worst_after=self.drop_worst_after,
+                    use_rdrop=self.use_rdrop,
+                    reg_alpha=self.reg_alpha,
+                    constraint_masks=constraint_masks_,
+                    constraint_start=self.constraint_start,
+                    constraint_end=self.constraint_end
+                )
+                # print(-lprobs.max(dim=-1)[0].squeeze(-1).sum(), loss_)
+                # loss = -lprobs.max(dim=-1)[0].squeeze(-1).sum()*self.lambda_reinforce + loss_
+
+                loss = loss*self.lambda_reinforce + loss_ # only supervised with more weights via reward
+
+            else:
+                loss = -lprobs.max(dim=-1)[0].squeeze(-1).sum()
+            
+        elif self.logprob:
+            loss = nll_loss
+
+        return loss, nll_loss, ntokens
+
+    def compute_accuracy(self, model, net_output, sample):
+        lprobs, target = self.get_lprobs_and_target(model, net_output, sample)
+        mask = target.ne(self.padding_idx)
+        n_correct = torch.sum(
+            lprobs.argmax(1).masked_select(mask).eq(target.masked_select(mask))
+        )
+        total = torch.sum(mask)
+        return n_correct, total
+
+    @classmethod
+    def reduce_metrics(cls, logging_outputs) -> None:
+        """Aggregate logging outputs from data parallel training."""
+        loss_sum = sum(log.get("loss", 0) for log in logging_outputs)
+        loss_sum_v1 = sum(log.get("loss_v1", 0) for log in logging_outputs)
+        loss_sum_v2 = sum(log.get("loss_v2", 0) for log in logging_outputs)
+        nll_loss_sum = sum(log.get("nll_loss", 0) for log in logging_outputs)
+        ntokens = sum(log.get("ntokens", 0) for log in logging_outputs)
+        nsentences = sum(log.get("nsentences", 0) for log in logging_outputs)
+        sample_size = sum(log.get("sample_size", 0) for log in logging_outputs)
+        sample_size_v1 = sum(log.get("sample_size_v1", 0) for log in logging_outputs)
+        sample_size_v2 = sum(log.get("sample_size_v2", 0) for log in logging_outputs)
+
+
+        reward = sum(log.get("reward", 0) for log in logging_outputs)
+
+        metrics.log_scalar(
+            "loss", loss_sum / sample_size, sample_size, round=3
+        )
+        metrics.log_scalar(
+            "loss_v1", loss_sum_v1 / max(sample_size_v1, 1), max(sample_size_v1, 1), round=3
+        )
+        metrics.log_scalar(
+            "loss_v2", loss_sum_v2 / max(sample_size_v2, 1), max(sample_size_v2, 1), round=3
+        )
+        metrics.log_scalar(
+            "nll_loss", nll_loss_sum / sample_size, ntokens, round=3
+        )
+        metrics.log_derived(
+            "ppl", lambda meters: utils.get_perplexity(meters["nll_loss"].avg)
+        )
+
+        metrics.log_scalar(
+            "ntokens", ntokens, 1, round=3
+        )
+        metrics.log_scalar(
+            "nsentences", nsentences, 1, round=3
+        )
+        metrics.log_scalar(
+            "sample_size", sample_size, 1, round=3
+        )
+        metrics.log_scalar(
+            "sample_size_v1", sample_size_v1, 1, round=3
+        )
+        metrics.log_scalar(
+            "sample_size_v2", sample_size_v2, 1, round=3
+        )
+
+        metrics.log_scalar(
+            "reward", reward / sample_size, sample_size, round=3
+        )
+
+        total = utils.item(sum(log.get("total", 0) for log in logging_outputs))
+        if total > 0:
+            metrics.log_scalar("total", total)
+            n_correct = utils.item(
+                sum(log.get("n_correct", 0) for log in logging_outputs)
+            )
+            metrics.log_scalar("n_correct", n_correct)
+            metrics.log_derived(
+                "accuracy",
+                lambda meters: round(
+                    meters["n_correct"].sum * 100.0 / meters["total"].sum, 3
+                )
+                if meters["total"].sum > 0
+                else float("nan"),
+            )
+
+    @staticmethod
+    def logging_outputs_can_be_summed() -> bool:
+        """
+        Whether the logging outputs returned by `forward` can be summed
+        across workers prior to calling `reduce_metrics`. Setting this
+        to True will improves distributed training speed.
+        """
+        return True

criterions/label_smoothed_encouraging_loss.py

@@ -0,0 +1,395 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import math
+from dataclasses import dataclass, field
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+import numpy as np
+from fairseq import metrics, utils
+from fairseq.criterions import FairseqCriterion, register_criterion
+from fairseq.dataclass import FairseqDataclass
+from omegaconf import II
+
+
+@dataclass
+class AdjustLabelSmoothedEncouragingLossConfig(FairseqDataclass):
+    label_smoothing: float = field(
+        default=0.0,
+        metadata={"help": "epsilon for label smoothing, 0 means no label smoothing"},
+    )
+    report_accuracy: bool = field(
+        default=False,
+        metadata={"help": "report accuracy metric"},
+    )
+    ignore_prefix_size: int = field(
+        default=0,
+        metadata={"help": "Ignore first N tokens"},
+    )
+    ignore_eos: bool = field(
+        default=False,
+        metadata={"help": "Ignore eos token"},
+    )
+    sentence_avg: bool = II("optimization.sentence_avg")
+    drop_worst_ratio: float = field(
+        default=0.0,
+        metadata={"help": "ratio for discarding bad samples"},
+    )
+    drop_worst_after: int = field(
+        default=0,
+        metadata={"help": "steps for discarding bad samples"},
+    )
+    use_rdrop: bool = field(
+        default=False, metadata={"help": "use R-Drop"}
+    )
+    reg_alpha: float = field(
+        default=1.0, metadata={"help": "weight for R-Drop"}
+    )
+    sample_patch_num: int = field(
+        default=196, metadata={"help": "sample patchs for v1"}
+    )
+    constraint_range: Optional[str] = field(
+        default=None,
+        metadata={"help": "constraint range"}
+    )
+    log_end: float = field(
+        default=0.75,
+        metadata={"help": "higher log_end is for cases with higher performance,"
+                          " we recommend 0.75 or 0.5 as your first try."}
+    )
+    drop_best_ratio: float = field(
+        default=0.0,
+        metadata={"help": "ratio for discarding best samples"},
+    )
+    drop_best_after: int = field(
+        default=0,
+        metadata={"help": "steps for discarding best samples"},
+    )
+
+
+
+def construct_rdrop_sample(x):
+    if isinstance(x, dict):
+        for key in x:
+            x[key] = construct_rdrop_sample(x[key])
+        return x
+    elif isinstance(x, torch.Tensor):
+        return x.repeat(2, *([1] * (x.dim()-1)))
+    elif isinstance(x, int):
+        return x * 2
+    elif isinstance(x, np.ndarray):
+        return x.repeat(2)
+    else:
+        raise NotImplementedError
+
+
+def kl_loss(p, q):
+    p_loss = F.kl_div(p, torch.exp(q), reduction='sum')
+    q_loss = F.kl_div(q, torch.exp(p), reduction='sum')
+    loss = (p_loss + q_loss) / 2
+    return loss
+
+
+def label_smoothed_nll_loss(
+        lprobs, target, epsilon, update_num, reduce=True,
+        drop_worst_ratio=0.0, drop_worst_after=0, use_rdrop=False, reg_alpha=1.0,
+        constraint_masks=None, constraint_start=None, constraint_end=None,        drop_best_ratio=0.0,
+        drop_best_after=0,
+):
+    if target.dim() == lprobs.dim() - 1:
+        target = target.unsqueeze(-1)
+    nll_loss = -lprobs.gather(dim=-1, index=target).squeeze(-1)
+    if constraint_masks is not None:
+        smooth_loss = -lprobs.masked_fill(~constraint_masks, 0).sum(dim=-1, keepdim=True).squeeze(-1)
+        eps_i = epsilon / (constraint_masks.sum(1) - 1 + 1e-6)
+    elif constraint_start is not None and constraint_end is not None:
+        constraint_range = [0, 1, 2, 3] + list(range(constraint_start, constraint_end))
+        smooth_loss = -lprobs[:, constraint_range].sum(dim=-1, keepdim=True).squeeze(-1)
+        eps_i = epsilon / (len(constraint_range) - 1 + 1e-6)
+    else:
+        smooth_loss = -lprobs.sum(dim=-1, keepdim=True).squeeze(-1)
+        eps_i = epsilon / (lprobs.size(-1) - 1)
+    loss = (1.0 - epsilon - eps_i) * nll_loss + eps_i * smooth_loss
+    if drop_worst_ratio > 0 and update_num > drop_worst_after:
+        if use_rdrop:
+            true_batch_size = loss.size(0) // 2
+            _, indices = torch.topk(loss[:true_batch_size], k=int(true_batch_size * (1 - drop_worst_ratio)), largest=False)
+            loss = torch.cat([loss[indices], loss[indices+true_batch_size]])
+            nll_loss = torch.cat([nll_loss[indices], nll_loss[indices+true_batch_size]])
+            lprobs = torch.cat([lprobs[indices], lprobs[indices+true_batch_size]])
+        else:
+            loss, indices = torch.topk(loss, k=int(loss.shape[0] * (1 - drop_worst_ratio)), largest=False)
+            nll_loss = nll_loss[indices]
+            lprobs = lprobs[indices]
+            target = target[indices]
+    if update_num > drop_best_after:
+        loss, indices = torch.topk(loss, k=int(loss.shape[0] * (1 - drop_best_ratio)), largest=True)
+        nll_loss = nll_loss[indices]
+        lprobs = lprobs[indices]
+        target = target[indices]
+
+    ntokens = loss.numel()
+    nll_loss = nll_loss.sum()
+    loss = loss.sum()
+    if use_rdrop:
+        true_batch_size = lprobs.size(0) // 2
+        p = lprobs[:true_batch_size]
+        q = lprobs[true_batch_size:]
+        if constraint_start is not None and constraint_end is not None:
+            constraint_range = [0, 1, 2, 3] + list(range(constraint_start, constraint_end))
+            p = p[:, constraint_range]
+            q = q[:, constraint_range]
+        loss += kl_loss(p, q) * reg_alpha
+
+    return loss, nll_loss, ntokens,lprobs,target
+
+
+@register_criterion(
+    "adjust_label_smoothed_encouraging_loss", dataclass=AdjustLabelSmoothedEncouragingLossConfig
+)
+class AdjustLabelSmoothedEncouragingLossCriterion(FairseqCriterion):
+    def __init__(
+        self,
+        task,
+        sentence_avg,
+        label_smoothing,
+        ignore_prefix_size=0,
+        ignore_eos=False,
+        report_accuracy=False,
+        drop_worst_ratio=0,
+        drop_worst_after=0,
+        use_rdrop=False,
+        reg_alpha=1.0,
+        sample_patch_num=196,
+        constraint_range=None,
+        log_end=0.75,
+        drop_best_ratio=0.0,
+        drop_best_after=0,
+    ):
+        super().__init__(task)
+        self.sentence_avg = sentence_avg
+        self.eps = label_smoothing
+        self.ignore_prefix_size = ignore_prefix_size
+        self.ignore_eos = ignore_eos
+        self.report_accuracy = report_accuracy
+        self.drop_worst_ratio = drop_worst_ratio
+        self.drop_worst_after = drop_worst_after
+        self.use_rdrop = use_rdrop
+        self.reg_alpha = reg_alpha
+        self.sample_patch_num = sample_patch_num
+
+        self.constraint_start = None
+        self.constraint_end = None
+        if constraint_range is not None:
+            constraint_start, constraint_end = constraint_range.split(',')
+            self.constraint_start = int(constraint_start)
+            self.constraint_end = int(constraint_end)
+        self.log_end = log_end
+        self.drop_best_ratio = drop_best_ratio
+        self.drop_best_after = drop_best_after
+        print('el, self.log_end=', self.log_end)
+    # @staticmethod
+    # def add_args(parser):
+    #     """Add criterion-specific arguments to the parser."""
+    #     # fmt: off
+    #     parser.add_argument('--log_end', type=float, default=1.0)
+
+    def forward(self, model, sample, update_num=0, reduce=True):
+        """Compute the loss for the given sample.
+
+        Returns a tuple with three elements:
+        1) the loss
+        2) the sample size, which is used as the denominator for the gradient
+        3) logging outputs to display while training
+        """
+        if isinstance(sample, list):
+            if self.sample_patch_num > 0:
+                sample[0]['net_input']['sample_patch_num'] = self.sample_patch_num
+            loss_v1, sample_size_v1, logging_output_v1 = self.forward(model, sample[0], update_num, reduce)
+            loss_v2, sample_size_v2, logging_output_v2 = self.forward(model, sample[1], update_num, reduce)
+            loss = loss_v1 / sample_size_v1 + loss_v2 / sample_size_v2
+            sample_size = 1
+            logging_output = {
+                "loss": loss.data,
+                "loss_v1": loss_v1.data,
+                "loss_v2": loss_v2.data,
+                "nll_loss": logging_output_v1["nll_loss"].data / sample_size_v1 + logging_output_v2["nll_loss"].data / sample_size_v2,
+                "ntokens": logging_output_v1["ntokens"] + logging_output_v2["ntokens"],
+                "nsentences": logging_output_v1["nsentences"] + logging_output_v2["nsentences"],
+                "sample_size": 1,
+                "sample_size_v1": sample_size_v1,
+                "sample_size_v2": sample_size_v2,
+            }
+            return loss, sample_size, logging_output
+
+        if self.use_rdrop:
+            construct_rdrop_sample(sample)
+
+        net_output = model(**sample["net_input"])
+        loss, nll_loss, ntokens = self.compute_loss(model, net_output, sample, update_num, reduce=reduce)
+        sample_size = (
+            sample["target"].size(0) if self.sentence_avg else ntokens
+        )
+        logging_output = {
+            "loss": loss.data,
+            "nll_loss": nll_loss.data,
+            "ntokens": sample["ntokens"],
+            "nsentences": sample["nsentences"],
+            "sample_size": sample_size,
+        }
+        if self.report_accuracy:
+            n_correct, total = self.compute_accuracy(model, net_output, sample)
+            logging_output["n_correct"] = utils.item(n_correct.data)
+            logging_output["total"] = utils.item(total.data)
+        return loss, sample_size, logging_output
+
+    def get_lprobs_and_target(self, model, net_output, sample):
+        conf = sample['conf'][:, None, None] if 'conf' in sample and sample['conf'] is not None else 1
+        constraint_masks = None
+        if "constraint_masks" in sample and sample["constraint_masks"] is not None:
+            constraint_masks = sample["constraint_masks"]
+            net_output[0].masked_fill_(~constraint_masks, -math.inf)
+        if self.constraint_start is not None and self.constraint_end is not None:
+            net_output[0][:, :, 4:self.constraint_start] = -math.inf
+            net_output[0][:, :, self.constraint_end:] = -math.inf
+        lprobs = model.get_normalized_probs(net_output, log_probs=True) * conf
+        target = model.get_targets(sample, net_output)
+        if self.ignore_prefix_size > 0:
+            lprobs = lprobs[:, self.ignore_prefix_size :, :].contiguous()
+            target = target[:, self.ignore_prefix_size :].contiguous()
+            if constraint_masks is not None:
+                constraint_masks = constraint_masks[:, self.ignore_prefix_size :, :].contiguous()
+        if self.ignore_eos:
+            bsz, seq_len, embed_dim = lprobs.size()
+            eos_indices = target.eq(self.task.tgt_dict.eos())
+            lprobs = lprobs[~eos_indices].reshape(bsz, seq_len-1, embed_dim)
+            target = target[~eos_indices].reshape(bsz, seq_len-1)
+            if constraint_masks is not None:
+                constraint_masks = constraint_masks[~eos_indices].reshape(bsz, seq_len-1, embed_dim)
+        if constraint_masks is not None:
+            constraint_masks = constraint_masks.view(-1, constraint_masks.size(-1))
+        return lprobs.view(-1, lprobs.size(-1)), target.view(-1), constraint_masks
+
+    def compute_loss(self, model, net_output, sample, update_num, reduce=True):
+        lprobs, target, constraint_masks = self.get_lprobs_and_target(model, net_output, sample)
+        if constraint_masks is not None:
+            constraint_masks = constraint_masks[target != self.padding_idx]
+        lprobs = lprobs[target != self.padding_idx]
+        target = target[target != self.padding_idx]
+        loss, nll_loss, ntokens, lprobs, target = label_smoothed_nll_loss(
+            lprobs,
+            target,
+            self.eps,
+            update_num,
+            reduce=reduce,
+            drop_worst_ratio=self.drop_worst_ratio,
+            drop_worst_after=self.drop_worst_after,
+            use_rdrop=self.use_rdrop,
+            reg_alpha=self.reg_alpha,
+            constraint_masks=constraint_masks,
+            constraint_start=self.constraint_start,
+            constraint_end=self.constraint_end
+        )
+        # for encouraging loss
+        probs = torch.exp(lprobs)
+        bonus = torch.log(torch.clamp((torch.ones_like(probs) - probs), min=1e-5))  # likelihood bonus
+        log_end = self.log_end
+        if log_end != 1.0:  # e.g. 0.9
+            y_log_end = torch.log(torch.ones_like(probs) - log_end)
+            bonus_after_log_end = 1 / (log_end - torch.ones_like(probs)) * (probs - log_end) + y_log_end
+            # x:log_end, y  torch.log(torch.clamp((torch.ones_like(probs) - probs), min=self.cl_eps))
+            bonus = torch.where(probs > log_end, bonus_after_log_end, bonus)
+        c_loss = F.nll_loss(
+            -bonus,
+            target.view(-1),
+            reduction='sum',
+        )
+        smoothing_c_loss = bonus.sum(dim=-1)
+        smoothing_c_loss = smoothing_c_loss.sum()
+        c_loss = c_loss * (1 - self.eps) + (self.eps / lprobs.size(-1)) * smoothing_c_loss
+        loss = loss + c_loss
+        # end for encouraging loss
+        return loss, nll_loss, ntokens
+
+    def compute_accuracy(self, model, net_output, sample):
+        lprobs, target = self.get_lprobs_and_target(model, net_output, sample)
+        mask = target.ne(self.padding_idx)
+        n_correct = torch.sum(
+            lprobs.argmax(1).masked_select(mask).eq(target.masked_select(mask))
+        )
+        total = torch.sum(mask)
+        return n_correct, total
+
+    @classmethod
+    def reduce_metrics(cls, logging_outputs) -> None:
+        """Aggregate logging outputs from data parallel training."""
+        loss_sum = sum(log.get("loss", 0) for log in logging_outputs)
+        loss_sum_v1 = sum(log.get("loss_v1", 0) for log in logging_outputs)
+        loss_sum_v2 = sum(log.get("loss_v2", 0) for log in logging_outputs)
+        nll_loss_sum = sum(log.get("nll_loss", 0) for log in logging_outputs)
+        ntokens = sum(log.get("ntokens", 0) for log in logging_outputs)
+        nsentences = sum(log.get("nsentences", 0) for log in logging_outputs)
+        sample_size = sum(log.get("sample_size", 0) for log in logging_outputs)
+        sample_size_v1 = sum(log.get("sample_size_v1", 0) for log in logging_outputs)
+        sample_size_v2 = sum(log.get("sample_size_v2", 0) for log in logging_outputs)
+
+        metrics.log_scalar(
+            "loss", loss_sum / sample_size, sample_size, round=3
+        )
+        metrics.log_scalar(
+            "loss_v1", loss_sum_v1 / max(sample_size_v1, 1), max(sample_size_v1, 1), round=3
+        )
+        metrics.log_scalar(
+            "loss_v2", loss_sum_v2 / max(sample_size_v2, 1), max(sample_size_v2, 1), round=3
+        )
+        metrics.log_scalar(
+            "nll_loss", nll_loss_sum / sample_size, ntokens, round=3
+        )
+        metrics.log_derived(
+            "ppl", lambda meters: utils.get_perplexity(meters["nll_loss"].avg)
+        )
+
+        metrics.log_scalar(
+            "ntokens", ntokens, 1, round=3
+        )
+        metrics.log_scalar(
+            "nsentences", nsentences, 1, round=3
+        )
+        metrics.log_scalar(
+            "sample_size", sample_size, 1, round=3
+        )
+        metrics.log_scalar(
+            "sample_size_v1", sample_size_v1, 1, round=3
+        )
+        metrics.log_scalar(
+            "sample_size_v2", sample_size_v2, 1, round=3
+        )
+
+        total = utils.item(sum(log.get("total", 0) for log in logging_outputs))
+        if total > 0:
+            metrics.log_scalar("total", total)
+            n_correct = utils.item(
+                sum(log.get("n_correct", 0) for log in logging_outputs)
+            )
+            metrics.log_scalar("n_correct", n_correct)
+            metrics.log_derived(
+                "accuracy",
+                lambda meters: round(
+                    meters["n_correct"].sum * 100.0 / meters["total"].sum, 3
+                )
+                if meters["total"].sum > 0
+                else float("nan"),
+            )
+
+    @staticmethod
+    def logging_outputs_can_be_summed() -> bool:
+        """
+        Whether the logging outputs returned by `forward` can be summed
+        across workers prior to calling `reduce_metrics`. Setting this
+        to True will improves distributed training speed.
+        """
+        return True

criterions/refcoco_scst_loss.py

@@ -0,0 +1,427 @@
+# Modified from OFA code.
+# 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 math
+import string
+from dataclasses import dataclass, field
+from collections import OrderedDict
+from typing import Optional
+
+import torch
+from fairseq import metrics, utils
+from fairseq.criterions import FairseqCriterion, register_criterion
+from fairseq.dataclass import FairseqDataclass
+from omegaconf import II
+
+from data import data_utils
+from utils.cider.pyciderevalcap.ciderD.ciderD import CiderD
+
+
+
+def scst_loss(lprobs, target, reward, ignore_index=None, reduce=True, ce=False):
+
+    if ce:
+        loss = -lprobs.gather(dim=-1, index=target.unsqueeze(-1)).squeeze(-1)
+    elif isinstance(reward, float):
+        loss = -lprobs.gather(dim=-1, index=target.unsqueeze(-1)).squeeze() * reward
+    else:
+        loss = -lprobs.gather(dim=-1, index=target.unsqueeze(-1)).squeeze() * reward.unsqueeze(-1)
+
+    if ignore_index is not None:
+        pad_mask = target.eq(ignore_index)
+        loss.masked_fill_(pad_mask, 0.0)
+        ntokens = (~pad_mask).sum()
+    else:
+        loss = loss.squeeze(-1)
+        ntokens = target.numel()
+    if reduce:
+        loss = loss.sum()
+    return loss, ntokens
+
+
+@dataclass
+class RefCOCOScstRewardCriterionConfig(FairseqDataclass):
+    scst_cider_cached_tokens: Optional[str] = field(
+        default="coco-train-words.p",
+        metadata={"help": "path to cached cPickle file used to calculate CIDEr scores"},
+    )
+    ignore_prefix_size: int = field(
+        default=0,
+        metadata={"help": "Ignore first N tokens"},
+    )
+    sentence_avg: bool = II("optimization.sentence_avg")
+    constraint_range: Optional[str] = field(
+        default=None,
+        metadata={"help": "constraint range"}
+    )
+
+
+    acc_thresh: Optional[float] = field(
+        default=None, metadata={"help": "acc thresh for refcoco"}
+    )
+    metric: Optional[str] = field(
+        default='acc',
+        metadata={"help": "metric"}
+    )
+
+    max_area_size: Optional[float] = field(
+        default=None, metadata={"help": "max_area_size"}
+    )
+
+    min_area_size: Optional[float] = field(
+        default=None, metadata={"help": "min_area_size"}
+    )
+    logprob: Optional[bool] = field(
+        default=False, metadata={"help": "maximise log prob"}
+    )
+
+    pos_reward: Optional[float] = field(
+        default=None, metadata={"help": "pos_reward"}
+    )
+
+    neg_reward: Optional[float] = field(
+        default=None, metadata={"help": "neg_reward"}
+    )
+
+    reinforce: Optional[bool] = field(
+        default=False, metadata={"help": "reinforce"}
+    )
+
+    lambda_reinforce: Optional[float] = field(
+        default=0, metadata={"help": "lambda_reinforce"}
+    )
+
+    medium_area: Optional[bool] = field(
+        default=False, metadata={"help": "reinforce"}
+    )
+
+@register_criterion(
+    "refcoco_scst_reward_criterion", dataclass=RefCOCOScstRewardCriterionConfig
+)
+class RefCOCOScstRewardCriterion(FairseqCriterion):
+    CIDER_REWARD_WEIGHT = 1
+
+    def __init__(
+        self,
+        task,
+        scst_cider_cached_tokens,
+        sentence_avg,
+        ignore_prefix_size=0,
+        constraint_range=None,
+        acc_thresh=None,
+        metric='acc',
+        max_area_size=None,
+        min_area_size=None,
+        logprob=False,
+        pos_reward=None,
+        neg_reward=None,
+        reinforce=False,
+        lambda_reinforce=0,
+        medium_area=False,
+    ):
+        super().__init__(task)
+        self.sentence_avg = sentence_avg
+        self.ignore_prefix_size = ignore_prefix_size
+        self.transtab = str.maketrans({key: None for key in string.punctuation})
+
+        self.constraint_start = None
+        self.constraint_end = None
+        if constraint_range is not None:
+            constraint_start, constraint_end = constraint_range.split(',')
+            self.constraint_start = int(constraint_start)
+            self.constraint_end = int(constraint_end)
+
+        self.metric = metric 
+        print("metric", metric)
+
+        self.acc_thresh = acc_thresh
+        self.metric = metric
+        self.min_area_size = min_area_size
+        self.max_area_size = max_area_size
+        self.logprob = logprob
+
+        self.pos_reward = pos_reward
+        self.neg_reward = neg_reward
+
+        self.reinforce = reinforce
+        self.lambda_reinforce = lambda_reinforce
+
+        self.medium_area = medium_area
+
+
+            
+
+    def forward(self, model, sample, update_num=0, reduce=True):
+        """Compute the loss for the given sample.
+
+        Returns a tuple with three elements:
+        1) the loss
+        2) the sample size, which is used as the denominator for the gradient
+        3) logging outputs to display while training
+        """
+        loss, score, ntokens, nsentences = self.compute_loss(model, sample, reduce=reduce)
+
+        sample_size = (
+            nsentences if self.sentence_avg else ntokens
+        )
+        logging_output = {
+            "loss": loss.data,
+            "score": score,
+            "ntokens": ntokens,
+            "nsentences": nsentences,
+            "sample_size": sample_size,
+        }
+        return loss, sample_size, logging_output
+
+    def _calculate_eval_scores(self, gen_res, gt_idx, gt_res):
+        '''
+        gen_res: generated captions, list of str
+        gt_idx: list of int, of the same length as gen_res
+        gt_res: ground truth captions, list of list of str.
+            gen_res[i] corresponds to gt_res[gt_idx[i]]
+            Each image can have multiple ground truth captions
+        '''
+
+        gen_res_size = len(gen_res)
+
+        res = OrderedDict()
+        for i in range(gen_res_size):
+            res[i] = [self._wrap_sentence(gen_res[i].strip().translate(self.transtab))]
+
+        gts = OrderedDict()
+        gt_res_ = [
+            [self._wrap_sentence(gt_res[i][j].strip().translate(self.transtab)) for j in range(len(gt_res[i]))]
+                for i in range(len(gt_res))
+        ]
+        for i in range(gen_res_size):
+            gts[i] = gt_res_[gt_idx[i]]
+
+        res_ = [{'image_id':i, 'caption': res[i]} for i in range(len(res))]
+
+        # replace with other metrics
+        if self.metric != 'cider':
+            predicts = [res[i][0] if isinstance(res[i], list) else res[i] for i in range(len(res))]
+
+            answers = [gts[i] for i in range(gen_res_size)]
+            
+            results = self.evaluator.run_evaluation(predicts, answers)
+            batch_cider_scores = results[self.metric]
+
+            batch_cider_scores = torch.tensor(batch_cider_scores).repeat(gen_res_size)
+        else:
+            _, batch_cider_scores = self.scst_cider_scorer.compute_score(gts, res_)
+
+        scores = self.CIDER_REWARD_WEIGHT * batch_cider_scores
+        return scores
+
+    @classmethod
+    def _wrap_sentence(self, s):
+        # ensure the sentence ends with <eos> token
+        # in order to keep consisitent with cider_cached_tokens
+        r = s.strip()
+        if r.endswith('.'):
+            r = r[:-1]
+        r += ' <eos>'
+        return r
+
+
+    def get_generator_out(self, model, sample):
+
+
+        model.eval()
+        with torch.no_grad():
+            self.task.scst_generator.model.eval()
+            gen_out = self.task.scst_generator.generate([model], sample)
+
+        gen_target = []
+        gen_res = []
+        gt_res = []
+        for i in range(len(gen_out)):
+            gen_res.append(gen_out[i][0]["tokens"][:-1] - len(self.task.src_dict) + self.task.cfg.num_bins)
+            gt_res.append(sample["target"][i][:-1] - len(self.task.src_dict) + self.task.cfg.num_bins)
+            gen_target.append(gen_out[i][0]["tokens"][:-1].int().cpu())
+
+        return gen_target, gen_res, gt_res
+
+    def _calculate_ap_score(self, hyps, refs, thresh=0.5, min_area_size=None, max_area_size=None, medium_area=False):
+        interacts = torch.cat(
+            [torch.where(hyps[:, :2] < refs[:, :2], refs[:, :2], hyps[:, :2]),
+             torch.where(hyps[:, 2:] < refs[:, 2:], hyps[:, 2:], refs[:, 2:])],
+            dim=1
+        )
+        area_predictions = (hyps[:, 2] - hyps[:, 0]) * (hyps[:, 3] - hyps[:, 1]) ## x1, y1, x2, y2, x1 < x2
+        area_targets = (refs[:, 2] - refs[:, 0]) * (refs[:, 3] - refs[:, 1])
+        interacts_w = interacts[:, 2] - interacts[:, 0]
+        interacts_h = interacts[:, 3] - interacts[:, 1]
+        area_interacts = interacts_w * interacts_h
+        ious = area_interacts / (area_predictions + area_targets - area_interacts + 1e-6)
+        
+
+        if max_area_size is not None and min_area_size is not None:
+            if medium_area:
+                ious =  ious * (torch.logical_and(area_targets > max_area_size, area_targets < min_area_size).float())
+
+            else:
+                ious =  ious * (torch.logical_or(area_targets < max_area_size, area_targets > min_area_size).float())
+
+        elif min_area_size is not None:
+            if medium_area:
+                ious =  ious * (area_targets < min_area_size).float() # as max areas
+            else:
+                ious =  ious * (area_targets > min_area_size).float()
+
+        elif max_area_size is not None:
+            if medium_area:
+                ious =  ious * (area_targets > max_area_size).float()
+            else:
+                ious =  ious * (area_targets < max_area_size).float()
+
+        if thresh is None:
+            return ious
+        else:
+            return ((ious >= thresh) & (interacts_w > 0) & (interacts_h > 0)).float()
+
+
+    def get_reward_and_scores(self, gen_res, gt_res, device, sample):
+        
+
+        hyps_, refs_ = torch.stack(gen_res, dim=0), torch.stack(gt_res, dim=0)
+
+        hyps = hyps_ / (self.task.cfg.num_bins - 1) * self.task.cfg.max_image_size
+        refs = refs_ / (self.task.cfg.num_bins - 1) * self.task.cfg.max_image_size
+        
+        hyps[:, ::2] /= sample['w_resize_ratios'].unsqueeze(1)
+        hyps[:, 1::2] /= sample['h_resize_ratios'].unsqueeze(1)
+        refs[:, ::2] /= sample['w_resize_ratios'].unsqueeze(1)
+        refs[:, 1::2] /= sample['h_resize_ratios'].unsqueeze(1)
+
+        if self.metric == 'acc':
+            scores = self._calculate_ap_score(hyps, sample['region_coords'].float(), thresh=self.acc_thresh, 
+                                              min_area_size=self.min_area_size, max_area_size=self.max_area_size, medium_area=self.medium_area)
+        else:
+            raise NotImplemented
+        
+
+        if self.pos_reward:
+            scores = torch.where(scores > 0, self.pos_reward, scores)
+        if self.neg_reward:
+            scores = torch.where(scores == 0, self.neg_reward, scores)
+
+        return scores, scores
+
+
+    def get_net_output(self, model, sample, gen_target):
+        def merge(sample_list, eos=self.task.tgt_dict.eos(), move_eos_to_beginning=False):
+            return data_utils.collate_tokens(
+                sample_list,
+                pad_idx=self.padding_idx,
+                eos_idx=eos,
+                left_pad=False,
+                move_eos_to_beginning=move_eos_to_beginning,
+            )
+
+        batch_size = len(sample["target"])
+        gen_target_size = len(gen_target)
+        seq_per_img = gen_target_size // batch_size
+
+        model.train()
+        sample_src_tokens = torch.repeat_interleave(
+            sample['net_input']['src_tokens'], seq_per_img, dim=0
+        )
+        sample_src_lengths = torch.repeat_interleave(
+            sample['net_input']['src_lengths'], seq_per_img, dim=0
+        )
+        sample_patch_images = torch.repeat_interleave(
+            sample['net_input']['patch_images'], seq_per_img, dim=0
+        )
+        sample_patch_masks = torch.repeat_interleave(
+            sample['net_input']['patch_masks'], seq_per_img, dim=0
+        )
+        gen_prev_output_tokens = torch.as_tensor(
+            merge(gen_target, eos=self.task.tgt_dict.bos(), move_eos_to_beginning=True),
+            device=sample["target"].device, dtype=torch.int64
+        )
+        gen_target_tokens = torch.as_tensor(
+            merge(gen_target), device=sample["target"].device, dtype=torch.int64
+        )
+
+        net_output = model(
+            src_tokens=sample_src_tokens, src_lengths=sample_src_lengths,
+            patch_images=sample_patch_images, patch_masks=sample_patch_masks,
+            prev_output_tokens=gen_prev_output_tokens
+        )
+
+        return net_output, gen_target_tokens
+
+    def get_lprobs_and_target(self, model, net_output, gen_target):
+        if self.constraint_start is not None and self.constraint_end is not None:
+            net_output[0][:, :, 4:self.constraint_start] = -math.inf
+            net_output[0][:, :, self.constraint_end:] = -math.inf
+        lprobs = model.get_normalized_probs(net_output, log_probs=True)
+        if self.ignore_prefix_size > 0:
+            if getattr(lprobs, "batch_first", False):
+                lprobs = lprobs[:, self.ignore_prefix_size :, :].contiguous()
+                gen_target = gen_target[:, self.ignore_prefix_size :].contiguous()
+            else:
+                lprobs = lprobs[self.ignore_prefix_size :, :, :].contiguous()
+                gen_target = gen_target[self.ignore_prefix_size :, :].contiguous()
+        return lprobs, gen_target
+
+    def compute_loss(self, model, sample, reduce=True):
+        gen_target, gen_res, gt_res = self.get_generator_out(model, sample)
+        reward, scores = self.get_reward_and_scores(gen_res, gt_res, device=sample["target"].device, sample=sample)
+
+        net_output, gen_target_tokens = self.get_net_output(model, sample, gen_target)
+
+        gen_lprobs, gen_target_tokens = self.get_lprobs_and_target(model, net_output, gen_target_tokens)
+        loss, ntokens = scst_loss(gen_lprobs, gen_target_tokens, reward, ignore_index=self.padding_idx, reduce=reduce)
+        nsentences = gen_target_tokens.size(0)
+
+        if self.lambda_reinforce > 0:
+            target = model.get_targets(sample, net_output)[:, :-1] # ignore eos token
+            if self.ignore_prefix_size > 0:
+                target = target[:, self.ignore_prefix_size :].contiguous()
+
+            loss_ce, ntokens_ = scst_loss(gen_lprobs, target, reward=1, ignore_index=self.padding_idx, reduce=reduce, ce=True)
+            
+            loss = loss_ce + self.lambda_reinforce*loss
+
+        return loss, scores.sum(), ntokens, nsentences
+
+    @classmethod
+    def reduce_metrics(cls, logging_outputs) -> None:
+        """Aggregate logging outputs from data parallel training."""
+        loss_sum = sum(log.get("loss", 0) for log in logging_outputs)
+        score_sum = sum(log.get("score", 0) for log in logging_outputs)
+        ntokens = sum(log.get("ntokens", 0) for log in logging_outputs)
+        nsentences = sum(log.get("nsentences", 0) for log in logging_outputs)
+        sample_size = sum(log.get("sample_size", 0) for log in logging_outputs)
+
+        metrics.log_scalar(
+            "loss", loss_sum / sample_size, sample_size, round=3
+        )
+        metrics.log_scalar(
+            "score", score_sum / nsentences, nsentences, round=3
+        )
+
+        metrics.log_scalar(
+            "ntokens", ntokens, 1, round=3
+        )
+        metrics.log_scalar(
+            "nsentences", nsentences, 1, round=3
+        )
+        metrics.log_scalar(
+            "sample_size", sample_size, 1, round=3
+        )
+
+    @staticmethod
+    def logging_outputs_can_be_summed() -> bool:
+        """
+        Whether the logging outputs returned by `forward` can be summed
+        across workers prior to calling `reduce_metrics`. Setting this
+        to True will improves distributed training speed.
+        """
+        return True

data/.ipynb_checkpoints/file_dataset-checkpoint.py

@@ -0,0 +1,107 @@
+# 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 os
+import torch
+import pickle
+
+
+class FileDataset:
+    def __init__(self, file_path, selected_col_ids=None, dtypes=None, separator="\t", cached_index=False):
+        self.file_path = file_path
+        assert os.path.exists(self.file_path), "Error: The local datafile {} not exists!".format(self.file_path)
+
+        self.separator = separator
+        if selected_col_ids is None:
+            # default to all fields
+            self.selected_col_ids = list(
+                range(len(open(self.file_path).readline().rstrip("\n").split(self.separator))))
+        else:
+            self.selected_col_ids = [int(col_id) for col_id in selected_col_ids.split(",")]
+        if dtypes is None:
+            # default to str
+            self.dtypes = [str for col_id in self.selected_col_ids]
+        else:
+            self.dtypes = [eval(col_dtype) for col_dtype in dtypes.split(",")]
+            assert len(self.dtypes) == len(self.selected_col_ids)
+
+        self.data_cnt = 0
+        try:
+            self.slice_id = torch.distributed.get_rank()
+            self.slice_count = torch.distributed.get_world_size()
+        except Exception:
+            self.slice_id = 0
+            self.slice_count = 1
+        self.cached_index = cached_index
+        self._init_seek_index()
+        self._reader = self._get_reader()
+        print("file {} slice_id {} row count {} total row count {}".format(
+            self.file_path, self.slice_id, self.row_count, self.total_row_count)
+        )
+
+    def _init_seek_index(self):
+        if self.cached_index:
+            cache_path = "{}.index".format(self.file_path)
+            assert os.path.exists(cache_path), "cache file {} not exists!".format(cache_path)
+            self.total_row_count, self.lineid_to_offset = pickle.load(open(cache_path, "rb"))
+            print("local datafile {} slice_id {} use cached row_count and line_idx-to-offset mapping".format(
+                self.file_path, self.slice_id))
+        else:
+            # make an iteration over the file to get row_count and line_idx-to-offset mapping
+            fp = open(self.file_path, "r")
+            print("local datafile {} slice_id {} begin to initialize row_count and line_idx-to-offset mapping".format(
+                self.file_path, self.slice_id))
+            self.total_row_count = 0
+            offset = 0
+            self.lineid_to_offset = []
+            for line in fp:
+                self.lineid_to_offset.append(offset)
+                self.total_row_count += 1
+                offset += len(line.encode('utf-8'))
+        self._compute_start_pos_and_row_count()
+        print("local datafile {} slice_id {} finished initializing row_count and line_idx-to-offset mapping".format(
+            self.file_path, self.slice_id))
+
+    def _compute_start_pos_and_row_count(self):
+        self.row_count = self.total_row_count // self.slice_count
+        if self.slice_id < self.total_row_count - self.row_count * self.slice_count:
+            self.row_count += 1
+            self.start_pos = self.row_count * self.slice_id
+        else:
+            self.start_pos = self.row_count * self.slice_id + (self.total_row_count - self.row_count * self.slice_count)
+
+    def _get_reader(self):
+        fp = open(self.file_path, "r")
+        fp.seek(self.lineid_to_offset[self.start_pos])
+        return fp
+
+    def _seek(self, offset=0):
+        try:
+            print("slice_id {} seek offset {}".format(self.slice_id, self.start_pos + offset))
+            self._reader.seek(self.lineid_to_offset[self.start_pos + offset])
+            self.data_cnt = offset
+        except Exception:
+            print("slice_id {} seek offset {}".format(self.slice_id, offset))
+            self._reader.seek(self.lineid_to_offset[offset])
+            self.data_cnt = offset
+
+    def __del__(self):
+        self._reader.close()
+
+    def __len__(self):
+        return self.row_count
+
+    def get_total_row_count(self):
+        return self.total_row_count
+
+    def __getitem__(self, index):
+        if self.data_cnt == self.row_count:
+            print("reach the end of datafile, start a new reader")
+            self.data_cnt = 0
+            self._reader = self._get_reader()
+        column_l = self._reader.readline().rstrip("\n").split(self.separator)
+        self.data_cnt += 1
+        column_l = [dtype(column_l[col_id]) for col_id, dtype in zip(self.selected_col_ids, self.dtypes)]
+        return column_l