change the service

LICENSE

@@ -1,201 +0,0 @@
-                                 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,27 +0,0 @@
----
-title: Chinese OCR
-emoji: 📖
-colorFrom: red
-colorTo: indigo
-sdk: gradio
-sdk_version: 3.9.1
-app_file: app.py
-pinned: false
----
-# Configuration
-`title`: _string_
-OFA Image Caption
-`emoji`: _string_
-🖼
-`colorFrom`: _string_
-red
-`colorTo`: _string_
-indigo
-`sdk`: _string_
-gradio
-`app_file`: _string_
-app.py
-
-
-`pinned`: _boolean_
-false

README_EncouragingLoss.md

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

app.py

@@ -193,8 +193,7 @@ description = "Gradio Demo for Chinese OCR based on OFA-Base. "\
 article = "<p style='text-align: center'><a href='https://github.com/OFA-Sys/OFA' target='_blank'>OFA Github " \
           "Repo</a></p> "
 examples = [['shupai.png'], ['chinese.jpg'], ['gaidao.jpeg'],
-            ['qiaodaima.png'], ['benpao.jpeg'],
-            ['xsd.jpg', 'General']]
+            ['qiaodaima.png'], ['xsd.jpg']]
 io = gr.Interface(fn=ocr, inputs=gr.inputs.Image(type='filepath', label='Image'),
                   outputs=[gr.outputs.Image(type='pil', label='Image'),
                            gr.outputs.Dataframe(headers=['Box ID', 'Text'], type='pandas', label='OCR Results')],

checkpoints.md

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

checkpoints_cn.md

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

colab.md

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

criterions/__init__.py

@@ -1,4 +0,0 @@
-from .scst_loss import ScstRewardCriterion
-from .label_smoothed_cross_entropy import AdjustLabelSmoothedCrossEntropyCriterion
-from .clip_scst_loss import ClipScstRewardCriterion
-from .label_smoothed_encouraging_loss import AdjustLabelSmoothedEncouragingLossCriterion

criterions/clip_scst_loss.py

@@ -1,277 +0,0 @@
-# Copyright 2022 The OFA-Sys Team. 
-# All rights reserved.
-# This source code is licensed under the Apache 2.0 license 
-# found in the LICENSE file in the root directory.
-
-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

@@ -1,343 +0,0 @@
-# Copyright 2022 The OFA-Sys Team. 
-# All rights reserved.
-# This source code is licensed under the Apache 2.0 license 
-# found in the LICENSE file in the root directory.
-
-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
-            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 = 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_encouraging_loss.py

@@ -1,395 +0,0 @@
-# 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/scst_loss.py

@@ -1,281 +0,0 @@
-# Copyright 2022 The OFA-Sys Team. 
-# All rights reserved.
-# This source code is licensed under the Apache 2.0 license 
-# found in the LICENSE file in the root directory.
-
-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):
-    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 ScstRewardCriterionConfig(FairseqDataclass):
-    scst_cider_cached_tokens: 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"}
-    )
-
-
-@register_criterion(
-    "scst_reward_criterion", dataclass=ScstRewardCriterionConfig
-)
-class ScstRewardCriterion(FairseqCriterion):
-    CIDER_REWARD_WEIGHT = 1
-
-    def __init__(
-        self,
-        task,
-        scst_cider_cached_tokens,
-        sentence_avg,
-        ignore_prefix_size=0,
-        constraint_range=None
-    ):
-        super().__init__(task)
-        self.scst_cider_scorer = CiderD(df=scst_cider_cached_tokens)
-        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)
-
-    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))]
-        _, 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):
-        def decode(toks):
-            hypo = toks.int().cpu()
-            hypo_str = self.task.tgt_dict.string(hypo)
-            hypo_str = self.task.bpe.decode(hypo_str).strip()
-            return hypo, hypo_str
-
-        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)):
-            for j in range(len(gen_out[i])):
-                hypo, hypo_str = decode(gen_out[i][j]["tokens"])
-                gen_target.append(hypo)
-                gen_res.append(hypo_str)
-            gt_res.append(
-                decode(utils.strip_pad(sample["target"][i], self.padding_idx))[1].split('&&')
-            )
-
-        return gen_target, gen_res, gt_res
-
-    def get_reward_and_scores(self, gen_res, gt_res, device):
-        batch_size = len(gt_res)
-        gen_res_size = len(gen_res)
-        seq_per_img = gen_res_size // batch_size
-
-        gt_idx = [i // seq_per_img for i in range(gen_res_size)]
-        scores = self._calculate_eval_scores(gen_res, gt_idx, gt_res)
-        sc_ = scores.reshape(batch_size, seq_per_img)
-        baseline = (sc_.sum(1, keepdims=True) - sc_) / (sc_.shape[1] - 1)
-        # sample - baseline
-        reward = scores.reshape(batch_size, seq_per_img)
-        reward = reward - baseline
-        reward = reward.reshape(gen_res_size)
-        reward = torch.as_tensor(reward, device=device, dtype=torch.float64)
-
-        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)
-        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)
-        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

data/cv_data/image_classify_dataset.py

@@ -1,196 +0,0 @@
-# Copyright 2022 The OFA-Sys Team. 
-# All rights reserved.
-# This source code is licensed under the Apache 2.0 license 
-# found in the LICENSE file in the root directory.
-
-from io import BytesIO
-
-import logging
-import warnings
-import functools
-
-import numpy as np
-import torch
-import base64
-from torchvision import transforms
-from timm.data import create_transform
-from utils.vision_helper import RandomAugment
-
-from PIL import Image, ImageFile
-
-from data import data_utils
-from data.ofa_dataset import OFADataset
-
-ImageFile.LOAD_TRUNCATED_IMAGES = True
-ImageFile.MAX_IMAGE_PIXELS = None
-Image.MAX_IMAGE_PIXELS = None
-
-logger = logging.getLogger(__name__)
-warnings.filterwarnings("ignore", "(Possibly )?corrupt EXIF data", UserWarning)
-
-IMAGENET_DEFAULT_MEAN = (0.485, 0.456, 0.406)
-IMAGENET_DEFAULT_STD = (0.229, 0.224, 0.225)
-
-def collate(samples, pad_idx, eos_idx):
-    if len(samples) == 0:
-        return {}
-
-    def merge(key):
-        return data_utils.collate_tokens(
-            [s[key] for s in samples],
-            pad_idx,
-            eos_idx=eos_idx,
-        )
-
-    id = np.array([s["id"] for s in samples])
-    src_tokens = merge("source")
-    src_lengths = torch.LongTensor([s["source"].ne(pad_idx).long().sum() for s in samples])
-
-    patch_images = torch.stack([sample['patch_image'] for sample in samples], dim=0)
-    patch_masks = torch.cat([sample['patch_mask'] for sample in samples])
-
-    conf = None
-    if samples[0].get("conf", None) is not None:
-        conf = torch.cat([s['conf'] for s in samples], dim=0)
-
-    ref_dict = None
-    if samples[0].get("ref_dict", None) is not None:
-        ref_dict = np.array([s['ref_dict'] for s in samples])
-
-    constraint_masks = None
-    if samples[0].get("constraint_mask", None) is not None:
-        constraint_masks = merge("constraint_mask")
-
-    prev_output_tokens = None
-    target = None
-    if samples[0].get("target", None) is not None:
-        target = merge("target")
-        tgt_lengths = torch.LongTensor(
-            [s["target"].ne(pad_idx).long().sum() for s in samples]
-        )
-        ntokens = tgt_lengths.sum().item()
-
-        if samples[0].get("prev_output_tokens", None) is not None:
-            prev_output_tokens = merge("prev_output_tokens")
-    else:
-        ntokens = src_lengths.sum().item()
-
-    batch = {
-        "id": id,
-        "nsentences": len(samples),
-        "ntokens": ntokens,
-        "net_input": {
-            "src_tokens": src_tokens,
-            "src_lengths": src_lengths,
-            "patch_images": patch_images,
-            "patch_masks": patch_masks,
-            "prev_output_tokens": prev_output_tokens
-        },
-        "conf": conf,
-        "ref_dict": ref_dict,
-        "constraint_masks": constraint_masks,
-        "target": target,
-    }
-
-    return batch
-
-
-class ImageClassifyDataset(OFADataset):
-    def __init__(
-        self,
-        split,
-        dataset,
-        bpe,
-        src_dict,
-        tgt_dict=None,
-        max_src_length=128,
-        max_tgt_length=30,
-        patch_image_size=224,
-        constraint_trie=None,
-        imagenet_default_mean_and_std=False
-    ):
-        super().__init__(split, dataset, bpe, src_dict, tgt_dict)
-        self.max_src_length = max_src_length
-        self.max_tgt_length = max_tgt_length
-        self.patch_image_size = patch_image_size
-
-        self.constraint_trie = constraint_trie
-
-        if imagenet_default_mean_and_std:
-            mean = IMAGENET_DEFAULT_MEAN
-            std = IMAGENET_DEFAULT_STD
-        else:
-            mean = [0.5, 0.5, 0.5]
-            std = [0.5, 0.5, 0.5]
-
-        if self.split != 'train':
-            self.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),
-            ])
-            logger.info("val split, do not use random augmentation.")
-        else:
-            self.patch_resize_transform = create_transform(
-                input_size=patch_image_size,
-                is_training=True,
-                color_jitter=0.4,
-                auto_augment='rand-m9-mstd0.5-inc1',
-                interpolation='bicubic',
-                re_prob=0.25,
-                re_mode='pixel',
-                re_count=1,
-                mean=mean,
-                std=std,
-            )
-            self.patch_resize_transform = transforms.Compose(functools.reduce(lambda x, y:x + y, [
-                [lambda image: image.convert("RGB"),],
-                self.patch_resize_transform.transforms[:2],
-                [self.patch_resize_transform.transforms[2]],
-                [RandomAugment(2, 7, isPIL=True, augs=['Identity', 'AutoContrast', 'Equalize', 'Brightness', 'Sharpness', 'ShearX', 'ShearY', 'TranslateX', 'TranslateY', 'Rotate']), ],
-                self.patch_resize_transform.transforms[3:],
-            ]))
-            logger.info("train split, use random augmentation.")
-
-    def __getitem__(self, index):
-        image, label_name = self.dataset[index]
-
-        image = Image.open(BytesIO(base64.urlsafe_b64decode(image)))
-        patch_image = self.patch_resize_transform(image)
-        patch_mask = torch.tensor([True])
-
-        src_item = self.encode_text(' what does the image describe?')
-        tgt_item = self.encode_text(" {}".format(label_name))
-        ref_dict = {label_name: 1.0}
-
-        src_item = torch.cat([self.bos_item, src_item, self.eos_item])
-        target_item = torch.cat([tgt_item, self.eos_item])
-        prev_output_item = torch.cat([self.bos_item, tgt_item])
-
-        example = {
-            "id": index,
-            "source": src_item,
-            "patch_image": patch_image,
-            "patch_mask": patch_mask,
-            "target": target_item,
-            "prev_output_tokens": prev_output_item,
-            "ref_dict": ref_dict,
-        }
-        if self.constraint_trie is not None:
-            constraint_mask = torch.zeros((len(prev_output_item), len(self.tgt_dict))).bool()
-            for i in range(len(prev_output_item)):
-                constraint_prefix_token = prev_output_item[:i+1].tolist()
-                constraint_nodes = self.constraint_trie.get_next_layer(constraint_prefix_token)
-                constraint_mask[i][constraint_nodes] = True
-            example["constraint_mask"] = constraint_mask
-        return example
-
-    def collater(self, samples, pad_to_length=None):
-        """Merge a list of samples to form a mini-batch.
-        Args:
-            samples (List[dict]): samples to collate
-        Returns:
-            dict: a mini-batch containing the data of the task
-        """
-        return collate(samples, pad_idx=self.pad, eos_idx=self.eos)

data/data_utils.py

@@ -1,601 +0,0 @@
-# Copyright 2022 The OFA-Sys Team. 
-# All rights reserved.
-# This source code is licensed under the Apache 2.0 license 
-# found in the LICENSE file in the root directory.
-
-try:
-    from collections.abc import Iterable
-except ImportError:
-    from collections import Iterable
-import contextlib
-import itertools
-import logging
-import re
-import warnings
-from typing import Optional, Tuple
-
-import numpy as np
-import torch
-
-from fairseq.file_io import PathManager
-from fairseq import utils
-import os
-
-logger = logging.getLogger(__name__)
-
-
-def infer_language_pair(path):
-    """Infer language pair from filename: <split>.<lang1>-<lang2>.(...).idx"""
-    src, dst = None, None
-    for filename in PathManager.ls(path):
-        parts = filename.split(".")
-        if len(parts) >= 3 and len(parts[1].split("-")) == 2:
-            return parts[1].split("-")
-    return src, dst
-
-
-def collate_tokens(
-    values,
-    pad_idx,
-    eos_idx=None,
-    left_pad=False,
-    move_eos_to_beginning=False,
-    pad_to_length=None,
-    pad_to_multiple=1,
-    pad_to_bsz=None,
-):
-    """Convert a list of 1d tensors into a padded 2d tensor."""
-    size = max(v.size(0) for v in values)
-    size = size if pad_to_length is None else max(size, pad_to_length)
-    if pad_to_multiple != 1 and size % pad_to_multiple != 0:
-        size = int(((size - 0.1) // pad_to_multiple + 1) * pad_to_multiple)
-
-    def copy_tensor(src, dst):
-        assert dst.numel() == src.numel()
-        if move_eos_to_beginning:
-            if eos_idx is None:
-                # if no eos_idx is specified, then use the last token in src
-                dst[0] = src[-1]
-            else:
-                dst[0] = eos_idx
-            dst[1:] = src[:-1]
-        else:
-            dst.copy_(src)
-
-    if values[0].dim() == 1:
-        res = values[0].new(len(values), size).fill_(pad_idx)
-    elif values[0].dim() == 2:
-        assert move_eos_to_beginning is False
-        res = values[0].new(len(values), size, values[0].size(1)).fill_(pad_idx)
-    else:
-        raise NotImplementedError
-
-    for i, v in enumerate(values):
-        copy_tensor(v, res[i][size - len(v) :] if left_pad else res[i][: len(v)])
-    return res
-
-
-def load_indexed_dataset(
-    path, dictionary=None, dataset_impl=None, combine=False, default="cached"
-):
-    """A helper function for loading indexed datasets.
-
-    Args:
-        path (str): path to indexed dataset (e.g., 'data-bin/train')
-        dictionary (~fairseq.data.Dictionary): data dictionary
-        dataset_impl (str, optional): which dataset implementation to use. If
-            not provided, it will be inferred automatically. For legacy indexed
-            data we use the 'cached' implementation by default.
-        combine (bool, optional): automatically load and combine multiple
-            datasets. For example, if *path* is 'data-bin/train', then we will
-            combine 'data-bin/train', 'data-bin/train1', ... and return a
-            single ConcatDataset instance.
-    """
-    import fairseq.data.indexed_dataset as indexed_dataset
-    from fairseq.data.concat_dataset import ConcatDataset
-
-    datasets = []
-    for k in itertools.count():
-        path_k = path + (str(k) if k > 0 else "")
-        try:
-            path_k = indexed_dataset.get_indexed_dataset_to_local(path_k)
-        except Exception as e:
-            if "StorageException: [404] Path not found" in str(e):
-                logger.warning(f"path_k: {e} not found")
-            else:
-                raise e
-
-        dataset_impl_k = dataset_impl
-        if dataset_impl_k is None:
-            dataset_impl_k = indexed_dataset.infer_dataset_impl(path_k)
-        dataset = indexed_dataset.make_dataset(
-            path_k,
-            impl=dataset_impl_k or default,
-            fix_lua_indexing=True,
-            dictionary=dictionary,
-        )
-        if dataset is None:
-            break
-        logger.info("loaded {:,} examples from: {}".format(len(dataset), path_k))
-        datasets.append(dataset)
-        if not combine:
-            break
-    if len(datasets) == 0:
-        return None
-    elif len(datasets) == 1:
-        return datasets[0]
-    else:
-        return ConcatDataset(datasets)
-
-
-@contextlib.contextmanager
-def numpy_seed(seed, *addl_seeds):
-    """Context manager which seeds the NumPy PRNG with the specified seed and
-    restores the state afterward"""
-    if seed is None:
-        yield
-        return
-    if len(addl_seeds) > 0:
-        seed = int(hash((seed, *addl_seeds)) % 1e6)
-    state = np.random.get_state()
-    np.random.seed(seed)
-    try:
-        yield
-    finally:
-        np.random.set_state(state)
-
-
-def collect_filtered(function, iterable, filtered):
-    """
-    Similar to :func:`filter` but collects filtered elements in ``filtered``.
-
-    Args:
-        function (callable): function that returns ``False`` for elements that
-            should be filtered
-        iterable (iterable): iterable to filter
-        filtered (list): list to store filtered elements
-    """
-    for el in iterable:
-        if function(el):
-            yield el
-        else:
-            filtered.append(el)
-
-
-def _filter_by_size_dynamic(indices, size_fn, max_positions, raise_exception=False):
-    def compare_leq(a, b):
-        return a <= b if not isinstance(a, tuple) else max(a) <= b
-
-    def check_size(idx):
-        if isinstance(max_positions, float) or isinstance(max_positions, int):
-            return size_fn(idx) <= max_positions
-        elif isinstance(max_positions, dict):
-            idx_size = size_fn(idx)
-            assert isinstance(idx_size, dict)
-            intersect_keys = set(max_positions.keys()) & set(idx_size.keys())
-            return all(
-                all(
-                    a is None or b is None or a <= b
-                    for a, b in zip(idx_size[key], max_positions[key])
-                )
-                for key in intersect_keys
-            )
-        else:
-            # For MultiCorpusSampledDataset, will generalize it later
-            if not isinstance(size_fn(idx), Iterable):
-                return all(size_fn(idx) <= b for b in max_positions)
-            return all(
-                a is None or b is None or a <= b
-                for a, b in zip(size_fn(idx), max_positions)
-            )
-
-    ignored = []
-    itr = collect_filtered(check_size, indices, ignored)
-    indices = np.fromiter(itr, dtype=np.int64, count=-1)
-    return indices, ignored
-
-
-def filter_by_size(indices, dataset, max_positions, raise_exception=False):
-    """
-    [deprecated] Filter indices based on their size.
-    Use `FairseqDataset::filter_indices_by_size` instead.
-
-    Args:
-        indices (List[int]): ordered list of dataset indices
-        dataset (FairseqDataset): fairseq dataset instance
-        max_positions (tuple): filter elements larger than this size.
-            Comparisons are done component-wise.
-        raise_exception (bool, optional): if ``True``, raise an exception if
-            any elements are filtered (default: False).
-    """
-    warnings.warn(
-        "data_utils.filter_by_size is deprecated. "
-        "Use `FairseqDataset::filter_indices_by_size` instead.",
-        stacklevel=2,
-    )
-    if isinstance(max_positions, float) or isinstance(max_positions, int):
-        if hasattr(dataset, "sizes") and isinstance(dataset.sizes, np.ndarray):
-            ignored = indices[dataset.sizes[indices] > max_positions].tolist()
-            indices = indices[dataset.sizes[indices] <= max_positions]
-        elif (
-            hasattr(dataset, "sizes")
-            and isinstance(dataset.sizes, list)
-            and len(dataset.sizes) == 1
-        ):
-            ignored = indices[dataset.sizes[0][indices] > max_positions].tolist()
-            indices = indices[dataset.sizes[0][indices] <= max_positions]
-        else:
-            indices, ignored = _filter_by_size_dynamic(
-                indices, dataset.size, max_positions
-            )
-    else:
-        indices, ignored = _filter_by_size_dynamic(indices, dataset.size, max_positions)
-
-    if len(ignored) > 0 and raise_exception:
-        raise Exception(
-            (
-                "Size of sample #{} is invalid (={}) since max_positions={}, "
-                "skip this example with --skip-invalid-size-inputs-valid-test"
-            ).format(ignored[0], dataset.size(ignored[0]), max_positions)
-        )
-    if len(ignored) > 0:
-        logger.warning(
-            (
-                "{} samples have invalid sizes and will be skipped, "
-                "max_positions={}, first few sample ids={}"
-            ).format(len(ignored), max_positions, ignored[:10])
-        )
-    return indices
-
-
-def filter_paired_dataset_indices_by_size(src_sizes, tgt_sizes, indices, max_sizes):
-    """Filter a list of sample indices. Remove those that are longer
-        than specified in max_sizes.
-
-    Args:
-        indices (np.array): original array of sample indices
-        max_sizes (int or list[int] or tuple[int]): max sample size,
-            can be defined separately for src and tgt (then list or tuple)
-
-    Returns:
-        np.array: filtered sample array
-        list: list of removed indices
-    """
-    if max_sizes is None:
-        return indices, []
-    if type(max_sizes) in (int, float):
-        max_src_size, max_tgt_size = max_sizes, max_sizes
-    else:
-        max_src_size, max_tgt_size = max_sizes
-    if tgt_sizes is None:
-        ignored = indices[src_sizes[indices] > max_src_size]
-    else:
-        ignored = indices[
-            (src_sizes[indices] > max_src_size) | (tgt_sizes[indices] > max_tgt_size)
-        ]
-    if len(ignored) > 0:
-        if tgt_sizes is None:
-            indices = indices[src_sizes[indices] <= max_src_size]
-        else:
-            indices = indices[
-                (src_sizes[indices] <= max_src_size)
-                & (tgt_sizes[indices] <= max_tgt_size)
-            ]
-    return indices, ignored.tolist()
-
-
-def batch_by_size(
-    indices,
-    num_tokens_fn,
-    num_tokens_vec=None,
-    max_tokens=None,
-    max_sentences=None,
-    required_batch_size_multiple=1,
-    fixed_shapes=None,
-):
-    """
-    Yield mini-batches of indices bucketed by size. Batches may contain
-    sequences of different lengths.
-
-    Args:
-        indices (List[int]): ordered list of dataset indices
-        num_tokens_fn (callable): function that returns the number of tokens at
-            a given index
-        num_tokens_vec (List[int], optional): precomputed vector of the number
-            of tokens for each index in indices (to enable faster batch generation)
-        max_tokens (int, optional): max number of tokens in each batch
-            (default: None).
-        max_sentences (int, optional): max number of sentences in each
-            batch (default: None).
-        required_batch_size_multiple (int, optional): require batch size to
-            be less than N or a multiple of N (default: 1).
-        fixed_shapes (List[Tuple[int, int]], optional): if given, batches will
-            only be created with the given shapes. *max_sentences* and
-            *required_batch_size_multiple* will be ignored (default: None).
-    """
-    try:
-        from fairseq.data.data_utils_fast import (
-            batch_by_size_fn,
-            batch_by_size_vec,
-            batch_fixed_shapes_fast,
-        )
-    except ImportError:
-        raise ImportError(
-            "Please build Cython components with: "
-            "`python setup.py build_ext --inplace`"
-        )
-    except ValueError:
-        raise ValueError(
-            "Please build (or rebuild) Cython components with `python setup.py build_ext --inplace`."
-        )
-
-    # added int() to avoid TypeError: an integer is required
-    max_tokens = (
-        int(max_tokens) if max_tokens is not None else -1
-    )
-    max_sentences = max_sentences if max_sentences is not None else -1
-    bsz_mult = required_batch_size_multiple
-
-    if not isinstance(indices, np.ndarray):
-        indices = np.fromiter(indices, dtype=np.int64, count=-1)
-
-    if num_tokens_vec is not None and not isinstance(num_tokens_vec, np.ndarray):
-        num_tokens_vec = np.fromiter(num_tokens_vec, dtype=np.int64, count=-1)
-
-    if fixed_shapes is None:
-        if num_tokens_vec is None:
-            return batch_by_size_fn(
-                indices,
-                num_tokens_fn,
-                max_tokens,
-                max_sentences,
-                bsz_mult,
-            )
-        else:
-            return batch_by_size_vec(
-                indices,
-                num_tokens_vec,
-                max_tokens,
-                max_sentences,
-                bsz_mult,
-            )
-
-    else:
-        fixed_shapes = np.array(fixed_shapes, dtype=np.int64)
-        sort_order = np.lexsort(
-            [
-                fixed_shapes[:, 1].argsort(),  # length
-                fixed_shapes[:, 0].argsort(),  # bsz
-            ]
-        )
-        fixed_shapes_sorted = fixed_shapes[sort_order]
-        return batch_fixed_shapes_fast(indices, num_tokens_fn, fixed_shapes_sorted)
-
-
-def post_process(sentence: str, symbol: str):
-    if symbol == "sentencepiece":
-        sentence = sentence.replace(" ", "").replace("\u2581", " ").strip()
-    elif symbol == "wordpiece":
-        sentence = sentence.replace(" ", "").replace("_", " ").strip()
-    elif symbol == "letter":
-        sentence = sentence.replace(" ", "").replace("|", " ").strip()
-    elif symbol == "silence":
-        import re
-        sentence = sentence.replace("<SIL>", "")
-        sentence = re.sub(' +', ' ', sentence).strip()
-    elif symbol == "_EOW":
-        sentence = sentence.replace(" ", "").replace("_EOW", " ").strip()
-    elif symbol in {"subword_nmt", "@@ ", "@@"}:
-        if symbol == "subword_nmt":
-            symbol = "@@ "
-        sentence = (sentence + " ").replace(symbol, "").rstrip()
-    elif symbol == "none":
-        pass
-    elif symbol is not None:
-        raise NotImplementedError(f"Unknown post_process option: {symbol}")
-    return sentence
-
-
-def compute_mask_indices(
-    shape: Tuple[int, int],
-    padding_mask: Optional[torch.Tensor],
-    mask_prob: float,
-    mask_length: int,
-    mask_type: str = "static",
-    mask_other: float = 0.0,
-    min_masks: int = 0,
-    no_overlap: bool = False,
-    min_space: int = 0,
-) -> np.ndarray:
-    """
-    Computes random mask spans for a given shape
-
-    Args:
-        shape: the the shape for which to compute masks.
-            should be of size 2 where first element is batch size and 2nd is timesteps
-        padding_mask: optional padding mask of the same size as shape, which will prevent masking padded elements
-        mask_prob: probability for each token to be chosen as start of the span to be masked. this will be multiplied by
-            number of timesteps divided by length of mask span to mask approximately this percentage of all elements.
-            however due to overlaps, the actual number will be smaller (unless no_overlap is True)
-        mask_type: how to compute mask lengths
-            static = fixed size
-            uniform = sample from uniform distribution [mask_other, mask_length*2]
-            normal = sample from normal distribution with mean mask_length and stdev mask_other. mask is min 1 element
-            poisson = sample from possion distribution with lambda = mask length
-        min_masks: minimum number of masked spans
-        no_overlap: if false, will switch to an alternative recursive algorithm that prevents spans from overlapping
-        min_space: only used if no_overlap is True, this is how many elements to keep unmasked between spans
-    """
-
-    bsz, all_sz = shape
-    mask = np.full((bsz, all_sz), False)
-
-    all_num_mask = int(
-        # add a random number for probabilistic rounding
-        mask_prob * all_sz / float(mask_length)
-        + np.random.rand()
-    )
-
-    all_num_mask = max(min_masks, all_num_mask)
-
-    mask_idcs = []
-    for i in range(bsz):
-        if padding_mask is not None:
-            sz = all_sz - padding_mask[i].long().sum().item()
-            num_mask = int(
-                # add a random number for probabilistic rounding
-                mask_prob * sz / float(mask_length)
-                + np.random.rand()
-            )
-            num_mask = max(min_masks, num_mask)
-        else:
-            sz = all_sz
-            num_mask = all_num_mask
-
-        if mask_type == "static":
-            lengths = np.full(num_mask, mask_length)
-        elif mask_type == "uniform":
-            lengths = np.random.randint(mask_other, mask_length * 2 + 1, size=num_mask)
-        elif mask_type == "normal":
-            lengths = np.random.normal(mask_length, mask_other, size=num_mask)
-            lengths = [max(1, int(round(x))) for x in lengths]
-        elif mask_type == "poisson":
-            lengths = np.random.poisson(mask_length, size=num_mask)
-            lengths = [int(round(x)) for x in lengths]
-        else:
-            raise Exception("unknown mask selection " + mask_type)
-
-        if sum(lengths) == 0:
-            lengths[0] = min(mask_length, sz - 1)
-
-        if no_overlap:
-            mask_idc = []
-
-            def arrange(s, e, length, keep_length):
-                span_start = np.random.randint(s, e - length)
-                mask_idc.extend(span_start + i for i in range(length))
-
-                new_parts = []
-                if span_start - s - min_space >= keep_length:
-                    new_parts.append((s, span_start - min_space + 1))
-                if e - span_start - keep_length - min_space > keep_length:
-                    new_parts.append((span_start + length + min_space, e))
-                return new_parts
-
-            parts = [(0, sz)]
-            min_length = min(lengths)
-            for length in sorted(lengths, reverse=True):
-                lens = np.fromiter(
-                    (e - s if e - s >= length + min_space else 0 for s, e in parts),
-                    np.int,
-                )
-                l_sum = np.sum(lens)
-                if l_sum == 0:
-                    break
-                probs = lens / np.sum(lens)
-                c = np.random.choice(len(parts), p=probs)
-                s, e = parts.pop(c)
-                parts.extend(arrange(s, e, length, min_length))
-            mask_idc = np.asarray(mask_idc)
-        else:
-            min_len = min(lengths)
-            if sz - min_len <= num_mask:
-                min_len = sz - num_mask - 1
-
-            mask_idc = np.random.choice(sz - min_len, num_mask, replace=False)
-
-            mask_idc = np.asarray(
-                [
-                    mask_idc[j] + offset
-                    for j in range(len(mask_idc))
-                    for offset in range(lengths[j])
-                ]
-            )
-
-        mask_idcs.append(np.unique(mask_idc[mask_idc < sz]))
-
-    min_len = min([len(m) for m in mask_idcs])
-    for i, mask_idc in enumerate(mask_idcs):
-        if len(mask_idc) > min_len:
-            mask_idc = np.random.choice(mask_idc, min_len, replace=False)
-        mask[i, mask_idc] = True
-
-    return mask
-
-
-def get_mem_usage():
-    try:
-        import psutil
-
-        mb = 1024 * 1024
-        return f"used={psutil.virtual_memory().used / mb}Mb; avail={psutil.virtual_memory().available / mb}Mb"
-    except ImportError:
-        return "N/A"
-
-
-# lens: torch.LongTensor
-# returns: torch.BoolTensor
-def lengths_to_padding_mask(lens):
-    bsz, max_lens = lens.size(0), torch.max(lens).item()
-    mask = torch.arange(max_lens).to(lens.device).view(1, max_lens)
-    mask = mask.expand(bsz, -1) >= lens.view(bsz, 1).expand(-1, max_lens)
-    return mask
-
-
-# lens: torch.LongTensor
-# returns: torch.BoolTensor
-def lengths_to_mask(lens):
-    return ~lengths_to_padding_mask(lens)
-
-
-def get_buckets(sizes, num_buckets):
-    buckets = np.unique(
-        np.percentile(
-            sizes,
-            np.linspace(0, 100, num_buckets + 1),
-            interpolation='lower',
-        )[1:]
-    )
-    return buckets
-
-
-def get_bucketed_sizes(orig_sizes, buckets):
-    sizes = np.copy(orig_sizes)
-    assert np.min(sizes) >= 0
-    start_val = -1
-    for end_val in buckets:
-        mask = (sizes > start_val) & (sizes <= end_val)
-        sizes[mask] = end_val
-        start_val = end_val
-    return sizes
-
-
-
-def _find_extra_valid_paths(dataset_path: str) -> set:
-    paths = utils.split_paths(dataset_path)
-    all_valid_paths = set()
-    for sub_dir in paths:
-        contents = PathManager.ls(sub_dir)
-        valid_paths = [c for c in contents if re.match("valid*[0-9].*", c) is not None]
-        all_valid_paths |= {os.path.basename(p) for p in valid_paths}
-    # Remove .bin, .idx etc
-    roots = {os.path.splitext(p)[0] for p in all_valid_paths}
-    return roots
-
-
-def raise_if_valid_subsets_unintentionally_ignored(train_cfg) -> None:
-    """Raises if there are paths matching 'valid*[0-9].*' which are not combined or ignored."""
-    if (
-        train_cfg.dataset.ignore_unused_valid_subsets
-        or train_cfg.dataset.combine_valid_subsets
-        or train_cfg.dataset.disable_validation
-        or not hasattr(train_cfg.task, "data")
-    ):
-        return
-    other_paths = _find_extra_valid_paths(train_cfg.task.data)
-    specified_subsets = train_cfg.dataset.valid_subset.split(",")
-    ignored_paths = [p for p in other_paths if p not in specified_subsets]
-    if ignored_paths:
-        advice = "Set --combine-val to combine them or --ignore-unused-valid-subsets to ignore them."
-        msg = f"Valid paths {ignored_paths} will be ignored. {advice}"
-        raise ValueError(msg)

data/file_dataset.py

@@ -1,107 +0,0 @@
-# Copyright 2022 The OFA-Sys Team. 
-# All rights reserved.
-# This source code is licensed under the Apache 2.0 license 
-# found in the LICENSE file in the root directory.
-
-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

data/mm_data/caption_dataset.py

@@ -1,160 +0,0 @@
-# Copyright 2022 The OFA-Sys Team. 
-# All rights reserved.
-# This source code is licensed under the Apache 2.0 license 
-# found in the LICENSE file in the root directory.
-
-from io import BytesIO
-
-import logging
-import warnings
-import string
-
-import numpy as np
-import torch
-import base64
-from torchvision import transforms
-
-from PIL import Image, ImageFile
-
-from data import data_utils
-from data.ofa_dataset import OFADataset
-
-ImageFile.LOAD_TRUNCATED_IMAGES = True
-ImageFile.MAX_IMAGE_PIXELS = None
-Image.MAX_IMAGE_PIXELS = None
-
-logger = logging.getLogger(__name__)
-warnings.filterwarnings("ignore", "(Possibly )?corrupt EXIF data", UserWarning)
-
-IMAGENET_DEFAULT_MEAN = (0.485, 0.456, 0.406)
-IMAGENET_DEFAULT_STD = (0.229, 0.224, 0.225)
-
-
-def collate(samples, pad_idx, eos_idx):
-    if len(samples) == 0:
-        return {}
-
-    def merge(key):
-        return data_utils.collate_tokens(
-            [s[key] for s in samples],
-            pad_idx,
-            eos_idx=eos_idx,
-        )
-
-    id = np.array([s["id"] for s in samples])
-    src_tokens = merge("source")
-    src_lengths = torch.LongTensor([s["source"].ne(pad_idx).long().sum() for s in samples])
-
-    patch_images = torch.stack([sample['patch_image'] for sample in samples], dim=0)
-    patch_masks = torch.cat([sample['patch_mask'] for sample in samples])
-
-    prev_output_tokens = None
-    target = None
-    if samples[0].get("target", None) is not None:
-        target = merge("target")
-        tgt_lengths = torch.LongTensor([s["target"].ne(pad_idx).long().sum() for s in samples])
-        ntokens = tgt_lengths.sum().item()
-
-        if samples[0].get("prev_output_tokens", None) is not None:
-            prev_output_tokens = merge("prev_output_tokens")
-    else:
-        ntokens = src_lengths.sum().item()
-
-    batch = {
-        "id": id,
-        "nsentences": len(samples),
-        "ntokens": ntokens,
-        "net_input": {
-            "src_tokens": src_tokens,
-            "src_lengths": src_lengths,
-            "patch_images": patch_images,
-            "patch_masks": patch_masks,
-            "prev_output_tokens": prev_output_tokens
-        },
-        "target": target,
-    }
-
-    return batch
-
-
-class CaptionDataset(OFADataset):
-    def __init__(
-        self,
-        split,
-        dataset,
-        bpe,
-        src_dict,
-        tgt_dict=None,
-        max_src_length=128,
-        max_tgt_length=30,
-        patch_image_size=224,
-        imagenet_default_mean_and_std=False,
-        scst=False
-    ):
-        super().__init__(split, dataset, bpe, src_dict, tgt_dict)
-        self.max_src_length = max_src_length
-        self.max_tgt_length = max_tgt_length
-        self.patch_image_size = patch_image_size
-        self.scst = scst
-
-        self.transtab = str.maketrans({key: None for key in string.punctuation})
-
-        if imagenet_default_mean_and_std:
-            mean = IMAGENET_DEFAULT_MEAN
-            std = IMAGENET_DEFAULT_STD
-        else:
-            mean = [0.5, 0.5, 0.5]
-            std = [0.5, 0.5, 0.5]
-
-        self.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),
-        ])
-
-        if type(bpe).__name__ == 'GPT2BPE':
-            self.prompt = " what does the image describe?"
-        elif type(bpe).__name__ == 'BertBPE':
-            self.prompt = "图片描述了什么内容?"
-
-    def __getitem__(self, index):
-        uniq_id, image, caption = self.dataset[index]
-
-        image = Image.open(BytesIO(base64.urlsafe_b64decode(image)))
-        patch_image = self.patch_resize_transform(image)
-        patch_mask = torch.tensor([True])
-
-        if self.split == 'train' and not self.scst:
-            caption = caption.translate(self.transtab).strip()
-            caption_token_list = caption.strip().split()
-            tgt_caption = ' '.join(caption_token_list[:self.max_tgt_length])
-        else:
-            caption = ' '.join(caption.strip().split())
-            caption_list = [cap.translate(self.transtab).strip() for cap in caption.strip().split('&&')]
-            tgt_caption = '&&'.join(caption_list)
-        src_item = self.encode_text(self.prompt)
-        tgt_item = self.encode_text(" {}".format(tgt_caption))
-
-        src_item = torch.cat([self.bos_item, src_item, self.eos_item])
-        target_item = torch.cat([tgt_item, self.eos_item])
-        prev_output_item = torch.cat([self.bos_item, tgt_item])
-
-        example = {
-            "id": uniq_id,
-            "source": src_item,
-            "patch_image": patch_image,
-            "patch_mask": patch_mask,
-            "target": target_item,
-            "prev_output_tokens": prev_output_item
-        }
-        return example
-
-    def collater(self, samples, pad_to_length=None):
-        """Merge a list of samples to form a mini-batch.
-        Args:
-            samples (List[dict]): samples to collate
-        Returns:
-            dict: a mini-batch containing the data of the task
-        """
-        return collate(samples, pad_idx=self.pad, eos_idx=self.eos)

data/mm_data/ocr_dataset.py

@@ -1,210 +0,0 @@
-# Copyright 2022 The OFA-Sys Team.
-# All rights reserved.
-# This source code is licensed under the Apache 2.0 license
-# found in the LICENSE file in the root directory.
-
-from io import BytesIO
-
-import logging
-import warnings
-import random
-import functools
-
-import torch
-import base64
-from torchvision import transforms
-from torchvision.transforms import InterpolationMode
-from torchvision.transforms import functional as F
-
-from PIL import Image, ImageFile
-
-from zhconv import convert
-import unicodedata
-
-from data import data_utils
-from data.ofa_dataset import OFADataset
-
-ImageFile.LOAD_TRUNCATED_IMAGES = True
-ImageFile.MAX_IMAGE_PIXELS = None
-Image.MAX_IMAGE_PIXELS = None
-
-logger = logging.getLogger(__name__)
-warnings.filterwarnings("ignore", "(Possibly )?corrupt EXIF data", UserWarning)
-
-IMAGENET_DEFAULT_MEAN = (0.485, 0.456, 0.406)
-IMAGENET_DEFAULT_STD = (0.229, 0.224, 0.225)
-
-
-def collate(samples, pad_idx, eos_idx):
-    if len(samples) == 0:
-        return {}
-
-    def merge(key):
-        return data_utils.collate_tokens(
-            [s[key] for s in samples],
-            pad_idx,
-            eos_idx=eos_idx,
-        )
-
-    id = np.array([s["id"] for s in samples])
-    src_tokens = merge("source")
-    src_lengths = torch.LongTensor([s["source"].ne(pad_idx).long().sum() for s in samples])
-
-    patch_images = torch.stack([sample['patch_image'] for sample in samples], dim=0)
-    patch_masks = torch.cat([sample['patch_mask'] for sample in samples])
-
-    prev_output_tokens = None
-    target = None
-    if samples[0].get("target", None) is not None:
-        target = merge("target")
-        tgt_lengths = torch.LongTensor([s["target"].ne(pad_idx).long().sum() for s in samples])
-        ntokens = tgt_lengths.sum().item()
-
-        if samples[0].get("prev_output_tokens", None) is not None:
-            prev_output_tokens = merge("prev_output_tokens")
-    else:
-        ntokens = src_lengths.sum().item()
-
-    batch = {
-        "id": id,
-        "nsentences": len(samples),
-        "ntokens": ntokens,
-        "net_input": {
-            "src_tokens": src_tokens,
-            "src_lengths": src_lengths,
-            "patch_images": patch_images,
-            "patch_masks": patch_masks,
-            "prev_output_tokens": prev_output_tokens
-        },
-        "target": target,
-    }
-
-    return batch
-
-
-def ocr_resize(img, patch_image_size, is_document=False, split='train'):
-    img = img.convert("RGB")
-    width, height = img.size
-
-    if is_document:
-        new_height, new_width = 64, 1920
-    else:
-        if width >= height:
-            new_width = max(64, patch_image_size)
-            new_height = max(64, int(patch_image_size * (height / width)))
-            if split != 'train':
-                top = int((patch_image_size - new_height) // 2)
-            else:
-                top = random.randint(0, patch_image_size - new_height)
-            bottom = patch_image_size - new_height - top
-            left, right = 0, 0
-        else:
-            new_height = max(64, patch_image_size)
-            new_width = max(64, int(patch_image_size * (width / height)))
-            if split != 'train':
-                left = int((patch_image_size - new_width) // 2)
-            else:
-                left = random.randint(0, patch_image_size - new_width)
-            right = patch_image_size - new_width - left
-            top, bottom = 0, 0
-
-    img_new = F.resize(
-        img,
-        [new_height, new_width],
-        interpolation=InterpolationMode.BICUBIC,
-    )
-
-    if is_document:
-        img_split = transforms.ToTensor()(img_new).chunk(4, dim=-1)
-        img_new = transforms.ToPILImage()(torch.cat(img_split, dim=-2))
-        new_width, new_height = img_new.size
-        top = random.randint(0, patch_image_size - new_height)
-        bottom = patch_image_size - new_height - top
-        left, right = 0, 0
-
-    img_new = F.pad(img_new, padding=[left, top, right, bottom], padding_mode="edge")
-    assert img_new.size == (patch_image_size, patch_image_size)
-
-    return img_new
-
-
-class OcrDataset(OFADataset):
-    def __init__(
-        self,
-        split,
-        dataset,
-        bpe,
-        src_dict,
-        tgt_dict=None,
-        max_src_length=80,
-        max_tgt_length=30,
-        patch_image_size=224,
-        imagenet_default_mean_and_std=False,
-        is_document=False,
-    ):
-        super().__init__(split, dataset, bpe, src_dict, tgt_dict)
-        self.max_src_length = max_src_length
-        self.max_tgt_length = max_tgt_length
-        self.patch_image_size = patch_image_size
-
-        if imagenet_default_mean_and_std:
-            mean = IMAGENET_DEFAULT_MEAN
-            std = IMAGENET_DEFAULT_STD
-        else:
-            mean = [0.5, 0.5, 0.5]
-            std = [0.5, 0.5, 0.5]
-
-        self.patch_resize_transform = transforms.Compose(
-            [
-                lambda image: ocr_resize(
-                    image, patch_image_size, is_document=is_document, split=split,
-                ),
-                transforms.ToTensor(),
-                transforms.Normalize(mean=mean, std=std),
-            ]
-        )
-
-        self.bpe = bpe
-        if type(bpe).__name__ == 'GPT2BPE':
-            self.prompt = " what are the texts on the image?"
-        elif type(bpe).__name__ == 'BertBPE':
-            self.prompt = "图片上的文字是什么?"
-
-    def __getitem__(self, index):
-        uniq_id, image, caption = self.dataset[index]
-
-        image = Image.open(BytesIO(base64.urlsafe_b64decode(image)))
-        patch_image = self.patch_resize_transform(image)
-        patch_mask = torch.tensor([True])
-
-        caption = unicodedata.normalize("NFKC", convert(caption, "zh-hans"))
-        if type(self.bpe).__name__ == 'GPT2BPE':
-            caption_token_list = caption.lower().strip().split()
-            tgt_caption = ' '.join(caption_token_list[:self.max_tgt_length])
-        elif type(self.bpe).__name__ == 'BertBPE':
-            tgt_caption = caption[: self.max_tgt_length].lower()
-        src_item = self.encode_text(self.prompt)
-        tgt_item = self.encode_text(" {}".format(tgt_caption))
-
-        src_item = torch.cat([self.bos_item, src_item, self.eos_item])
-        target_item = torch.cat([tgt_item, self.eos_item])
-        prev_output_item = torch.cat([self.bos_item, tgt_item])
-
-        example = {
-            "id": uniq_id,
-            "source": src_item,
-            "patch_image": patch_image,
-            "patch_mask": patch_mask,
-            "target": target_item,
-            "prev_output_tokens": prev_output_item,
-        }
-        return example
-
-    def collater(self, samples, pad_to_length=None):
-        """Merge a list of samples to form a mini-batch.
-        Args:
-            samples (List[dict]): samples to collate
-        Returns:
-            dict: a mini-batch containing the data required for the task
-        """
-        return collate(samples, pad_idx=self.pad, eos_idx=self.eos)

data/mm_data/refcoco_dataset.py

@@ -1,174 +0,0 @@
-# Copyright 2022 The OFA-Sys Team. 
-# All rights reserved.
-# This source code is licensed under the Apache 2.0 license 
-# found in the LICENSE file in the root directory.
-
-from io import BytesIO
-
-import logging
-import warnings
-
-import numpy as np
-import torch
-import base64
-import utils.transforms as T
-
-from PIL import Image, ImageFile
-
-from data import data_utils
-from data.ofa_dataset import OFADataset
-
-ImageFile.LOAD_TRUNCATED_IMAGES = True
-ImageFile.MAX_IMAGE_PIXELS = None
-Image.MAX_IMAGE_PIXELS = None
-
-logger = logging.getLogger(__name__)
-warnings.filterwarnings("ignore", "(Possibly )?corrupt EXIF data", UserWarning)
-
-IMAGENET_DEFAULT_MEAN = (0.485, 0.456, 0.406)
-IMAGENET_DEFAULT_STD = (0.229, 0.224, 0.225)
-
-
-def collate(samples, pad_idx, eos_idx):
-    if len(samples) == 0:
-        return {}
-
-    def merge(key):
-        return data_utils.collate_tokens(
-            [s[key] for s in samples],
-            pad_idx,
-            eos_idx=eos_idx,
-        )
-
-    id = np.array([s["id"] for s in samples])
-    src_tokens = merge("source")
-    src_lengths = torch.LongTensor([s["source"].ne(pad_idx).long().sum() for s in samples])
-
-    patch_images = torch.stack([sample['patch_image'] for sample in samples], dim=0)
-    patch_masks = torch.cat([sample['patch_mask'] for sample in samples])
-
-    w_resize_ratios = torch.stack([s["w_resize_ratio"] for s in samples], dim=0)
-    h_resize_ratios = torch.stack([s["h_resize_ratio"] for s in samples], dim=0)
-    region_coords = torch.stack([s['region_coord'] for s in samples], dim=0)
-
-    prev_output_tokens = None
-    target = None
-    if samples[0].get("target", None) is not None:
-        target = merge("target")
-        tgt_lengths = torch.LongTensor([s["target"].ne(pad_idx).long().sum() for s in samples])
-        ntokens = tgt_lengths.sum().item()
-
-        if samples[0].get("prev_output_tokens", None) is not None:
-            prev_output_tokens = merge("prev_output_tokens")
-    else:
-        ntokens = src_lengths.sum().item()
-
-    batch = {
-        "id": id,
-        "nsentences": len(samples),
-        "ntokens": ntokens,
-        "net_input": {
-            "src_tokens": src_tokens,
-            "src_lengths": src_lengths,
-            "patch_images": patch_images,
-            "patch_masks": patch_masks,
-            "prev_output_tokens": prev_output_tokens
-        },
-        "target": target,
-        "w_resize_ratios": w_resize_ratios,
-        "h_resize_ratios": h_resize_ratios,
-        "region_coords": region_coords
-    }
-
-    return batch
-
-
-class RefcocoDataset(OFADataset):
-    def __init__(
-        self,
-        split,
-        dataset,
-        bpe,
-        src_dict,
-        tgt_dict=None,
-        max_src_length=80,
-        max_tgt_length=30,
-        patch_image_size=512,
-        imagenet_default_mean_and_std=False,
-        num_bins=1000,
-        max_image_size=512
-    ):
-        super().__init__(split, dataset, bpe, src_dict, tgt_dict)
-        self.max_src_length = max_src_length
-        self.max_tgt_length = max_tgt_length
-        self.patch_image_size = patch_image_size
-        self.num_bins = num_bins
-
-        if imagenet_default_mean_and_std:
-            mean = IMAGENET_DEFAULT_MEAN
-            std = IMAGENET_DEFAULT_STD
-        else:
-            mean = [0.5, 0.5, 0.5]
-            std = [0.5, 0.5, 0.5]
-
-        # for positioning
-        self.positioning_transform = T.Compose([
-            T.RandomResize([patch_image_size], max_size=patch_image_size),
-            T.ToTensor(),
-            T.Normalize(mean=mean, std=std, max_image_size=max_image_size)
-        ])
-
-        if type(bpe).__name__ == 'GPT2BPE':
-            self.prompt = ' which region does the text " {} " describe?'
-        elif type(bpe).__name__ == 'BertBPE':
-            self.prompt = '这段文字" {} "描述的是哪个区域？'
-
-    def __getitem__(self, index):
-        uniq_id, base64_str, text, region_coord = self.dataset[index]
-
-        image = Image.open(BytesIO(base64.urlsafe_b64decode(base64_str))).convert("RGB")
-        w, h = image.size
-        boxes_target = {"boxes": [], "labels": [], "area": [], "size": torch.tensor([h, w])}
-        x0, y0, x1, y1 = region_coord.strip().split(',')
-        region = torch.tensor([float(x0), float(y0), float(x1), float(y1)])
-        boxes_target["boxes"] = torch.tensor([[float(x0), float(y0), float(x1), float(y1)]])
-        boxes_target["labels"] = np.array([0])
-        boxes_target["area"] = torch.tensor([(float(x1) - float(x0)) * (float(y1) - float(y0))])
-
-        patch_image, patch_boxes = self.positioning_transform(image, boxes_target)
-        resize_h, resize_w = patch_boxes["size"][0], patch_boxes["size"][1]
-        patch_mask = torch.tensor([True])
-        quant_x0 = "<bin_{}>".format(int((patch_boxes["boxes"][0][0] * (self.num_bins - 1)).round()))
-        quant_y0 = "<bin_{}>".format(int((patch_boxes["boxes"][0][1] * (self.num_bins - 1)).round()))
-        quant_x1 = "<bin_{}>".format(int((patch_boxes["boxes"][0][2] * (self.num_bins - 1)).round()))
-        quant_y1 = "<bin_{}>".format(int((patch_boxes["boxes"][0][3] * (self.num_bins - 1)).round()))
-        region_coord = "{} {} {} {}".format(quant_x0, quant_y0, quant_x1, quant_y1)
-        src_caption = self.pre_caption(text, self.max_src_length)
-        src_item = self.encode_text(self.prompt.format(src_caption))
-        tgt_item = self.encode_text(region_coord, use_bpe=False)
-
-        src_item = torch.cat([self.bos_item, src_item, self.eos_item])
-        target_item = torch.cat([tgt_item, self.eos_item])
-        prev_output_item = torch.cat([self.bos_item, tgt_item])
-
-        example = {
-            "id": uniq_id,
-            "source": src_item,
-            "patch_image": patch_image,
-            "patch_mask": patch_mask,
-            "target": target_item,
-            "prev_output_tokens": prev_output_item,
-            "w_resize_ratio": resize_w / w,
-            "h_resize_ratio": resize_h / h,
-            "region_coord": region
-        }
-        return example
-
-    def collater(self, samples, pad_to_length=None):
-        """Merge a list of samples to form a mini-batch.
-        Args:
-            samples (List[dict]): samples to collate
-        Returns:
-            dict: a mini-batch containing the data of the task
-        """
-        return collate(samples, pad_idx=self.pad, eos_idx=self.eos)

data/mm_data/snli_ve_dataset.py

@@ -1,203 +0,0 @@
-# Copyright 2022 The OFA-Sys Team. 
-# All rights reserved.
-# This source code is licensed under the Apache 2.0 license 
-# found in the LICENSE file in the root directory.
-
-from io import BytesIO
-
-import logging
-import warnings
-
-import numpy as np
-import torch
-import base64
-from torchvision import transforms
-
-from PIL import Image, ImageFile
-
-from data import data_utils
-from data.ofa_dataset import OFADataset
-
-ImageFile.LOAD_TRUNCATED_IMAGES = True
-ImageFile.MAX_IMAGE_PIXELS = None
-Image.MAX_IMAGE_PIXELS = None
-
-logger = logging.getLogger(__name__)
-warnings.filterwarnings("ignore", "(Possibly )?corrupt EXIF data", UserWarning)
-
-IMAGENET_DEFAULT_MEAN = (0.485, 0.456, 0.406)
-IMAGENET_DEFAULT_STD = (0.229, 0.224, 0.225)
-
-
-def collate(samples, pad_idx, eos_idx):
-    if len(samples) == 0:
-        return {}
-
-    def merge(key):
-        return data_utils.collate_tokens(
-            [s[key] for s in samples],
-            pad_idx,
-            eos_idx=eos_idx,
-        )
-
-    id = np.array([s["id"] for s in samples])
-    src_tokens = merge("source")
-    src_lengths = torch.LongTensor([s["source"].ne(pad_idx).long().sum() for s in samples])
-
-    patch_images = torch.stack([sample['patch_image'] for sample in samples], dim=0)
-    patch_masks = torch.cat([sample['patch_mask'] for sample in samples])
-
-    ref_dict = None
-    if samples[0].get("ref_dict", None) is not None:
-        ref_dict = np.array([s['ref_dict'] for s in samples])
-
-    constraint_masks = None
-    if samples[0].get("constraint_mask", None) is not None:
-        constraint_masks = merge("constraint_mask")
-
-    decoder_prompts = None
-    if samples[0].get("decoder_prompt", None) is not None:
-        decoder_prompts = np.array([s['decoder_prompt'].tolist() for s in samples])
-
-    prev_output_tokens = None
-    target = None
-    if samples[0].get("target", None) is not None:
-        target = merge("target")
-        tgt_lengths = torch.LongTensor(
-            [s["target"].ne(pad_idx).long().sum() for s in samples]
-        )
-        ntokens = tgt_lengths.sum().item()
-
-        if samples[0].get("prev_output_tokens", None) is not None:
-            prev_output_tokens = merge("prev_output_tokens")
-    else:
-        ntokens = src_lengths.sum().item()
-
-    batch = {
-        "id": id,
-        "nsentences": len(samples),
-        "ntokens": ntokens,
-        "net_input": {
-            "src_tokens": src_tokens,
-            "src_lengths": src_lengths,
-            "patch_images": patch_images,
-            "patch_masks": patch_masks,
-            "prev_output_tokens": prev_output_tokens
-        },
-        "ref_dict": ref_dict,
-        "constraint_masks": constraint_masks,
-        "decoder_prompts": decoder_prompts,
-        "target": target
-    }
-
-    return batch
-
-
-class SnliVeDataset(OFADataset):
-    def __init__(
-        self,
-        split,
-        dataset,
-        bpe,
-        src_dict,
-        tgt_dict=None,
-        max_src_length=80,
-        max_tgt_length=30,
-        patch_image_size=224,
-        add_caption=False,
-        constraint_trie=None,
-        imagenet_default_mean_and_std=False,
-        prompt_type="none"
-    ):
-        super().__init__(split, dataset, bpe, src_dict, tgt_dict)
-        self.max_src_length = max_src_length
-        self.max_tgt_length = max_tgt_length
-        self.patch_image_size = patch_image_size
-
-        self.add_caption = add_caption
-        self.constraint_trie = constraint_trie
-        self.prompt_type = prompt_type
-
-        if imagenet_default_mean_and_std:
-            mean = IMAGENET_DEFAULT_MEAN
-            std = IMAGENET_DEFAULT_STD
-        else:
-            mean = [0.5, 0.5, 0.5]
-            std = [0.5, 0.5, 0.5]
-
-        self.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),
-        ])
-
-    def __getitem__(self, index):
-        uniq_id, image, hypothesis, caption, label = self.dataset[index]
-        if label == 'contradiction':
-            label = 'no'
-        elif label == 'entailment':
-            label = 'yes'
-        elif label == 'neutral':
-            label = 'maybe'
-        else:
-            raise NotImplementedError
-
-        image = Image.open(BytesIO(base64.urlsafe_b64decode(image)))
-        patch_image = self.patch_resize_transform(image)
-        patch_mask = torch.tensor([True])
-
-        hypothesis = self.pre_caption(hypothesis, self.max_src_length)
-        src_item = self.encode_text(' does the image describe " {} "?'.format(hypothesis))
-        tgt_item = self.encode_text(" {}".format(label))
-        ref_dict = {label: 1.0}
-
-        if self.add_caption:
-            caption = self.pre_caption(caption, self.max_src_length)
-            src_item = self.encode_text(' can image and text1 " {} " imply text2 " {} "?'.format(caption, hypothesis))
-
-        src_item = torch.cat([self.bos_item, src_item, self.eos_item])
-        if self.prompt_type == 'none':
-            prev_output_item = torch.cat([self.bos_item, tgt_item])
-            target_item = torch.cat([prev_output_item[1:], self.eos_item])
-            decoder_prompt = self.bos_item
-        elif self.prompt_type == 'src':
-            prev_output_item = torch.cat([src_item, tgt_item])
-            target_item = torch.cat([prev_output_item[1:], self.eos_item])
-            decoder_prompt = src_item
-        elif self.prompt_type == 'prev_output':
-            prev_output_item = torch.cat([src_item[:-1], tgt_item])
-            target_item = torch.cat([prev_output_item[1:], self.eos_item])
-            decoder_prompt = src_item[:-1]
-        else:
-            raise NotImplementedError
-        target_item[:-len(tgt_item)-1] = self.tgt_dict.pad()
-
-        example = {
-            "id": uniq_id,
-            "source": src_item,
-            "patch_image": patch_image,
-            "patch_mask": patch_mask,
-            "target": target_item,
-            "prev_output_tokens": prev_output_item,
-            "decoder_prompt": decoder_prompt,
-            "ref_dict": ref_dict,
-        }
-        if self.constraint_trie is not None:
-            constraint_mask = torch.zeros((len(target_item), len(self.tgt_dict))).bool()
-            start_idx = len(target_item) - len(tgt_item) - 1
-            for i in range(len(target_item)-len(tgt_item)-1, len(target_item)):
-                constraint_prefix_token = [self.tgt_dict.bos()] + target_item[start_idx:i].tolist()
-                constraint_nodes = self.constraint_trie.get_next_layer(constraint_prefix_token)
-                constraint_mask[i][constraint_nodes] = True
-            example["constraint_mask"] = constraint_mask
-        return example
-
-    def collater(self, samples, pad_to_length=None):
-        """Merge a list of samples to form a mini-batch.
-        Args:
-            samples (List[dict]): samples to collate
-        Returns:
-            dict: a mini-batch containing the data of the task
-        """
-        return collate(samples, pad_idx=self.pad, eos_idx=self.eos)

data/mm_data/vqa_gen_dataset.py

@@ -1,218 +0,0 @@
-# Copyright 2022 The OFA-Sys Team. 
-# All rights reserved.
-# This source code is licensed under the Apache 2.0 license 
-# found in the LICENSE file in the root directory.
-
-from io import BytesIO
-
-import logging
-import warnings
-
-import numpy as np
-import torch
-import base64
-from torchvision import transforms
-
-from PIL import Image, ImageFile
-
-from data import data_utils
-from data.ofa_dataset import OFADataset
-
-ImageFile.LOAD_TRUNCATED_IMAGES = True
-ImageFile.MAX_IMAGE_PIXELS = None
-Image.MAX_IMAGE_PIXELS = None
-
-logger = logging.getLogger(__name__)
-warnings.filterwarnings("ignore", "(Possibly )?corrupt EXIF data", UserWarning)
-
-IMAGENET_DEFAULT_MEAN = (0.485, 0.456, 0.406)
-IMAGENET_DEFAULT_STD = (0.229, 0.224, 0.225)
-
-
-def collate(samples, pad_idx, eos_idx):
-    if len(samples) == 0:
-        return {}
-
-    def merge(key):
-        return data_utils.collate_tokens(
-            [s[key] for s in samples],
-            pad_idx,
-            eos_idx=eos_idx,
-        )
-
-    id = np.array([s["id"] for s in samples])
-    src_tokens = merge("source")
-    src_lengths = torch.LongTensor([s["source"].ne(pad_idx).long().sum() for s in samples])
-
-    patch_images = torch.stack([sample['patch_image'] for sample in samples], dim=0)
-    patch_masks = torch.cat([sample['patch_mask'] for sample in samples])
-
-    conf = None
-    if samples[0].get("conf", None) is not None:
-        conf = torch.cat([s['conf'] for s in samples], dim=0)
-
-    ref_dict = None
-    if samples[0].get("ref_dict", None) is not None:
-        ref_dict = np.array([s['ref_dict'] for s in samples])
-
-    constraint_masks = None
-    if samples[0].get("constraint_mask", None) is not None:
-        constraint_masks = merge("constraint_mask")
-
-    decoder_prompts = None
-    if samples[0].get("decoder_prompt", None) is not None:
-        decoder_prompts = np.array([s['decoder_prompt'].tolist() for s in samples])
-
-    prefix_tokens = None
-    if samples[0].get("decoder_prompt", None) is not None:
-        prefix_tokens = merge("decoder_prompt")
-        prefix_tokens = prefix_tokens[:, 1:]
-
-    prev_output_tokens = None
-    target = None
-    if samples[0].get("target", None) is not None:
-        target = merge("target")
-        tgt_lengths = torch.LongTensor(
-            [s["target"].ne(pad_idx).long().sum() for s in samples]
-        )
-        ntokens = tgt_lengths.sum().item()
-
-        if samples[0].get("prev_output_tokens", None) is not None:
-            prev_output_tokens = merge("prev_output_tokens")
-    else:
-        ntokens = src_lengths.sum().item()
-
-    batch = {
-        "id": id,
-        "nsentences": len(samples),
-        "ntokens": ntokens,
-        "net_input": {
-            "src_tokens": src_tokens,
-            "src_lengths": src_lengths,
-            "patch_images": patch_images,
-            "patch_masks": patch_masks,
-            "prev_output_tokens": prev_output_tokens
-        },
-        "conf": conf,
-        "ref_dict": ref_dict,
-        "constraint_masks": constraint_masks,
-        "decoder_prompts": decoder_prompts,
-        "target": target,
-        "prefix_tokens": prefix_tokens
-    }
-
-    return batch
-
-
-class VqaGenDataset(OFADataset):
-    def __init__(
-        self,
-        split,
-        dataset,
-        bpe,
-        src_dict,
-        tgt_dict=None,
-        max_src_length=128,
-        max_object_length=30,
-        max_tgt_length=30,
-        patch_image_size=224,
-        add_object=False,
-        constraint_trie=None,
-        imagenet_default_mean_and_std=False,
-        prompt_type="none"
-    ):
-        super().__init__(split, dataset, bpe, src_dict, tgt_dict)
-        self.max_src_length = max_src_length
-        self.max_object_length = max_object_length
-        self.max_tgt_length = max_tgt_length
-        self.patch_image_size = patch_image_size
-
-        self.add_object = add_object
-        self.constraint_trie = constraint_trie
-        self.prompt_type = prompt_type
-
-        if imagenet_default_mean_and_std:
-            mean = IMAGENET_DEFAULT_MEAN
-            std = IMAGENET_DEFAULT_STD
-        else:
-            mean = [0.5, 0.5, 0.5]
-            std = [0.5, 0.5, 0.5]
-
-        self.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),
-        ])
-
-    def __getitem__(self, index):
-        item = self.dataset[index]
-        if len(item) == 5:
-            uniq_id, image, question, ref, predict_objects = item
-        else:
-            uniq_id, image, question, ref, predict_objects, caption = item
-
-        image = Image.open(BytesIO(base64.urlsafe_b64decode(image)))
-        patch_image = self.patch_resize_transform(image)
-        patch_mask = torch.tensor([True])
-
-        question = self.pre_question(question, self.max_src_length)
-        question = question + '?' if not question.endswith('?') else question
-        src_item = self.encode_text(' {}'.format(question))
-
-        ref_dict = {item.split('|!+')[1]: float(item.split('|!+')[0]) for item in ref.split('&&')}
-        answer = max(ref_dict, key=ref_dict.get)
-        conf = torch.tensor([ref_dict[answer]])
-        tgt_item = self.encode_text(" {}".format(answer))
-
-        if self.add_object and predict_objects is not None:
-            predict_object_seq = ' '.join(predict_objects.strip().split('&&')[:self.max_object_length])
-            predict_object_item = self.encode_text(" object: {}".format(predict_object_seq))
-            src_item = torch.cat([src_item, predict_object_item])
-
-        src_item = torch.cat([self.bos_item, src_item, self.eos_item])
-        if self.prompt_type == 'none':
-            prev_output_item = torch.cat([self.bos_item, tgt_item])
-            target_item = torch.cat([prev_output_item[1:], self.eos_item])
-            decoder_prompt = self.bos_item
-        elif self.prompt_type == 'src':
-            prev_output_item = torch.cat([src_item, tgt_item])
-            target_item = torch.cat([prev_output_item[1:], self.eos_item])
-            decoder_prompt = src_item
-        elif self.prompt_type == 'prev_output':
-            prev_output_item = torch.cat([src_item[:-1], tgt_item])
-            target_item = torch.cat([prev_output_item[1:], self.eos_item])
-            decoder_prompt = src_item[:-1]
-        else:
-            raise NotImplementedError
-        target_item[:-len(tgt_item)-1] = self.tgt_dict.pad()
-
-        example = {
-            "id": uniq_id,
-            "source": src_item,
-            "patch_image": patch_image,
-            "patch_mask": patch_mask,
-            "target": target_item,
-            "prev_output_tokens": prev_output_item,
-            "decoder_prompt": decoder_prompt,
-            "ref_dict": ref_dict,
-            "conf": conf,
-        }
-        if self.constraint_trie is not None:
-            constraint_mask = torch.zeros((len(target_item), len(self.tgt_dict))).bool()
-            start_idx = len(target_item) - len(tgt_item) - 1
-            for i in range(len(target_item)-len(tgt_item)-1, len(target_item)):
-                constraint_prefix_token = [self.tgt_dict.bos()] + target_item[start_idx:i].tolist()
-                constraint_nodes = self.constraint_trie.get_next_layer(constraint_prefix_token)
-                constraint_mask[i][constraint_nodes] = True
-            example["constraint_mask"] = constraint_mask
-        return example
-
-    def collater(self, samples, pad_to_length=None):
-        """Merge a list of samples to form a mini-batch.
-        Args:
-            samples (List[dict]): samples to collate
-        Returns:
-            dict: a mini-batch containing the data of the task
-        """
-        return collate(samples, pad_idx=self.pad, eos_idx=self.eos)

data/nlg_data/summary_dataset.py

@@ -1,131 +0,0 @@
-# Copyright 2022 The OFA-Sys Team.
-# All rights reserved.
-# This source code is licensed under the Apache 2.0 license
-# found in the LICENSE file in the root directory.
-
-import logging
-import warnings
-import torch
-import numpy as np
-
-from data import data_utils
-from data.ofa_dataset import OFADataset
-
-logger = logging.getLogger(__name__)
-warnings.filterwarnings("ignore", "(Possibly )?corrupt EXIF data", UserWarning)
-
-
-def collate(samples, pad_idx, eos_idx):
-    if len(samples) == 0:
-        return {}
-
-    def merge(key):
-        return data_utils.collate_tokens(
-            [s[key] for s in samples],
-            pad_idx,
-            eos_idx=eos_idx,
-        )
-
-    src_tokens = merge("source")
-    src_lengths = torch.LongTensor([s["source"].ne(pad_idx).long().sum() for s in samples])
-
-    prev_output_tokens = None
-    target = None
-    if samples[0].get("target", None) is not None:
-        target = merge("target")
-        tgt_lengths = torch.LongTensor(
-            [s["target"].ne(pad_idx).long().sum() for s in samples]
-        )
-        ntokens = tgt_lengths.sum().item()
-
-        if samples[0].get("prev_output_tokens", None) is not None:
-            prev_output_tokens = merge("prev_output_tokens")
-    else:
-        ntokens = src_lengths.sum().item()
-
-    target_strs = np.array([s["target_str"] for s in samples])
-
-    batch = {
-        "nsentences": len(samples),
-        "ntokens": ntokens,
-        "net_input": {
-            "src_tokens": src_tokens,
-            "src_lengths": src_lengths,
-            "prev_output_tokens": prev_output_tokens
-        },
-        "target": target,
-        "target_strs": target_strs
-    }
-
-    return batch
-
-
-class SummaryDataset(OFADataset):
-    def __init__(
-        self,
-        split,
-        dataset,
-        bpe,
-        src_dict,
-        tgt_dict=None,
-        code_dict_size=8192,
-        num_bins=1000,
-        max_src_length=512,
-        max_tgt_length=128,
-        noise_ratio=0.0
-    ):
-        super().__init__(split, dataset, bpe, src_dict, tgt_dict)
-        self.max_src_length = max_src_length
-        self.max_tgt_length = max_tgt_length
-        self.code_dict_size = code_dict_size
-        self.num_bins = num_bins
-        self.noise_ratio = noise_ratio
-
-        if type(bpe).__name__ == 'GPT2BPE':
-            self.prompt = ' what is the summary of article " {} "?'
-        elif type(bpe).__name__ == 'BertBPE':
-            self.prompt = "{} 请用一个句子简单总结上文："
-
-    def __getitem__(self, index):
-        source, target = self.dataset[index]
-        target_str = target.lower()
-
-        source = self.pre_caption(source, max_words=self.max_src_length)
-        target = self.pre_caption(target, max_words=self.max_tgt_length)
-        source = source.replace('<unk>', 'unk')
-        target = target.replace('<unk>', 'unk')
-
-        src_item = self.encode_text(
-            self.prompt.format(source),
-            length=self.max_src_length
-        )
-        tgt_item = self.encode_text('{}'.format(target))
-        noise_tgt_item = self.add_noise_to_tgt(tgt_item.clone(), self.noise_ratio)
-
-        src_item = torch.cat([self.bos_item, src_item, self.eos_item])
-        target_item = torch.cat([tgt_item, self.eos_item])
-        prev_output_item = torch.cat([self.bos_item, noise_tgt_item])
-
-        example = {
-            "source": src_item,
-            "target": target_item,
-            "prev_output_tokens": prev_output_item,
-            "target_str": target_str
-        }
-        return example
-
-    def add_noise_to_tgt(self, target, p):
-        noise_indices = torch.FloatTensor(target.size(0)).uniform_() < p
-        target[noise_indices] = torch.randint(
-            4, len(self.src_dict) - self.code_dict_size - self.num_bins, size=(noise_indices.sum(),)
-        )
-        return target
-
-    def collater(self, samples, pad_to_length=None):
-        """Merge a list of samples to form a mini-batch.
-        Args:
-            samples (List[dict]): samples to collate
-        Returns:
-            dict: a mini-batch containing the data of the task
-        """
-        return collate(samples, pad_idx=self.pad, eos_idx=self.eos)

data/nlu_data/cola_dataset.py

@@ -1,138 +0,0 @@
-# Copyright 2022 The OFA-Sys Team. 
-# All rights reserved.
-# This source code is licensed under the Apache 2.0 license 
-# found in the LICENSE file in the root directory.
-
-import logging
-import warnings
-import torch
-import numpy as np
-
-from data import data_utils
-from data.ofa_dataset import OFADataset
-
-logger = logging.getLogger(__name__)
-warnings.filterwarnings("ignore", "(Possibly )?corrupt EXIF data", UserWarning)
-
-
-def collate(samples, pad_idx, eos_idx):
-    if len(samples) == 0:
-        return {}
-
-    def merge(key):
-        return data_utils.collate_tokens(
-            [s[key] for s in samples],
-            pad_idx,
-            eos_idx=eos_idx,
-        )
-
-    src_tokens = merge("source")
-    src_lengths = torch.LongTensor([s["source"].ne(pad_idx).long().sum() for s in samples])
-
-    ref_dict = None
-    if samples[0].get("ref_dict", None) is not None:
-        ref_dict = np.array([s['ref_dict'] for s in samples])
-
-    constraint_masks = None
-    if samples[0].get("constraint_mask", None) is not None:
-        constraint_masks = merge("constraint_mask")
-
-    prev_output_tokens = None
-    target = None
-    if samples[0].get("target", None) is not None:
-        target = merge("target")
-        tgt_lengths = torch.LongTensor(
-            [s["target"].ne(pad_idx).long().sum() for s in samples]
-        )
-        ntokens = tgt_lengths.sum().item()
-
-        if samples[0].get("prev_output_tokens", None) is not None:
-            prev_output_tokens = merge("prev_output_tokens")
-    else:
-        ntokens = src_lengths.sum().item()
-
-    batch = {
-        "nsentences": len(samples),
-        "ntokens": ntokens,
-        "net_input": {
-            "src_tokens": src_tokens,
-            "src_lengths": src_lengths,
-            "prev_output_tokens": prev_output_tokens
-        },
-        "ref_dict": ref_dict,
-        "constraint_masks": constraint_masks,
-        "target": target,
-    }
-
-    return batch
-
-
-class COLADataset(OFADataset):
-    def __init__(
-        self,
-        split,
-        dataset,
-        bpe,
-        src_dict,
-        tgt_dict=None,
-        max_src_length=512,
-        max_tgt_length=30,
-        constraint_trie=None,
-        prompt_type="none"
-    ):
-        super().__init__(split, dataset, bpe, src_dict, tgt_dict)
-        self.max_src_length = max_src_length
-        self.max_tgt_length = max_tgt_length
-        self.constraint_trie = constraint_trie
-        self.prompt_type = prompt_type
-
-    def __getitem__(self, index):
-        sentence, label = self.dataset[index]
-        if label == '0':
-            label = 'no'
-        elif label == '1':
-            label = 'yes'
-        else:
-            raise NotImplementedError
-
-        sentence = ' '.join(sentence.lower().strip().split()[:self.max_src_length])
-        src_item = self.encode_text(' is the text " {} " grammatically correct?'.format(sentence))
-        tgt_item = self.encode_text(" {}".format(label))
-        assert tgt_item.size(0) == 1
-        ref_dict = {label: 1.0}
-
-        src_item = torch.cat([self.bos_item, src_item, self.eos_item])
-        if self.prompt_type == 'none':
-            prev_output_item = self.bos_item
-            target_item = tgt_item
-        elif self.prompt_type == 'src':
-            prev_output_item = src_item.clone()
-            target_item = torch.cat([prev_output_item[1:], tgt_item])
-        elif self.prompt_type == 'prev_output':
-            prev_output_item = src_item[:-1].clone()
-            target_item = torch.cat([prev_output_item[1:], tgt_item])
-        else:
-            raise NotImplementedError
-        target_item[:-1] = self.tgt_dict.pad()
-
-        example = {
-            "source": src_item,
-            "target": target_item,
-            "prev_output_tokens": prev_output_item,
-            "ref_dict": ref_dict,
-        }
-        if self.constraint_trie is not None:
-            constraint_mask = torch.zeros((len(prev_output_item), len(self.tgt_dict))).bool()
-            constraint_nodes = self.constraint_trie.get_next_layer(self.bos_item.tolist())
-            constraint_mask[-1][constraint_nodes] = True
-            example["constraint_mask"] = constraint_mask
-        return example
-
-    def collater(self, samples, pad_to_length=None):
-        """Merge a list of samples to form a mini-batch.
-        Args:
-            samples (List[dict]): samples to collate
-        Returns:
-            dict: a mini-batch containing the data of the task
-        """
-        return collate(samples, pad_idx=self.pad, eos_idx=self.eos)

data/nlu_data/mnli_dataset.py

@@ -1,143 +0,0 @@
-# Copyright 2022 The OFA-Sys Team. 
-# All rights reserved.
-# This source code is licensed under the Apache 2.0 license 
-# found in the LICENSE file in the root directory.
-
-import logging
-import warnings
-import torch
-import numpy as np
-
-from data import data_utils
-from data.ofa_dataset import OFADataset
-
-logger = logging.getLogger(__name__)
-warnings.filterwarnings("ignore", "(Possibly )?corrupt EXIF data", UserWarning)
-
-
-def collate(samples, pad_idx, eos_idx):
-    if len(samples) == 0:
-        return {}
-
-    def merge(key):
-        return data_utils.collate_tokens(
-            [s[key] for s in samples],
-            pad_idx,
-            eos_idx=eos_idx,
-        )
-
-    src_tokens = merge("source")
-    src_lengths = torch.LongTensor([s["source"].ne(pad_idx).long().sum() for s in samples])
-
-    ref_dict = None
-    if samples[0].get("ref_dict", None) is not None:
-        ref_dict = np.array([s['ref_dict'] for s in samples])
-
-    constraint_masks = None
-    if samples[0].get("constraint_mask", None) is not None:
-        constraint_masks = merge("constraint_mask")
-
-    prev_output_tokens = None
-    target = None
-    if samples[0].get("target", None) is not None:
-        target = merge("target")
-        tgt_lengths = torch.LongTensor(
-            [s["target"].ne(pad_idx).long().sum() for s in samples]
-        )
-        ntokens = tgt_lengths.sum().item()
-
-        if samples[0].get("prev_output_tokens", None) is not None:
-            prev_output_tokens = merge("prev_output_tokens")
-    else:
-        ntokens = src_lengths.sum().item()
-
-    batch = {
-        "nsentences": len(samples),
-        "ntokens": ntokens,
-        "net_input": {
-            "src_tokens": src_tokens,
-            "src_lengths": src_lengths,
-            "prev_output_tokens": prev_output_tokens
-        },
-        "ref_dict": ref_dict,
-        "constraint_masks": constraint_masks,
-        "target": target,
-    }
-
-    return batch
-
-
-class MNLIDataset(OFADataset):
-    def __init__(
-        self,
-        split,
-        dataset,
-        bpe,
-        src_dict,
-        tgt_dict=None,
-        max_src_length=512,
-        max_tgt_length=30,
-        constraint_trie=None,
-        prompt_type="none"
-    ):
-        super().__init__(split, dataset, bpe, src_dict, tgt_dict)
-        self.max_src_length = max_src_length
-        self.max_tgt_length = max_tgt_length
-        self.constraint_trie = constraint_trie
-        self.prompt_type = prompt_type
-
-    def __getitem__(self, index):
-        sentence1, sentence2, label = self.dataset[index]
-        if label == '0':
-            label = 'maybe'
-        elif label == '1':
-            label = 'yes'
-        elif label == '2':
-            label = 'no'
-        else:
-            raise NotImplementedError
-
-        sentence1 = ' '.join(sentence1.lower().strip().split()[:self.max_src_length])
-        sentence2 = ' '.join(sentence2.lower().strip().split()[:self.max_src_length])
-        src_item = self.encode_text(
-            ' can text1 " {} " imply text2 " {} "?'.format(sentence1, sentence2)
-        )
-        tgt_item = self.encode_text(" {}".format(label))
-        assert tgt_item.size(0) == 1
-        ref_dict = {label: 1.0}
-
-        src_item = torch.cat([self.bos_item, src_item, self.eos_item])
-        if self.prompt_type == 'none':
-            prev_output_item = self.bos_item
-            target_item = tgt_item
-        elif self.prompt_type == 'src':
-            prev_output_item = src_item.clone()
-            target_item = torch.cat([prev_output_item[1:], tgt_item])
-        elif self.prompt_type == 'prev_output':
-            prev_output_item = src_item[:-1].clone()
-            target_item = torch.cat([prev_output_item[1:], tgt_item])
-        else:
-            raise NotImplementedError
-        target_item[:-1] = self.tgt_dict.pad()
-
-        example = {
-            "source": src_item,
-            "target": target_item,
-            "prev_output_tokens": prev_output_item,
-            "ref_dict": ref_dict,
-        }
-        if self.constraint_trie is not None:
-            constraint_mask = torch.zeros((len(prev_output_item), len(self.tgt_dict))).bool()
-            constraint_nodes = self.constraint_trie.get_next_layer(self.bos_item.tolist())
-            constraint_mask[-1][constraint_nodes] = True
-            example["constraint_mask"] = constraint_mask
-        return example
-
-    def collater(self, samples, pad_to_length=None):
-        """Merge a list of samples to form a mini-batch.
-        Args:
-            samples (List[dict]): samples to collate
-        Returns:
-            dict: a mini-batch containing the data of the task
-        """
-        return collate(samples, pad_idx=self.pad, eos_idx=self.eos)

data/nlu_data/mrpc_dataset.py

@@ -1,141 +0,0 @@
-# Copyright 2022 The OFA-Sys Team. 
-# All rights reserved.
-# This source code is licensed under the Apache 2.0 license 
-# found in the LICENSE file in the root directory.
-
-import logging
-import warnings
-import torch
-import numpy as np
-
-from data import data_utils
-from data.ofa_dataset import OFADataset
-
-logger = logging.getLogger(__name__)
-warnings.filterwarnings("ignore", "(Possibly )?corrupt EXIF data", UserWarning)
-
-
-def collate(samples, pad_idx, eos_idx):
-    if len(samples) == 0:
-        return {}
-
-    def merge(key):
-        return data_utils.collate_tokens(
-            [s[key] for s in samples],
-            pad_idx,
-            eos_idx=eos_idx,
-        )
-
-    src_tokens = merge("source")
-    src_lengths = torch.LongTensor([s["source"].ne(pad_idx).long().sum() for s in samples])
-
-    ref_dict = None
-    if samples[0].get("ref_dict", None) is not None:
-        ref_dict = np.array([s['ref_dict'] for s in samples])
-
-    constraint_masks = None
-    if samples[0].get("constraint_mask", None) is not None:
-        constraint_masks = merge("constraint_mask")
-
-    prev_output_tokens = None
-    target = None
-    if samples[0].get("target", None) is not None:
-        target = merge("target")
-        tgt_lengths = torch.LongTensor(
-            [s["target"].ne(pad_idx).long().sum() for s in samples]
-        )
-        ntokens = tgt_lengths.sum().item()
-
-        if samples[0].get("prev_output_tokens", None) is not None:
-            prev_output_tokens = merge("prev_output_tokens")
-    else:
-        ntokens = src_lengths.sum().item()
-
-    batch = {
-        "nsentences": len(samples),
-        "ntokens": ntokens,
-        "net_input": {
-            "src_tokens": src_tokens,
-            "src_lengths": src_lengths,
-            "prev_output_tokens": prev_output_tokens
-        },
-        "ref_dict": ref_dict,
-        "constraint_masks": constraint_masks,
-        "target": target,
-    }
-
-    return batch
-
-
-class MRPCDataset(OFADataset):
-    def __init__(
-        self,
-        split,
-        dataset,
-        bpe,
-        src_dict,
-        tgt_dict=None,
-        max_src_length=512,
-        max_tgt_length=30,
-        constraint_trie=None,
-        prompt_type="none"
-    ):
-        super().__init__(split, dataset, bpe, src_dict, tgt_dict)
-        self.max_src_length = max_src_length
-        self.max_tgt_length = max_tgt_length
-        self.constraint_trie = constraint_trie
-        self.prompt_type = prompt_type
-
-    def __getitem__(self, index):
-        sentence1, sentence2, label = self.dataset[index]
-        if label == '0':
-            label = 'no'
-        elif label == '1':
-            label = 'yes'
-        else:
-            raise NotImplementedError
-
-        sentence1 = ' '.join(sentence1.lower().strip().split()[:self.max_src_length])
-        sentence2 = ' '.join(sentence2.lower().strip().split()[:self.max_src_length])
-        src_item = self.encode_text(
-            ' does text1 " {} " and text2 " {} " have the same semantics?'.format(sentence1, sentence2),
-        )
-        tgt_item = self.encode_text(" {}".format(label))
-        assert tgt_item.size(0) == 1
-        ref_dict = {label: 1.0}
-
-        src_item = torch.cat([self.bos_item, src_item, self.eos_item])
-        if self.prompt_type == 'none':
-            prev_output_item = self.bos_item
-            target_item = tgt_item
-        elif self.prompt_type == 'src':
-            prev_output_item = src_item.clone()
-            target_item = torch.cat([prev_output_item[1:], tgt_item])
-        elif self.prompt_type == 'prev_output':
-            prev_output_item = src_item[:-1].clone()
-            target_item = torch.cat([prev_output_item[1:], tgt_item])
-        else:
-            raise NotImplementedError
-        target_item[:-1] = self.tgt_dict.pad()
-
-        example = {
-            "source": src_item,
-            "target": target_item,
-            "prev_output_tokens": prev_output_item,
-            "ref_dict": ref_dict,
-        }
-        if self.constraint_trie is not None:
-            constraint_mask = torch.zeros((len(prev_output_item), len(self.tgt_dict))).bool()
-            constraint_nodes = self.constraint_trie.get_next_layer(self.bos_item.tolist())
-            constraint_mask[-1][constraint_nodes] = True
-            example["constraint_mask"] = constraint_mask
-        return example
-
-    def collater(self, samples, pad_to_length=None):
-        """Merge a list of samples to form a mini-batch.
-        Args:
-            samples (List[dict]): samples to collate
-        Returns:
-            dict: a mini-batch containing the data of the task
-        """
-        return collate(samples, pad_idx=self.pad, eos_idx=self.eos)

data/nlu_data/qnli_dataset.py

@@ -1,141 +0,0 @@
-# Copyright 2022 The OFA-Sys Team. 
-# All rights reserved.
-# This source code is licensed under the Apache 2.0 license 
-# found in the LICENSE file in the root directory.
-
-import logging
-import warnings
-import torch
-import numpy as np
-
-from data import data_utils
-from data.ofa_dataset import OFADataset
-
-logger = logging.getLogger(__name__)
-warnings.filterwarnings("ignore", "(Possibly )?corrupt EXIF data", UserWarning)
-
-
-def collate(samples, pad_idx, eos_idx):
-    if len(samples) == 0:
-        return {}
-
-    def merge(key):
-        return data_utils.collate_tokens(
-            [s[key] for s in samples],
-            pad_idx,
-            eos_idx=eos_idx,
-        )
-
-    src_tokens = merge("source")
-    src_lengths = torch.LongTensor([s["source"].ne(pad_idx).long().sum() for s in samples])
-
-    ref_dict = None
-    if samples[0].get("ref_dict", None) is not None:
-        ref_dict = np.array([s['ref_dict'] for s in samples])
-
-    constraint_masks = None
-    if samples[0].get("constraint_mask", None) is not None:
-        constraint_masks = merge("constraint_mask")
-
-    prev_output_tokens = None
-    target = None
-    if samples[0].get("target", None) is not None:
-        target = merge("target")
-        tgt_lengths = torch.LongTensor(
-            [s["target"].ne(pad_idx).long().sum() for s in samples]
-        )
-        ntokens = tgt_lengths.sum().item()
-
-        if samples[0].get("prev_output_tokens", None) is not None:
-            prev_output_tokens = merge("prev_output_tokens")
-    else:
-        ntokens = src_lengths.sum().item()
-
-    batch = {
-        "nsentences": len(samples),
-        "ntokens": ntokens,
-        "net_input": {
-            "src_tokens": src_tokens,
-            "src_lengths": src_lengths,
-            "prev_output_tokens": prev_output_tokens
-        },
-        "ref_dict": ref_dict,
-        "constraint_masks": constraint_masks,
-        "target": target,
-    }
-
-    return batch
-
-
-class QNLIDataset(OFADataset):
-    def __init__(
-        self,
-        split,
-        dataset,
-        bpe,
-        src_dict,
-        tgt_dict=None,
-        max_src_length=512,
-        max_tgt_length=30,
-        constraint_trie=None,
-        prompt_type="none"
-    ):
-        super().__init__(split, dataset, bpe, src_dict, tgt_dict)
-        self.max_src_length = max_src_length
-        self.max_tgt_length = max_tgt_length
-        self.constraint_trie = constraint_trie
-        self.prompt_type = prompt_type
-
-    def __getitem__(self, index):
-        question, sentence, label = self.dataset[index]
-        if label == '0' or label == 'not_entailment':
-            label = 'no'
-        elif label == '1' or label == 'entailment':
-            label = 'yes'
-        else:
-            raise NotImplementedError
-
-        question = ' '.join(question.lower().strip().split()[:self.max_src_length])
-        sentence = ' '.join(sentence.lower().strip().split()[:self.max_src_length])
-        src_item = self.encode_text(
-            ' does " {} " contain the answer to question " {} "?'.format(sentence, question)
-        )
-        tgt_item = self.encode_text(" {}".format(label))
-        assert tgt_item.size(0) == 1
-        ref_dict = {label: 1.0}
-
-        src_item = torch.cat([self.bos_item, src_item, self.eos_item])
-        if self.prompt_type == 'none':
-            prev_output_item = self.bos_item
-            target_item = tgt_item
-        elif self.prompt_type == 'src':
-            prev_output_item = src_item.clone()
-            target_item = torch.cat([prev_output_item[1:], tgt_item])
-        elif self.prompt_type == 'prev_output':
-            prev_output_item = src_item[:-1].clone()
-            target_item = torch.cat([prev_output_item[1:], tgt_item])
-        else:
-            raise NotImplementedError
-        target_item[:-1] = self.tgt_dict.pad()
-
-        example = {
-            "source": src_item,
-            "target": target_item,
-            "prev_output_tokens": prev_output_item,
-            "ref_dict": ref_dict,
-        }
-        if self.constraint_trie is not None:
-            constraint_mask = torch.zeros((len(prev_output_item), len(self.tgt_dict))).bool()
-            constraint_nodes = self.constraint_trie.get_next_layer(self.bos_item.tolist())
-            constraint_mask[-1][constraint_nodes] = True
-            example["constraint_mask"] = constraint_mask
-        return example
-
-    def collater(self, samples, pad_to_length=None):
-        """Merge a list of samples to form a mini-batch.
-        Args:
-            samples (List[dict]): samples to collate
-        Returns:
-            dict: a mini-batch containing the data of the task
-        """
-        return collate(samples, pad_idx=self.pad, eos_idx=self.eos)

data/nlu_data/qqp_dataset.py

@@ -1,141 +0,0 @@
-# Copyright 2022 The OFA-Sys Team. 
-# All rights reserved.
-# This source code is licensed under the Apache 2.0 license 
-# found in the LICENSE file in the root directory.
-
-import logging
-import warnings
-import torch
-import numpy as np
-
-from data import data_utils
-from data.ofa_dataset import OFADataset
-
-logger = logging.getLogger(__name__)
-warnings.filterwarnings("ignore", "(Possibly )?corrupt EXIF data", UserWarning)
-
-
-def collate(samples, pad_idx, eos_idx):
-    if len(samples) == 0:
-        return {}
-
-    def merge(key):
-        return data_utils.collate_tokens(
-            [s[key] for s in samples],
-            pad_idx,
-            eos_idx=eos_idx,
-        )
-
-    src_tokens = merge("source")
-    src_lengths = torch.LongTensor([s["source"].ne(pad_idx).long().sum() for s in samples])
-
-    ref_dict = None
-    if samples[0].get("ref_dict", None) is not None:
-        ref_dict = np.array([s['ref_dict'] for s in samples])
-
-    constraint_masks = None
-    if samples[0].get("constraint_mask", None) is not None:
-        constraint_masks = merge("constraint_mask")
-
-    prev_output_tokens = None
-    target = None
-    if samples[0].get("target", None) is not None:
-        target = merge("target")
-        tgt_lengths = torch.LongTensor(
-            [s["target"].ne(pad_idx).long().sum() for s in samples]
-        )
-        ntokens = tgt_lengths.sum().item()
-
-        if samples[0].get("prev_output_tokens", None) is not None:
-            prev_output_tokens = merge("prev_output_tokens")
-    else:
-        ntokens = src_lengths.sum().item()
-
-    batch = {
-        "nsentences": len(samples),
-        "ntokens": ntokens,
-        "net_input": {
-            "src_tokens": src_tokens,
-            "src_lengths": src_lengths,
-            "prev_output_tokens": prev_output_tokens
-        },
-        "ref_dict": ref_dict,
-        "constraint_masks": constraint_masks,
-        "target": target,
-    }
-
-    return batch
-
-
-class QQPDataset(OFADataset):
-    def __init__(
-        self,
-        split,
-        dataset,
-        bpe,
-        src_dict,
-        tgt_dict=None,
-        max_src_length=512,
-        max_tgt_length=30,
-        constraint_trie=None,
-        prompt_type="none"
-    ):
-        super().__init__(split, dataset, bpe, src_dict, tgt_dict)
-        self.max_src_length = max_src_length
-        self.max_tgt_length = max_tgt_length
-        self.constraint_trie = constraint_trie
-        self.prompt_type = prompt_type
-
-    def __getitem__(self, index):
-        question1, question2, label = self.dataset[index]
-        if label == '0':
-            label = 'no'
-        elif label == '1':
-            label = 'yes'
-        else:
-            raise NotImplementedError
-
-        question1 = ' '.join(question1.lower().strip().split()[:self.max_src_length])
-        question2 = ' '.join(question2.lower().strip().split()[:self.max_src_length])
-        src_item = self.encode_text(
-            ' is question " {} " and question " {} " equivalent?'.format(question1, question2)
-        )
-        tgt_item = self.encode_text(" {}".format(label))
-        assert tgt_item.size(0) == 1
-        ref_dict = {label: 1.0}
-
-        src_item = torch.cat([self.bos_item, src_item, self.eos_item])
-        if self.prompt_type == 'none':
-            prev_output_item = self.bos_item
-            target_item = tgt_item
-        elif self.prompt_type == 'src':
-            prev_output_item = src_item.clone()
-            target_item = torch.cat([prev_output_item[1:], tgt_item])
-        elif self.prompt_type == 'prev_output':
-            prev_output_item = src_item[:-1].clone()
-            target_item = torch.cat([prev_output_item[1:], tgt_item])
-        else:
-            raise NotImplementedError
-        target_item[:-1] = self.tgt_dict.pad()
-
-        example = {
-            "source": src_item,
-            "target": target_item,
-            "prev_output_tokens": prev_output_item,
-            "ref_dict": ref_dict,
-        }
-        if self.constraint_trie is not None:
-            constraint_mask = torch.zeros((len(prev_output_item), len(self.tgt_dict))).bool()
-            constraint_nodes = self.constraint_trie.get_next_layer(self.bos_item.tolist())
-            constraint_mask[-1][constraint_nodes] = True
-            example["constraint_mask"] = constraint_mask
-        return example
-
-    def collater(self, samples, pad_to_length=None):
-        """Merge a list of samples to form a mini-batch.
-        Args:
-            samples (List[dict]): samples to collate
-        Returns:
-            dict: a mini-batch containing the data of the task
-        """
-        return collate(samples, pad_idx=self.pad, eos_idx=self.eos)

data/nlu_data/rte_dataset.py

@@ -1,141 +0,0 @@
-# Copyright 2022 The OFA-Sys Team. 
-# All rights reserved.
-# This source code is licensed under the Apache 2.0 license 
-# found in the LICENSE file in the root directory.
-
-import logging
-import warnings
-import torch
-import numpy as np
-
-from data import data_utils
-from data.ofa_dataset import OFADataset
-
-logger = logging.getLogger(__name__)
-warnings.filterwarnings("ignore", "(Possibly )?corrupt EXIF data", UserWarning)
-
-
-def collate(samples, pad_idx, eos_idx):
-    if len(samples) == 0:
-        return {}
-
-    def merge(key):
-        return data_utils.collate_tokens(
-            [s[key] for s in samples],
-            pad_idx,
-            eos_idx=eos_idx,
-        )
-
-    src_tokens = merge("source")
-    src_lengths = torch.LongTensor([s["source"].ne(pad_idx).long().sum() for s in samples])
-
-    ref_dict = None
-    if samples[0].get("ref_dict", None) is not None:
-        ref_dict = np.array([s['ref_dict'] for s in samples])
-
-    constraint_masks = None
-    if samples[0].get("constraint_mask", None) is not None:
-        constraint_masks = merge("constraint_mask")
-
-    prev_output_tokens = None
-    target = None
-    if samples[0].get("target", None) is not None:
-        target = merge("target")
-        tgt_lengths = torch.LongTensor(
-            [s["target"].ne(pad_idx).long().sum() for s in samples]
-        )
-        ntokens = tgt_lengths.sum().item()
-
-        if samples[0].get("prev_output_tokens", None) is not None:
-            prev_output_tokens = merge("prev_output_tokens")
-    else:
-        ntokens = src_lengths.sum().item()
-
-    batch = {
-        "nsentences": len(samples),
-        "ntokens": ntokens,
-        "net_input": {
-            "src_tokens": src_tokens,
-            "src_lengths": src_lengths,
-            "prev_output_tokens": prev_output_tokens
-        },
-        "ref_dict": ref_dict,
-        "constraint_masks": constraint_masks,
-        "target": target,
-    }
-
-    return batch
-
-
-class RTEDataset(OFADataset):
-    def __init__(
-        self,
-        split,
-        dataset,
-        bpe,
-        src_dict,
-        tgt_dict=None,
-        max_src_length=512,
-        max_tgt_length=30,
-        constraint_trie=None,
-        prompt_type="none"
-    ):
-        super().__init__(split, dataset, bpe, src_dict, tgt_dict)
-        self.max_src_length = max_src_length
-        self.max_tgt_length = max_tgt_length
-        self.constraint_trie = constraint_trie
-        self.prompt_type = prompt_type
-
-    def __getitem__(self, index):
-        sentence1, sentence2, label = self.dataset[index]
-        if label == 'not_entailment':
-            label = 'no'
-        elif label == 'entailment':
-            label = 'yes'
-        else:
-            raise NotImplementedError
-
-        sentence1 = ' '.join(sentence1.lower().strip().split()[:self.max_src_length])
-        sentence2 = ' '.join(sentence2.lower().strip().split()[:self.max_src_length])
-        src_item = self.encode_text(
-            ' can text1 " {} " imply text2 " {} "?'.format(sentence1, sentence2),
-        )
-        tgt_item = self.encode_text(" {}".format(label))
-        assert tgt_item.size(0) == 1
-        ref_dict = {label: 1.0}
-
-        src_item = torch.cat([self.bos_item, src_item, self.eos_item])
-        if self.prompt_type == 'none':
-            prev_output_item = self.bos_item
-            target_item = tgt_item
-        elif self.prompt_type == 'src':
-            prev_output_item = src_item.clone()
-            target_item = torch.cat([prev_output_item[1:], tgt_item])
-        elif self.prompt_type == 'prev_output':
-            prev_output_item = src_item[:-1].clone()
-            target_item = torch.cat([prev_output_item[1:], tgt_item])
-        else:
-            raise NotImplementedError
-        target_item[:-1] = self.tgt_dict.pad()
-
-        example = {
-            "source": src_item,
-            "target": target_item,
-            "prev_output_tokens": prev_output_item,
-            "ref_dict": ref_dict,
-        }
-        if self.constraint_trie is not None:
-            constraint_mask = torch.zeros((len(prev_output_item), len(self.tgt_dict))).bool()
-            constraint_nodes = self.constraint_trie.get_next_layer(self.bos_item.tolist())
-            constraint_mask[-1][constraint_nodes] = True
-            example["constraint_mask"] = constraint_mask
-        return example
-
-    def collater(self, samples, pad_to_length=None):
-        """Merge a list of samples to form a mini-batch.
-        Args:
-            samples (List[dict]): samples to collate
-        Returns:
-            dict: a mini-batch containing the data of the task
-        """
-        return collate(samples, pad_idx=self.pad, eos_idx=self.eos)

data/nlu_data/sst2_dataset.py

@@ -1,138 +0,0 @@
-# Copyright 2022 The OFA-Sys Team. 
-# All rights reserved.
-# This source code is licensed under the Apache 2.0 license 
-# found in the LICENSE file in the root directory.
-
-import logging
-import warnings
-import torch
-import numpy as np
-
-from data import data_utils
-from data.ofa_dataset import OFADataset
-
-logger = logging.getLogger(__name__)
-warnings.filterwarnings("ignore", "(Possibly )?corrupt EXIF data", UserWarning)
-
-
-def collate(samples, pad_idx, eos_idx):
-    if len(samples) == 0:
-        return {}
-
-    def merge(key):
-        return data_utils.collate_tokens(
-            [s[key] for s in samples],
-            pad_idx,
-            eos_idx=eos_idx,
-        )
-
-    src_tokens = merge("source")
-    src_lengths = torch.LongTensor([s["source"].ne(pad_idx).long().sum() for s in samples])
-
-    ref_dict = None
-    if samples[0].get("ref_dict", None) is not None:
-        ref_dict = np.array([s['ref_dict'] for s in samples])
-
-    constraint_masks = None
-    if samples[0].get("constraint_mask", None) is not None:
-        constraint_masks = merge("constraint_mask")
-
-    prev_output_tokens = None
-    target = None
-    if samples[0].get("target", None) is not None:
-        target = merge("target")
-        tgt_lengths = torch.LongTensor(
-            [s["target"].ne(pad_idx).long().sum() for s in samples]
-        )
-        ntokens = tgt_lengths.sum().item()
-
-        if samples[0].get("prev_output_tokens", None) is not None:
-            prev_output_tokens = merge("prev_output_tokens")
-    else:
-        ntokens = src_lengths.sum().item()
-
-    batch = {
-        "nsentences": len(samples),
-        "ntokens": ntokens,
-        "net_input": {
-            "src_tokens": src_tokens,
-            "src_lengths": src_lengths,
-            "prev_output_tokens": prev_output_tokens
-        },
-        "ref_dict": ref_dict,
-        "constraint_masks": constraint_masks,
-        "target": target,
-    }
-
-    return batch
-
-
-class SST2Dataset(OFADataset):
-    def __init__(
-        self,
-        split,
-        dataset,
-        bpe,
-        src_dict,
-        tgt_dict=None,
-        max_src_length=512,
-        max_tgt_length=30,
-        constraint_trie=None,
-        prompt_type="none"
-    ):
-        super().__init__(split, dataset, bpe, src_dict, tgt_dict)
-        self.max_src_length = max_src_length
-        self.max_tgt_length = max_tgt_length
-        self.constraint_trie = constraint_trie
-        self.prompt_type = prompt_type
-
-    def __getitem__(self, index):
-        sentence, label = self.dataset[index]
-        if label == '0':
-            label = 'negative'
-        elif label == '1':
-            label = 'positive'
-        else:
-            raise NotImplementedError
-
-        sentence = ' '.join(sentence.lower().strip().split()[:self.max_src_length])
-        src_item = self.encode_text(' is the sentiment of text " {} " positive or negative?'.format(sentence))
-        tgt_item = self.encode_text(" {}".format(label))
-        assert tgt_item.size(0) == 1
-        ref_dict = {label: 1.0}
-
-        src_item = torch.cat([self.bos_item, src_item, self.eos_item])
-        if self.prompt_type == 'none':
-            prev_output_item = self.bos_item
-            target_item = tgt_item
-        elif self.prompt_type == 'src':
-            prev_output_item = src_item.clone()
-            target_item = torch.cat([prev_output_item[1:], tgt_item])
-        elif self.prompt_type == 'prev_output':
-            prev_output_item = src_item[:-1].clone()
-            target_item = torch.cat([prev_output_item[1:], tgt_item])
-        else:
-            raise NotImplementedError
-        target_item[:-1] = self.tgt_dict.pad()
-
-        example = {
-            "source": src_item,
-            "target": target_item,
-            "prev_output_tokens": prev_output_item,
-            "ref_dict": ref_dict,
-        }
-        if self.constraint_trie is not None:
-            constraint_mask = torch.zeros((len(prev_output_item), len(self.tgt_dict))).bool()
-            constraint_nodes = self.constraint_trie.get_next_layer(self.bos_item.tolist())
-            constraint_mask[-1][constraint_nodes] = True
-            example["constraint_mask"] = constraint_mask
-        return example
-
-    def collater(self, samples, pad_to_length=None):
-        """Merge a list of samples to form a mini-batch.
-        Args:
-            samples (List[dict]): samples to collate
-        Returns:
-            dict: a mini-batch containing the data of the task
-        """
-        return collate(samples, pad_idx=self.pad, eos_idx=self.eos)

data/ofa_dataset.py

@@ -1,79 +0,0 @@
-# Copyright 2022 The OFA-Sys Team. 
-# All rights reserved.
-# This source code is licensed under the Apache 2.0 license 
-# found in the LICENSE file in the root directory.
-
-import logging
-import re
-import torch.utils.data
-from fairseq.data import FairseqDataset
-
-logger = logging.getLogger(__name__)
-
-
-class OFADataset(FairseqDataset):
-    def __init__(self, split, dataset, bpe, src_dict, tgt_dict):
-        self.split = split
-        self.dataset = dataset
-        self.bpe = bpe
-        self.src_dict = src_dict
-        self.tgt_dict = tgt_dict
-
-        self.bos = src_dict.bos()
-        self.eos = src_dict.eos()
-        self.pad = src_dict.pad()
-        self.bos_item = torch.LongTensor([self.bos])
-        self.eos_item = torch.LongTensor([self.eos])
-
-    def __len__(self):
-        return len(self.dataset)
-
-    def encode_text(self, text, length=None, append_bos=False, append_eos=False, use_bpe=True):
-        s = self.tgt_dict.encode_line(
-            line=self.bpe.encode(text) if use_bpe else text,
-            add_if_not_exist=False,
-            append_eos=False
-        ).long()
-        if length is not None:
-            s = s[:length]
-        if append_bos:
-            s = torch.cat([self.bos_item, s])
-        if append_eos:
-            s = torch.cat([s, self.eos_item])
-        return s
-
-    def pre_question(self, question, max_ques_words=None):
-        question = question.lower().lstrip(",.!?*#:;~").replace('-', ' ').replace('/', ' ')
-
-        question = re.sub(
-            r"\s{2,}",
-            ' ',
-            question,
-        )
-        question = question.rstrip('\n')
-        question = question.strip(' ')
-
-        # truncate question
-        question_words = question.split(' ')
-        if max_ques_words is not None and len(question_words) > max_ques_words:
-            question = ' '.join(question_words[:max_ques_words])
-
-        return question
-
-    def pre_caption(self, caption, max_words=None):
-        caption = caption.lower().lstrip(",.!?*#:;~").replace('-', ' ').replace('/', ' ').replace('<person>', 'person')
-
-        caption = re.sub(
-            r"\s{2,}",
-            ' ',
-            caption,
-        )
-        caption = caption.rstrip('\n')
-        caption = caption.strip(' ')
-
-        # truncate caption
-        caption_words = caption.split(' ')
-        if max_words is not None and len(caption_words) > max_words:
-            caption = ' '.join(caption_words[:max_words])
-
-        return caption

data/pretrain_data/unify_dataset.py

@@ -1,636 +0,0 @@
-# Copyright 2022 The OFA-Sys Team.
-# All rights reserved.
-# This source code is licensed under the Apache 2.0 license
-# found in the LICENSE file in the root directory.
-
-from io import BytesIO
-
-import math
-import logging
-import random
-import warnings
-
-import numpy as np
-import torch
-import base64
-from torchvision import transforms
-
-from PIL import Image, ImageFile
-
-from data import data_utils
-from data.ofa_dataset import OFADataset
-from utils.vision_helper import RandomAugment
-import utils.transforms as T
-
-ImageFile.LOAD_TRUNCATED_IMAGES = True
-ImageFile.MAX_IMAGE_PIXELS = None
-Image.MAX_IMAGE_PIXELS = None
-
-logger = logging.getLogger(__name__)
-warnings.filterwarnings("ignore", "(Possibly )?corrupt EXIF data", UserWarning)
-
-
-def get_whole_word_mask(bpe, dictionary):
-    if bpe is not None:
-
-        def is_beginning_of_word(i):
-            if i < dictionary.nspecial:
-                # special elements are always considered beginnings
-                return True
-            tok = dictionary[i]
-            if tok.startswith("madeupword"):
-                return True
-            try:
-                return bpe.is_beginning_of_word(tok)
-            except ValueError:
-                return True
-
-        mask_whole_words = torch.ByteTensor(
-            list(map(is_beginning_of_word, range(len(dictionary))))
-        )
-        return mask_whole_words
-    return None
-
-
-def collate(samples, pad_idx, eos_idx):
-    if len(samples) == 0:
-        return {}
-
-    def merge(key):
-        return data_utils.collate_tokens(
-            [s[key] for s in samples],
-            pad_idx,
-            eos_idx=eos_idx,
-        )
-
-    id = np.array([s["id"] for s in samples])
-    src_tokens = merge("source")
-    src_lengths = torch.LongTensor([s["source"].ne(pad_idx).long().sum() for s in samples])
-
-    patch_images = torch.stack([sample['patch_image'] for sample in samples], dim=0)
-    patch_masks = torch.cat([sample['patch_mask'] for sample in samples])
-
-    code_masks = None
-    if samples[0].get("code_mask", None) is not None:
-        code_masks = torch.cat([sample['code_mask'] for sample in samples])
-
-    conf = torch.cat([s['conf'] for s in samples], dim=0)
-
-    prev_output_tokens = None
-    target = None
-    if samples[0].get("target", None) is not None:
-        target = merge("target")
-        tgt_lengths = torch.LongTensor([s["target"].ne(pad_idx).long().sum() for s in samples])
-        ntokens = tgt_lengths.sum().item()
-
-        if samples[0].get("prev_output_tokens", None) is not None:
-            prev_output_tokens = merge("prev_output_tokens")
-    else:
-        ntokens = src_lengths.sum().item()
-
-    batch = {
-        "id": id,
-        "nsentences": len(samples),
-        "ntokens": ntokens,
-        "net_input": {
-            "src_tokens": src_tokens,
-            "src_lengths": src_lengths,
-            "patch_images": patch_images,
-            "patch_masks": patch_masks,
-            "code_masks": code_masks,
-            "prev_output_tokens": prev_output_tokens
-        },
-        "target": target,
-        "conf": conf
-    }
-
-    return batch
-
-
-class UnifyDataset(OFADataset):
-    def __init__(
-        self,
-        split,
-        dataset,
-        bpe,
-        src_dict,
-        tgt_dict=None,
-        max_src_length=128,
-        max_tgt_length=30,
-        seed=7,
-        code_dict_size=8192,
-        num_bins=1000,
-        patch_image_size=384,
-        code_image_size=128,
-        pure_text_dataset=None,
-        pure_image_dataset=None,
-        detection_dataset=None,
-        all_object_list=None,
-        all_caption_list=None,
-        type2ans_dict=None,
-        ans2type_dict=None,
-        max_image_size=512,
-        mask_ratio=0.3,
-        random_ratio=0.0,
-        keep_ratio=0.0,
-        mask_length="span-poisson",
-        poisson_lambda=3.0,
-        replace_length=1
-    ):
-        super().__init__(split, dataset, bpe, src_dict, tgt_dict)
-        self.max_src_length = max_src_length
-        self.max_tgt_length = max_tgt_length
-        self.seed = seed
-        self.code_dict_size = code_dict_size
-        self.num_bins = num_bins
-        self.patch_image_size = patch_image_size
-        self.code_image_size = code_image_size
-
-        self.pure_text_dataset = pure_text_dataset
-        self.pure_image_dataset = pure_image_dataset
-        self.detection_dataset = detection_dataset
-        self.epoch = 0
-
-        self.all_object_list = all_object_list
-        self.all_caption_list = all_caption_list
-        self.type2ans_dict = type2ans_dict
-        self.ans2type_dict = ans2type_dict
-
-        self.mask_ratio = mask_ratio
-        self.random_ratio = random_ratio
-        self.keep_ratio = keep_ratio
-        self.mask_length = mask_length
-        self.poisson_lambda = poisson_lambda
-        self.replace_length = replace_length
-        if self.replace_length not in [-1, 0, 1]:
-            raise ValueError(f"invalid arg: replace_length={self.replace_length}")
-        if self.mask_length not in ["subword", "word", "span-poisson"]:
-            raise ValueError(f"invalid arg: mask-length={self.mask_length}")
-        if self.mask_length == "subword" and self.replace_length not in [0, 1]:
-            raise ValueError(f"if using subwords, use replace-length=1 or 0")
-
-        self.mask_idx = src_dict.index("<mask>")
-        self.mask_whole_word = (
-            get_whole_word_mask(self.bpe, self.src_dict)
-            if self.mask_length != "subword"
-            else None
-        )
-        self.mask_span_distribution = None
-        if self.mask_length == "span-poisson":
-            _lambda = self.poisson_lambda
-            lambda_to_the_k = 1
-            e_to_the_minus_lambda = math.exp(-_lambda)
-            k_factorial = 1
-            ps = []
-            for k in range(0, 128):
-                ps.append(e_to_the_minus_lambda * lambda_to_the_k / k_factorial)
-                lambda_to_the_k *= _lambda
-                k_factorial *= k + 1
-                if ps[-1] < 0.0000001:
-                    break
-            ps = torch.FloatTensor(ps)
-            self.mask_span_distribution = torch.distributions.Categorical(ps)
-
-        self.pos_tgt_item = self.encode_text(" yes")
-        self.neg_tgt_item = self.encode_text(" no")
-
-        self.mask_left = self.mask_top = int(0.5 * self.code_image_size)
-        self.mask_right = self.mask_bottom = int(1.5 * self.code_image_size)
-        self.mask_ids = [
-            i*self.code_image_size*2+j
-            for i in range(self.code_image_size*2) for j in range(self.code_image_size*2)
-            if not (self.mask_left <= i < self.mask_right and self.mask_top <= j < self.mask_bottom)
-        ]
-
-        scales = np.arange(patch_image_size, 481).tolist()
-
-        # for image-text pair
-        self.patch_resize_transform = transforms.Compose([
-            T.RandomResize(scales, max_size=672),
-            transforms.CenterCrop(patch_image_size),
-            RandomAugment(2, 7, isPIL=True, augs=['Identity', 'AutoContrast', 'Equalize', 'Brightness', 'Sharpness',
-                                                  'ShearX', 'ShearY', 'TranslateX', 'TranslateY', 'Rotate']),
-            transforms.ToTensor(),
-            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
-        ])
-        # for pure image
-        self.patch_crop_transform = transforms.Compose([
-            transforms.ToTensor(),
-            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
-        ])
-        # for detection
-        self.detection_transform = T.Compose([
-            T.RandomHorizontalFlip(),
-            T.LargeScaleJitter(output_size=self.code_image_size*2, aug_scale_min=1.0, aug_scale_max=1.5),
-            T.ToTensor(),
-            T.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], max_image_size=max_image_size)
-        ])
-        # for visual grounding
-        self.visual_grounding_transform = T.Compose([
-            T.RandomResize(scales, max_size=672),
-            T.ObjectCenterCrop((patch_image_size, patch_image_size)),
-            T.ToTensor(),
-            T.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], max_image_size=max_image_size)
-        ])
-
-    def set_epoch(self, epoch, **unused):
-        self.epoch = epoch
-
-    def get_negative_caption(self, caption, gt_objects):
-        prob = random.random()
-        if gt_objects is not None and gt_objects != '' and prob > 0.6:
-            gt_object = random.choice(gt_objects.strip().split('&&'))
-            negative_object = random.choice(self.all_object_list[:-1])
-            negative_object = self.all_object_list[-1] if negative_object == gt_object else negative_object
-            negative_caption = caption.replace(gt_object, negative_object)
-        else:
-            negative_caption = random.choice(self.all_caption_list)
-        return negative_caption
-
-    def get_negative_answer(self, answer, conf):
-        prob = random.random()
-        if conf > (prob + 0.1) and answer in self.ans2type_dict:
-            negative_answer_type = self.ans2type_dict[answer]
-            if negative_answer_type == 'how many' and answer.isdigit() and prob > 0.5:
-                negative_answer = int(answer) + random.choice([-1, 1]) if answer != 0 else 1
-            else:
-                negative_answer_list = self.type2ans_dict[negative_answer_type]
-                negative_answer = random.choice(negative_answer_list[:-1])
-                negative_answer = negative_answer_list[-1] if negative_answer == answer else negative_answer
-            return negative_answer
-
-        negative_answer_list = self.type2ans_dict['other']
-        negative_answer = random.choice(negative_answer_list[:-1])
-        negative_answer = negative_answer_list[-1] if negative_answer == answer else negative_answer
-        return negative_answer
-
-    def process_image_text_pair(self, index):
-        uniq_id, image, caption, question, refs, gt_objects, dataset_name, type = self.dataset[index]
-
-        image = Image.open(BytesIO(base64.urlsafe_b64decode(image))).convert("RGB")
-        patch_image = self.patch_resize_transform(image) if type != 'visual_grounding' else None
-        patch_mask = torch.tensor([True])
-        conf = torch.tensor([1.0])
-        if type == 'caption':
-            tgt_caption = self.pre_caption(caption, self.max_tgt_length)
-            pos_src_caption = self.pre_caption(caption, self.max_src_length)
-            neg_src_caption = self.pre_caption(self.get_negative_caption(caption, gt_objects), self.max_src_length)
-            src_item = self.encode_text(" what does the image describe?")
-            tgt_item = self.encode_text(" {}".format(tgt_caption))
-            pos_src_item = self.encode_text(' does the image describe " {} "?'.format(pos_src_caption))
-            neg_src_item = self.encode_text(' does the image describe " {} "?'.format(neg_src_caption))
-        elif type == 'qa':
-            question = self.pre_question(question, self.max_src_length)
-            ref_dict = {item.split('|!+')[1]: float(item.split('|!+')[0]) for item in refs.split('&&')}
-            answer = max(ref_dict, key=ref_dict.get)
-            conf = ref_dict[answer]
-            src_item = self.encode_text(" {}".format(question))
-            tgt_item = self.encode_text(" {}".format(answer))
-            conf = torch.tensor([conf])
-            pos_src_item = self.encode_text(' what is the answer to question " {} ". is " {} "?'.format(question, answer))
-            neg_src_item = self.encode_text(
-                ' what is the answer to question " {} ". is " {} "?'.format(question, self.get_negative_answer(answer, conf))
-            )
-        elif type == 'visual_grounding':
-            conf = torch.tensor([1.0])
-            w, h = image.size
-            boxes_target = {"boxes": [], "labels": [], "area": [], "size": torch.tensor([h, w])}
-            x0, y0, x1, y1 = refs.strip().split(',')
-            boxes_target["boxes"] = torch.tensor([[float(x0), float(y0), float(x1), float(y1)]])
-            boxes_target["labels"] = np.array([0])
-            boxes_target["area"] = torch.tensor([(float(x1) - float(x0)) * (float(y1) - float(y0))])
-            patch_image, boxes_target = self.visual_grounding_transform(image, boxes_target)
-            quant_x0 = "<bin_{}>".format(int((boxes_target["boxes"][0][0] * (self.num_bins - 1)).round()))
-            quant_y0 = "<bin_{}>".format(int((boxes_target["boxes"][0][1] * (self.num_bins - 1)).round()))
-            quant_x1 = "<bin_{}>".format(int((boxes_target["boxes"][0][2] * (self.num_bins - 1)).round()))
-            quant_y1 = "<bin_{}>".format(int((boxes_target["boxes"][0][3] * (self.num_bins - 1)).round()))
-            region_coord = "{} {} {} {}".format(quant_x0, quant_y0, quant_x1, quant_y1)
-            src_caption = self.pre_caption(caption, self.max_src_length)
-            src_item = self.encode_text(' which region does the text " {} " describe?'.format(src_caption))
-            tgt_item = self.encode_text(region_coord, use_bpe=False)
-        else:
-            logger.info('type {} is not implemented'.format(type))
-            raise NotImplementedError
-
-        src_item = torch.cat([self.bos_item, src_item, self.eos_item])
-        target_item = torch.cat([tgt_item, self.eos_item])
-        prev_output_item = torch.cat([self.bos_item, tgt_item])
-        pos_src_item = torch.cat([self.bos_item, pos_src_item, self.eos_item]) if type != 'visual_grounding' else None
-        neg_src_item = torch.cat([self.bos_item, neg_src_item, self.eos_item]) if type != 'visual_grounding' else None
-
-        if type == 'caption' and dataset_name == 'cc12m':
-            target_item[:2] = self.src_dict.pad()
-            target_item[-1] = self.eos_item
-
-        example = {
-            "id": uniq_id,
-            "source": src_item,
-            "patch_image": patch_image,
-            "patch_mask": patch_mask,
-            "target": target_item,
-            "prev_output_tokens": prev_output_item,
-            "conf": conf,
-        }
-
-        examples = [example]
-        prob = random.random()
-        if type == 'visual_grounding':
-            region_example = example.copy()
-            region_prefix_item = self.encode_text('  what does the region describe? region:')
-            region_coord_item = self.encode_text('{}'.format(region_coord), use_bpe=False)
-            region_src_item = torch.cat([region_prefix_item, region_coord_item])
-            region_tgt_item = self.encode_text(' {}'.format(self.pre_caption(caption, self.max_tgt_length)))
-            region_example["source"] = torch.cat([self.bos_item, region_src_item, self.eos_item])
-            region_example["target"] = torch.cat([region_tgt_item, self.eos_item])
-            region_example["prev_output_tokens"] = torch.cat([self.bos_item, region_tgt_item])
-            region_example["conf"] = torch.tensor([1.0])
-            examples.append(region_example)
-        elif prob >= 0.5 and self.split == 'train':
-            pos_example = example.copy()
-            pos_example["source"] = pos_src_item
-            pos_example["target"] = torch.cat([self.pos_tgt_item, self.eos_item])
-            pos_example["prev_output_tokens"] = torch.cat([self.bos_item, self.pos_tgt_item])
-            examples.append(pos_example)
-        elif self.split == 'train':
-            neg_example = example.copy()
-            neg_example["source"] = neg_src_item
-            neg_example["target"] = torch.cat([self.neg_tgt_item, self.eos_item])
-            neg_example["prev_output_tokens"] = torch.cat([self.bos_item, self.neg_tgt_item])
-            examples.append(neg_example)
-        return examples
-
-    def process_pure_text(self, index):
-        patch_image = torch.zeros((3, self.code_image_size*2, self.code_image_size*2))
-        patch_mask = torch.tensor([False])
-        code_mask = torch.tensor([False])
-        conf = torch.tensor([2.0])
-
-        examples = []
-        for _ in range(2):
-            uniq_id, text = self.pure_text_dataset[index]
-            text = text.strip().lower()
-            text_item = self.encode_text(" {}".format(text), length=512)
-            text_item = text_item[-256:]
-            text_item = torch.cat([self.bos_item, text_item, self.eos_item])
-            mask_text_item = self.add_whole_word_mask(text_item.clone(), self.mask_ratio)
-            prefix_item = self.encode_text(' what is the complete text of " "?')
-            src_item = torch.cat([prefix_item[:-2], mask_text_item[1:-1], prefix_item[-2:]])
-            tgt_item = text_item[1:-1]
-            src_item = torch.cat([self.bos_item, src_item, self.eos_item])
-            target_item = torch.cat([tgt_item, self.eos_item])
-            prev_output_item = torch.cat([self.bos_item, tgt_item])
-            example = {
-                "id": uniq_id,
-                "source": src_item,
-                "patch_image": patch_image,
-                "patch_mask": patch_mask,
-                "code_mask": code_mask,
-                "target": target_item,
-                "prev_output_tokens": prev_output_item,
-                "conf": conf,
-            }
-            examples.append(example)
-
-        return examples
-
-    def process_pure_image(self, index):
-        image_id, image, code = self.pure_image_dataset[index]
-        image = Image.open(BytesIO(base64.urlsafe_b64decode(image))).convert("RGB")
-        patch_image = self.patch_crop_transform(image)
-        patch_image[:, self.mask_top:self.mask_bottom, self.mask_left:self.mask_right] = 0
-        patch_mask = torch.tensor([True])
-        src_item = self.encode_text(" what is the image in the middle part?")
-        image_code = torch.LongTensor([int(num) for num in code.strip().split()])
-        tgt_item = image_code + len(self.src_dict) - self.code_dict_size - self.num_bins
-        code_mask = torch.tensor([True])
-        conf = torch.tensor([2.0])
-
-        src_item = torch.cat([self.bos_item, src_item, self.eos_item])
-        target_item = torch.cat([tgt_item, self.eos_item])
-        prev_output_item = torch.cat([self.bos_item, tgt_item])
-
-        example = {
-            "id": image_id,
-            "source": src_item,
-            "patch_image": patch_image,
-            "patch_mask": patch_mask,
-            "code_mask": code_mask,
-            "target": target_item,
-            "prev_output_tokens": prev_output_item,
-            "conf": conf,
-        }
-        return [example]
-
-    def process_detection(self, index):
-        image_id, image, label = self.detection_dataset[index]
-        image = Image.open(BytesIO(base64.urlsafe_b64decode(image))).convert("RGB")
-
-        w, h = image.size
-        boxes_target = {"boxes": [], "labels": [], "area": [], "size": torch.tensor([h, w])}
-        label_list = label.strip().split('&&')
-        for label in label_list:
-            x0, y0, x1, y1, cat_id, cat = label.strip().split(',', 5)
-            boxes_target["boxes"].append([float(x0), float(y0), float(x1), float(y1)])
-            boxes_target["labels"].append(cat)
-            boxes_target["area"].append((float(x1) - float(x0)) * (float(y1) - float(y0)))
-        boxes_target["boxes"] = torch.tensor(boxes_target["boxes"])
-        boxes_target["labels"] = np.array(boxes_target["labels"])
-        boxes_target["area"] = torch.tensor(boxes_target["area"])
-
-        patch_image, boxes_target = self.detection_transform(image, boxes_target)
-        patch_mask = torch.tensor([True])
-        code_mask = torch.tensor([False])
-        conf = torch.tensor([2.0])
-
-        quant_boxes = []
-        for i, box in enumerate(boxes_target["boxes"]):
-            quant_boxes.extend(["<bin_{}>".format(int((pos * (self.num_bins - 1)).round())) for pos in box[:4]])
-            quant_boxes.append(self.bpe.encode(' {}'.format(boxes_target["labels"][i])))
-        src_item = self.encode_text(' what are the objects in the image?')
-        tgt_item = self.encode_text(' '.join(quant_boxes), use_bpe=False)
-
-        src_item = torch.cat([self.bos_item, src_item, self.eos_item])
-        target_item = torch.cat([tgt_item, self.eos_item])
-        prev_output_item = torch.cat([self.bos_item, tgt_item])
-
-        example = {
-            "id": image_id,
-            "source": src_item,
-            "patch_image": patch_image,
-            "patch_mask": patch_mask,
-            "code_mask": code_mask,
-            "target": target_item,
-            "prev_output_tokens": prev_output_item,
-            "conf": conf,
-        }
-        return [example]
-
-    def __getitem__(self, index):
-        with data_utils.numpy_seed(self.seed, self.epoch):
-            pair_samples = self.process_image_text_pair(index)
-            extra_samples = []
-            if self.split == 'train' and self.dataset.data_cnt % 8 == 0:
-                extra_samples += self.process_pure_text(0) if self.pure_text_dataset else []
-                extra_samples += self.process_pure_image(0) if self.pure_image_dataset else []
-                extra_samples += self.process_detection(0) if self.detection_dataset else []
-        return pair_samples, extra_samples
-
-    def word_starts(self, source):
-        if self.mask_whole_word is not None:
-            is_word_start = self.mask_whole_word.gather(0, source)
-        else:
-            is_word_start = torch.ones(source.size())
-        is_word_start[0] = 0
-        is_word_start[-1] = 0
-        return is_word_start
-
-    def add_whole_word_mask(self, source, p):
-        is_word_start = self.word_starts(source)
-        num_to_mask = int(math.ceil(is_word_start.float().sum() * p))
-        num_inserts = 0
-        if num_to_mask == 0:
-            return source
-
-        if self.mask_span_distribution is not None:
-            lengths = self.mask_span_distribution.sample(sample_shape=(num_to_mask,))
-
-            # Make sure we have enough to mask
-            cum_length = torch.cumsum(lengths, 0)
-            while cum_length[-1] < num_to_mask:
-                lengths = torch.cat(
-                    [
-                        lengths,
-                        self.mask_span_distribution.sample(sample_shape=(num_to_mask,)),
-                    ],
-                    dim=0,
-                )
-                cum_length = torch.cumsum(lengths, 0)
-
-            # Trim to masking budget
-            i = 0
-            while cum_length[i] < num_to_mask:
-                i += 1
-            lengths[i] = num_to_mask - (0 if i == 0 else cum_length[i - 1])
-            num_to_mask = i + 1
-            lengths = lengths[:num_to_mask]
-
-            # Handle 0-length mask (inserts) separately
-            lengths = lengths[lengths > 0]
-            num_inserts = num_to_mask - lengths.size(0)
-            num_to_mask -= num_inserts
-            if num_to_mask == 0:
-                return self.add_insertion_noise(source, num_inserts / source.size(0))
-
-            assert (lengths > 0).all()
-        else:
-            lengths = torch.ones((num_to_mask,)).long()
-        assert is_word_start[-1] == 0
-        word_starts = is_word_start.nonzero(as_tuple=False)
-        indices = word_starts[
-            torch.randperm(word_starts.size(0))[:num_to_mask]
-        ].squeeze(1)
-        mask_random = torch.FloatTensor(num_to_mask).uniform_() < self.random_ratio
-
-        source_length = source.size(0)
-        assert source_length - 1 not in indices
-        to_keep = torch.ones(source_length, dtype=torch.bool)
-        is_word_start[
-            -1
-        ] = 255  # acts as a long length, so spans don't go over the end of doc
-        if self.replace_length == 0:
-            to_keep[indices] = 0
-        else:
-            # keep index, but replace it with [MASK]
-            source[indices] = self.mask_idx
-            source[indices[mask_random]] = torch.randint(
-                4, len(self.tgt_dict) - self.code_dict_size - self.num_bins, size=(mask_random.sum(),)
-            )
-
-        if self.mask_span_distribution is not None:
-            assert len(lengths.size()) == 1
-            assert lengths.size() == indices.size()
-            lengths -= 1
-            while indices.size(0) > 0:
-                assert lengths.size() == indices.size()
-                lengths -= is_word_start[indices + 1].long()
-                uncompleted = lengths >= 0
-                indices = indices[uncompleted] + 1
-                mask_random = mask_random[uncompleted]
-                lengths = lengths[uncompleted]
-                if self.replace_length != -1:
-                    # delete token
-                    to_keep[indices] = 0
-                else:
-                    # keep index, but replace it with [MASK]
-                    source[indices] = self.mask_idx
-                    source[indices[mask_random]] = torch.randint(
-                        4, len(self.tgt_dict) - self.code_dict_size - self.num_bins, size=(mask_random.sum(),)
-                    )
-        else:
-            # A bit faster when all lengths are 1
-            while indices.size(0) > 0:
-                uncompleted = is_word_start[indices + 1] == 0
-                indices = indices[uncompleted] + 1
-                mask_random = mask_random[uncompleted]
-                if self.replace_length != -1:
-                    # delete token
-                    to_keep[indices] = 0
-                else:
-                    # keep index, but replace it with [MASK]
-                    source[indices] = self.mask_idx
-                    source[indices[mask_random]] = torch.randint(
-                        4, len(self.tgt_dict) - self.code_dict_size - self.num_bins, size=(mask_random.sum(),)
-                    )
-
-                assert source_length - 1 not in indices
-
-        source = source[to_keep]
-
-        if num_inserts > 0:
-            source = self.add_insertion_noise(source, num_inserts / source.size(0))
-
-        return source
-
-    def add_insertion_noise(self, tokens, p):
-        if p == 0.0:
-            return tokens
-
-        num_tokens = len(tokens)
-        n = int(math.ceil(num_tokens * p))
-
-        noise_indices = torch.randperm(num_tokens + n - 2)[:n] + 1
-        noise_mask = torch.zeros(size=(num_tokens + n,), dtype=torch.bool)
-        noise_mask[noise_indices] = 1
-        result = torch.LongTensor(n + len(tokens)).fill_(-1)
-
-        num_random = int(math.ceil(n * self.random_ratio))
-        result[noise_indices[num_random:]] = self.mask_idx
-        result[noise_indices[:num_random]] = torch.randint(
-            low=4, high=len(self.tgt_dict)-self.code_dict_size-self.num_bins, size=(num_random,)
-        )
-
-        result[~noise_mask] = tokens
-
-        assert (result >= 0).all()
-        return result
-
-    def collater(self, samples, pad_to_length=None):
-        """Merge samples of different tasks to form two mini-batches.
-        Args:
-            samples (List[Tuple]): samples to collate
-        Returns:
-            Tuple[dict]: two mini-batch containing the data of different tasks
-        """
-
-        samples_v1 = []   # containing image-text pairs
-        samples_v2 = []   # containing detection data, text data and image data
-        for sample_tuple in samples:
-            samples_v1 += sample_tuple[0]
-            samples_v2 += sample_tuple[1]
-        if samples_v2 != []:
-            res_v1 = collate(samples_v1, pad_idx=self.src_dict.pad(), eos_idx=self.eos)
-            res_v2 = collate(samples_v2, pad_idx=self.src_dict.pad(), eos_idx=self.eos)
-            return res_v1, res_v2
-        else:
-            res_v1 = collate(samples_v1, pad_idx=self.src_dict.pad(), eos_idx=self.eos)
-            return res_v1

datasets.md

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

evaluate.py

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

ezocr/LICENSE

@@ -1,201 +0,0 @@
-                                 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 [yyyy] [name of copyright owner]
-
-   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.

ezocr/README.md

@@ -1,49 +0,0 @@
-# EasyOCR Lite
-
-从EasyOCR提取文本定位有关代码，进一步适配中文，修正缺陷
-
-## 安装
-
-Python版本至少为3.8。
-
-
-首先按照PyTorch官方说明安装PyTorch。
-
-```
-pip install -e .
-```
-
-## 使用
-
-``` python3
-from easyocrlite import ReaderLite
-
-reader = ReaderLite()
-results = reader.process('my_awesome_handwriting.png')
-```
-
-返回的内容为边界框和对应的图像区域的列表。
-其它说明见[demo](./demo.ipynb)。
-
-
-## 致谢
-
-基于[EasyOCR](https://github.com/JaidedAI/EasyOCR)修改实现。以下为EasyOCR致谢：
-
-This project is based on research and code from several papers and open-source repositories.
-
-All deep learning execution is based on [Pytorch](https://pytorch.org). :heart:
-
-Detection execution uses the CRAFT algorithm from this [official repository](https://github.com/clovaai/CRAFT-pytorch) and their [paper](https://arxiv.org/abs/1904.01941) (Thanks @YoungminBaek from [@clovaai](https://github.com/clovaai)). We also use their pretrained model. Training script is provided by [@gmuffiness](https://github.com/gmuffiness).
-
-The recognition model is a CRNN ([paper](https://arxiv.org/abs/1507.05717)). It is composed of 3 main components: feature extraction (we are currently using [Resnet](https://arxiv.org/abs/1512.03385)) and VGG, sequence labeling ([LSTM](https://www.bioinf.jku.at/publications/older/2604.pdf)) and decoding ([CTC](https://www.cs.toronto.edu/~graves/icml_2006.pdf)). The training pipeline for recognition execution is a modified version of the [deep-text-recognition-benchmark](https://github.com/clovaai/deep-text-recognition-benchmark) framework. (Thanks [@ku21fan](https://github.com/ku21fan) from [@clovaai](https://github.com/clovaai)) This repository is a gem that deserves more recognition.
-
-Beam search code is based on this [repository](https://github.com/githubharald/CTCDecoder) and his [blog](https://towardsdatascience.com/beam-search-decoding-in-ctc-trained-neural-networks-5a889a3d85a7). (Thanks [@githubharald](https://github.com/githubharald))
-
-Data synthesis is based on [TextRecognitionDataGenerator](https://github.com/Belval/TextRecognitionDataGenerator). (Thanks [@Belval](https://github.com/Belval))
-
-And a good read about CTC from distill.pub [here](https://distill.pub/2017/ctc/).
-
-
-## 许可证 (注意！)
-Apache 2.0

ezocr/build/lib/easyocrlite/__init__.py

@@ -1 +0,0 @@
-from easyocrlite.reader import ReaderLite

ezocr/build/lib/easyocrlite/reader.py

@@ -1,272 +0,0 @@
-from __future__ import annotations
-
-import logging
-import os
-from pathlib import Path
-from typing import Tuple
-
-import cv2
-import numpy as np
-import torch
-from PIL import Image, ImageEnhance
-
-from easyocrlite.model import CRAFT
-
-from easyocrlite.utils.download_utils import prepare_model
-from easyocrlite.utils.image_utils import (
-    adjust_result_coordinates,
-    boxed_transform,
-    normalize_mean_variance,
-    resize_aspect_ratio,
-)
-from easyocrlite.utils.detect_utils import (
-    extract_boxes,
-    extract_regions_from_boxes,
-    box_expand,
-    greedy_merge,
-)
-from easyocrlite.types import BoxTuple, RegionTuple
-import easyocrlite.utils.utils as utils
-
-logger = logging.getLogger(__name__)
-
-MODULE_PATH = (
-    os.environ.get("EASYOCR_MODULE_PATH")
-    or os.environ.get("MODULE_PATH")
-    or os.path.expanduser("~/.EasyOCR/")
-)
-
-
-class ReaderLite(object):
-    def __init__(
-        self,
-        gpu=True,
-        model_storage_directory=None,
-        download_enabled=True,
-        verbose=True,
-        quantize=True,
-        cudnn_benchmark=False,
-    ):
-
-        self.verbose = verbose
-
-        model_storage_directory = Path(
-            model_storage_directory
-            if model_storage_directory
-            else MODULE_PATH + "/model"
-        )
-        self.detector_path = prepare_model(
-            model_storage_directory, download_enabled, verbose
-        )
-
-        self.quantize = quantize
-        self.cudnn_benchmark = cudnn_benchmark
-        if gpu is False:
-            self.device = "cpu"
-            if verbose:
-                logger.warning(
-                    "Using CPU. Note: This module is much faster with a GPU."
-                )
-        elif not torch.cuda.is_available():
-            self.device = "cpu"
-            if verbose:
-                logger.warning(
-                    "CUDA not available - defaulting to CPU. Note: This module is much faster with a GPU."
-                )
-        elif gpu is True:
-            self.device = "cuda"
-        else:
-            self.device = gpu
-
-        self.detector = CRAFT()
-
-        state_dict = torch.load(self.detector_path, map_location=self.device)
-        if list(state_dict.keys())[0].startswith("module"):
-            state_dict = {k[7:]: v for k, v in state_dict.items()}
-
-        self.detector.load_state_dict(state_dict)
-
-        if self.device == "cpu":
-            if self.quantize:
-                try:
-                    torch.quantization.quantize_dynamic(
-                        self.detector, dtype=torch.qint8, inplace=True
-                    )
-                except:
-                    pass
-        else:
-            self.detector = torch.nn.DataParallel(self.detector).to(self.device)
-            import torch.backends.cudnn as cudnn
-
-            cudnn.benchmark = self.cudnn_benchmark
-
-        self.detector.eval()
-
-    def process(
-        self,
-        image_path: str,
-        max_size: int = 960,
-        expand_ratio: float = 1.0,
-        sharp: float = 1.0,
-        contrast: float = 1.0,
-        text_confidence: float = 0.7,
-        text_threshold: float = 0.4,
-        link_threshold: float = 0.4,
-        slope_ths: float = 0.1,
-        ratio_ths: float = 0.5,
-        center_ths: float = 0.5,
-        dim_ths: float = 0.5,
-        space_ths: float = 1.0,
-        add_margin: float = 0.1,
-        min_size: float = 0.01,
-    ) -> Tuple[BoxTuple, list[np.ndarray]]:
-
-        image = Image.open(image_path).convert('RGB')
-
-        tensor, inverse_ratio = self.preprocess(
-            image, max_size, expand_ratio, sharp, contrast
-        )
-
-        scores = self.forward_net(tensor)
-
-        boxes = self.detect(scores, text_confidence, text_threshold, link_threshold)
-
-        image = np.array(image)
-        region_list, box_list = self.postprocess(
-            image,
-            boxes,
-            inverse_ratio,
-            slope_ths,
-            ratio_ths,
-            center_ths,
-            dim_ths,
-            space_ths,
-            add_margin,
-            min_size,
-        )
-
-        # get cropped image
-        image_list = []
-        for region in region_list:
-            x_min, x_max, y_min, y_max = region
-            crop_img = image[y_min:y_max, x_min:x_max, :]
-            image_list.append(
-                (
-                    ((x_min, y_min), (x_max, y_min), (x_max, y_max), (x_min, y_max)),
-                    crop_img,
-                )
-            )
-
-        for box in box_list:
-            transformed_img = boxed_transform(image, np.array(box, dtype="float32"))
-            image_list.append((box, transformed_img))
-
-        # sort by top left point
-        image_list = sorted(image_list, key=lambda x: (x[0][0][1], x[0][0][0]))
-
-        return image_list
-
-    def preprocess(
-        self,
-        image: Image.Image,
-        max_size: int,
-        expand_ratio: float = 1.0,
-        sharp: float = 1.0,
-        contrast: float = 1.0,
-    ) -> torch.Tensor:
-        if sharp != 1:
-            enhancer = ImageEnhance.Sharpness(image)
-            image = enhancer.enhance(sharp)
-        if contrast != 1:
-            enhancer = ImageEnhance.Contrast(image)
-            image = enhancer.enhance(contrast)
-
-        image = np.array(image)
-
-        image, target_ratio = resize_aspect_ratio(
-            image, max_size, interpolation=cv2.INTER_LINEAR, expand_ratio=expand_ratio
-        )
-        inverse_ratio = 1 / target_ratio
-
-        x = np.transpose(normalize_mean_variance(image), (2, 0, 1))
-
-        x = torch.tensor(np.array([x]), device=self.device)
-
-        return x, inverse_ratio
-
-    @torch.no_grad()
-    def forward_net(self, tensor: torch.Tensor) -> torch.Tensor:
-        scores, feature = self.detector(tensor)
-        return scores[0]
-
-    def detect(
-        self,
-        scores: torch.Tensor,
-        text_confidence: float = 0.7,
-        text_threshold: float = 0.4,
-        link_threshold: float = 0.4,
-    ) -> list[BoxTuple]:
-        # make score and link map
-        score_text = scores[:, :, 0].cpu().data.numpy()
-        score_link = scores[:, :, 1].cpu().data.numpy()
-        # extract box
-        boxes, _ = extract_boxes(
-            score_text, score_link, text_confidence, text_threshold, link_threshold
-        )
-        return boxes
-
-    def postprocess(
-        self,
-        image: np.ndarray,
-        boxes: list[BoxTuple],
-        inverse_ratio: float,
-        slope_ths: float = 0.1,
-        ratio_ths: float = 0.5,
-        center_ths: float = 0.5,
-        dim_ths: float = 0.5,
-        space_ths: float = 1.0,
-        add_margin: float = 0.1,
-        min_size: int = 0,
-    ) -> Tuple[list[RegionTuple], list[BoxTuple]]:
-
-        # coordinate adjustment
-        boxes = adjust_result_coordinates(boxes, inverse_ratio)
-
-        max_y, max_x, _ = image.shape
-
-        # extract region and merge
-        region_list, box_list = extract_regions_from_boxes(boxes, slope_ths)
-
-        region_list = greedy_merge(
-            region_list,
-            ratio_ths=ratio_ths,
-            center_ths=center_ths,
-            dim_ths=dim_ths,
-            space_ths=space_ths,
-            verbose=0
-        )
-
-        # add margin
-        region_list = [
-            region.expand(add_margin, (max_x, max_y)).as_tuple()
-            for region in region_list
-        ]
-
-        box_list = [box_expand(box, add_margin, (max_x, max_y)) for box in box_list]
-
-        # filter by size
-        if min_size:
-            if min_size < 1:
-                min_size = int(min(max_y, max_x) * min_size)
-
-            region_list = [
-                i for i in region_list if max(i[1] - i[0], i[3] - i[2]) > min_size
-            ]
-            box_list = [
-                i
-                for i in box_list
-                if max(utils.diff([c[0] for c in i]), utils.diff([c[1] for c in i]))
-                > min_size
-            ]
-
-        return region_list, box_list

ezocr/build/lib/easyocrlite/types.py

@@ -1,5 +0,0 @@
-from typing import Tuple
-
-Point = Tuple[int, int]
-BoxTuple = Tuple[Point, Point, Point, Point]
-RegionTuple = Tuple[int, int, int, int]

ezocr/easyocrlite.egg-info/PKG-INFO

@@ -1,8 +0,0 @@
-Metadata-Version: 2.1
-Name: easyocrlite
-Version: 0.0.1
-License: Apache License 2.0
-Keywords: ocr optical character recognition deep learning neural network
-Classifier: Development Status :: 5 - Production/Stable
-Requires-Python: >=3.7
-License-File: LICENSE

ezocr/easyocrlite.egg-info/SOURCES.txt

@@ -1,11 +0,0 @@
-LICENSE
-README.md
-setup.py
-easyocrlite/__init__.py
-easyocrlite/reader.py
-easyocrlite/types.py
-easyocrlite.egg-info/PKG-INFO
-easyocrlite.egg-info/SOURCES.txt
-easyocrlite.egg-info/dependency_links.txt
-easyocrlite.egg-info/requires.txt
-easyocrlite.egg-info/top_level.txt

ezocr/easyocrlite.egg-info/dependency_links.txt

@@ -1 +0,0 @@
-

ezocr/easyocrlite.egg-info/requires.txt

@@ -1,5 +0,0 @@
-torch
-torchvision>=0.5
-opencv-python-headless<=4.5.4.60
-numpy
-Pillow

ezocr/easyocrlite.egg-info/top_level.txt

@@ -1 +0,0 @@
-easyocrlite

ezocr/easyocrlite/__init__.py

@@ -1 +0,0 @@
-from easyocrlite.reader import ReaderLite

ezocr/easyocrlite/model/__init__.py

@@ -1 +0,0 @@
-from .craft import CRAFT

ezocr/easyocrlite/model/craft.py

@@ -1,174 +0,0 @@
-"""
-Copyright (c) 2019-present NAVER Corp.
-MIT License
-"""
-from __future__ import annotations
-
-from collections import namedtuple
-from typing import Iterable, Tuple
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import torchvision
-from packaging import version
-from torchvision import models
-
-VGGOutputs = namedtuple(
-    "VggOutputs", ["fc7", "relu5_3", "relu4_3", "relu3_2", "relu2_2"]
-)
-
-def init_weights(modules: Iterable[nn.Module]):
-    for m in modules:
-        if isinstance(m, nn.Conv2d):
-            nn.init.xavier_uniform_(m.weight)
-            if m.bias is not None:
-                nn.init.zeros_(m.bias)
-        elif isinstance(m, nn.BatchNorm2d):
-            nn.init.constant_(m.weight, 1.0)
-            nn.init.zeros_(m.bias)
-        elif isinstance(m, nn.Linear):
-            nn.init.normal_(m.weight, 0, 0.01)
-            nn.init.zeros_(m.bias)
-
-
-class VGG16_BN(nn.Module):
-    def __init__(self, pretrained: bool=True, freeze: bool=True):
-        super().__init__()
-        if version.parse(torchvision.__version__) >= version.parse("0.13"):
-            vgg_pretrained_features = models.vgg16_bn(
-                weights=models.VGG16_BN_Weights.DEFAULT if pretrained else None
-            ).features
-        else:  # torchvision.__version__ < 0.13
-            models.vgg.model_urls["vgg16_bn"] = models.vgg.model_urls[
-                "vgg16_bn"
-            ].replace("https://", "http://")
-            vgg_pretrained_features = models.vgg16_bn(pretrained=pretrained).features
-
-        self.slice1 = torch.nn.Sequential()
-        self.slice2 = torch.nn.Sequential()
-        self.slice3 = torch.nn.Sequential()
-        self.slice4 = torch.nn.Sequential()
-        self.slice5 = torch.nn.Sequential()
-        for x in range(12):  # conv2_2
-            self.slice1.add_module(str(x), vgg_pretrained_features[x])
-        for x in range(12, 19):  # conv3_3
-            self.slice2.add_module(str(x), vgg_pretrained_features[x])
-        for x in range(19, 29):  # conv4_3
-            self.slice3.add_module(str(x), vgg_pretrained_features[x])
-        for x in range(29, 39):  # conv5_3
-            self.slice4.add_module(str(x), vgg_pretrained_features[x])
-
-        # fc6, fc7 without atrous conv
-        self.slice5 = torch.nn.Sequential(
-            nn.MaxPool2d(kernel_size=3, stride=1, padding=1),
-            nn.Conv2d(512, 1024, kernel_size=3, padding=6, dilation=6),
-            nn.Conv2d(1024, 1024, kernel_size=1),
-        )
-
-        if not pretrained:
-            init_weights(self.slice1.modules())
-            init_weights(self.slice2.modules())
-            init_weights(self.slice3.modules())
-            init_weights(self.slice4.modules())
-
-        init_weights(self.slice5.modules())  # no pretrained model for fc6 and fc7
-
-        if freeze:
-            for param in self.slice1.parameters():  # only first conv
-                param.requires_grad = False
-
-    def forward(self, x: torch.Tensor) -> VGGOutputs:
-        h = self.slice1(x)
-        h_relu2_2 = h
-        h = self.slice2(h)
-        h_relu3_2 = h
-        h = self.slice3(h)
-        h_relu4_3 = h
-        h = self.slice4(h)
-        h_relu5_3 = h
-        h = self.slice5(h)
-        h_fc7 = h
-        
-        out = VGGOutputs(h_fc7, h_relu5_3, h_relu4_3, h_relu3_2, h_relu2_2)
-        return out
-
-
-class DoubleConv(nn.Module):
-    def __init__(self, in_ch: int, mid_ch: int, out_ch: int):
-        super().__init__()
-        self.conv = nn.Sequential(
-            nn.Conv2d(in_ch + mid_ch, mid_ch, kernel_size=1),
-            nn.BatchNorm2d(mid_ch),
-            nn.ReLU(inplace=True),
-            nn.Conv2d(mid_ch, out_ch, kernel_size=3, padding=1),
-            nn.BatchNorm2d(out_ch),
-            nn.ReLU(inplace=True),
-        )
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = self.conv(x)
-        return x
-
-
-class CRAFT(nn.Module):
-    def __init__(self, pretrained: bool=False, freeze: bool=False):
-        super(CRAFT, self).__init__()
-
-        """ Base network """
-        self.basenet = VGG16_BN(pretrained, freeze)
-
-        """ U network """
-        self.upconv1 = DoubleConv(1024, 512, 256)
-        self.upconv2 = DoubleConv(512, 256, 128)
-        self.upconv3 = DoubleConv(256, 128, 64)
-        self.upconv4 = DoubleConv(128, 64, 32)
-
-        num_class = 2
-        self.conv_cls = nn.Sequential(
-            nn.Conv2d(32, 32, kernel_size=3, padding=1),
-            nn.ReLU(inplace=True),
-            nn.Conv2d(32, 32, kernel_size=3, padding=1),
-            nn.ReLU(inplace=True),
-            nn.Conv2d(32, 16, kernel_size=3, padding=1),
-            nn.ReLU(inplace=True),
-            nn.Conv2d(16, 16, kernel_size=1),
-            nn.ReLU(inplace=True),
-            nn.Conv2d(16, num_class, kernel_size=1),
-        )
-
-        init_weights(self.upconv1.modules())
-        init_weights(self.upconv2.modules())
-        init_weights(self.upconv3.modules())
-        init_weights(self.upconv4.modules())
-        init_weights(self.conv_cls.modules())
-
-    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
-        """Base network"""
-        sources = self.basenet(x)
-
-        """ U network """
-        y = torch.cat([sources[0], sources[1]], dim=1)
-        y = self.upconv1(y)
-
-        y = F.interpolate(
-            y, size=sources[2].size()[2:], mode="bilinear", align_corners=False
-        )
-        y = torch.cat([y, sources[2]], dim=1)
-        y = self.upconv2(y)
-
-        y = F.interpolate(
-            y, size=sources[3].size()[2:], mode="bilinear", align_corners=False
-        )
-        y = torch.cat([y, sources[3]], dim=1)
-        y = self.upconv3(y)
-
-        y = F.interpolate(
-            y, size=sources[4].size()[2:], mode="bilinear", align_corners=False
-        )
-        y = torch.cat([y, sources[4]], dim=1)
-        feature = self.upconv4(y)
-
-        y = self.conv_cls(feature)
-
-        return y.permute(0, 2, 3, 1), feature

ezocr/easyocrlite/reader.py

@@ -1,271 +0,0 @@
-from __future__ import annotations
-
-import logging
-import os
-from pathlib import Path
-from typing import Tuple
-
-import cv2
-import numpy as np
-import torch
-from PIL import Image, ImageEnhance
-
-from easyocrlite.model import CRAFT
-
-from easyocrlite.utils.download_utils import prepare_model
-from easyocrlite.utils.image_utils import (
-    adjust_result_coordinates,
-    boxed_transform,
-    normalize_mean_variance,
-    resize_aspect_ratio,
-)
-from easyocrlite.utils.detect_utils import (
-    extract_boxes,
-    extract_regions_from_boxes,
-    box_expand,
-    greedy_merge,
-)
-from easyocrlite.types import BoxTuple, RegionTuple
-import easyocrlite.utils.utils as utils
-
-logger = logging.getLogger(__name__)
-
-MODULE_PATH = (
-    os.environ.get("EASYOCR_MODULE_PATH")
-    or os.environ.get("MODULE_PATH")
-    or os.path.expanduser("~/.EasyOCR/")
-)
-
-
-class ReaderLite(object):
-    def __init__(
-        self,
-        gpu=True,
-        model_storage_directory=None,
-        download_enabled=True,
-        verbose=True,
-        quantize=True,
-        cudnn_benchmark=False,
-    ):
-
-        self.verbose = verbose
-
-        model_storage_directory = Path(
-            model_storage_directory
-            if model_storage_directory
-            else MODULE_PATH + "/model"
-        )
-        self.detector_path = prepare_model(
-            model_storage_directory, download_enabled, verbose
-        )
-
-        self.quantize = quantize
-        self.cudnn_benchmark = cudnn_benchmark
-        if gpu is False:
-            self.device = "cpu"
-            if verbose:
-                logger.warning(
-                    "Using CPU. Note: This module is much faster with a GPU."
-                )
-        elif not torch.cuda.is_available():
-            self.device = "cpu"
-            if verbose:
-                logger.warning(
-                    "CUDA not available - defaulting to CPU. Note: This module is much faster with a GPU."
-                )
-        elif gpu is True:
-            self.device = "cuda"
-        else:
-            self.device = gpu
-
-        self.detector = CRAFT()
-
-        state_dict = torch.load(self.detector_path, map_location=self.device)
-        if list(state_dict.keys())[0].startswith("module"):
-            state_dict = {k[7:]: v for k, v in state_dict.items()}
-
-        self.detector.load_state_dict(state_dict)
-
-        if self.device == "cpu":
-            if self.quantize:
-                try:
-                    torch.quantization.quantize_dynamic(
-                        self.detector, dtype=torch.qint8, inplace=True
-                    )
-                except:
-                    pass
-        else:
-            self.detector = torch.nn.DataParallel(self.detector).to(self.device)
-            import torch.backends.cudnn as cudnn
-
-            cudnn.benchmark = self.cudnn_benchmark
-
-        self.detector.eval()
-
-    def process(
-        self,
-        image_path: str,
-        max_size: int = 960,
-        expand_ratio: float = 1.0,
-        sharp: float = 1.0,
-        contrast: float = 1.0,
-        text_confidence: float = 0.7,
-        text_threshold: float = 0.4,
-        link_threshold: float = 0.4,
-        slope_ths: float = 0.1,
-        ratio_ths: float = 0.5,
-        center_ths: float = 0.5,
-        dim_ths: float = 0.5,
-        space_ths: float = 1.0,
-        add_margin: float = 0.1,
-        min_size: float = 0.01,
-    ) -> Tuple[BoxTuple, list[np.ndarray]]:
-
-        image = Image.open(image_path).convert('RGB')
-        tensor, inverse_ratio = self.preprocess(
-            image, max_size, expand_ratio, sharp, contrast
-        )
-
-        scores = self.forward_net(tensor)
-
-        boxes = self.detect(scores, text_confidence, text_threshold, link_threshold)
-
-        image = np.array(image)
-        region_list, box_list = self.postprocess(
-            image,
-            boxes,
-            inverse_ratio,
-            slope_ths,
-            ratio_ths,
-            center_ths,
-            dim_ths,
-            space_ths,
-            add_margin,
-            min_size,
-        )
-
-        # get cropped image
-        image_list = []
-        for region in region_list:
-            x_min, x_max, y_min, y_max = region
-            crop_img = image[y_min:y_max, x_min:x_max, :]
-            image_list.append(
-                (
-                    ((x_min, y_min), (x_max, y_min), (x_max, y_max), (x_min, y_max)),
-                    crop_img,
-                )
-            )
-
-        for box in box_list:
-            transformed_img = boxed_transform(image, np.array(box, dtype="float32"))
-            image_list.append((box, transformed_img))
-
-        # sort by top left point
-        image_list = sorted(image_list, key=lambda x: (x[0][0][1], x[0][0][0]))
-
-        return image_list
-
-    def preprocess(
-        self,
-        image: Image.Image,
-        max_size: int,
-        expand_ratio: float = 1.0,
-        sharp: float = 1.0,
-        contrast: float = 1.0,
-    ) -> torch.Tensor:
-        if sharp != 1:
-            enhancer = ImageEnhance.Sharpness(image)
-            image = enhancer.enhance(sharp)
-        if contrast != 1:
-            enhancer = ImageEnhance.Contrast(image)
-            image = enhancer.enhance(contrast)
-
-        image = np.array(image)
-
-        image, target_ratio = resize_aspect_ratio(
-            image, max_size, interpolation=cv2.INTER_LINEAR, expand_ratio=expand_ratio
-        )
-        inverse_ratio = 1 / target_ratio
-
-        x = np.transpose(normalize_mean_variance(image), (2, 0, 1))
-
-        x = torch.tensor(np.array([x]), device=self.device)
-
-        return x, inverse_ratio
-
-    @torch.no_grad()
-    def forward_net(self, tensor: torch.Tensor) -> torch.Tensor:
-        scores, feature = self.detector(tensor)
-        return scores[0]
-
-    def detect(
-        self,
-        scores: torch.Tensor,
-        text_confidence: float = 0.7,
-        text_threshold: float = 0.4,
-        link_threshold: float = 0.4,
-    ) -> list[BoxTuple]:
-        # make score and link map
-        score_text = scores[:, :, 0].cpu().data.numpy()
-        score_link = scores[:, :, 1].cpu().data.numpy()
-        # extract box
-        boxes, _ = extract_boxes(
-            score_text, score_link, text_confidence, text_threshold, link_threshold
-        )
-        return boxes
-
-    def postprocess(
-        self,
-        image: np.ndarray,
-        boxes: list[BoxTuple],
-        inverse_ratio: float,
-        slope_ths: float = 0.1,
-        ratio_ths: float = 0.5,
-        center_ths: float = 0.5,
-        dim_ths: float = 0.5,
-        space_ths: float = 1.0,
-        add_margin: float = 0.1,
-        min_size: int = 0,
-    ) -> Tuple[list[RegionTuple], list[BoxTuple]]:
-
-        # coordinate adjustment
-        boxes = adjust_result_coordinates(boxes, inverse_ratio)
-
-        max_y, max_x, _ = image.shape
-
-        # extract region and merge
-        region_list, box_list = extract_regions_from_boxes(boxes, slope_ths)
-
-        region_list = greedy_merge(
-            region_list,
-            ratio_ths=ratio_ths,
-            center_ths=center_ths,
-            dim_ths=dim_ths,
-            space_ths=space_ths,
-            verbose=0
-        )
-
-        # add margin
-        region_list = [
-            region.expand(add_margin, (max_x, max_y)).as_tuple()
-            for region in region_list
-        ]
-
-        box_list = [box_expand(box, add_margin, (max_x, max_y)) for box in box_list]
-
-        # filter by size
-        if min_size:
-            if min_size < 1:
-                min_size = int(min(max_y, max_x) * min_size)
-
-            region_list = [
-                i for i in region_list if max(i[1] - i[0], i[3] - i[2]) > min_size
-            ]
-            box_list = [
-                i
-                for i in box_list
-                if max(utils.diff([c[0] for c in i]), utils.diff([c[1] for c in i]))
-                > min_size
-            ]
-
-        return region_list, box_list

ezocr/easyocrlite/types.py

@@ -1,5 +0,0 @@
-from typing import Tuple
-
-Point = Tuple[int, int]
-BoxTuple = Tuple[Point, Point, Point, Point]
-RegionTuple = Tuple[int, int, int, int]