From 72f33be6f22af7fc03f8f5d7e31db523d2a3f293 Mon Sep 17 00:00:00 2001 From: Quentin Fuxa Date: Sun, 7 Sep 2025 15:25:14 +0200 Subject: [PATCH] translation: use of get_nllb_code --- whisperlivekit/core.py | 4 +++- whisperlivekit/translation/mapping_languages.py | 2 +- whisperlivekit/translation/translation.py | 13 +++++++------ 3 files changed, 11 insertions(+), 8 deletions(-) diff --git a/whisperlivekit/core.py b/whisperlivekit/core.py index 3a6b3c1..bc07aa8 100644 --- a/whisperlivekit/core.py +++ b/whisperlivekit/core.py @@ -133,12 +133,14 @@ class TranscriptionEngine: self.diarization_model = SortformerDiarization() else: raise ValueError(f"Unknown diarization backend: {self.args.diarization_backend}") - + + self.translation_model = None if self.args.target_language: if self.args.language == 'auto': raise Exception('Translation cannot be set with language auto') else: from whisperlivekit.translation.translation import load_model + self.translation_model = load_model() TranscriptionEngine._initialized = True diff --git a/whisperlivekit/translation/mapping_languages.py b/whisperlivekit/translation/mapping_languages.py index da6881d..becd657 100644 --- a/whisperlivekit/translation/mapping_languages.py +++ b/whisperlivekit/translation/mapping_languages.py @@ -132,7 +132,7 @@ NLLB_TO_NAME = {lang["nllb"]: lang["name"] for lang in LANGUAGES} def get_nllb_code(crowdin_code): - return CROWDIN_TO_NLLB.get(crowdin_code, crowdin_code) + return CROWDIN_TO_NLLB.get(crowdin_code, None) def get_crowdin_code(nllb_code): diff --git a/whisperlivekit/translation/translation.py b/whisperlivekit/translation/translation.py index 91c76fb..ccfedb5 100644 --- a/whisperlivekit/translation/translation.py +++ b/whisperlivekit/translation/translation.py @@ -2,8 +2,7 @@ import ctranslate2 import transformers from dataclasses import dataclass import huggingface_hub - -src_lang = "eng_Latn" +from .mapping_languages import get_nllb_code @dataclass class TranslationModel(): @@ -30,8 +29,10 @@ def translate(input, translation_model, tgt_lang): if __name__ == '__main__': - tgt_lang = "fra_Latn" - src_lang = "eng_Latn" - translation_model = load_model(src_lang) - result = translate('Hello world', translation_model=translation_model, tgt_lang=tgt_lang) + tgt_lang = 'fr' + src_lang = "en" + nllb_tgt_lang = get_nllb_code(tgt_lang) + nllb_src_lang = get_nllb_code(src_lang) + translation_model = load_model(nllb_src_lang) + result = translate('Hello world', translation_model=translation_model, tgt_lang=nllb_tgt_lang) print(result) \ No newline at end of file