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     |  10000   |     iwslt2017    |40|   bg   |   bulgarian    |  10000   |       xnli       |41|   bn   |    bengali     |  10000   |  open_subtitles  |42|   bs   |    bosnian     |  10000   |  open_subtitles  |43|   cs   |     czech      |  10000   |       ecb        |44|   da   |     danish     |  10000   |  open_subtitles  |45|   de   |     german     |  10000   |       ecb        |46|   el   |  modern greek  |  10000   |       ecb        |47|   en   |    english     |  10000   |       ecb        |48|   eo   |   esperanto    |  10000   |     tatoeba      |49|   es   |    spanish     |  10000   |     tatoeba      |50|   et   |    estonian    |  10000   |      emea        |51|   fi   |    finnish     |  10000   |       ecb        |52|   fo   |    faroese     |  10000   |  nordic_langid   |53|   fr   |    french      |  10000   |     iwslt2017    |54|   ga   |    irish       |  10000   | multi_para_crawl |55|   gl   |    galician    |  3096    |     tatoeba      |56|   hi   |     hindi      |  10000   |  open_subtitles  |57|  hi_en |     hindi      |  7180    | cmu_hinglish_dog |58|   hr   |    croatian    |  10000   |   hrenwac_para   |59|   hu   |    hungarian   |  3801    |   europa_ecdc_tm; europa_eac_tm   |60|   hy   |    armenian    |   660    |  open_subtitles  |61|   id   |   indonesian   |  10000   |   id_panl_bppt   |62|   is   |   icelandic    |  2973    |   europa_ecdc_tm; europa_eac_tm   |63|   it   |    italian     |  10000   |     iwslt2017    |64|   ja   |    japanese    |  10000   |     iwslt2017    |65|   ko   |    korean      |  10000   |     iwslt2017    |66|   lt   |   lithuanian   |  10000   |       emea       |67|   lv   |    latvian     |  4595    |   europa_ecdc_tm; europa_eac_tm   |68|   mr   |    marathi     |  10000   |     tatoeba      |69|   mt   |    maltese     |  10000   | multi_para_crawl |70|   nl   |    dutch       |  10000   |       kde4       |71|   no   |   norwegian    |  10000   | multi_para_crawl |72|   pl   |    polish      |  10000   |       ecb        |73|   pt   |   portuguese   |  10000   |     tatoeba      |74|   ro   |    romanian    |  10000   |       kde4       |75|   ru   |    russian     |  10000   |       xnli       |76|   sk   |    slovak      |  10000   | multi_para_crawl |77|   sl   |   slovenian    |  4589    |   europa_ecdc_tm; europa_eac_tm   |78|   sw   |    swahili     |  10000   |       xnli       |79|   sv   |    swedish     |  10000   |       kde4       |80|   th   |     thai       |  10000   |       xnli       |81|   tl   |    tagalog     |  10000   | multi_para_crawl |82|   tn   |    serpeti     |  10000   |   autshumato     |83|   tr   |    turkish     |  10000   |       xnli       |84|   ts   |    dzonga      |  10000   |    autshumato    |85|   ur   |     urdu       |  10000   |       xnli       |86|   vi   |   vietnamese   |  10000   |       xnli       |87|   yo   |     yoruba     |  9970    |    menyo20k_mt   |88|   zh   |    chinese     |  10000   |       xnli       |89|   zu   |  zulu, south africa  |  10000   |    autshumato    |90"""91 92 93def main():94    args = get_args()95 96    subset_dataset_dict = dict()97 98    lines = s.strip().split("\n")99 100    with open(args.train_subset, "w", encoding="utf-8") as ftrain, open(args.valid_subset, "w", encoding="utf-8") as fvalid:101        for line in lines:102            row = str(line).split("|")103            row = [col.strip() for col in row if len(col) != 0]104 105            if len(row) != 4:106                raise AssertionError("not 4 item, line: {}".format(line))107 108            abbr = row[0]109            full = row[1]110            total = int(row[2])111            subsets = [e.strip() for e in row[3].split(";")]112 113            count = 0114            for subset in subsets:115                if subset in subset_dataset_dict.keys():116                    dataset_dict = subset_dataset_dict[subset]117                else:118                    dataset_dict = load_dataset(119                        "qgyd2021/language_identification",120                        name=subset,121                        cache_dir=args.dataset_cache_dir,122                        # download_mode=DownloadMode.FORCE_REDOWNLOAD123                    )124                    subset_dataset_dict[subset] = dataset_dict125 126                train_dataset = dataset_dict["train"]127                for sample in train_dataset:128                    text = sample["text"]129                    language = sample["language"]130                    data_source = sample["data_source"]131 132                    if count > total:133                        break134 135                    if language != abbr:136                        continue137 138                    split = "train" if random.random() < 0.8 else "valid"139 140                    row_ = {141                        "text": text,142                        "label": language,143                        "language": full,144                        "data_source": data_source,145                        "split": split,146                    }147                    row_ = json.dumps(row_, ensure_ascii=False)148 149                    if split == "train":150                        ftrain.write("{}\n".format(row_))151                    elif split == "valid":152                        fvalid.write("{}\n".format(row_))153                    else:154                        raise AssertionError155 156                    count += 1157 158    return159 160 161if __name__ == "__main__":162    main()163