Build options validation

machine/jobs/build_clearml_helper.py

@@ -2,15 +2,18 @@
 import logging
 import os
 from datetime import datetime, timezone
-from typing import Callable, Optional, Union, cast
+from typing import Callable, Optional, TypeVar, Union, cast
 
 from clearml import Task
 from dynaconf.base import Settings
+from pydantic import BaseModel, ValidationError
 
 from ..utils.canceled_error import CanceledError
 from ..utils.phased_progress_reporter import PhaseProgressStatus
 from ..utils.progress_status import ProgressStatus
 
+T = TypeVar("T", bound=BaseModel)
+
 
 class ProgressInfo:
     last_percent_completed: Union[int, None] = 0
@@ -119,7 +122,7 @@ def local_progress(progress_status: ProgressStatus) -> None:
     return local_progress
 
 
-def update_settings(settings: Settings, args: dict, task: Optional[Task], logger: logging.Logger):
+def update_settings(settings: Settings, args: dict, task: Optional[Task], logger: logging.Logger, model: type[T]):
     settings.update(args)
     settings.model_type = cast(str, settings.model_type).lower()
     if "build_options" in settings:
@@ -129,6 +132,10 @@ def update_settings(settings: Settings, args: dict, task: Optional[Task], logger
             raise ValueError("Build options could not be parsed: Invalid JSON") from e
         except TypeError as e:
             raise TypeError(f"Build options could not be parsed: {e}") from e
+        try:
+            model.model_validate(build_options)
+        except ValidationError as e:
+            raise ValueError(f"Invalid build options: {e}") from e
         settings.update({settings.model_type: build_options})
         if "align_pretranslations" in build_options:
             settings.update({"align_pretranslations": build_options["align_pretranslations"]})

machine/jobs/build_nmt_engine.py

@@ -8,6 +8,7 @@
 from ..utils.progress_status import ProgressStatus
 from .build_clearml_helper import report_clearml_progress, update_settings
 from .config import SETTINGS
+from .nmt_build_options import NmtBuildOptions
 from .nmt_engine_build_job import NmtEngineBuildJob
 from .nmt_model_factory import NmtModelFactory
 from .shared_file_service_factory import SharedFileServiceType
@@ -45,7 +46,7 @@ def clearml_progress(status: ProgressStatus) -> None:
 
     try:
         logger.info("NMT Engine Build Job started")
-        update_settings(SETTINGS, args, task, logger)
+        update_settings(SETTINGS, args, task, logger, NmtBuildOptions)
 
         translation_file_service = TranslationFileService(SharedFileServiceType.CLEARML, SETTINGS)
         nmt_model_factory: NmtModelFactory

machine/jobs/build_smt_engine.py

@@ -14,6 +14,7 @@
 )
 from .config import SETTINGS
 from .shared_file_service_factory import SharedFileServiceType
+from .smt_build_options import SmtBuildOptions
 from .smt_engine_build_job import SmtEngineBuildJob
 from .smt_model_factory import SmtModelFactory
 from .translation_file_service import TranslationFileService
@@ -46,7 +47,7 @@ def run(args: dict) -> None:
 
     try:
         logger.info("SMT Engine Build Job started")
-        update_settings(SETTINGS, args, task, logger)
+        update_settings(SETTINGS, args, task, logger, SmtBuildOptions)
 
         shared_file_service = TranslationFileService(SharedFileServiceType.CLEARML, SETTINGS)
         smt_model_factory: SmtModelFactory

machine/jobs/build_word_alignment_model.py

@@ -16,6 +16,7 @@
 from .shared_file_service_factory import SharedFileServiceType
 from .thot.thot_word_alignment_model_factory import ThotWordAlignmentModelFactory
 from .word_alignment_build_job import WordAlignmentBuildJob
+from .word_alignment_build_options import WordAlignmentBuildOptions
 from .word_alignment_file_service import WordAlignmentFileService
 from .word_alignment_model_factory import WordAlignmentModelFactory
 
@@ -47,7 +48,7 @@ def run(args: dict):
 
     try:
         logger.info("Word Alignment Build Job started")
-        update_settings(SETTINGS, args, task, logger)
+        update_settings(SETTINGS, args, task, logger, WordAlignmentBuildOptions)
 
         word_alignment_file_service = WordAlignmentFileService(SharedFileServiceType.CLEARML, SETTINGS)
         word_alignment_model_factory: WordAlignmentModelFactory

machine/jobs/nmt_build_options.py

@@ -0,0 +1,71 @@
+from typing import Literal
+
+from pydantic import BaseModel, ConfigDict
+
+
+class TrainParams(BaseModel):
+    model_config = ConfigDict(
+        extra="forbid",
+        strict=True,
+    )
+
+    do_train: bool | None = None
+    optim: str | None = None
+    warmup_steps: int | None = None
+    per_device_train_batch_size: int | None = None
+    gradient_accumulation_steps: int | None = None
+    label_smoothing_factor: float | None = None
+    group_by_length: bool | None = None
+    gradient_checkpointing: bool | None = None
+    lr_scheduler_type: str | None = None
+    learning_rate: float | None = None
+    fp16: bool | None = None
+    tf32: bool | None = None
+    save_strategy: str | None = None
+    max_steps: int | None = None
+
+
+class GenerateParams(BaseModel):
+    model_config = ConfigDict(
+        extra="forbid",
+        strict=True,
+    )
+
+    device: int | None = None
+    num_beams: int | None = 2
+    batch_size: int | None = None
+    oom_batch_size_backoff_mult: float | None = None
+
+
+class TokenizerConfig(BaseModel):
+    model_config = ConfigDict(
+        extra="forbid",
+        strict=True,
+    )
+
+    add_unk_src_tokens: bool | None = None
+    add_unk_trg_tokens: bool | None = None
+
+
+class NmtBuildOptions(BaseModel):
+    model_config = ConfigDict(
+        extra="forbid",
+        strict=True,
+    )
+
+    align_pretranslations: bool | None = None
+    tags: list[str] | str | None = None
+    use_key_terms: bool | None = None
+    parent_model_name: (
+        Literal[
+            "facebook/nllb-200-distilled-1.3B",
+            "facebook/nllb-200-distilled-600M",
+            "hf-internal-testing/tiny-random-nllb",
+        ]
+        | None
+    ) = None
+    train_params: TrainParams | None = None
+    generate_params: GenerateParams | None = None
+    tokenizer: TokenizerConfig | None = None
+    attn_implementation: str | None = None
+    max_steps: int | None = None

machine/jobs/smt_build_options.py

@@ -0,0 +1,22 @@
+from pydantic import BaseModel, ConfigDict
+
+
+class ThotMt(BaseModel):
+    model_config = ConfigDict(
+        extra="forbid",
+        strict=True,
+    )
+
+    word_alignment_model_type: str | None = None
+    tokenizer: str | None = None
+
+
+class SmtBuildOptions(BaseModel):
+    model_config = ConfigDict(
+        extra="forbid",
+        strict=True,
+    )
+
+    tags: list[str] | str | None = None
+    use_key_terms: bool | None = None
+    thot_mt: ThotMt | None = None

machine/jobs/word_alignment_build_options.py

@@ -0,0 +1,23 @@
+from pydantic import BaseModel, ConfigDict
+
+
+class ThotAlign(BaseModel):
+    model_config = ConfigDict(
+        extra="forbid",
+        strict=True,
+    )
+
+    word_alignment_heuristic: str | None = None
+    model_type: str | None = None
+    tokenizer: str | None = None
+
+
+class WordAlignmentBuildOptions(BaseModel):
+    model_config = ConfigDict(
+        extra="forbid",
+        strict=True,
+    )
+
+    tags: list[str] | str | None = None
+    use_key_terms: bool | None = None
+    thot_align: ThotAlign | None = None