utils_preprocess.py

# coding=utf-8
# Copyright 2020 Google and DeepMind.
#
# 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.
from __future__ import absolute_import, division, print_function

import os
import ast
import csv
import json
import shutil
import random
import argparse

from collections import defaultdict
from transformers import BertTokenizer, XLMTokenizer, XLMRobertaTokenizer

TOKENIZERS = {
    "bert": BertTokenizer,
    "xlm": XLMTokenizer,
    "xlmr": XLMRobertaTokenizer,
    "syntax-xlmr": XLMRobertaTokenizer,
}


def panx_tokenize_preprocess(args):
    def _preprocess_one_file(infile, outfile, idxfile, tokenizer, max_len):
        if not os.path.exists(infile):
            print(f"{infile} not exists")
            return 0
        special_tokens_count = 3 if isinstance(tokenizer, XLMRobertaTokenizer) else 2
        max_seq_len = max_len - special_tokens_count
        subword_len_counter = idx = 0
        exception_counter = 0
        with open(infile, "rt") as fin, open(outfile, "w") as fout, open(
                idxfile, "w"
        ) as fidx:
            for line in fin:
                line = line.strip()
                if not line:
                    fout.write("\n")
                    fidx.write("\n")
                    idx += 1
                    subword_len_counter = 0
                    continue

                items = line.split()
                token = items[0].strip()
                label = items[1].strip()
                head = -1
                if len(items) == 3:
                    head = items[2].strip()
                current_subwords_len = len(tokenizer.tokenize(token))

                if (
                        current_subwords_len == 0 or current_subwords_len > max_seq_len
                ) and len(token) != 0:
                    token = tokenizer.unk_token
                    current_subwords_len = 1

                # adding current token exceeds max sequence length
                # so start a new sequence
                if (subword_len_counter + current_subwords_len) > max_seq_len:
                    fout.write(f"\n{token}\t{label}\t{head}\n")
                    fidx.write(f"\n{idx}\n")
                    subword_len_counter = current_subwords_len
                    exception_counter += 1
                else:
                    fout.write(f"{token}\t{label}\t{head}\n")
                    fidx.write(f"{idx}\n")
                    subword_len_counter += current_subwords_len

        print(f"{exception_counter} examples are truncated.")
        return 1

    model_type = args.model_type
    tokenizer = TOKENIZERS[model_type].from_pretrained(
        args.model_name_or_path,
        do_lower_case=args.do_lower_case,
        cache_dir=args.cache_dir if args.cache_dir else None,
    )
    for lang in args.languages.split(","):
        out_dir = os.path.join(args.output_dir, lang)
        if not os.path.exists(out_dir):
            os.makedirs(out_dir)
        if lang == "en":
            files = ["dev", "test", "train"]
        else:
            files = ["dev", "test"]
        suffix = args.model_name_or_path
        if os.path.isdir(suffix):
            suffix = list(filter(None, suffix.split("/"))).pop()
        for file in files:
            infile = os.path.join(args.data_dir, f"{file}-{lang}.tsv")
            outfile = os.path.join(out_dir, "{}.{}".format(file, suffix))
            idxfile = os.path.join(out_dir, "{}.{}.idx".format(file, suffix))
            if os.path.exists(outfile) and os.path.exists(idxfile):
                print(f"{outfile} and {idxfile} exist")
            else:
                code = _preprocess_one_file(
                    infile, outfile, idxfile, tokenizer, args.max_len
                )
                if code > 0:
                    print(f"finish preprocessing {outfile}")


def panx_preprocess(args):
    def _process_one_file(infile, outfile):
        lines = open(infile, "r").readlines()
        if lines[-1].strip() == "":
            lines = lines[:-1]
        with open(outfile, "w") as fout:
            for l in lines:
                items = l.strip().split("\t")
                if len(items) == 2:
                    label = items[1].strip()
                    idx = items[0].find(":")
                    if idx != -1:
                        token = items[0][idx + 1:].strip()
                        fout.write(f"{token}\t{label}\n")
                        # if 'test' in infile:
                        #   fout.write(f'{token}\n')
                        # else:
                        #   fout.write(f'{token}\t{label}\n')
                else:
                    fout.write("\n")

    if not os.path.exists(args.output_dir):
        os.makedirs(args.output_dir)
    langs = "ar he vi id jv ms tl eu ml ta te af nl en de el bn hi mr ur fa fr it pt es bg ru ja ka ko th sw yo my zh kk tr et fi hu".split(
        " "
    )
    for lg in langs:
        for split in ["train", "test", "dev"]:
            infile = os.path.join(args.data_dir, f"{lg}-{split}")
            outfile = os.path.join(args.output_dir, f"{split}-{lg}.tsv")
            _process_one_file(infile, outfile)


def udpos_tokenize_preprocess(args):
    def _preprocess_one_file(infile, outfile, idxfile, tokenizer, max_len):
        if not os.path.exists(infile):
            print(f"{infile} does not exist")
            return
        subword_len_counter = idx = 0
        special_tokens_count = 3 if isinstance(tokenizer, XLMRobertaTokenizer) else 2
        max_seq_len = max_len - special_tokens_count
        with open(infile, "rt") as fin, open(outfile, "w") as fout, open(
                idxfile, "w"
        ) as fidx:
            for line in fin:
                line = line.strip()
                if len(line) == 0 or line == "":
                    fout.write("\n")
                    fidx.write("\n")
                    idx += 1
                    subword_len_counter = 0
                    continue

                items = line.split()
                if len(items) == 2:
                    label = items[1].strip()
                else:
                    label = "X"
                token = items[0].strip()
                current_subwords_len = len(tokenizer.tokenize(token))

                if (
                        current_subwords_len == 0 or current_subwords_len > max_seq_len
                ) and len(token) != 0:
                    token = tokenizer.unk_token
                    current_subwords_len = 1

                if (subword_len_counter + current_subwords_len) > max_seq_len:
                    fout.write(f"\n{token}\t{label}\n")
                    fidx.write(f"\n{idx}\n")
                    subword_len_counter = current_subwords_len
                else:
                    fout.write(f"{token}\t{label}\n")
                    fidx.write(f"{idx}\n")
                    subword_len_counter += current_subwords_len

    model_type = args.model_type
    tokenizer = TOKENIZERS[model_type].from_pretrained(
        args.model_name_or_path,
        do_lower_case=args.do_lower_case,
        cache_dir=args.cache_dir if args.cache_dir else None,
    )
    for lang in args.languages.split(","):
        out_dir = os.path.join(args.output_dir, lang)
        if not os.path.exists(out_dir):
            os.makedirs(out_dir)
        if lang == "en":
            files = ["dev", "test", "train"]
        else:
            files = ["dev", "test"]
        suffix = args.model_name_or_path
        if os.path.isdir(suffix):
            suffix = list(filter(None, suffix.split("/"))).pop()
        for file in files:
            infile = os.path.join(args.data_dir, "{}-{}.tsv".format(file, lang))
            outfile = os.path.join(out_dir, "{}.{}".format(file, suffix))
            idxfile = os.path.join(out_dir, "{}.{}.idx".format(file, suffix))
            if os.path.exists(outfile) and os.path.exists(idxfile):
                print(f"{outfile} and {idxfile} exist")
            else:
                _preprocess_one_file(infile, outfile, idxfile, tokenizer, args.max_len)
                print(f"finish preprocessing {outfile}")


def udpos_preprocess(args):
    def _read_one_file(file):
        data = []
        sent, tag, lines, head, deptag = [], [], [], [], []
        for line in open(file, "r"):
            items = line.strip().split("\t")
            if len(items) != 10:
                empty = all(w == "_" for w in sent)
                if not empty:
                    data.append((sent, tag, head, deptag, lines))
                sent, tag, head, deptag, lines = [], [], [], [], []
            else:
                sent.append(items[1].strip())
                tag.append(items[3].strip())
                deptag.append(items[7].strip())
                head.append(int(items[6].strip()))
                lines.append(line.strip())
                assert len(sent) == int(items[0]), \
                    "line={}, sent={}, tag={}".format(line, sent, tag)
        return data

    def isfloat(value):
        try:
            float(value)
            return True
        except ValueError:
            return False

    def remove_empty_space(data):
        new_data = {}
        for split in data:
            new_data[split] = []
            for sent, tag, head, deptag, lines in data[split]:
                new_sent = ["".join(w.replace("\u200c", "").split(" ")) for w in sent]
                lines = [line.replace("\u200c", "") for line in lines]
                assert len(" ".join(new_sent).split(" ")) == len(tag)
                new_data[split].append((new_sent, tag, head, deptag, lines))
        return new_data

    def check_file(file):
        for i, l in enumerate(open(file)):
            items = l.strip().split("\t")
            assert len(items[0].split(" ")) == len(
                items[1].split(" ")
            ), "idx={}, line={}".format(i, l)

    def _write_files(data, output_dir, lang, suffix):
        for split in data:
            if len(data[split]) > 0:
                prefix = os.path.join(output_dir, f"{split}-{lang}")
                if suffix == "mt":
                    with open(prefix + ".mt.tsv", "w") as fout:
                        for idx, (sent, tag, _) in enumerate(data[split]):
                            newline = "\n" if idx != len(data[split]) - 1 else ""
                            fout.write(
                                "{}\t{}{}".format(
                                    " ".join(sent), " ".join(tag), newline
                                )
                            )
                            # if split == 'test':
                            #   fout.write('{}{}'.format(' '.join(sent, newline)))
                            # else:
                            #   fout.write('{}\t{}{}'.format(' '.join(sent), ' '.join(tag), newline))
                    check_file(prefix + ".mt.tsv")
                    print("  - finish checking " + prefix + ".mt.tsv")
                elif suffix == "tsv":
                    with open(prefix + ".tsv", "w") as fout:
                        for sidx, (sent, tag, head, deptag, _) in enumerate(data[split]):
                            for widx, (w, t, h, d) in enumerate(zip(sent, tag, head, deptag)):
                                newline = (
                                    ""
                                    if (sidx == len(data[split]) - 1)
                                       and (widx == len(sent) - 1)
                                    else "\n"
                                )
                                fout.write("{}\t{}\t{}\t{}{}".format(w, t, h, d, newline))
                                # if split == 'test':
                                #   fout.write('{}{}'.format(w, newline))
                                # else:
                                #   fout.write('{}\t{}{}'.format(w, t, newline))
                            fout.write("\n")
                elif suffix == "conll":
                    with open(prefix + ".conll", "w") as fout:
                        for _, _, lines in data[split]:
                            for l in lines:
                                fout.write(l.strip() + "\n")
                            fout.write("\n")
                print(f"finish writing file to {prefix}.{suffix}")

    if not os.path.exists(args.output_dir):
        os.makedirs(args.output_dir)

    languages = \
        "af ar bg de el en es et eu fa fi fr he hi hu id it ja kk ko mr nl pt ru ta te th tl tr ur vi yo zh".split(" ")
    for root, dirs, files in os.walk(args.data_dir):
        lg = root.strip().split("/")[-1]
        if root == args.data_dir or lg not in languages:
            continue

        data = {k: [] for k in ["train", "dev", "test"]}
        for f in sorted(files):
            if f.endswith("conll"):
                file = os.path.join(root, f)
                examples = _read_one_file(file)
                if "train" in f:
                    data["train"].extend(examples)
                elif "dev" in f:
                    data["dev"].extend(examples)
                elif "test" in f:
                    data["test"].extend(examples)
                else:
                    print("split not found: ", file)
                print(
                    " - finish reading {}, {}".format(
                        file, [(k, len(v)) for k, v in data.items()]
                    )
                )

        data = remove_empty_space(data)
        for sub in ["tsv"]:
            _write_files(data, args.output_dir, lg, sub)


def pawsx_preprocess(args):
    def _preprocess_one_file(infile, outfile, remove_label=False):
        data = []
        for i, line in enumerate(open(infile, "r")):
            if i == 0:
                continue
            items = line.strip().split("\t")
            sent1 = " ".join(items[1].strip().split(" "))
            sent2 = " ".join(items[2].strip().split(" "))
            label = items[3]
            data.append([sent1, sent2, label])

        with open(outfile, "w") as fout:
            writer = csv.writer(fout, delimiter="\t")
            for sent1, sent2, label in data:
                if remove_label:
                    writer.writerow([sent1, sent2])
                else:
                    writer.writerow([sent1, sent2, label])

    if not os.path.exists(args.output_dir):
        os.makedirs(args.output_dir)

    split2file = {"train": "train", "test": "test_2k", "dev": "dev_2k"}
    for lang in ["en", "de", "es", "fr", "ja", "ko", "zh"]:
        for split in ["train", "test", "dev"]:
            if split == "train" and lang != "en":
                # continue
                file = 'translated_train'
            else:
                file = split2file[split]
            infile = os.path.join(args.data_dir, lang, "{}.tsv".format(file))
            outfile = os.path.join(args.output_dir, "{}-{}.tsv".format(split, lang))
            # _preprocess_one_file(infile, outfile, remove_label=(split == 'test'))
            _preprocess_one_file(infile, outfile)
            print(f"finish preprocessing {outfile}")


def xnli_preprocess(args):
    def _preprocess_file(infile, output_dir, split):
        all_langs = defaultdict(list)
        for i, line in enumerate(open(infile, "r")):
            if i == 0:
                continue

            items = line.strip().split("\t")
            lang = items[0].strip()
            label = (
                "contradiction"
                if items[1].strip() == "contradictory"
                else items[1].strip()
            )
            sent1 = " ".join(items[6].strip().split(" "))
            sent2 = " ".join(items[7].strip().split(" "))
            all_langs[lang].append((sent1, sent2, label))
        print(f"# langs={len(all_langs)}")
        for lang, pairs in all_langs.items():
            outfile = os.path.join(output_dir, "{}-{}.tsv".format(split, lang))
            with open(outfile, "w") as fout:
                writer = csv.writer(fout, delimiter="\t")
                for (sent1, sent2, label) in pairs:
                    writer.writerow([sent1, sent2, label])
                    # if split == 'test':
                    #   writer.writerow([sent1, sent2])
                    # else:
                    #   writer.writerow([sent1, sent2, label])
            print(f"finish preprocess {outfile}")

    def _preprocess_train_file(infile, outfile):
        with open(outfile, "w") as fout:
            writer = csv.writer(fout, delimiter="\t")
            for i, line in enumerate(open(infile, "r")):
                if i == 0:
                    continue

                items = line.strip().split("\t")
                sent1 = " ".join(items[0].strip().split(" "))
                sent2 = " ".join(items[1].strip().split(" "))
                label = (
                    "contradiction"
                    if items[2].strip() == "contradictory"
                    else items[2].strip()
                )
                writer.writerow([sent1, sent2, label])
        print(f"finish preprocess {outfile}")

    train_langs = 'ar,bg,de,el,en,es,fr,hi,ru,tr,ur,vi,zh'.split(',')
    if not os.path.exists(args.output_dir):
        os.makedirs(args.output_dir)
    for l in train_langs:
        infile = os.path.join(args.data_dir, "XNLI-MT-1.0/multinli/multinli.train.{}.tsv".format(l))
        outfile = os.path.join(args.output_dir, "train-{}.tsv".format(l))
        _preprocess_train_file(infile, outfile)

    for split in ["test", "dev"]:
        infile = os.path.join(args.data_dir, "XNLI-1.0/xnli.{}.tsv".format(split))
        print(f"reading file {infile}")
        _preprocess_file(infile, args.output_dir, split)


def tatoeba_preprocess(args):
    lang3_dict = {
        "afr": "af",
        "ara": "ar",
        "bul": "bg",
        "ben": "bn",
        "deu": "de",
        "ell": "el",
        "spa": "es",
        "est": "et",
        "eus": "eu",
        "pes": "fa",
        "fin": "fi",
        "fra": "fr",
        "heb": "he",
        "hin": "hi",
        "hun": "hu",
        "ind": "id",
        "ita": "it",
        "jpn": "ja",
        "jav": "jv",
        "kat": "ka",
        "kaz": "kk",
        "kor": "ko",
        "mal": "ml",
        "mar": "mr",
        "nld": "nl",
        "por": "pt",
        "rus": "ru",
        "swh": "sw",
        "tam": "ta",
        "tel": "te",
        "tha": "th",
        "tgl": "tl",
        "tur": "tr",
        "urd": "ur",
        "vie": "vi",
        "cmn": "zh",
        "eng": "en",
    }
    if not os.path.exists(args.output_dir):
        os.makedirs(args.output_dir)
    for sl3, sl2 in lang3_dict.items():
        if sl3 != "eng":
            src_file = f"{args.data_dir}/tatoeba.{sl3}-eng.{sl3}"
            tgt_file = f"{args.data_dir}/tatoeba.{sl3}-eng.eng"
            src_out = f"{args.output_dir}/{sl2}-en.{sl2}"
            tgt_out = f"{args.output_dir}/{sl2}-en.en"
            gold_out = f"{args.output_dir}/{sl2}-en.en.gold"
            shutil.copy(src_file, src_out)
            tgts = [l.strip() for l in open(tgt_file)]
            idx = range(len(tgts))
            data = zip(tgts, idx)
            with open(tgt_out, "w") as ftgt, open(gold_out, "w") as fgold:
                for t, i in sorted(data, key=lambda x: x[0]):
                    ftgt.write(f"{t}\n")
                    fgold.write(f"{i}\n")


def xquad_preprocess(args):
    pass
    # Remove the test annotations to prevent accidental cheating
    # remove_qa_test_annotations(args.data_dir)


def mlqa_preprocess(args):
    pass
    # Remove the test annotations to prevent accidental cheating
    # remove_qa_test_annotations(args.data_dir)


def tydiqa_preprocess(args):
    LANG2ISO = {
        "arabic": "ar",
        "bengali": "bn",
        "english": "en",
        "finnish": "fi",
        "indonesian": "id",
        "korean": "ko",
        "russian": "ru",
        "swahili": "sw",
        "telugu": "te",
    }
    assert os.path.exists(args.data_dir)
    train_file = os.path.join(args.data_dir, "tydiqa-goldp-v1.1-train.json")
    os.makedirs(args.output_dir, exist_ok=True)

    # Split the training file into language-specific files
    lang2data = defaultdict(list)
    with open(train_file, "r") as f_in:
        data = json.load(f_in)
        version = data["version"]
        for doc in data["data"]:
            for par in doc["paragraphs"]:
                context = par["context"]
                for qa in par["qas"]:
                    question = qa["question"]
                    question_id = qa["id"]
                    example_lang = question_id.split("-")[0]
                    q_id = question_id.split("-")[-1]
                    for answer in qa["answers"]:
                        a_start, a_text = answer["answer_start"], answer["text"]
                        a_end = a_start + len(a_text)
                        assert context[a_start:a_end] == a_text
                    lang2data[example_lang].append(
                        {
                            "paragraphs": [
                                {
                                    "context": context,
                                    "qas": [
                                        {
                                            "answers": qa["answers"],
                                            "question": question,
                                            "id": q_id,
                                        }
                                    ],
                                }
                            ]
                        }
                    )

    for lang, data in lang2data.items():
        out_file = os.path.join(
            args.output_dir, "tydiqa.%s.train.json" % LANG2ISO[lang]
        )
        with open(out_file, "w") as f:
            json.dump({"data": data, "version": version}, f)

    # Rename the dev files
    dev_dir = os.path.join(args.data_dir, "tydiqa-goldp-v1.1-dev")
    assert os.path.exists(dev_dir)
    for lang, iso in LANG2ISO.items():
        src_file = os.path.join(dev_dir, "tydiqa-goldp-dev-%s.json" % lang)
        dst_file = os.path.join(dev_dir, "tydiqa.%s.dev.json" % iso)
        os.rename(src_file, dst_file)

    # Remove the test annotations to prevent accidental cheating
    # remove_qa_test_annotations(dev_dir)


def remove_qa_test_annotations(test_dir):
    assert os.path.exists(test_dir)
    for file_name in os.listdir(test_dir):
        new_data = []
        test_file = os.path.join(test_dir, file_name)
        with open(test_file, "r") as f:
            data = json.load(f)
            version = data["version"]
            for doc in data["data"]:
                for par in doc["paragraphs"]:
                    context = par["context"]
                    for qa in par["qas"]:
                        question = qa["question"]
                        question_id = qa["id"]
                        for answer in qa["answers"]:
                            a_start, a_text = answer["answer_start"], answer["text"]
                            a_end = a_start + len(a_text)
                            assert context[a_start:a_end] == a_text
                        new_data.append(
                            {
                                "paragraphs": [
                                    {
                                        "context": context,
                                        "qas": [
                                            {
                                                "answers": [
                                                    {"answer_start": 0, "text": ""}
                                                ],
                                                "question": question,
                                                "id": question_id,
                                            }
                                        ],
                                    }
                                ]
                            }
                        )
        with open(test_file, "w") as f:
            json.dump({"data": new_data, "version": version}, f)


def mtop_preprocess(args):
    DOMAINS = [
        'alarm', 'calling', 'event', 'messaging', 'music', 'news',
        'people', 'recipes', 'reminder', 'timer', 'watcher'
    ]
    FIELD_NAMES = ['id', 'intent', 'slots', 'utterance', 'domain', 'language', 'compositional_rep', 'tokens']

    def convert(infile, outfile):
        with open(outfile, 'w', encoding='utf-8') as fw:
            with open(infile, 'r') as f:
                for line in f:
                    line = line.strip()
                    if not line:
                        continue
                    fields = line.split('\t')
                    intent = fields[1].split(':')[1]
                    slots = []
                    # slots field can be empty string
                    if len(fields[2]) > 0:
                        # slots are comma separated
                        # 4:20:SL:CONTACT,21:26:SL:TYPE_CONTENT,36:38:SL:RECIPIENT
                        _slots = fields[2].split(',')
                        for s in _slots:
                            if s:
                                # each slot has 4 fields, separated by ':'
                                # start_offset:end_offset:SL:slot_tag
                                _s = s.split(':')
                                #assert len(_s) == 4, (s, _s)
                                slots.append([int(_s[0]), int(_s[1]), _s[3]])

                    assert len(fields) == len(FIELD_NAMES)
                    try:
                        # obj = json.loads(fields[7])
                        obj = ast.literal_eval(fields[7])
                    except ValueError as e:
                        print(fields[7])
                        print(e)

                    tokens = []
                    for j, (w, span) in enumerate(zip(obj['tokens'], obj['tokenSpans'])):
                        tokens.append({
                            'token': w,
                            'token_range': (span['start'], span['start'] + span['length'])
                        })
                    # print(tokens)

                    words = []
                    bio_tags = ['O'] * len(tokens)
                    total_slots = 0
                    for j, tok in enumerate(tokens):
                        words.append(tok['token'])
                        tok_start, tok_end = tok['token_range']
                        # we match word offset with slot offsets to find proper tag
                        for slot_start, slot_end, slot_tag in slots:
                            # if a token starts at the slot start offset
                            if tok_start == slot_start:
                                bio_tags[j] = 'B-{}'.format(slot_tag)
                                total_slots += 1
                                break
                            # if a token starts after a slot's start offset
                            # but the token ends at or before the slot ends
                            elif tok_start > slot_start and tok_end <= slot_end:
                                # an I-{} tag must be preceded by a B-{} or I-{} tag
                                if bio_tags[j - 1] in ['B-{}'.format(slot_tag), 'I-{}'.format(slot_tag)]:
                                    bio_tags[j] = 'I-{}'.format(slot_tag)
                                break

                    assert len(words) == len(bio_tags)
                    fw.write(json.dumps({
                        'tokens': words,
                        'slot_labels': bio_tags,
                        'intent_label': intent
                    }) + '\n')

    # infile contains one example per line (tab separated fields)
    # there must be len(FIELD_NAMES) fields in every line
    def vocab_process(data_dir):
        slot_label_vocab = 'slot_label.txt'
        intent_label_vocab = 'intent_label.txt'

        intent_vocab = set()
        slot_vocab = set()
        for lang in ['en', 'de', 'es', 'fr', 'hi', 'th']:
            for split in ['train', 'dev', 'test']:
                with open(os.path.join(data_dir, '{}-{}.jsonl'.format(split, lang)),
                          'r', encoding='utf-8') as f_r:
                    for line in f_r:
                        line = line.strip()
                        ex = json.loads(line)
                        intent_vocab.add(ex['intent_label'])
                        for slot in ex['slot_labels']:
                            slot_vocab.add(slot)

        with open(os.path.join(data_dir, intent_label_vocab), 'w', encoding='utf-8') as f_w1, \
                open(os.path.join(data_dir, slot_label_vocab), 'w', encoding='utf-8') as f_w2:
            intent_vocab = sorted(list(intent_vocab))
            for intent in intent_vocab:
                f_w1.write(intent + '\n')

            slot_vocab = sorted(list(slot_vocab), key=lambda x: (x[2:], x[:2]))
            for slot in slot_vocab:
                f_w2.write(slot + '\n')

    for lang in ['en', 'de', 'es', 'fr', 'hi', 'th']:
        for split in ['train', 'eval', 'test']:
            convert(
                os.path.join(args.data_dir, lang, '{}.txt'.format(split)),
                os.path.join(
                    args.output_dir,
                    '{}-{}.jsonl'.format('dev' if split == 'eval' else split, lang)
                )
            )
    vocab_process(args.output_dir)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    ## Required parameters
    parser.add_argument(
        "--data_dir",
        default=None,
        type=str,
        required=True,
        help="The input data dir. Should contain the .tsv files (or other data files) for the task.",
    )
    parser.add_argument(
        "--output_dir",
        default=None,
        type=str,
        required=True,
        help="The output data dir where any processed files will be written to.",
    )
    parser.add_argument(
        "--task", default="panx", type=str, required=True, help="The task name"
    )
    parser.add_argument(
        "--model_name_or_path",
        default="bert-base-multilingual-cased",
        type=str,
        help="The pre-trained model",
    )
    parser.add_argument("--model_type", default="bert", type=str, help="model type")
    parser.add_argument(
        "--max_len", default=512, type=int, help="the maximum length of sentences"
    )
    parser.add_argument(
        "--do_lower_case", action="store_true", help="whether to do lower case"
    )
    parser.add_argument("--cache_dir", default=None, type=str, help="cache directory")
    parser.add_argument("--languages", default="en", type=str, help="process language")
    parser.add_argument(
        "--remove_last_token",
        action="store_true",
        help="whether to remove the last token",
    )
    parser.add_argument(
        "--remove_test_label",
        action="store_true",
        help="whether to remove test set label",
    )
    args = parser.parse_args()

    if args.task == "panx_tokenize":
        panx_tokenize_preprocess(args)
    if args.task == "panx":
        panx_preprocess(args)
    if args.task == "udpos_tokenize":
        udpos_tokenize_preprocess(args)
    if args.task == "udpos":
        udpos_preprocess(args)
    if args.task == "pawsx":
        pawsx_preprocess(args)
    if args.task == "xnli":
        xnli_preprocess(args)
    if args.task == "tatoeba":
        tatoeba_preprocess(args)
    if args.task == "xquad":
        xquad_preprocess(args)
    if args.task == "mlqa":
        mlqa_preprocess(args)
    if args.task == "tydiqa":
        tydiqa_preprocess(args)
    if args.task == "mtop":
        mtop_preprocess(args)