# Copyright 2022 The TensorFlow Authors. All Rights Reserved.

#

# 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.

"""BERT finetuning task dataset generator."""

import functools

import json

import os

# Import libraries

from absl import app

from absl import flags

import tensorflow as tf

from official.nlp.data import classifier_data_lib

from official.nlp.data import sentence_retrieval_lib

# word-piece tokenizer based squad_lib

from official.nlp.data import squad_lib as squad_lib_wp

# sentence-piece tokenizer based squad_lib

from official.nlp.data import squad_lib_sp

from official.nlp.data import tagging_data_lib

from official.nlp.tools import tokenization

FLAGS = flags.FLAGS

flags.DEFINE_enum(

"fine_tuning_task_type", "classification",

["classification", "regression", "squad", "retrieval", "tagging"],

"The name of the BERT fine tuning task for which data "

"will be generated.")

# BERT classification specific flags.

flags.DEFINE_string(

"input_data_dir", None,

"The input data dir. Should contain the .tsv files (or other data files) "

"for the task.")

flags.DEFINE_enum(

"classification_task_name", "MNLI", [

"AX", "COLA", "IMDB", "MNLI", "MRPC", "PAWS-X", "QNLI", "QQP", "RTE",

"SST-2", "STS-B", "WNLI", "XNLI", "XTREME-XNLI", "XTREME-PAWS-X",

"AX-g", "SUPERGLUE-RTE", "CB", "BoolQ", "WIC"

], "The name of the task to train BERT classifier. The "

"difference between XTREME-XNLI and XNLI is: 1. the format "

"of input tsv files; 2. the dev set for XTREME is english "

"only and for XNLI is all languages combined. Same for "

"PAWS-X.")

# MNLI task-specific flag.

flags.DEFINE_enum("mnli_type", "matched", ["matched", "mismatched"],

"The type of MNLI dataset.")

# XNLI task-specific flag.

flags.DEFINE_string(

"xnli_language", "en",

"Language of training data for XNLI task. If the value is 'all', the data "

"of all languages will be used for training.")

# PAWS-X task-specific flag.

flags.DEFINE_string(

"pawsx_language", "en",

"Language of training data for PAWS-X task. If the value is 'all', the data "

"of all languages will be used for training.")

# XTREME classification specific flags. Only used in XtremePawsx and XtremeXnli.

flags.DEFINE_string(

"translated_input_data_dir", None,

"The translated input data dir. Should contain the .tsv files (or other "

"data files) for the task.")

# Retrieval task-specific flags.

flags.DEFINE_enum("retrieval_task_name", "bucc", ["bucc", "tatoeba"],

"The name of sentence retrieval task for scoring")

# Tagging task-specific flags.

flags.DEFINE_enum("tagging_task_name", "panx", ["panx", "udpos"],

"The name of BERT tagging (token classification) task.")

flags.DEFINE_bool("tagging_only_use_en_train", True,

"Whether only use english training data in tagging.")

# BERT Squad task-specific flags.

flags.DEFINE_string(

"squad_data_file", None,

"The input data file in for generating training data for BERT squad task.")

flags.DEFINE_string(

"translated_squad_data_folder", None,

"The translated data folder for generating training data for BERT squad "

"task.")

flags.DEFINE_integer(

"doc_stride", 128,

"When splitting up a long document into chunks, how much stride to "

"take between chunks.")

flags.DEFINE_integer(

"max_query_length", 64,

"The maximum number of tokens for the question. Questions longer than "

"this will be truncated to this length.")

flags.DEFINE_bool(

"version_2_with_negative", False,

"If true, the SQuAD examples contain some that do not have an answer.")

flags.DEFINE_bool(

"xlnet_format", False,

"If true, then data will be preprocessed in a paragraph, query, class order"

" instead of the BERT-style class, paragraph, query order.")

# XTREME specific flags.

flags.DEFINE_bool("only_use_en_dev", True, "Whether only use english dev data.")

# Shared flags across BERT fine-tuning tasks.

flags.DEFINE_string("vocab_file", None,

"The vocabulary file that the BERT model was trained on.")

flags.DEFINE_string(

"train_data_output_path", None,

"The path in which generated training input data will be written as tf"

" records.")

flags.DEFINE_string(

"eval_data_output_path", None,

"The path in which generated evaluation input data will be written as tf"

" records.")

flags.DEFINE_string(

"test_data_output_path", None,

"The path in which generated test input data will be written as tf"

" records. If None, do not generate test data. Must be a pattern template"

" as test_{}.tfrecords if processor has language specific test data.")

flags.DEFINE_string("meta_data_file_path", None,

"The path in which input meta data will be written.")

flags.DEFINE_bool(

"do_lower_case", True,

"Whether to lower case the input text. Should be True for uncased "

"models and False for cased models.")

flags.DEFINE_integer(

"max_seq_length", 128,

"The maximum total input sequence length after WordPiece tokenization. "

"Sequences longer than this will be truncated, and sequences shorter "

"than this will be padded.")

flags.DEFINE_string("sp_model_file", "",

"The path to the model used by sentence piece tokenizer.")

flags.DEFINE_enum(

"tokenization", "WordPiece", ["WordPiece", "SentencePiece"],

"Specifies the tokenizer implementation, i.e., whether to use WordPiece "

"or SentencePiece tokenizer. Canonical BERT uses WordPiece tokenizer, "

"while ALBERT uses SentencePiece tokenizer.")

flags.DEFINE_string(

"tfds_params", "", "Comma-separated list of TFDS parameter assigments for "

"generic classfication data import (for more details "

"see the TfdsProcessor class documentation).")

def generate_classifier_dataset():

max_seq_length=FLAGS.max_seq_length)

def generate_regression_dataset():

raise ValueError("No data processor found for the given regression task.")

def generate_squad_dataset():

version_2_with_negative=FLAGS.version_2_with_negative)

def generate_retrieval_dataset():

FLAGS.test_data_output_path, FLAGS.max_seq_length)

def generate_tagging_dataset():

FLAGS.test_data_output_path, processor_text_fn)

def main(_):

writer.write(json.dumps(input_meta_data, indent=4) + "\n")

if __name__ == "__main__":

flags.mark_flag_as_required("meta_data_file_path")

app.run(main)