mirror of https://github.com/razor-ai/soup.git
52 lines
1.6 KiB
Python
52 lines
1.6 KiB
Python
"""MLX backend trainer registry — routes task → MLX trainer class.
|
|
|
|
All trainer imports are deferred to function bodies so the MLX classes don't
|
|
load on non-Apple Silicon machines during CLI startup (project convention:
|
|
lazy heavy imports).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
|
|
class _MLXRegistry(dict):
|
|
"""Dict-like registry that lazy-imports MLX trainer classes."""
|
|
|
|
def _resolve(self, task: str):
|
|
if task == "sft":
|
|
from soup_cli.trainer.mlx_sft import MLXSFTTrainerWrapper
|
|
|
|
return MLXSFTTrainerWrapper
|
|
if task == "dpo":
|
|
from soup_cli.trainer.mlx_dpo import MLXDPOTrainerWrapper
|
|
|
|
return MLXDPOTrainerWrapper
|
|
if task == "grpo":
|
|
from soup_cli.trainer.mlx_grpo import MLXGRPOTrainerWrapper
|
|
|
|
return MLXGRPOTrainerWrapper
|
|
raise KeyError(task)
|
|
|
|
def __contains__(self, key: object) -> bool:
|
|
return key in ("sft", "dpo", "grpo")
|
|
|
|
def __getitem__(self, key: str):
|
|
return self._resolve(key)
|
|
|
|
def keys(self): # noqa: D401
|
|
return ("sft", "dpo", "grpo")
|
|
|
|
|
|
MLX_TRAINER_REGISTRY: _MLXRegistry = _MLXRegistry()
|
|
|
|
|
|
def get_mlx_trainer(task: str):
|
|
"""Return the MLX trainer class for a task, or raise ValueError."""
|
|
try:
|
|
return MLX_TRAINER_REGISTRY[task]
|
|
except KeyError as exc:
|
|
supported = ", ".join(MLX_TRAINER_REGISTRY.keys())
|
|
raise ValueError(
|
|
f"MLX backend does not support task '{task}'. "
|
|
f"Supported: {supported}. Use backend=transformers for full task coverage."
|
|
) from exc
|