# coding=utf-8
# Copyright 2019 The TensorFlow Datasets Authors.
#
# 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.

"""The General Language Understanding Evaluation (GLUE) benchmark."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import csv
import os
import numpy as np
import six

import tensorflow as tf
from tensorflow_datasets.core import api_utils
import tensorflow_datasets.public_api as tfds

_GLUE_CITATION = """\
@inproceedings{wang2019glue,
  title={{GLUE}: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding},
  author={Wang, Alex and Singh, Amanpreet and Michael, Julian and Hill, Felix and Levy, Omer and Bowman, Samuel R.},
  note={In the Proceedings of ICLR.},
  year={2019}
}

Note that each GLUE dataset has its own citation. Please see the source to see
the correct citation for each contained dataset."""

_MRPC_DEV_IDS = "https://firebasestorage.googleapis.com/v0/b/mtl-sentence-representations.appspot.com/o/data%2Fmrpc_dev_ids.tsv?alt=media&token=ec5c0836-31d5-48f4-b431-7480817f1adc"
_MRPC_TRAIN = "https://dl.fbaipublicfiles.com/senteval/senteval_data/msr_paraphrase_train.txt"
_MRPC_TEST = "https://dl.fbaipublicfiles.com/senteval/senteval_data/msr_paraphrase_test.txt"

_MNLI_BASE_KWARGS = dict(
    text_features={
        "premise": "sentence1",
        "hypothesis": "sentence2",
    },
    label_classes=["entailment", "neutral", "contradiction"],
    label_column="gold_label",
    data_url="https://firebasestorage.googleapis.com/v0/b/mtl-sentence-representations.appspot.com/o/data%2FMNLI.zip?alt=media&token=50329ea1-e339-40e2-809c-10c40afff3ce",
    data_dir="MNLI",
    citation="""\
      @InProceedings{N18-1101,
        author = "Williams, Adina
                  and Nangia, Nikita
                  and Bowman, Samuel",
        title = "A Broad-Coverage Challenge Corpus for
                 Sentence Understanding through Inference",
        booktitle = "Proceedings of the 2018 Conference of
                     the North American Chapter of the
                     Association for Computational Linguistics:
                     Human Language Technologies, Volume 1 (Long
                     Papers)",
        year = "2018",
        publisher = "Association for Computational Linguistics",
        pages = "1112--1122",
        location = "New Orleans, Louisiana",
        url = "http://aclweb.org/anthology/N18-1101"
      }
      @article{bowman2015large,
        title={A large annotated corpus for learning natural language inference},
        author={Bowman, Samuel R and Angeli, Gabor and Potts, Christopher and Manning, Christopher D},
        journal={arXiv preprint arXiv:1508.05326},
        year={2015}
      }""",
    url="http://www.nyu.edu/projects/bowman/multinli/",
    train_shards=2)


class GlueConfig(tfds.core.BuilderConfig):
  """BuilderConfig for GLUE."""

  @api_utils.disallow_positional_args
  def __init__(self,
               text_features,
               label_column,
               data_url,
               data_dir,
               citation,
               url,
               label_classes=None,
               train_shards=1,
               process_label=lambda x: x,
               **kwargs):
    """BuilderConfig for GLUE.

    Args:
      text_features: `dict[string, string]`, map from the name of the feature
        dict for each text field to the name of the column in the tsv file
      label_column: `string`, name of the column in the tsv file corresponding
        to the label
      data_url: `string`, url to download the zip file from
      data_dir: `string`, the path to the folder containing the tsv files in the
        downloaded zip
      citation: `string`, citation for the data set
      url: `string`, url for information about the data set
      label_classes: `list[string]`, the list of classes if the label is
        categorical. If not provided, then the label will be of type
        `tf.float32`.
      train_shards: `int`, number of shards for the train data set
      process_label: `Function[string, any]`, function taking in the raw value
        of the label and processing it to the form required by the label feature
      **kwargs: keyword arguments forwarded to super.
    """
    # Version history:
    # 1.0.0: S3 (new shuffling, sharding and slicing mechanism).
    # 0.0.2: Initial version.
    super(GlueConfig, self).__init__(
        version=tfds.core.Version(
            "0.0.2", experiments={tfds.core.Experiment.S3: False}),
        supported_versions=[
            tfds.core.Version("1.0.0"),
        ],
        **kwargs)
    self.text_features = text_features
    self.label_column = label_column
    self.label_classes = label_classes
    self.data_url = data_url
    self.data_dir = data_dir
    self.citation = citation
    self.url = url
    self.train_shards = train_shards
    self.process_label = process_label


class Glue(tfds.core.GeneratorBasedBuilder):
  """The General Language Understanding Evaluation (GLUE) benchmark."""
  BUILDER_CONFIGS = [
      GlueConfig(
          name="cola",
          description="""\
            The Corpus of Linguistic Acceptability consists of English
            acceptability judgments drawn from books and journal articles on
            linguistic theory. Each example is a sequence of words annotated
            with whether it is a grammatical English sentence.""",
          text_features={"sentence": "sentence"},
          label_classes=["unacceptable", "acceptable"],
          label_column="is_acceptable",
          data_url="https://firebasestorage.googleapis.com/v0/b/mtl-sentence-representations.appspot.com/o/data%2FCoLA.zip?alt=media&token=46d5e637-3411-4188-bc44-5809b5bfb5f4",
          data_dir="CoLA",
          citation="""\
            @article{warstadt2018neural,
              title={Neural Network Acceptability Judgments},
              author={Warstadt, Alex and Singh, Amanpreet and Bowman, Samuel R},
              journal={arXiv preprint arXiv:1805.12471},
              year={2018}
            }""",
          url="https://nyu-mll.github.io/CoLA/"),
      GlueConfig(
          name="sst2",
          description="""\
            The Stanford Sentiment Treebank consists of sentences from movie reviews and
            human annotations of their sentiment. The task is to predict the sentiment of a
            given sentence. We use the two-way (positive/negative) class split, and use only
            sentence-level labels.""",
          text_features={"sentence": "sentence"},
          label_classes=["negative", "positive"],
          label_column="label",
          data_url="https://firebasestorage.googleapis.com/v0/b/mtl-sentence-representations.appspot.com/o/data%2FSST-2.zip?alt=media&token=aabc5f6b-e466-44a2-b9b4-cf6337f84ac8",
          data_dir="SST-2",
          citation="""\
            @inproceedings{socher2013recursive,
              title={Recursive deep models for semantic compositionality over a sentiment treebank},
              author={Socher, Richard and Perelygin, Alex and Wu, Jean and Chuang, Jason and Manning, Christopher D and Ng, Andrew and Potts, Christopher},
              booktitle={Proceedings of the 2013 conference on empirical methods in natural language processing},
              pages={1631--1642},
              year={2013}
            }""",
          url="https://nlp.stanford.edu/sentiment/index.html"),
      GlueConfig(
          name="mrpc",
          description="""\
            The Microsoft Research Paraphrase Corpus (Dolan & Brockett, 2005) is a corpus of
            sentence pairs automatically extracted from online news sources, with human annotations
            for whether the sentences in the pair are semantically equivalent.""",
          text_features={
              "sentence1": "",
              "sentence2": ""
          },
          label_classes=["not_equivalent", "equivalent"],
          label_column="Quality",
          data_url="",  # MRPC isn't hosted by GLUE.
          data_dir="MRPC",
          citation="""\
            @inproceedings{dolan2005automatically,
              title={Automatically constructing a corpus of sentential paraphrases},
              author={Dolan, William B and Brockett, Chris},
              booktitle={Proceedings of the Third International Workshop on Paraphrasing (IWP2005)},
              year={2005}
            }""",
          url="https://www.microsoft.com/en-us/download/details.aspx?id=52398"
      ),
      GlueConfig(
          name="qqp",
          description="""\
            The Quora Question Pairs2 dataset is a collection of question pairs from the
            community question-answering website Quora. The task is to determine whether a
            pair of questions are semantically equivalent.""",
          text_features={
              "question1": "question1",
              "question2": "question2",
          },
          label_classes=["not_duplicate", "duplicate"],
          label_column="is_duplicate",
          data_url="https://firebasestorage.googleapis.com/v0/b/mtl-sentence-representations.appspot.com/o/data%2FQQP.zip?alt=media&token=700c6acf-160d-4d89-81d1-de4191d02cb5",
          data_dir="QQP",
          citation="""\
          @online{WinNT,
            author = {Iyer, Shankar and Dandekar, Nikhil and Csernai, Kornel},
            title = {First Quora Dataset Release: Question Pairs},
            year = 2017,
            url = {https://data.quora.com/First-Quora-Dataset-Release-Question-Pairs},
            urldate = {2019-04-03}
          }""",
          url="https://data.quora.com/First-Quora-Dataset-Release-Question-Pairs"
      ),
      GlueConfig(
          name="stsb",
          description="""\
            The Semantic Textual Similarity Benchmark (Cer et al., 2017) is a collection of
            sentence pairs drawn from news headlines, video and image captions, and natural
            language inference data. Each pair is human-annotated with a similarity score
            from 1 to 5.""",
          text_features={
              "sentence1": "sentence1",
              "sentence2": "sentence2",
          },
          label_column="score",
          data_url="https://firebasestorage.googleapis.com/v0/b/mtl-sentence-representations.appspot.com/o/data%2FSTS-B.zip?alt=media&token=bddb94a7-8706-4e0d-a694-1109e12273b5",
          data_dir="STS-B",
          citation="""\
            @article{cer2017semeval,
              title={Semeval-2017 task 1: Semantic textual similarity-multilingual and cross-lingual focused evaluation},
              author={Cer, Daniel and Diab, Mona and Agirre, Eneko and Lopez-Gazpio, Inigo and Specia, Lucia},
              journal={arXiv preprint arXiv:1708.00055},
              year={2017}
            }""",
          url="http://ixa2.si.ehu.es/stswiki/index.php/STSbenchmark",
          process_label=np.float32),
      GlueConfig(
          name="mnli",
          description="""\
            The Multi-Genre Natural Language Inference Corpusn is a crowdsourced
            collection of sentence pairs with textual entailment annotations. Given a premise sentence
            and a hypothesis sentence, the task is to predict whether the premise entails the hypothesis
            (entailment), contradicts the hypothesis (contradiction), or neither (neutral). The premise sentences are
            gathered from ten different sources, including transcribed speech, fiction, and government reports.
            We use the standard test set, for which we obtained private labels from the authors, and evaluate
            on both the matched (in-domain) and mismatched (cross-domain) section. We also use and recommend
            the SNLI corpus as 550k examples of auxiliary training data.""",
          **_MNLI_BASE_KWARGS),
      GlueConfig(
          name="mnli_mismatched",
          description="""\
          The mismatched validation and test splits from MNLI.
          See the "mnli" BuilderConfig for additional information.""",
          **_MNLI_BASE_KWARGS),
      GlueConfig(
          name="mnli_matched",
          description="""\
          The matched validation and test splits from MNLI.
          See the "mnli" BuilderConfig for additional information.""",
          **_MNLI_BASE_KWARGS),
      GlueConfig(
          name="qnli",
          description="""\
            The Stanford Question Answering Dataset is a question-answering
            dataset consisting of question-paragraph pairs, where one of the sentences in the paragraph (drawn
            from Wikipedia) contains the answer to the corresponding question (written by an annotator). We
            convert the task into sentence pair classification by forming a pair between each question and each
            sentence in the corresponding context, and filtering out pairs with low lexical overlap between the
            question and the context sentence. The task is to determine whether the context sentence contains
            the answer to the question. This modified version of the original task removes the requirement that
            the model select the exact answer, but also removes the simplifying assumptions that the answer
            is always present in the input and that lexical overlap is a reliable cue.""",
          text_features={
              "question": "question",
              "sentence": "sentence",
          },
          label_classes=["entailment", "not_entailment"],
          label_column="label",
          data_url="https://firebasestorage.googleapis.com/v0/b/mtl-sentence-representations.appspot.com/o/data%2FQNLIv2.zip?alt=media&token=6fdcf570-0fc5-4631-8456-9505272d1601",
          data_dir="QNLI",
          citation="""\
            @article{rajpurkar2016squad,
              title={Squad: 100,000+ questions for machine comprehension of text},
              author={Rajpurkar, Pranav and Zhang, Jian and Lopyrev, Konstantin and Liang, Percy},
              journal={arXiv preprint arXiv:1606.05250},
              year={2016}
            }""",
          url="https://rajpurkar.github.io/SQuAD-explorer/"),
      GlueConfig(
          name="rte",
          description="""\
            The Recognizing Textual Entailment (RTE) datasets come from a series of annual textual
            entailment challenges. We combine the data from RTE1 (Dagan et al., 2006), RTE2 (Bar Haim
            et al., 2006), RTE3 (Giampiccolo et al., 2007), and RTE5 (Bentivogli et al., 2009).4 Examples are
            constructed based on news and Wikipedia text. We convert all datasets to a two-class split, where
            for three-class datasets we collapse neutral and contradiction into not entailment, for consistency.""",
          text_features={
              "sentence1": "sentence1",
              "sentence2": "sentence2",
          },
          label_classes=["entailment", "not_entailment"],
          label_column="label",
          data_url="https://firebasestorage.googleapis.com/v0/b/mtl-sentence-representations.appspot.com/o/data%2FRTE.zip?alt=media&token=5efa7e85-a0bb-4f19-8ea2-9e1840f077fb",
          data_dir="RTE",
          citation="""\
            @inproceedings{dagan2005pascal,
              title={The PASCAL recognising textual entailment challenge},
              author={Dagan, Ido and Glickman, Oren and Magnini, Bernardo},
              booktitle={Machine Learning Challenges Workshop},
              pages={177--190},
              year={2005},
              organization={Springer}
            }
            @inproceedings{bar2006second,
              title={The second pascal recognising textual entailment challenge},
              author={Bar-Haim, Roy and Dagan, Ido and Dolan, Bill and Ferro, Lisa and Giampiccolo, Danilo and Magnini, Bernardo and Szpektor, Idan},
              booktitle={Proceedings of the second PASCAL challenges workshop on recognising textual entailment},
              volume={6},
              number={1},
              pages={6--4},
              year={2006},
              organization={Venice}
            }
            @inproceedings{giampiccolo2007third,
              title={The third pascal recognizing textual entailment challenge},
              author={Giampiccolo, Danilo and Magnini, Bernardo and Dagan, Ido and Dolan, Bill},
              booktitle={Proceedings of the ACL-PASCAL workshop on textual entailment and paraphrasing},
              pages={1--9},
              year={2007},
              organization={Association for Computational Linguistics}
            }
            @inproceedings{bentivogli2009fifth,
              title={The Fifth PASCAL Recognizing Textual Entailment Challenge.},
              author={Bentivogli, Luisa and Clark, Peter and Dagan, Ido and Giampiccolo, Danilo},
              booktitle={TAC},
              year={2009}
            }""",
          url="https://aclweb.org/aclwiki/Recognizing_Textual_Entailment"
      ),
      GlueConfig(
          name="wnli",
          description="""\
            The Winograd Schema Challenge (Levesque et al., 2011) is a reading comprehension task
            in which a system must read a sentence with a pronoun and select the referent of that pronoun from
            a list of choices. The examples are manually constructed to foil simple statistical methods: Each
            one is contingent on contextual information provided by a single word or phrase in the sentence.
            To convert the problem into sentence pair classification, we construct sentence pairs by replacing
            the ambiguous pronoun with each possible referent. The task is to predict if the sentence with the
            pronoun substituted is entailed by the original sentence. We use a small evaluation set consisting of
            new examples derived from fiction books that was shared privately by the authors of the original
            corpus. While the included training set is balanced between two classes, the test set is imbalanced
            between them (65% not entailment). Also, due to a data quirk, the development set is adversarial:
            hypotheses are sometimes shared between training and development examples, so if a model memorizes the
            training examples, they will predict the wrong label on corresponding development set
            example. As with QNLI, each example is evaluated separately, so there is not a systematic correspondence
            between a model's score on this task and its score on the unconverted original task. We
            call converted dataset WNLI (Winograd NLI).""",
          text_features={
              "sentence1": "sentence1",
              "sentence2": "sentence2",
          },
          label_classes=["not_entailment", "entailment"],
          label_column="label",
          data_url="https://firebasestorage.googleapis.com/v0/b/mtl-sentence-representations.appspot.com/o/data%2FWNLI.zip?alt=media&token=068ad0a0-ded7-4bd7-99a5-5e00222e0faf",
          data_dir="WNLI",
          citation="""\
            @inproceedings{levesque2012winograd,
              title={The winograd schema challenge},
              author={Levesque, Hector and Davis, Ernest and Morgenstern, Leora},
              booktitle={Thirteenth International Conference on the Principles of Knowledge Representation and Reasoning},
              year={2012}
            }""",
          url="https://cs.nyu.edu/faculty/davise/papers/WinogradSchemas/WS.html"
      ),
  ]

  def _info(self):
    features = {
        text_feature: tfds.features.Text()
        for text_feature in six.iterkeys(self.builder_config.text_features)
    }
    if self.builder_config.label_classes:
      features["label"] = tfds.features.ClassLabel(
          names=self.builder_config.label_classes)
    else:
      features["label"] = tf.float32
    features["idx"] = tf.int32
    return tfds.core.DatasetInfo(
        builder=self,
        description=self.builder_config.description,
        features=tfds.features.FeaturesDict(features),
        urls=[
            self.builder_config.url,
            "https://gluebenchmark.com/",
        ],
        citation=self.builder_config.citation + "\n" + _GLUE_CITATION,
    )

  def _split_generators(self, dl_manager):
    if self.builder_config.name == "mrpc":
      data_dir = None
      mrpc_files = dl_manager.download({
          "dev_ids": _MRPC_DEV_IDS,
          "train": _MRPC_TRAIN,
          "test": _MRPC_TEST,
      })
    else:
      dl_dir = dl_manager.download_and_extract(self.builder_config.data_url)
      data_dir = os.path.join(dl_dir, self.builder_config.data_dir)
      mrpc_files = None
    train_split = tfds.core.SplitGenerator(
        name=tfds.Split.TRAIN,
        num_shards=self.builder_config.train_shards,
        gen_kwargs={
            "data_file": os.path.join(data_dir or "", "train.tsv"),
            "split": "train",
            "mrpc_files": mrpc_files,
        })
    if self.builder_config.name == "mnli":
      return [
          train_split,
          _mnli_split_generator(
              "validation_matched", data_dir, "dev", matched=True),
          _mnli_split_generator(
              "validation_mismatched", data_dir, "dev", matched=False),
          _mnli_split_generator("test_matched", data_dir, "test", matched=True),
          _mnli_split_generator(
              "test_mismatched", data_dir, "test", matched=False)
      ]
    elif self.builder_config.name == "mnli_matched":
      return [
          _mnli_split_generator("validation", data_dir, "dev", matched=True),
          _mnli_split_generator("test", data_dir, "test", matched=True)
      ]
    elif self.builder_config.name == "mnli_mismatched":
      return [
          _mnli_split_generator("validation", data_dir, "dev", matched=False),
          _mnli_split_generator("test", data_dir, "test", matched=False)
      ]
    else:
      return [
          train_split,
          tfds.core.SplitGenerator(
              name=tfds.Split.VALIDATION,
              num_shards=1,
              gen_kwargs={
                  "data_file": os.path.join(data_dir or "", "dev.tsv"),
                  "split": "dev",
                  "mrpc_files": mrpc_files,
              }),
          tfds.core.SplitGenerator(
              name=tfds.Split.TEST,
              num_shards=1,
              gen_kwargs={
                  "data_file": os.path.join(data_dir or "", "test.tsv"),
                  "split": "test",
                  "mrpc_files": mrpc_files,
              }),
      ]

  def _generate_examples(self, data_file, split, mrpc_files):
    if self.builder_config.name == "mrpc":
      # We have to prepare the MRPC dataset from the original sources ourselves.
      examples = self._generate_example_mrpc_files(
          mrpc_files=mrpc_files, split=split)
      for example in examples:
        yield example["idx"], example
    else:
      process_label = self.builder_config.process_label
      label_classes = self.builder_config.label_classes

      # The train and dev files for CoLA are the only tsv files without a
      # header.
      is_cola_non_test = self.builder_config.name == "cola" and split != "test"

      with tf.io.gfile.GFile(data_file) as f:
        reader = csv.DictReader(f, delimiter="\t", quoting=csv.QUOTE_NONE)
        if is_cola_non_test:
          reader = csv.reader(f, delimiter="\t", quoting=csv.QUOTE_NONE)

        for n, row in enumerate(reader):
          if is_cola_non_test:
            row = {
                "sentence": row[3],
                "is_acceptable": row[1],
            }

          example = {
              feat: row[col]
              for feat, col in six.iteritems(self.builder_config.text_features)
          }
          example["idx"] = n

          if self.builder_config.label_column in row:
            label = row[self.builder_config.label_column]
            # For some tasks, the label is represented as 0 and 1 in the tsv
            # files and needs to be cast to integer to work with the feature.
            if label_classes and label not in label_classes:
              label = int(label) if label else None
            example["label"] = process_label(label)
          else:
            example["label"] = process_label(-1)

          # Filter out corrupted rows.
          for value in six.itervalues(example):
            if value is None:
              break
          else:
            yield example["idx"], example

  def _generate_example_mrpc_files(self, mrpc_files, split):
    if split == "test":
      with tf.io.gfile.GFile(mrpc_files["test"]) as f:
        reader = csv.DictReader(f, delimiter="\t", quoting=csv.QUOTE_NONE)
        for n, row in enumerate(reader):
          yield {
              "sentence1": row["#1 String"],
              "sentence2": row["#2 String"],
              "label": -1,
              "idx": n,
          }
    else:
      with tf.io.gfile.GFile(mrpc_files["dev_ids"]) as f:
        reader = csv.reader(f, delimiter="\t", quoting=csv.QUOTE_NONE)
        dev_ids = [[row[0], row[1]] for row in reader]
      with tf.io.gfile.GFile(mrpc_files["train"]) as f:
        # The first 3 bytes are the utf-8 BOM \xef\xbb\xbf, which messes with
        # the Quality key.
        f.seek(3)
        reader = csv.DictReader(f, delimiter="\t", quoting=csv.QUOTE_NONE)
        for n, row in enumerate(reader):
          is_row_in_dev = [row["#1 ID"], row["#2 ID"]] in dev_ids
          if is_row_in_dev == (split == "dev"):
            yield {
                "sentence1": row["#1 String"],
                "sentence2": row["#2 String"],
                "label": int(row["Quality"]),
                "idx": n,
            }


def _mnli_split_generator(name, data_dir, split, matched):
  return tfds.core.SplitGenerator(
      name=name,
      num_shards=1,
      gen_kwargs={
          "data_file": os.path.join(
              data_dir,
              "%s_%s.tsv" % (split, "matched" if matched else "mismatched")),
          "split": split,
          "mrpc_files": None,
      })
