Repository: mozilla-ai/lumigator
Pull request #1178: Support for language code based HF translation models like facebook/m2m100* family
Author: login: HareeshBahuleyan
Created: 2025-03-12T09:26:50Z | Merged: 2025-03-14T10:13:26Z

# What's changing

Add `HuggingFaceLanguageCodeTranslationClient` to support m2m100-* models that support translation where the language pairs are specified in the pipeline as iso codes. Example:
```
pipeline('translation', 'facebook/m2m100_418M', src_lang='en', tgt_lang="de")
```


Note, previously we added `HuggingFacePrefixTranslationClient` to support models that work with prefix before the input sentence in #1147 

Closes https://github.com/mozilla-ai/lumigator/issues/1192

# How to test it

Steps to test the changes:

1. Upload en to de sample dataset
2. Create an experiment:
```
curl -X 'POST' \
  'http://localhost:8000/api/v1/experiments/' \
  -H 'accept: application/json' \
  -H 'Content-Type: application/json' \
  -d '{
  "name": "translation",
  "dataset": "82096978-17a6-40e7-8698-67fc7338320c",
  "max_samples": 2,
  "task_definition": {
    "task": "translation",
    "source_language": "English",
    "target_language": "German"
  }
}'
```

3. Create and run workflow:
```
curl -X 'POST' \
  'http://localhost:8000/api/v1/workflows/' \
  -H 'accept: application/json' \
  -H 'Content-Type: application/json' \
  -d '{
  "name": "tranz",
  "experiment_id": "24",
  "task_definition": {
    "task": "translation",
    "source_language": "English",
    "target_language": "German"
  },
  "model": "facebook/m2m100_418M",
  "provider": "hf",
  "dataset": "82096978-17a6-40e7-8698-67fc7338320c",
  "max_samples": 2
}'
```

4. Check UI to see if the generation is as expected.
![image](https://github.com/user-attachments/assets/9e457c91-3a83-405c-ab7f-9f6cfd42820b)


# Additional notes for reviewers

Currently, we don't have a way to check if the model specified supports the language that the user has provided. 

TODO:
- [x] Add jobs unit tests 

# I already...

- [x] Tested the changes in a working environment to ensure they work as expected
- [x] Added some tests for any new functionality
- [ ] Updated the documentation (both comments in code and [product documentation](https://mozilla-ai.github.io/lumigator) under `/docs`)
- [ ] Checked if a (backend) DB migration step was required and included it if required



Diff:
diff --git a/lumigator/jobs/inference/model_clients/huggingface_clients.py b/lumigator/jobs/inference/model_clients/huggingface_clients.py
index 79356a959..7ce66546c 100644
--- a/lumigator/jobs/inference/model_clients/huggingface_clients.py
+++ b/lumigator/jobs/inference/model_clients/huggingface_clients.py
@@ -5,7 +5,8 @@
 from model_clients.mixins.generation_config_mixin import GenerationConfigMixin
 from model_clients.mixins.huggingface_model_mixin import HuggingFaceModelMixin
 from model_clients.mixins.huggingface_seq2seq_pipeline_mixin import HuggingFaceSeq2SeqPipelineMixin
-from model_clients.translation_utils import load_translation_config, resolve_user_input_language
+from model_clients.mixins.language_code_mixin import LanguageCodesSetupMixin
+from model_clients.translation_utils import load_translation_config
 from transformers import AutoConfig, pipeline
 
 from schemas import PredictionResult
@@ -36,10 +37,12 @@ def create(config: InferenceJobConfig, api_key: str | None = None) -> BaseModelC
             # Load the supported translation model families configuration
             translation_config = load_translation_config()
 
-            prefix_models = translation_config.get("prefix_translation_models", [])
-            if model_name in prefix_models:
+            if model_name in translation_config["prefix_based"]:
                 logger.info(f"Running inference with HuggingFacePrefixTranslationClient for {model_name}")
-                return HuggingFacePrefixTranslationClient(config)
+                return HuggingFacePrefixTranslationClient(config, api_key)
+            elif model_name in translation_config["language_code_based"]:
+                logger.info(f"Running inference with HuggingFaceLanguageCodeTranslationClient for {model_name}")
+                return HuggingFaceLanguageCodeTranslationClient(config, api_key)
             else:
                 logger.error(
                     f"Translation task with HF seq2seq models: {model_name} is not supported. "
@@ -69,22 +72,9 @@ def __init__(self, config: InferenceJobConfig, api_key: str | None = None):
         self.api_key = api_key
         self.model = self.initialize_model(self.config.hf_pipeline)
         self.tokenizer = self.initialize_tokenizer(self.config.hf_pipeline)
-        self.pipeline = self.initialize_pipeline(self.config.hf_pipeline, self.model, self.tokenizer, api_key)
+        self.pipeline = self.initialize_pipeline(self.config.hf_pipeline, self.model, self.tokenizer, api_key=api_key)
         self.set_seq2seq_max_length()
 
-    def set_seq2seq_max_length(self):
-        """Set the maximum sequence length for the seq2seq model.
-
-        This method ensures that the tokenizer and model have the same maximum position embeddings
-        and adjusts the generation configuration accordingly.
-        """
-        # If the model is of the HF Hub the odds of this being wrong are low, but it's still good to check that the
-        # tokenizer model and the model have the same max_position_embeddings.
-        max_pos_emb = self.get_max_position_embeddings(self.pipeline.model)
-        self.adjust_tokenizer_max_length(self.pipeline, max_pos_emb)
-        # Adjust output sequence generation max tokens.
-        self.adjust_config_max_new_tokens(self.config.generation_config, max_pos_emb)
-
     def predict(self, examples: list) -> list[PredictionResult]:
         generations = self.pipeline(
             examples, max_new_tokens=self.config.generation_config.max_new_tokens, truncation=True
@@ -138,53 +128,35 @@ def predict(self, examples: list[list[dict[str, str]]]) -> list[PredictionResult
 
 
 class HuggingFacePrefixTranslationClient(
-    BaseModelClient, HuggingFaceModelMixin, HuggingFaceSeq2SeqPipelineMixin, GenerationConfigMixin
+    BaseModelClient,
+    HuggingFaceModelMixin,
+    HuggingFaceSeq2SeqPipelineMixin,
+    GenerationConfigMixin,
+    LanguageCodesSetupMixin,
 ):
     """Client for T5-style models that use prefixes for translation"""
 
     def __init__(self, config: InferenceJobConfig, api_key: str | None = None):
         self.config = config
         self.api_key = api_key
+        self.setup_translation_languages(config.task_definition)
         self.model = self.initialize_model(self.config.hf_pipeline)
         self.tokenizer = self.initialize_tokenizer(self.config.hf_pipeline)
 
-        source_language_user_input = getattr(config.task_definition, "source_language", None)
-        target_language_user_input = getattr(config.task_definition, "target_language", None)
-
-        if not source_language_user_input or not target_language_user_input:
-            raise ValueError("Source and target languages must be provided for translation task.")
-
-        source_language_info = resolve_user_input_language(source_language_user_input)
-        target_language_info = resolve_user_input_language(target_language_user_input)
-
-        self.source_language_iso_code = source_language_info["iso_code"]  # e.g. "en"
-        self.source_language = source_language_info["full_name"]  # e.g. "English"
-        self.target_language_iso_code = target_language_info["iso_code"]  # e.g. "de"
-        self.target_language = target_language_info["full_name"]  # e.g. "German"
-
         # Modify the task to include the the source and target languages
         specific_pipeline_task = (
             f"{TaskType.TRANSLATION.value}_{self.source_language_iso_code}_to_{self.target_language_iso_code}"
         )
         self.pipeline = self.initialize_pipeline(
-            self.config.hf_pipeline, self.model, self.tokenizer, specific_pipeline_task, api_key
+            self.config.hf_pipeline,
+            self.model,
+            self.tokenizer,
+            specific_pipeline_task=specific_pipeline_task,
+            api_key=api_key,
         )
         self.set_seq2seq_max_length()
         self.prefix = f"translate {self.source_language} to {self.target_language}: "
 
-    def set_seq2seq_max_length(self):
-        """Set the maximum sequence length for the seq2seq model.
-
-        This method ensures that the tokenizer and model have the same maximum position embeddings
-        and adjusts the generation configuration accordingly.
-        """
-        # If the model is of the HF Hub the odds of this being wrong are low, but it's still good to check that the
-        # tokenizer model and the model have the same max_position_embeddings.
-        max_pos_emb = self.get_max_position_embeddings(self.pipeline.model)
-        self.adjust_tokenizer_max_length(self.pipeline, max_pos_emb)
-        # Adjust output sequence generation max tokens.
-        self.adjust_config_max_new_tokens(self.config.generation_config, max_pos_emb)
-
     def predict(self, examples: list) -> list[PredictionResult]:
         prefixed_examples = [self.prefix + example for example in examples]
 
@@ -198,3 +170,39 @@ def predict(self, examples: list) -> list[PredictionResult]:
             prediction_results.append(prediction_result)
 
         return prediction_results
+
+
+class HuggingFaceLanguageCodeTranslationClient(
+    BaseModelClient,
+    HuggingFaceModelMixin,
+    HuggingFaceSeq2SeqPipelineMixin,
+    GenerationConfigMixin,
+    LanguageCodesSetupMixin,
+):
+    """Client for translation models that require language codes (mBART, NLLB, M2M)"""
+
+    def __init__(self, config: InferenceJobConfig, api_key: str | None = None):
+        self.config = config
+        self.api_key = api_key
+        self.setup_translation_languages(config.task_definition)
+        self.model = self.initialize_model(self.config.hf_pipeline)
+        self.tokenizer = self.initialize_tokenizer(self.config.hf_pipeline)
+
+        self.pipeline = self.initialize_pipeline(self.config.hf_pipeline, self.model, self.tokenizer, api_key=api_key)
+        self.set_seq2seq_max_length()
+
+    def predict(self, examples: str | list[list[dict[str, str]]]) -> list[PredictionResult]:
+        generations = self.pipeline(
+            examples,
+            max_new_tokens=self.config.generation_config.max_new_tokens,
+            truncation=True,
+            src_lang=self.source_language_iso_code,
+            tgt_lang=self.target_language_iso_code,
+        )
+
+        prediction_results = []
+        for generation in generations:
+            prediction_result = PredictionResult(prediction=generation["translation_text"])
+            prediction_results.append(prediction_result)
+
+        return prediction_results
diff --git a/lumigator/jobs/inference/model_clients/mixins/huggingface_seq2seq_pipeline_mixin.py b/lumigator/jobs/inference/model_clients/mixins/huggingface_seq2seq_pipeline_mixin.py
index 363faebde..10faa3eb7 100644
--- a/lumigator/jobs/inference/model_clients/mixins/huggingface_seq2seq_pipeline_mixin.py
+++ b/lumigator/jobs/inference/model_clients/mixins/huggingface_seq2seq_pipeline_mixin.py
@@ -119,3 +119,27 @@ def adjust_tokenizer_max_length(self, pipeline: Pipeline, max_pos_emb: int | Non
             f"is bigger than the model's max_position_embeddings ({max_pos_emb})."
         )
         pipeline.tokenizer.model_max_length = max_pos_emb
+
+    def set_seq2seq_max_length(self) -> None:
+        """Set the maximum sequence length for input and generation for the seq2seq model.
+
+        This method ensures that the tokenizer and model have the same maximum position embeddings
+        and adjusts the generation configuration accordingly.
+
+        :raises TypeError: If the pipeline, model, or config is None.
+        """
+        if not hasattr(self, "pipeline") or self.pipeline is None:
+            raise TypeError("The pipeline cannot be None")
+        if not hasattr(self.pipeline, "model") or self.pipeline.model is None:
+            raise TypeError("The pipeline's model cannot be None")
+        if not hasattr(self, "config") or self.config is None:
+            raise TypeError("The config cannot be None")
+        if not hasattr(self.config, "generation_config") or self.config.generation_config is None:
+            raise TypeError("The generation_config cannot be None")
+
+        # If the model is of the HF Hub the odds of this being wrong are low, but it's still good to check that the
+        # tokenizer model and the model have the same max_position_embeddings.
+        max_pos_emb = self.get_max_position_embeddings(self.pipeline.model)
+        self.adjust_tokenizer_max_length(self.pipeline, max_pos_emb)
+        # Adjust output sequence generation max tokens.
+        self.adjust_config_max_new_tokens(self.config.generation_config, max_pos_emb)
diff --git a/lumigator/jobs/inference/model_clients/mixins/language_code_mixin.py b/lumigator/jobs/inference/model_clients/mixins/language_code_mixin.py
new file mode 100644
index 000000000..ff0a2b195
--- /dev/null
+++ b/lumigator/jobs/inference/model_clients/mixins/language_code_mixin.py
@@ -0,0 +1,22 @@
+from lumigator_schemas.tasks import TaskDefinition
+from model_clients.translation_utils import resolve_user_input_language
+
+
+class LanguageCodesSetupMixin:
+    """Mixin with common functionality for translation clients"""
+
+    def setup_translation_languages(self, task_definition: TaskDefinition):
+        """Initialize source and target language information for translation"""
+        source_language_user_input = getattr(task_definition, "source_language", None)
+        target_language_user_input = getattr(task_definition, "target_language", None)
+
+        if not source_language_user_input or not target_language_user_input:
+            raise ValueError("Source and target languages must be provided for translation task.")
+
+        source_language_info = resolve_user_input_language(source_language_user_input)
+        target_language_info = resolve_user_input_language(target_language_user_input)
+
+        self.source_language_iso_code = source_language_info["iso_code"]  # e.g. "en"
+        self.source_language = source_language_info["full_name"]  # e.g. "English"
+        self.target_language_iso_code = target_language_info["iso_code"]  # e.g. "de"
+        self.target_language = target_language_info["full_name"]  # e.g. "German"
diff --git a/lumigator/jobs/inference/model_clients/translation_models.yaml b/lumigator/jobs/inference/model_clients/translation_models.yaml
index 3af8fee76..0ffd5c1ee 100644
--- a/lumigator/jobs/inference/model_clients/translation_models.yaml
+++ b/lumigator/jobs/inference/model_clients/translation_models.yaml
@@ -1,7 +1,7 @@
 # Configuration file enumerating supported translation models
 
 # Models that use the "translate {source_language} to {target_language}: " prefix format
-prefix_translation_models:
+prefix_based:
   # Google's Flan-T5 family
   - google/flan-t5-small
   - google/flan-t5-base
@@ -15,3 +15,9 @@ prefix_translation_models:
   - bigscience/mt0-large
   - bigscience/mt0-xl
   - bigscience/mt0-xxl
+
+# Models that require the source and target language codes to be specified
+# in the pipeline configuration for example src_lang="en", tgt_lang="de"
+language_code_based:
+  - facebook/m2m100_418M
+  - facebook/m2m100_1.2B
diff --git a/lumigator/jobs/inference/model_clients/translation_utils.py b/lumigator/jobs/inference/model_clients/translation_utils.py
index aeef6b41c..5cb5fec1d 100644
--- a/lumigator/jobs/inference/model_clients/translation_utils.py
+++ b/lumigator/jobs/inference/model_clients/translation_utils.py
@@ -3,6 +3,7 @@
 import yaml
 from langcodes import Language, find, tag_is_valid
 from langcodes.tag_parser import LanguageTagError
+from loguru import logger
 
 
 def load_translation_config():
@@ -11,12 +12,20 @@ def load_translation_config():
         # Get the current file's directory and construct path to config file
         current_file = Path(__file__).resolve()
         config_path = current_file.parent / "translation_models.yaml"
-
-        # Open and load the yaml file
         with config_path.open() as file:
             translation_config = yaml.safe_load(file)
 
-        return translation_config
+        # Define default configuration
+        default_config = {"prefix_based": [], "language_code_based": []}
+
+        # Merge loaded config with default config
+        validated_config = {**default_config, **translation_config}
+
+        for key in default_config:
+            if key not in translation_config:
+                logger.warning(f"Missing required key in translation_models.yaml: {key}.")
+
+        return validated_config
     except Exception as e:
         raise Exception("Error loading translation models configuration.") from e
 
diff --git a/lumigator/jobs/inference/tests/test_huggingface_clients.py b/lumigator/jobs/inference/tests/test_huggingface_clients.py
index 0d1f48d84..f2d0d4ab3 100644
--- a/lumigator/jobs/inference/tests/test_huggingface_clients.py
+++ b/lumigator/jobs/inference/tests/test_huggingface_clients.py
@@ -5,6 +5,7 @@
 from lumigator_schemas.tasks import TaskDefinition, TaskType
 from model_clients.huggingface_clients import (
     HuggingFaceCausalLMClient,
+    HuggingFaceLanguageCodeTranslationClient,
     HuggingFacePrefixTranslationClient,
     HuggingFaceSeq2SeqSummarizationClient,
 )
@@ -231,7 +232,7 @@ def mock_config(self):
         config.generation_config.max_new_tokens = 100
         return config
 
-    @patch("model_clients.huggingface_clients.resolve_user_input_language")
+    @patch("model_clients.translation_utils.resolve_user_input_language")
     @patch("model_clients.mixins.huggingface_seq2seq_pipeline_mixin.AutoTokenizer")
     @patch("model_clients.mixins.huggingface_seq2seq_pipeline_mixin.AutoModelForSeq2SeqLM")
     @patch("model_clients.mixins.huggingface_seq2seq_pipeline_mixin.pipeline")
@@ -277,7 +278,7 @@ def test_initialization(self, mock_pipeline, mock_automodel, mock_tokenizer, moc
         assert pipeline_args["task"] == "translation_en_to_fr"
         assert client.pipeline == mock_pipeline_instance
 
-    @patch("model_clients.huggingface_clients.resolve_user_input_language")
+    @patch("model_clients.translation_utils.resolve_user_input_language")
     @patch("model_clients.mixins.huggingface_seq2seq_pipeline_mixin.AutoTokenizer")
     @patch("model_clients.mixins.huggingface_seq2seq_pipeline_mixin.AutoModelForSeq2SeqLM")
     @patch("model_clients.mixins.huggingface_seq2seq_pipeline_mixin.pipeline")
@@ -326,3 +327,124 @@ def test_predict(self, mock_pipeline, mock_automodel, mock_tokenizer, mock_resol
             "translate English to French: I am a translation model.",
         ]
         mock_pipeline_instance.assert_called_once_with(expected_prefixed_inputs, max_new_tokens=100, truncation=True)
+
+
+class TestHuggingFaceLanguageCodeTranslationClient:
+    @pytest.fixture
+    def mock_config(self):
+        """Create a mock InferenceJobConfig for testing language code translation client."""
+        config = MagicMock(spec=InferenceJobConfig)
+        config.task_definition = MagicMock(spec=TaskDefinition)
+        config.task_definition.source_language = "en"
+        config.task_definition.target_language = "fr"
+
+        config.hf_pipeline = MagicMock(spec=HfPipelineConfig)
+        config.hf_pipeline.model_name_or_path = "mock-language-code-model"
+        config.hf_pipeline.task = TaskType.TRANSLATION
+        config.hf_pipeline.trust_remote_code = False
+        config.hf_pipeline.use_fast = True
+        config.hf_pipeline.torch_dtype = "float32"
+        config.hf_pipeline.revision = "main"
+        config.hf_pipeline.device = "cpu"
+
+        config.generation_config = MagicMock(spec=GenerationConfig)
+        config.generation_config.max_new_tokens = 100
+        return config
+
+    @patch("model_clients.translation_utils.resolve_user_input_language")
+    @patch("model_clients.mixins.huggingface_seq2seq_pipeline_mixin.AutoTokenizer")
+    @patch("model_clients.mixins.huggingface_seq2seq_pipeline_mixin.AutoModelForSeq2SeqLM")
+    @patch("model_clients.mixins.huggingface_seq2seq_pipeline_mixin.pipeline")
+    def test_initialization(self, mock_pipeline, mock_automodel, mock_tokenizer, mock_resolve_lang, mock_config):
+        """Test initialization of the language code translation client."""
+        # Setup mocks
+        mock_model = MagicMock()
+        mock_model.config.max_position_embeddings = 512
+        mock_automodel.from_pretrained.return_value = mock_model
+
+        mock_tokenizer_instance = MagicMock(spec=PreTrainedTokenizer)
+        mock_tokenizer_instance.model_max_length = 512
+        mock_tokenizer.from_pretrained.return_value = mock_tokenizer_instance
+
+        mock_pipeline_instance = MagicMock(spec=Pipeline)
+        mock_pipeline_instance.model = mock_model
+        mock_pipeline_instance.tokenizer = mock_tokenizer_instance
+        mock_pipeline.return_value = mock_pipeline_instance
+
+        # Mock language resolution
+        mock_resolve_lang.side_effect = [
+            {"iso_code": "en", "full_name": "English"},
+            {"iso_code": "fr", "full_name": "French"},
+        ]
+
+        # Initialize client
+        client = HuggingFaceLanguageCodeTranslationClient(mock_config)
+
+        # Verify initialization
+        mock_tokenizer.from_pretrained.assert_called_once()
+        mock_automodel.from_pretrained.assert_called_once()
+        mock_pipeline.assert_called_once()
+
+        # Check that correct language parameters are set
+        assert client.source_language_iso_code == "en"
+        assert client.source_language == "English"
+        assert client.target_language_iso_code == "fr"
+        assert client.target_language == "French"
+
+        # Verify pipeline task is correctly set for translation
+        mock_pipeline.assert_called_once()
+        pipeline_args = mock_pipeline.call_args[1]
+        assert pipeline_args["task"] == TaskType.TRANSLATION
+        assert client.pipeline == mock_pipeline_instance
+
+    @patch("model_clients.translation_utils.resolve_user_input_language")
+    @patch("model_clients.mixins.huggingface_seq2seq_pipeline_mixin.AutoTokenizer")
+    @patch("model_clients.mixins.huggingface_seq2seq_pipeline_mixin.AutoModelForSeq2SeqLM")
+    @patch("model_clients.mixins.huggingface_seq2seq_pipeline_mixin.pipeline")
+    def test_predict(self, mock_pipeline, mock_automodel, mock_tokenizer, mock_resolve_lang, mock_config):
+        """Test the predict method of the language code translation client."""
+        # Setup mocks
+        mock_model = MagicMock()
+        mock_model.config.max_position_embeddings = 512
+        mock_automodel.from_pretrained.return_value = mock_model
+
+        mock_tokenizer_instance = MagicMock(spec=PreTrainedTokenizer)
+        mock_tokenizer_instance.model_max_length = 512
+        mock_tokenizer.from_pretrained.return_value = mock_tokenizer_instance
+
+        mock_pipeline_instance = MagicMock(spec=Pipeline)
+        mock_pipeline_instance.model = mock_model
+        mock_pipeline_instance.tokenizer = mock_tokenizer_instance
+        mock_pipeline_instance.return_value = [
+            {"translation_text": "Bonjour, comment ça va?"},
+            {"translation_text": "Je suis un modèle de traduction."},
+        ]
+        mock_pipeline.return_value = mock_pipeline_instance
+
+        # Set max_new_tokens in config
+        mock_config.generation_config.max_new_tokens = 100
+
+        # Mock language resolution
+        mock_resolve_lang.side_effect = [
+            {"iso_code": "en", "full_name": "English"},
+            {"iso_code": "fr", "full_name": "French"},
+        ]
+
+        # Initialize client and call predict
+        client = HuggingFaceLanguageCodeTranslationClient(mock_config)
+        result = client.predict(["Hello, how are you?", "I am a translation model."])
+
+        # Verify prediction
+        assert len(result) == 2
+        assert all(isinstance(r, PredictionResult) for r in result)
+        assert result[0].prediction == "Bonjour, comment ça va?"
+        assert result[1].prediction == "Je suis un modèle de traduction."
+
+        # Verify pipeline was called with correct language codes
+        mock_pipeline_instance.assert_called_once_with(
+            ["Hello, how are you?", "I am a translation model."],
+            max_new_tokens=100,
+            truncation=True,
+            src_lang="en",
+            tgt_lang="fr",
+        )
diff --git a/lumigator/jobs/inference/tests/test_language_code_mixin.py b/lumigator/jobs/inference/tests/test_language_code_mixin.py
new file mode 100644
index 000000000..b2fd8920a
--- /dev/null
+++ b/lumigator/jobs/inference/tests/test_language_code_mixin.py
@@ -0,0 +1,67 @@
+from unittest.mock import MagicMock, patch
+
+import pytest
+from lumigator_schemas.tasks import TaskDefinition
+from model_clients.mixins.language_code_mixin import LanguageCodesSetupMixin
+
+
+class TestLanguageCodesSetupMixin:
+    @pytest.fixture
+    def language_mixin(self):
+        """Return a simple instance of the LanguageCodesSetupMixin."""
+        return LanguageCodesSetupMixin()
+
+    @patch("model_clients.mixins.language_code_mixin.resolve_user_input_language")
+    def test_setup_translation_languages_success(self, mock_resolve, language_mixin):
+        """Test successful setup of source and target languages."""
+        # Setup mock task definition with source and target languages
+        task_def = MagicMock(spec=TaskDefinition)
+        task_def.source_language = "english"
+        task_def.target_language = "french"
+
+        # Mock the resolve_user_input_language function
+        mock_resolve.side_effect = [
+            {"iso_code": "en", "full_name": "English"},
+            {"iso_code": "fr", "full_name": "French"},
+        ]
+
+        # Call the method
+        language_mixin.setup_translation_languages(task_def)
+
+        # Verify the attributes were set correctly
+        assert language_mixin.source_language_iso_code == "en"
+        assert language_mixin.source_language == "English"
+        assert language_mixin.target_language_iso_code == "fr"
+        assert language_mixin.target_language == "French"
+
+        # Verify resolve_user_input_language was called twice with correct arguments
+        assert mock_resolve.call_count == 2
+        mock_resolve.assert_any_call("english")
+        mock_resolve.assert_any_call("french")
+
+    @patch("model_clients.mixins.language_code_mixin.resolve_user_input_language")
+    def test_setup_translation_languages_with_iso_codes(self, mock_resolve, language_mixin):
+        """Test setup with ISO codes instead of full language names."""
+        # Setup mock task definition with ISO code languages
+        task_def = MagicMock(spec=TaskDefinition)
+        task_def.source_language = "en"
+        task_def.target_language = "fr"
+
+        # Mock the resolve_user_input_language function
+        mock_resolve.side_effect = [
+            {"iso_code": "en", "full_name": "English"},
+            {"iso_code": "fr", "full_name": "French"},
+        ]
+
+        # Call the method
+        language_mixin.setup_translation_languages(task_def)
+
+        # Verify the attributes were set correctly
+        assert language_mixin.source_language_iso_code == "en"
+        assert language_mixin.source_language == "English"
+        assert language_mixin.target_language_iso_code == "fr"
+        assert language_mixin.target_language == "French"
+
+        # Verify resolve_user_input_language was called with ISO codes
+        mock_resolve.assert_any_call("en")
+        mock_resolve.assert_any_call("fr")


Reviews:
user: login: aittalam
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-13T14:57:02Z
user: login: HareeshBahuleyan
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-13T15:58:05Z
user: login: aittalam
body: Thanks @HareeshBahuleyan! Looks good to me, I left a comment re: trapping potentially unsupported languages early but I think it's out of scope wrt this PR. Feel free to merge this when all checks are good and choose how to move forward.
state: APPROVED
author_association: MEMBER
submitted_at: 2025-03-14T09:37:34Z

Review comments:
diff_hunk: @@ -36,10 +37,15 @@ def create(config: InferenceJobConfig, api_key: str | None = None) -> BaseModelC
             # Load the supported translation model families configuration
             translation_config = load_translation_config()
 
-            prefix_models = translation_config.get("prefix_translation_models", [])
+            prefix_models = translation_config.get("prefix_based", [])
path: lumigator/jobs/inference/model_clients/huggingface_clients.py
user: login: aittalam
body: nit: I understand the `get` to make sure there is something returned even if the keys disappear from the input file. WDYT about bringing this logic in the `load_translation_config` method instead?
One advantage is that you would not need for the presence of those keys in here, and at the same time you could issue some kind of warning if the file does not contain the keys you are expecting (or error, if you think the application should break if those keys are not present)
created_at: 2025-03-13T14:57:02Z
author_association: MEMBER
diff_hunk: @@ -36,10 +37,15 @@ def create(config: InferenceJobConfig, api_key: str | None = None) -> BaseModelC
             # Load the supported translation model families configuration
             translation_config = load_translation_config()
 
-            prefix_models = translation_config.get("prefix_translation_models", [])
+            prefix_models = translation_config.get("prefix_based", [])
path: lumigator/jobs/inference/model_clients/huggingface_clients.py
user: login: HareeshBahuleyan
body: Good idea, I moved the logic to load_translation_config() and included issuing a warning 👍 
created_at: 2025-03-13T15:58:05Z
author_association: CONTRIBUTOR

PR comments:
user: login: aittalam
created_at: 2025-03-13T15:02:16Z
body: Hi Hareesh, many thanks for this PR! 
Re: "we don't have a way to check if the model specified supports the language that the user has provided", is it something we just need to implement or we do not have a way to get it from HF? What happens if the user provides the wrong pair? Alternatively, is there information we can gather about the models we are supporting that we could encode in a structured way (e.g. in the yaml file holding the translation models) explicitly stating which language pairs are available?
author_association: MEMBER
user: login: HareeshBahuleyan
created_at: 2025-03-13T16:06:47Z
body: Yes Davide, regarding the language compatibility one way we can achieve this is by calling the hf_api to get the model tags
https://huggingface.co/docs/huggingface_hub/en/package_reference/hf_api#huggingface_hub.HfApi.get_model_tags

And validate if the language is present in the list (not sure if it is a fool-proof way to check for all models, since it depends on the contributors adding those tags): 
![image](https://github.com/user-attachments/assets/1fc643e1-bebb-47e1-b6a0-accdfcaa8b84)

We could build a functionality for this, at the moment I think words from unsupported languages may be treated as unknown tokens.
author_association: CONTRIBUTOR
user: login: aittalam
created_at: 2025-03-14T09:35:44Z
body: > Yes Davide, regarding the language compatibility one way we can achieve this is by calling the hf_api to get the model tags https://huggingface.co/docs/huggingface_hub/en/package_reference/hf_api#huggingface_hub.HfApi.get_model_tags
> 
> And validate if the language is present in the list (not sure if it is a fool-proof way to check for all models, since it depends on the contributors adding those tags)

That's a good point! Still better than nothing though, perhaps instead of having a breaking error we could give a warning to the user that the language is not officially supported so they can decide... The only thing I'd like to avoid is people using unsupported languages from the UI and just seeing bad results with no apparent reason (which btw might happen even if the languages are supported 😅). WDYT?

If that sounds good to you, I think this could be a follow-up so we can keep this PR small and easy. We might even choose to capture this in the UI directly to allow users make a decision before the job is submitted, but I'll let you and @khaledosman decide whether it makes sense and how to implement it (e.g. direct call to HF API vs storing languages info ourselves in the yaml file and making it available via the API).
author_association: MEMBER
user: login: HareeshBahuleyan
created_at: 2025-03-14T10:12:27Z
body: @aittalam Sounds good. Created an issue to track this: https://github.com/mozilla-ai/lumigator/issues/1207

> The only thing I'd like to avoid is people using unsupported languages from the UI and just seeing bad results with no apparent reason (which btw might happen even if the languages are supported 😅).

That is quite true!
For now, we have language selection from the drop-down (just 5 options): 
![image](https://github.com/user-attachments/assets/c6b8b47c-8e2b-49ae-92b0-9825f49d6230)
I ll ensure that the HF models that are displayed in the UI support them (according to their model cards on HF)

And I think this problem I think goes beyond HF models - since it also happen with less powerful LLMs (for example local models):
- User may ask to translate to a very-low resource language and the model is unable to do so

And also applies to summarization models. User may upload a summarization dataset in German and try to use it along with BART! (less likely but plausible scenario 😅)   



author_association: CONTRIBUTOR