intelli-zen/language_identification

1#!/usr/bin/python32# -*- coding: utf-8 -*-3import argparse4import json5import os6import random7import sys8 9pwd = os.path.abspath(os.path.dirname(__file__))10sys.path.append(os.path.join(pwd, "../../"))11 12from datasets import load_dataset, DownloadMode13 14from project_settings import project_path15 16 17def get_args():18    parser = argparse.ArgumentParser()19    parser.add_argument(20        "--dataset_cache_dir",21        default=(project_path / "hub_datasets").as_posix(),22        type=str23    )24    parser.add_argument(25        "--train_subset",26        default="train.jsonl",27        type=str28    )29    parser.add_argument(30        "--valid_subset",31        default="valid.jsonl",32        type=str33    )34    args = parser.parse_args()35    return args36 37 38s = """39|   ar   |     arabic     |  100000   |     iwslt2017    |40|   bg   |   bulgarian    |  100000   |       xnli       |41|   bn   |    bengali     |  36064    |  open_subtitles  |42|   bs   |    bosnian     |  10212    |  open_subtitles  |43|   cs   |     czech      |  100000   |       emea       |44|   da   |     danish     |  100000   |  open_subtitles  |45|   de   |     german     |  100000   |     iwslt2017    |46|   el   |  modern greek  |  100000   |       emea       |47|   en   |    english     |  100000   |     iwslt2017    |48|   eo   |   esperanto    |  94101    | tatoeba; open_subtitles |49|   es   |    spanish     |  100000   |       xnli       |50|   et   |    estonian    |  100000   |       emea       |51|   fi   |    finnish     |  100000   |    ecb; kde4     |52|   fo   |    faroese     |  23807    |  nordic_langid   |53|   fr   |    french      |  100000   |     iwslt2017    |54|   ga   |    irish       |  100000   | multi_para_crawl |55|   gl   |    galician    |  3096     |     tatoeba      |56|   hi   |     hindi      |  100000   |       xnli       |57|  hi_en |     hindi      |  7180     | cmu_hinglish_dog |58|   hr   |    croatian    |  95844    |   hrenwac_para   |59|   hu   |    hungarian   |  3801     |   europa_ecdc_tm; europa_eac_tm   |60|   hy   |    armenian    |   660     |  open_subtitles  |61|   id   |   indonesian   |  23940    |   id_panl_bppt   |62|   is   |   icelandic    |  100000   | multi_para_crawl |63|   it   |    italian     |  100000   |     iwslt2017    |64|   ja   |    japanese    |  100000   |     iwslt2017    |65|   ko   |    korean      |  100000   |     iwslt2017    |66|   lt   |   lithuanian   |  100000   |       emea       |67|   lv   |    latvian     |  100000   | multi_para_crawl |68|   mr   |    marathi     |  51807    |     tatoeba      |69|   mt   |    maltese     |  100000   | multi_para_crawl |70|   nl   |    dutch       |  100000   |       kde4       |71|   no   |   norwegian    |  100000   | multi_para_crawl |72|   pl   |    polish      |  100000   | para_crawl_en_pl |73|   pt   |   portuguese   |  100000   | para_crawl_en_pt |74|   ro   |    romanian    |  100000   |     iwslt2017    |75|   ru   |    russian     |  100000   |       xnli       |76|   sk   |    slovak      |  100000   | multi_para_crawl |77|   sl   |   slovenian    |  100000   | para_crawl_en_sl |78|   sw   |    swahili     |  100000   |       xnli       |79|   sv   |    swedish     |  100000   |       kde4       |80|   th   |     thai       |  100000   |       xnli       |81|   tl   |    tagalog     |  97241    | multi_para_crawl |82|   tn   |    serpeti     |  100000   |   autshumato     |83|   tr   |    turkish     |  100000   |       xnli       |84|   ts   |    dzonga      |  100000   |    autshumato    |85|   uk   |    ukrainian   |  88533    |  para_pat_en_uk  |86|   ur   |     urdu       |  100000   |       xnli       |87|   vi   |   vietnamese   |  100000   |       xnli       |88|   yo   |     yoruba     |  9970    |    menyo20k_mt   |89|   zh   |    chinese     |  100000   |       xnli       |90|   zu   |  zulu, south africa  |  26801   |    autshumato    |91"""92 93 94def main():95    args = get_args()96 97    subset_dataset_dict = dict()98 99    lines = s.strip().split("\n")100 101    with open(args.train_subset, "w", encoding="utf-8") as ftrain, open(args.valid_subset, "w", encoding="utf-8") as fvalid:102        for line in lines:103            row = str(line).split("|")104            row = [col.strip() for col in row if len(col) != 0]105 106            if len(row) != 4:107                raise AssertionError("not 4 item, line: {}".format(line))108 109            abbr = row[0]110            full = row[1]111            total = int(row[2])112            subsets = [e.strip() for e in row[3].split(";")]113 114            count = 0115            for subset in subsets:116                if subset in subset_dataset_dict.keys():117                    dataset_dict = subset_dataset_dict[subset]118                else:119                    dataset_dict = load_dataset(120                        "qgyd2021/language_identification",121                        name=subset,122                        cache_dir=args.dataset_cache_dir,123                        # download_mode=DownloadMode.FORCE_REDOWNLOAD124                    )125                    subset_dataset_dict[subset] = dataset_dict126 127                train_dataset = dataset_dict["train"]128                for sample in train_dataset:129                    text = sample["text"]130                    language = sample["language"]131                    data_source = sample["data_source"]132 133                    if count > total:134                        break135 136                    if language != abbr:137                        continue138 139                    split = "train" if random.random() < 0.8 else "valid"140 141                    row_ = {142                        "text": text,143                        "label": language,144                        "language": full,145                        "data_source": data_source,146                        "split": split,147                    }148                    row_ = json.dumps(row_, ensure_ascii=False)149 150                    if split == "train":151                        ftrain.write("{}\n".format(row_))152                    elif split == "valid":153                        fvalid.write("{}\n".format(row_))154                    else:155                        raise AssertionError156 157                    count += 1158 159    return160 161 162if __name__ == "__main__":163    main()164