mirror of
https://github.com/chanzuckerberg/cellxgene.git
synced 2026-09-24 06:58:12 +08:00
Andrew Tolopko
30e19e47c6
add `cellxgene annotate` subcommand for invoking MLflow model to generate new `obs` annotations, initially intended for cell type annotations.
21 lines
780 B
Python
21 lines
780 B
Python
import shutil
|
|
from tempfile import TemporaryDirectory, mkstemp
|
|
|
|
import mlflow
|
|
|
|
|
|
def write_model(model) -> str:
|
|
with TemporaryDirectory() as mlflow_model_dir:
|
|
mlflow.pyfunc.save_model(mlflow_model_dir, python_model=model)
|
|
return shutil.make_archive(mkstemp()[1], "zip", mlflow_model_dir)
|
|
|
|
|
|
class FakeModel(mlflow.pyfunc.PythonModel):
|
|
def __init__(self, input_to_output: dict = {}):
|
|
self.input_to_output = input_to_output
|
|
|
|
def predict(self, context, model_input) -> None:
|
|
# this stdout output is useful for validating the input in a test, noting that this model will be invoked in a
|
|
# subprocess, so stdout is one means of communicating information back to the test code
|
|
print(f"__MODEL_INPUT__={model_input.iloc[0][0]}")
|