subtify / translate_transcriptions.py
Maximofn's picture
Fix bugs
61afcdf
raw
history blame
7.06 kB
import torch
from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
from lang_list import LANGUAGE_NAME_TO_CODE, WHISPER_LANGUAGES
import argparse
import re
from tqdm import tqdm
MAX_LENGTH = 500
MAGIC_STRING = "[$&]"
DEBUG = False
language_dict = {}
# Iterate over the LANGUAGE_NAME_TO_CODE dictionary
for language_name, language_code in LANGUAGE_NAME_TO_CODE.items():
# Extract the language code (the first two characters before the underscore)
lang_code = language_code.split('_')[0].lower()
# Check if the language code is present in WHISPER_LANGUAGES
if lang_code in WHISPER_LANGUAGES:
# Construct the entry for the resulting dictionary
language_dict[language_name] = {
"transcriber": lang_code,
"translator": language_code
}
def translate(transcribed_text, source_languaje, target_languaje, translate_model, translate_tokenizer, device="cpu"):
# Get source and target languaje codes
source_languaje_code = language_dict[source_languaje]["translator"]
target_languaje_code = language_dict[target_languaje]["translator"]
encoded = translate_tokenizer(transcribed_text, return_tensors="pt").to(device)
generated_tokens = translate_model.generate(
**encoded,
forced_bos_token_id=translate_tokenizer.lang_code_to_id[target_languaje_code]
)
translated = translate_tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
return translated
def main(transcription_file, source_languaje, target_languaje, translate_model, translate_tokenizer, device):
output_folder = "translated_transcriptions"
_, transcription_file_name = transcription_file.split("/")
transcription_file_name, _ = transcription_file_name.split(".")
# Read transcription
with open(transcription_file, "r") as f:
transcription = f.read().splitlines()
# Concatenate transcriptions
raw_transcription = ""
progress_bar = tqdm(total=len(transcription), desc='Concatenate transcriptions progress')
for line in transcription:
if re.match(r"\d+$", line):
pass
elif re.match(r"\d\d:\d\d:\d\d,\d+ --> \d\d:\d\d:\d\d,\d+", line):
pass
elif re.match(r"^$", line):
pass
else:
line = re.sub(r"\[SPEAKER_\d\d\]:", MAGIC_STRING, line)
raw_transcription += f"{line} "
progress_bar.update(1)
progress_bar.close()
# Save raw transcription
if DEBUG:
output_file = f"{output_folder}/{transcription_file_name}_raw.srt"
with open(output_file, "w") as f:
f.write(raw_transcription)
# Split raw transcription
raw_transcription_list = raw_transcription.split(MAGIC_STRING)
if raw_transcription_list[0] == "":
raw_transcription_list = raw_transcription_list[1:]
# Concatenate transcripts and translate when length is less than MAX_LENGTH
translated_transcription = ""
concatenate_transcription = raw_transcription_list[0] + MAGIC_STRING
progress_bar = tqdm(total=len(raw_transcription_list), desc='Translate transcriptions progress')
progress_bar.update(1)
if len(raw_transcription_list) > 1:
for transcription in raw_transcription_list[1:]:
if len(concatenate_transcription) + len(transcription) < MAX_LENGTH:
concatenate_transcription += transcription + MAGIC_STRING
else:
translation = translate(concatenate_transcription, source_languaje, target_languaje, translate_model, translate_tokenizer, device)
translated_transcription += translation
concatenate_transcription = transcription + MAGIC_STRING
progress_bar.update(1)
# Translate last part
translation = translate(concatenate_transcription, source_languaje, target_languaje, translate_model, translate_tokenizer, device)
translated_transcription += translation
else:
translated_transcription = translate(concatenate_transcription, source_languaje, target_languaje, translate_model, translate_tokenizer, device)
progress_bar.close()
# Save translated transcription raw
if DEBUG:
output_file = f"{output_folder}/{transcription_file_name}_{target_languaje}_raw.srt"
with open(output_file, "w") as f:
f.write(translated_transcription)
# Read transcription
with open(transcription_file, "r") as f:
transcription = f.read().splitlines()
# Add time stamps
translated_transcription_time_stamps = ""
translated_transcription_list = translated_transcription.split(MAGIC_STRING)
progress_bar = tqdm(total=len(translated_transcription_list), desc='Add time stamps to translated transcriptions progress')
i = 0
for line in transcription:
if re.match(r"\d+$", line):
translated_transcription_time_stamps += f"{line}\n"
elif re.match(r"\d\d:\d\d:\d\d,\d+ --> \d\d:\d\d:\d\d,\d+", line):
translated_transcription_time_stamps += f"{line}\n"
elif re.match(r"^$", line):
translated_transcription_time_stamps += f"{line}\n"
else:
if (i < len(translated_transcription_list)):
if len(translated_transcription_list[i]) > 0:
if translated_transcription_list[i][0] == " ": # Remove space at the beginning
translated_transcription_list[i] = translated_transcription_list[i][1:]
speaker = ""
if re.match(r"\[SPEAKER_\d\d\]:", line):
speaker = re.match(r"\[SPEAKER_\d\d\]:", line).group(0)
translated_transcription_time_stamps += f"{speaker} {translated_transcription_list[i]}\n"
i += 1
progress_bar.update(1)
progress_bar.close()
# Save translated transcription
output_file = f"{output_folder}/{transcription_file_name}_{target_languaje}.srt"
with open(output_file, "w") as f:
f.write(translated_transcription_time_stamps)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("transcription_file", help="Transcribed text")
parser.add_argument("--source_languaje", type=str, required=True)
parser.add_argument("--target_languaje", type=str, required=True)
parser.add_argument("--device", type=str, default="cpu")
args = parser.parse_args()
transcription_file = args.transcription_file
source_languaje = args.source_languaje
target_languaje = args.target_languaje
device = args.device
# model
print("Loading translation model")
translate_model = MBartForConditionalGeneration.from_pretrained("facebook/mbart-large-50-many-to-many-mmt").to(device)
translate_tokenizer = MBart50TokenizerFast.from_pretrained("facebook/mbart-large-50-many-to-many-mmt")
print("Translation model loaded")
main(transcription_file, source_languaje, target_languaje, translate_model, translate_tokenizer, device)