def resolve(self, registry_class: str, uri: str, name: str,
options: JavaMap):
"""Resolve a ML model.
Parameters
----------
registry_class : str
The full class name for the registry.
uri : str
The model URI.
name : str
Mode name
options : dict[str, str]
Options passed to the model.
"""
if registry_class not in self.registry_map:
from rikai.internal.reflection import find_class
from rikai.spark.sql.codegen.base import Registry
cls = find_class(registry_class)
if not issubclass(cls, Registry):
raise ValueError(
f"Class '{registry_class}' is not a Registry'")
self.registry_map[registry_class] = cls(self.spark)
registry = self.registry_map[registry_class]
# Convert JavaMap to dict
options = {key: options[key] for key in options.keys()}
return registry.resolve(uri, name, options)
def resolve(self, registry_class: str, spec):
"""Resolve a ML model.
Parameters
----------
registry_class : str
The full class name for the registry.
uri : str
The model URI.
name : str
Mode name
options : dict[str, str]
Options passed to the model.
"""
if registry_class not in self.registry_map:
from rikai.internal.reflection import find_class
from rikai.spark.sql.codegen.base import Registry
cls = find_class(registry_class)
if not issubclass(cls, Registry):
raise ValueError(
f"Class '{registry_class}' is not a Registry'")
self.registry_map[registry_class] = cls(self.spark)
registry = self.registry_map[registry_class]
try:
return registry.resolve(spec)
except Exception as e:
raise RuntimeError(
f"could not resolve {spec} from {registry}") from e
def post_processing(self) -> Optional[Callable]:
"""Return post-processing transform if exists"""
if ("transforms" not in self._spec
or "post" not in self._spec["transforms"]):
return None
f = find_class(self._spec["transforms"]["post"])
return f
def post_processing(self) -> Optional[Callable]:
"""Return post-processing transform if exists"""
if ("transforms" not in self._spec
or "post" not in self._spec["transforms"]):
# Passthrough
return lambda x: x
f = find_class(self._spec["transforms"]["post"])
return f(self.options)
