opsml.model.interfaces.vowpal

  1from pathlib import Path
  2from typing import Any, Dict, Optional
  3
  4from pydantic import model_validator
  5
  6from opsml.helpers.utils import get_class_name
  7from opsml.model.interfaces.base import ModelInterface
  8from opsml.types import CommonKwargs, ModelReturn, Suffix, TrainedModelType
  9
 10try:
 11    import vowpalwabbit as vw
 12
 13    class VowpalWabbitModel(ModelInterface):
 14        """Model interface for VowPal Wabbit model.
 15
 16        Args:
 17            model:
 18                vowpal wabbit workspace
 19            sample_data:
 20                Sample data to be used for type inference.
 21                For vowpal wabbit models this should be a string.
 22            arguments:
 23                Vowpal Wabbit arguments. This will be inferred automatically from the workspace
 24
 25        Returns:
 26            VowpalWabbitModel
 27        """
 28
 29        model: Optional[vw.Workspace] = None
 30        sample_data: Optional[str] = None
 31        arguments: str = ""
 32
 33        @property
 34        def model_class(self) -> str:
 35            return TrainedModelType.VOWPAL.value
 36
 37        @classmethod
 38        def _get_sample_data(cls, sample_data: str) -> str:
 39            """Check sample data and returns one record to be used
 40            during type inference and ONNX conversion/validation.
 41
 42            Returns:
 43                Sample data with only one record
 44            """
 45
 46            assert isinstance(sample_data, str), "Sample data must be a string"
 47            return sample_data
 48
 49        @model_validator(mode="before")
 50        @classmethod
 51        def check_model(cls, model_args: Dict[str, Any]) -> Dict[str, Any]:
 52            model = model_args.get("model")
 53
 54            if model_args.get("modelcard_uid", False):
 55                return model_args
 56
 57            assert model is not None, "Model must not be None"
 58
 59            sample_data = cls._get_sample_data(sample_data=model_args[CommonKwargs.SAMPLE_DATA.value])
 60            model_args[CommonKwargs.SAMPLE_DATA.value] = sample_data
 61            model_args[CommonKwargs.DATA_TYPE.value] = get_class_name(sample_data)
 62            model_args[CommonKwargs.VOWPAL_ARGS.value] = model.get_arguments()
 63
 64            return model_args
 65
 66        def save_model(self, path: Path) -> None:
 67            """Saves vowpal model to ".model" file.
 68
 69            Args:
 70                path:
 71                    base path to save model to
 72            """
 73            assert self.model is not None, "No model found"
 74            self.model.save(path.as_posix())
 75
 76        def load_model(self, path: Path, **kwargs: Any) -> None:
 77            """Loads a vowpal model from ".model" file with arguments.
 78
 79            Args:
 80                path:
 81                    base path to load from
 82                **kwargs:
 83                    Additional arguments to be passed to the workspace. This is supplied via
 84                    "arguments" key.
 85
 86                    Example:
 87                        arguments="--cb 4"
 88                        or
 89                        arguments=["--cb", "4"]
 90
 91                    There is no need to specify "-i" argument. This is done automatically during loading.
 92            """
 93            args = kwargs.get("arguments", self.arguments)
 94
 95            if args is None:
 96                args = self.arguments
 97
 98            elif isinstance(args, str):
 99                args = args + f" -i {path.as_posix()}"
100
101            elif isinstance(args, list):
102                args.append(f"-i {path.as_posix()}")
103                args = " ".join(args)
104
105            else:
106                raise ValueError("Arguments must be a string or a list")
107
108            self.model = vw.Workspace(args)
109
110        def save_onnx(self, path: Path) -> ModelReturn:  # pylint: disable=redundant-returns-doc
111            """Onnx is not supported for vowpal wabbit models.
112
113            Args:
114                path:
115                    Path to save
116
117            Returns:
118                ModelReturn
119            """
120            raise ValueError("Onnx is not supported for vowpal wabbit models")
121
122        @property
123        def model_suffix(self) -> str:
124            """Returns suffix for model storage"""
125            return Suffix.MODEL.value
126
127        @staticmethod
128        def name() -> str:
129            return VowpalWabbitModel.__name__
130
131except ModuleNotFoundError:
132    from opsml.model.interfaces.backups import (
133        VowpalWabbitModelNoModule as VowpalWabbitModel,
134    )