app.py

# -*- coding: utf-8 -*-

"""
@Author     : Rong Ye
@Time       : May 2022
@Contact    : yerong@bytedance
@Description:
"""

import os
import traceback
import shutil
import yaml
import re
from pydub import AudioSegment
import gradio as gr
from huggingface_hub import snapshot_download


LANGUAGE_CODES = {
    "German": "de",
    "Spanish": "es",
    "French": "fr",
    "Italian": "it",
    "Netherlands": "nl",
    "Portuguese": "pt",
    "Romanian": "ro",
    "Russian": "ru",
}

LANG_GEN_SETUPS = {
    "de": {"beam": 10, "lenpen": 0.7},
    "es": {"beam": 10, "lenpen": 0.1},
    "fr": {"beam": 10, "lenpen": 1.0},
    "it": {"beam": 10, "lenpen": 0.5},
    "nl": {"beam": 10, "lenpen": 0.4},
    "pt": {"beam": 10, "lenpen": 0.9},
    "ro": {"beam": 10, "lenpen": 1.0},
    "ru": {"beam": 10, "lenpen": 0.3},
}

os.system("git clone https://github.com/ReneeYe/ConST")
os.system("mv ConST ConST_git")
os.system('mv -n ConST_git/* ./')
os.system("rm -rf ConST_git")
os.system("pip3 install --editable ./")
os.system("mkdir -p data checkpoint")


huggingface_model_dir = snapshot_download(repo_id="ReneeYe/ConST_en2x_models")
print(huggingface_model_dir)


def convert_audio_to_16k_wav(audio_input):
    sound = AudioSegment.from_file(audio_input)
    sample_rate = sound.frame_rate
    num_channels = sound.channels
    num_frames = int(sound.frame_count())
    filename = audio_input.split("/")[-1]
    print("original file is at:", audio_input)
    if (num_channels > 1) or (sample_rate != 16000): # convert to mono-channel 16k wav
        if num_channels > 1:
            sound = sound.set_channels(1)
        if sample_rate != 16000:
            sound = sound.set_frame_rate(16000)
        num_frames = int(sound.frame_count())
        filename = filename.replace(".wav", "") + "_16k.wav"
        sound.export(f"data/{filename}", format="wav")
    else:
        shutil.copy(audio_input, f'data/{filename}')
    return filename, num_frames


def prepare_tsv(file_name, n_frame, language, task="ST"):
    tgt_lang = LANGUAGE_CODES[language]
    with open("data/test_case.tsv", "w") as f:
        f.write("id\taudio\tn_frames\ttgt_text\tspeaker\tsrc_lang\ttgt_lang\tsrc_text\n")
        f.write(f"sample\t{file_name}\t{n_frame}\tThis is in {tgt_lang}.\tspk.1\ten\t{tgt_lang}\tThis is English.\n")


def get_vocab_and_yaml(language):
    tgt_lang = LANGUAGE_CODES[language]
    # get: spm_ende.model and spm_ende.txt, and save to data/xxx
    # if exist, no need to download
    shutil.copy(os.path.join(huggingface_model_dir, f"vocabulary/spm_en{tgt_lang}.model"), "./data")
    shutil.copy(os.path.join(huggingface_model_dir, f"vocabulary/spm_en{tgt_lang}.txt"), "./data")

    # write yaml file
    abs_path = os.popen("pwd").read().strip()
    yaml_dict = LANG_GEN_SETUPS[tgt_lang]
    yaml_dict["input_channels"] = 1
    yaml_dict["use_audio_input"] = True
    yaml_dict["prepend_tgt_lang_tag"] = True
    yaml_dict["prepend_src_lang_tag"] = True
    yaml_dict["audio_root"] = os.path.join(abs_path, "data")
    yaml_dict["vocab_filename"] = f"spm_en{tgt_lang}.txt"
    yaml_dict["bpe_tokenizer"] = {"bpe": "sentencepiece",
                                  "sentencepiece_model": os.path.join(abs_path, f"data/spm_en{tgt_lang}.model")}
    with open("data/config.yaml", "w") as f:
        yaml.dump(yaml_dict, f)


def get_model(language):
    # download models to checkpoint/xxx
    return os.path.join(huggingface_model_dir, f"models/const_en{LANGUAGE_CODES[language]}.pt")


def generate(model_path):
    os.system(f"python3 fairseq_cli/generate.py data/ --gen-subset test_case --task speech_to_text --prefix-size 1 \
                --max-tokens 4000000 --max-source-positions 4000000 \
                --config-yaml config.yaml  --path {model_path} | tee temp.txt")
    output = os.popen("grep ^D temp.txt | sort -n -k 2 -t '-' | cut -f 3")
    return output.read().strip()


def post_processing(raw_sentence):
    output_sentence = raw_sentence
    if ":" in raw_sentence:
        splited_sent = raw_sentence.split(":")
        if len(splited_sent) == 2:
            prefix = splited_sent[0].strip()
            if len(prefix) <= 3:
                output_sentence = splited_sent[1].strip()
            elif ("(" in prefix) and (")" in prefix):
                bgm = re.findall(r"\(.*?\)", prefix)[0]
                if len(prefix.replace(bgm, "").strip()) <= 3:
                    output_sentence = splited_sent[1].strip()
                elif len(splited_sent[1].strip()) > 8:
                    output_sentence = splited_sent[1].strip()

    elif ("(" in raw_sentence) and (")" in raw_sentence):
        bgm_list = re.findall(r"\(.*?\)", raw_sentence)
        for bgm in bgm_list:
            if len(raw_sentence.replace(bgm, "").strip()) > 5:
                output_sentence = output_sentence.replace(bgm, "").strip()
        if len(output_sentence) <= 5:
            output_sentence = raw_sentence
    return output_sentence


def remove_temp_files(audio_file):
    os.remove("temp.txt")
    os.remove("data/test_case.tsv")
    os.remove(f"data/{audio_file}")


def run(audio_file, language):
    try:
        converted_audio_file, n_frame = convert_audio_to_16k_wav(audio_file)
        prepare_tsv(converted_audio_file, n_frame, language)
        get_vocab_and_yaml(language)
        model_path = get_model(language)
        generated_output = post_processing(generate(model_path))
        remove_temp_files(converted_audio_file)
        return generated_output
    except:
        traceback.print_exc()
        return error_output(language)


def error_output(language):
    return f"Fail to translate the audio into {language}, you may use the examples I provide."


inputs = [
        gr.inputs.Audio(source="microphone", type="filepath", label="Record something (in English)..."),
        gr.inputs.Dropdown(list(LANGUAGE_CODES.keys()), default="German", label="From English to Languages X..."),
    ]

iface = gr.Interface(
    fn=run,
    inputs=inputs,
    outputs=[gr.outputs.Textbox(label="The translation")],
    examples=[['short-case.wav', "German"], ['long-case.wav', "German"]],
    title="ConST: an end-to-end speech translator",
    description='ConST is an end-to-end speech-to-text translation model, whose algorithm corresponds to the '
                'NAACL 2022 paper *"Cross-modal Contrastive Learning for Speech Translation"* (see the paper at https://arxiv.org/abs/2205.02444 for more details). '
                'This is a live demo for ConST, to translate English into eight European languages. \n'
                'p.s. For better experience, we recommend using **Chrome** to record audio.',
    article="- The motivation of the ConST model is to use the contrastive learning method to learn similar representations for semantically similar speech and text, " \
            "thus leveraging MT to help improve ST performance. \n"
            "- The models you are experiencing are trained based on the MuST-C dataset (https://ict.fbk.eu/must-c/), " \
            "which only contains about 250k parallel data at each translation direction. "
            "The translation performance of these language directions varies from 20-30+ BLEU, "
            "so it is normal to find some flaws in the translation, and we are trying to improve the models, "
            "such as training on larger datasets and developing more advanced algorithms.\n"
            "- If you want to know how to train the models, you may refer to https://github.com/ReneeYe/ConST.",
    theme="peach",
)
iface.launch()