models/__init__.py

"""Model loading and inference for A1-Max MuQ LoRA."""

from models.loader import (
    A1MaxInferenceHead,
    ModelCache,
    get_model_cache,
)
from models.inference import (
    extract_muq_embeddings,
    predict_with_ensemble,
)
from models.calibration import (
    calibrate_predictions,
    predictions_to_calibrated_dict,
    get_calibration_context,
)

__all__ = [
    "A1MaxInferenceHead",
    "ModelCache",
    "get_model_cache",
    "extract_muq_embeddings",
    "predict_with_ensemble",
    "calibrate_predictions",
    "predictions_to_calibrated_dict",
    "get_calibration_context",
]