Refactor language validation logic

modules/translation/nllb_inference.py

@@ -37,6 +37,17 @@ class NLLBInference(TranslationBase):
                      tgt_lang: str,
                      progress: gr.Progress = gr.Progress()
                      ):
+        def validate_language(lang: str) -> str:
+            if lang in NLLB_AVAILABLE_LANGS:
+                return NLLB_AVAILABLE_LANGS[lang]
+            elif lang not in NLLB_AVAILABLE_LANGS.values():
+                raise ValueError(
+                    f"Language '{lang}' is not supported. Use one of: {list(NLLB_AVAILABLE_LANGS.keys())}")
+            return lang
+
+        src_lang = validate_language(src_lang)
+        tgt_lang = validate_language(tgt_lang)
+
         if model_size != self.current_model_size or self.model is None:
             print("\nInitializing NLLB Model..\n")
             progress(0, desc="Initializing NLLB Model..")
@@ -48,8 +59,7 @@ class NLLBInference(TranslationBase):
             self.tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path=model_size,
                                                            cache_dir=os.path.join(self.model_dir, "tokenizers"),
                                                            local_files_only=local_files_only)
-        src_lang = NLLB_AVAILABLE_LANGS[src_lang]
-        tgt_lang = NLLB_AVAILABLE_LANGS[tgt_lang]
+
         self.pipeline = pipeline("translation",
                                  model=self.model,
                                  tokenizer=self.tokenizer,