Skip to content

Generalize model type #13

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Sep 27, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@

setup(
name="unboxapi",
version="0.0.1",
version="0.0.3",
description="The official Python client library for Unbox AI, the Testing and Debugging Platform for AI",
url="https://github.com/unboxai/unboxapi-python-client",
author="Unbox AI",
Expand Down
64 changes: 40 additions & 24 deletions unboxapi/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@
from .api import Api
from .datasets import Dataset
from .exceptions import UnboxException
from .models import Model, ModelType, create_template_model, create_rasa_model
from .models import Model, ModelType, create_template_model


class DeploymentType(Enum):
Expand Down Expand Up @@ -47,6 +47,8 @@ def add_model(
description: str = None,
requirements_txt_file: Optional[str] = None,
setup_script: Optional[str] = None,
custom_model_code: Optional[str] = None,
dependent_dir: Optional[List[str]] = None,
**kwargs,
) -> Model:
"""Uploads a model.
Expand All @@ -69,50 +71,64 @@ def add_model(
Path to a requirements file containing dependencies needed by the predict function
setup_script (Optional[str]):
Path to a setup script executing any commands necessary to run before loading the model
custom_model_code (Optional[str]):
Custom code needed to initialize the model. Model object must be none in this case.
dependent_dir (Optional[str]):
Path to a dir of file dependencies needed to load the model

Returns:
Model:
Returns uploaded model
"""
if custom_model_code:
assert (
model_type is ModelType.custom
), "model_type must be ModelType.custom if specifying custom_model_code"
with TempDirectory() as dir:

if model_type == ModelType.rasa:
bento_service = create_rasa_model(
dir, requirements_txt_file, setup_script
)
else:
bento_service = create_template_model(
model_type, dir, requirements_txt_file, setup_script
)
if model_type == ModelType.transformers:
if "tokenizer" not in kwargs:
raise UnboxException(
"Must specify tokenizer in kwargs when using a transformers model"
)
bento_service.pack(
"model", {"model": model, "tokenizer": kwargs["tokenizer"]}
bento_service = create_template_model(
model_type,
dir,
requirements_txt_file,
setup_script,
custom_model_code,
)
if model_type is ModelType.transformers:
if "tokenizer" not in kwargs:
raise UnboxException(
"Must specify tokenizer in kwargs when using a transformers model"
)
kwargs.pop("tokenizer")
else:
bento_service.pack("model", model)
bento_service.pack(
"model", {"model": model, "tokenizer": kwargs["tokenizer"]}
)
kwargs.pop("tokenizer")
elif model_type not in [ModelType.custom, ModelType.rasa]:
bento_service.pack("model", model)

bento_service.pack("function", function)
bento_service.pack("kwargs", kwargs)

with TempDirectory() as temp_dir:
_write_bento_content_to_dir(bento_service, temp_dir)

if model_type == ModelType.rasa:
destination_path = os.path.join(temp_dir, "TemplateModel/nlu")
shutil.copytree(model.model_metadata.model_dir, destination_path)
if model_type is ModelType.rasa:
dependent_dir = model.model_metadata.model_dir

if dependent_dir is not None:
dependent_dir = os.path.abspath(dependent_dir)
shutil.copytree(
dependent_dir,
os.path.join(
temp_dir,
f"TemplateModel/{os.path.basename(dependent_dir)}",
),
)

with TempDirectory() as tarfile_dir:
tarfile_path = f"{tarfile_dir}/model"

with tarfile.open(tarfile_path, mode="w:gz") as tar:
tar.add(temp_dir, arcname=bento_service.name)

print("Connecting to Unbox server")
endpoint = "models"
payload = dict(
name=name,
Expand Down
130 changes: 45 additions & 85 deletions unboxapi/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,12 +8,13 @@
class ModelType(Enum):
"""Task Type List"""

custom = "Custom"
fasttext = "FasttextModelArtifact"
sklearn = "SklearnModelArtifact"
pytorch = "PytorchModelArtifact"
tensorflow = "TensorflowSavedModelArtifact"
transformers = "TransformersModelArtifact"
rasa = ""
rasa = "Rasa"


class Model:
Expand Down Expand Up @@ -46,16 +47,36 @@ def to_dict(self):


def _predict_function(model_type: ModelType):
if model_type == ModelType.transformers:
if model_type is ModelType.transformers:
return "results = self.artifacts.function(self.artifacts.model.get('model'), text, tokenizer=self.artifacts.model.get('tokenizer'), **self.artifacts.kwargs)"
elif model_type in [ModelType.rasa, ModelType.custom]:
return "results = self.artifacts.function(model, text, **self.artifacts.kwargs)"
else:
return "results = self.artifacts.function(self.artifacts.model, text, **self.artifacts.kwargs)"


def _model(model_type: ModelType):
if model_type is ModelType.custom or model_type is ModelType.rasa:
return ""
else:
return f"from bentoml.frameworks.{model_type.name} import {model_type.value}"


def _artifacts(model_type: ModelType):
if model_type is ModelType.custom or model_type is ModelType.rasa:
return "@artifacts([PickleArtifact('function'), PickleArtifact('kwargs')])"
else:
return f"@artifacts([{model_type.value}('model'), PickleArtifact('function'), PickleArtifact('kwargs')])"


def _format_custom_code(custom_model_code: Optional[str]):
if custom_model_code is None:
return ""
return textwrap.indent(textwrap.dedent("\n" + custom_model_code), prefix=" ")


def _env_dependencies(
tmp_dir: str,
requirements_txt_file: Optional[str] = None,
setup_script: Optional[str] = None,
tmp_dir: str, requirements_txt_file: Optional[str], setup_script: Optional[str]
):
unbox_req_file = f"{tmp_dir}/requirements.txt"
env_wrapper_str = ""
Expand All @@ -81,26 +102,38 @@ def _env_dependencies(
def create_template_model(
model_type: ModelType,
tmp_dir: str,
requirements_txt_file: Optional[str] = None,
setup_script: Optional[str] = None,
requirements_txt_file: Optional[str],
setup_script: Optional[str],
custom_model_code: Optional[str],
):
if model_type is ModelType.rasa:
custom_model_code = """
from rasa.nlu.model import Interpreter
model = Interpreter.load("nlu")
"""
if custom_model_code:
assert (
"model = " in custom_model_code
), "custom_model_code must intialize a `model` var"
with open(f"template_model.py", "w") as python_file:
file_contents = f"""\
import json
import os
import pandas as pd
from typing import List

from bentoml import env, artifacts, api, BentoService
from bentoml.frameworks.{model_type.name} import {model_type.value}
{_model(model_type)}
from bentoml.service.artifacts.common import PickleArtifact
from bentoml.adapters import JsonInput
from bentoml.types import JsonSerializable
from bentoml.utils.tempdir import TempDirectory


cwd = os.getcwd()
os.chdir(os.path.dirname(os.path.abspath(__file__)))
{_format_custom_code(custom_model_code)}
os.chdir(cwd)

{_env_dependencies(tmp_dir, requirements_txt_file, setup_script)}
@artifacts([{model_type.value}('model'), PickleArtifact('function'), PickleArtifact('kwargs')])
{_artifacts(model_type)}
class TemplateModel(BentoService):
@api(input=JsonInput())
def predict(self, parsed_json: JsonSerializable):
Expand Down Expand Up @@ -148,76 +181,3 @@ def tokenize_from_path(self, parsed_json: JsonSerializable):
from template_model import TemplateModel

return TemplateModel()


def create_rasa_model(
tmp_dir: str,
requirements_txt_file: Optional[str] = None,
setup_script: Optional[str] = None,
):
with open(f"template_model.py", "w") as python_file:
file_contents = f"""\
import json
import os
import pandas as pd
from typing import List

from bentoml import env, artifacts, api, BentoService
from bentoml.service.artifacts.common import PickleArtifact
from bentoml.adapters import JsonInput
from bentoml.types import JsonSerializable
from bentoml.utils.tempdir import TempDirectory
from rasa.nlu.model import Interpreter

model = Interpreter.load(f"{{os.path.dirname(os.path.abspath(__file__))}}/nlu")


{_env_dependencies(tmp_dir, requirements_txt_file, setup_script)}
@artifacts([PickleArtifact('function'), PickleArtifact('kwargs')])
class TemplateModel(BentoService):
@api(input=JsonInput())
def predict(self, parsed_json: JsonSerializable):
text = parsed_json['text']
results = self.artifacts.function(model, text, **self.artifacts.kwargs)
return results

@api(input=JsonInput())
def predict_from_path(self, parsed_json: JsonSerializable):
input_path = parsed_json["input_path"]
output_path = parsed_json["output_path"]
text = pd.read_csv(input_path)['text'].tolist()
results = self.artifacts.function(model, text, **self.artifacts.kwargs)
with open(output_path, 'w') as f:
if type(results) == list:
json.dump(results, f)
else:
json.dump(results.tolist(), f)
return "Success"

@api(input=JsonInput())
def tokenize(self, parsed_json: JsonSerializable):
text = parsed_json['text']
results = None
if "tokenizer" in self.artifacts.kwargs:
results = self.artifacts.kwargs["tokenizer"](text)
return results

@api(input=JsonInput())
def tokenize_from_path(self, parsed_json: JsonSerializable):
input_path = parsed_json["input_path"]
output_path = parsed_json["output_path"]
text = pd.read_csv(input_path)['text'].tolist()
if "tokenizer" in self.artifacts.kwargs:
results = self.artifacts.kwargs["tokenizer"](text)
with open(output_path, 'w') as f:
if type(results) == list:
json.dump(results, f)
else:
json.dump(results.tolist(), f)
return "Success"
"""
python_file.write(textwrap.dedent(file_contents))

from template_model import TemplateModel

return TemplateModel()
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy