sumanthd/IndicTrans-MultilingualTranslation

1import os2from tqdm import tqdm3import sys4 5LANGS = [6    "as",7    "bn",8    "gu",9    "hi",10    "kn",11    "ml",12    "mr",13    "or",14    "pa",15    "ta",16    "te",17    #"ur"18]19 20 21def add_token(sent, tag_infos):22    """ add special tokens specified by tag_infos to each element in list23 24    tag_infos: list of tuples (tag_type,tag)25 26    each tag_info results in a token of the form: __{tag_type}__{tag}__27 28    """29 30    tokens = []31    for tag_type, tag in tag_infos:32        token = '__' + tag_type + '__' + tag + '__'33        tokens.append(token)34 35    return ' '.join(tokens) + ' ' + sent36 37 38def concat_data(data_dir, outdir, lang_pair_list,39                out_src_lang='SRC', out_trg_lang='TGT', split='train'):40    """41    data_dir: input dir, contains directories for language pairs named l1-l242    """43    os.makedirs(outdir, exist_ok=True)44 45    out_src_fname = '{}/{}.{}'.format(outdir, split, out_src_lang)46    out_trg_fname = '{}/{}.{}'.format(outdir, split, out_trg_lang)47#     out_meta_fname='{}/metadata.txt'.format(outdir)48 49    print()50    print(out_src_fname)51    print(out_trg_fname)52#     print(out_meta_fname)53 54    # concatenate train data55    if os.path.isfile(out_src_fname):56        os.unlink(out_src_fname)57    if os.path.isfile(out_trg_fname):58        os.unlink(out_trg_fname)59#     if os.path.isfile(out_meta_fname):60#         os.unlink(out_meta_fname)61 62    for src_lang, trg_lang in tqdm(lang_pair_list):63        print('src: {}, tgt:{}'.format(src_lang, trg_lang))64 65        in_src_fname = '{}/{}-{}/{}.{}'.format(66            data_dir, src_lang, trg_lang, split, src_lang)67        in_trg_fname = '{}/{}-{}/{}.{}'.format(68            data_dir, src_lang, trg_lang, split, trg_lang)69        70        if not os.path.exists(in_src_fname):71            continue72        if not os.path.exists(in_trg_fname):73            continue74 75        print(in_src_fname)76        os.system('cat {} >> {}'.format(in_src_fname, out_src_fname))77 78        print(in_trg_fname)79        os.system('cat {} >> {}'.format(in_trg_fname, out_trg_fname))80 81 82#     with open('{}/lang_pairs.txt'.format(outdir),'w',encoding='utf-8') as lpfile:83#         lpfile.write('\n'.join( [ '-'.join(x) for x in lang_pair_list ] ))84 85    corpus_stats(data_dir, outdir, lang_pair_list, split)86 87 88def corpus_stats(data_dir, outdir, lang_pair_list, split):89    """90    data_dir: input dir, contains directories for language pairs named l1-l291    """92 93    with open('{}/{}_lang_pairs.txt'.format(outdir, split), 'w', encoding='utf-8') as lpfile:94 95        for src_lang, trg_lang in tqdm(lang_pair_list):96            print('src: {}, tgt:{}'.format(src_lang, trg_lang))97 98            in_src_fname = '{}/{}-{}/{}.{}'.format(99                data_dir, src_lang, trg_lang, split, src_lang)100    #         in_trg_fname='{}/{}-{}/train.{}'.format(data_dir,src_lang,trg_lang,trg_lang)101            if not os.path.exists(in_src_fname):102                continue103 104            print(in_src_fname)105            corpus_size = 0106            with open(in_src_fname, 'r', encoding='utf-8') as infile:107                corpus_size = sum(map(lambda x: 1, infile))108 109            lpfile.write('{}\t{}\t{}\n'.format(110                src_lang, trg_lang, corpus_size))111 112 113if __name__ == '__main__':114 115    in_dir = sys.argv[1]116    out_dir = sys.argv[2]117    src_lang = sys.argv[3]118    tgt_lang = sys.argv[4]119    split = sys.argv[5]120    lang_pair_list = []121 122    if src_lang == 'en':123        for lang in LANGS:124            lang_pair_list.append(['en', lang])125    else:126        for lang in LANGS:127            lang_pair_list.append([lang, 'en'])128 129    concat_data(in_dir, out_dir, lang_pair_list, split=split)130 131