<filename>workspace/asr/rewrite_noise_label.py
import json
import os
import six
from camel_tools.utils.stringutils import force_encoding, force_unicode
from camel_tools.utils.charmap import CharMapper

src_fname = "tarred_audio_manifest.json"
tgt_fname = "tarred_audio_manifest_rn.json"
src_manifests = ["data/arabic/train/tarred/small/bucket0/{}",
                 "data/arabic/train/tarred/small/bucket1/{}",
                 "data/arabic/train/tarred/small/bucket2/{}",
                 "data/arabic/train/tarred/small/bucket3/{}",
                 "data/arabic/train/tarred/small/bucket4/{}",
                 "data/arabic/train/tarred/medium/bucket0/{}",
                 "data/arabic/train/tarred/medium/bucket1/{}",
                 "data/arabic/train/tarred/medium/bucket2/{}",
                 "data/arabic/train/tarred/medium/bucket3/{}",
                 "data/arabic/train/tarred/medium/bucket4/{}",
                 "data/arabic/train/tarred/medium/bucket5/{}",
                 "data/arabic/train/tarred/large/bucket0/{}",
                 "data/arabic/train/tarred/large/bucket1/{}",
                 "data/arabic/train/tarred/large/bucket2/{}",
                 "data/arabic/train/tarred/large/bucket3/{}",
                 "data/arabic/train/tarred/large/bucket4/{}",
                 "data/arabic/train/tarred/extra_large/bucket0/{}",
                 "data/arabic/train/tarred/extra_large/bucket1/{}",
                 "data/arabic/train/tarred/extra_large/bucket2/{}",
                 "data/arabic/train/tarred/extra_large/bucket3/{}",
                 "data/arabic/train/tarred/extra_large/bucket4/{}"]

cat_lang_map = {
    "@@aircraft@@": "طائرات",
    "@@animal@@": "حيوان",
    "@@cough@@": "سعال",
    "@@door@@": "باب",
    "@@explosion@@": "انفجار",
    "@@food@@": "غذاء",
    "@@footstep@@": "وقع أقدام",
    "@@gunshot@@": "طلق ناري",
    "@@laugh@@": "يضحك",
    "@@music@@": "موسيقى",
    "@@noise@@": "ضوضاء",
    "@@phone@@": "هاتف",
    "@@siren@@": "صفارة إنذار",
    "@@sneeze@@": "عطس",
    "@@throat@@": "حلق",
    "@@thunder@@": "رعد",
    "@@tools@@": "أدوات",
    "@@train@@": "قطار",
    "@@vehicle@@": "مركبة",
    "@@water@@": "ماء",
    "@@wind@@": "ريح",
    "@@sniff@@": "شم"
}

count = 0
pre_count = 0
clean_mapper = CharMapper.builtin_mapper('arclean')
for src in src_manifests:
    with open(src.format(src_fname), encoding="utf-8") as src_json:
        with open(src.format(tgt_fname), encoding="utf-8", mode="w") as tgt_json:
            print("Started rewrite {} -> {}".format(os.path.join(src.format(src_fname)),
                                                    os.path.join(src.format(tgt_fname))))
            for line in src_json:
                jd = json.loads(line)
                duration = float(jd["duration"])
                text = str(jd["text"]).strip().replace("  ", " ").strip()
                audio_filepath = str(jd["audio_filepath"])
                shard_id = int(jd["shard_id"])

                clean_text = force_unicode(text)
                if six.PY3:
                    clean_text = clean_mapper.map_string(clean_text)
                else:
                    clean_text = force_encoding(clean_text)
                    clean_text = clean_mapper.map_string(clean_text)

                for key, ar in cat_lang_map.items():
                    noise_ar_val = "[" + str(ar) + "]"
                    if noise_ar_val in clean_text:
                        clean_text = clean_text.replace(noise_ar_val, key)
                        pre_count = pre_count + 1
                        if pre_count % 1000 == 0:
                            print("Found {} lines with noise.".format(pre_count))

                if len(audio_filepath) > 0 and len(text) > 0 and duration > 0.0:
                    metadata = {
                        "audio_filepath": audio_filepath,
                        "duration": duration,
                        "text": clean_text,
                        "shard_id": shard_id
                    }

                    json.dump(metadata, tgt_json)
                    tgt_json.write("\n")

                count = count + 1
                if count % 100000 == 0:
                    print("Completed checking {} lines.".format(count))

            tgt_json.close()
        src_json.close()
        print("Finished rewrite {} -> {}\n\n".format(os.path.join(src.format(src_fname)),
                                                     os.path.join(src.format(tgt_fname))))

count = 0
pre_count = 0
src_dev = "data/arabic/dev/manifest_{}.json"
with open(src_dev.format("dev"), encoding="utf-8") as src_json:
    with open(src_dev.format("dev_rn"), encoding="utf-8", mode="w") as tgt_json:
        print("Started rewrite {} -> {}".format(os.path.join(src_dev.format("dev")),
                                                os.path.join(src_dev.format("dev_rn"))))
        for line in src_json:
            jd = json.loads(line)
            duration = float(jd["duration"])
            text = str(jd["text"]).strip().replace("  ", " ").strip()
            audio_filepath = str(jd["audio_filepath"])

            clean_text = force_unicode(text)
            if six.PY3:
                clean_text = clean_mapper.map_string(clean_text)
            else:
                clean_text = force_encoding(clean_text)
                clean_text = clean_mapper.map_string(clean_text)

            for key, ar in cat_lang_map.items():
                noise_ar_val = "[" + str(ar) + "]"
                if noise_ar_val in clean_text:
                    clean_text = clean_text.replace(noise_ar_val, key)
                    pre_count = pre_count + 1
                    if pre_count % 1000 == 0:
                        print("Found {} lines with noise.".format(pre_count))

            if len(audio_filepath) > 0 and len(text) > 0 and duration > 0.0:
                metadata = {
                    "audio_filepath": audio_filepath,
                    "duration": duration,
                    "text": clean_text,
                }

                json.dump(metadata, tgt_json)
                tgt_json.write("\n")

            count = count + 1
            if count % 100000 == 0:
                print("Completed checking {} lines.".format(count))

        tgt_json.close()
    src_json.close()
    print("Finished rewrite {} -> {}\n\n".format(os.path.join(src_dev.format("dev")),
                                                 os.path.join(src_dev.format("dev_rn"))))
