evaluator / models /model_selector.py
kleervoyans's picture
Create model_selector.py
08c5787 verified
raw
history blame contribute delete
976 Bytes
import logging
from models.model_loader import ModelLoader
logger = logging.getLogger(__name__)
class ModelSelector:
"""
Given a list of candidate model names, attempts to load each in turn
and returns the first that succeeds.
"""
def __init__(self, candidates=None, quantize=True):
# Default candidates: NLLB distilled, then M2M100
self.candidates = candidates or [
"facebook/nllb-200-distilled-600M",
"facebook/m2m100_418M"
]
self.quantize = quantize
def select(self) -> str:
loader = ModelLoader(quantize=self.quantize)
for name in self.candidates:
try:
_, _ = loader.load(name)
logger.info(f"✔️ Selected translation model: {name}")
return name
except Exception as e:
logger.warning(f"Failed to load {name}: {e}")
raise RuntimeError("No translation model could be loaded.")