Armazenamento de modelos personalizados no registro do modelo Snowflake¶

Definição do contexto do modelo¶

Os modelos geralmente exigem um ou mais arquivos estáticos, como arquivos de configuração ou pesos de modelo, além do código. Os modelos personalizados devem fornecer informações sobre todos esses artefatos para que o registro saiba como compactá-los junto com o modelo. Os artefatos podem ser declarados usando a classe ModelContext da seguinte forma.

from snowflake.ml.model import custom_model

mc = custom_model.ModelContext(
    artifacts={
        'config': 'local_model_dir/config.json'
    },
)

Os caminhos para arquivos de artefatos no contexto do modelo são relativos ao diretório atual no ambiente do qual o modelo é registrado. O Registro do modelo Snowflake usa essas informações para garantir que todo o código e dados necessários sejam implantados no warehouse onde seu modelo será executado. No tempo de execução, o modelo pode encontrar esses artefatos chamando self.context.path('config') (o valor 'config' é o mesmo que a chave no dicionário passada para ModelContext).

Além de arquivos estáticos, um modelo pode compor outros modelos ou pipelines de um tipo suportado (por exemplo, um snowflake.ml.modeling.pipeline ou um modelo scikit-learn). O registro já sabe como registrar em log esses tipos de objetos, então você pode passar os objetos Python diretamente no contexto do modelo usando model_refs, como mostrado aqui. Você não precisa colocar esses objetos em um pacote por conta própria. Isso pode ser útil para agregação bootstrap (bagging) ou para pré-processamento ou pós-processamento.

mc = custom_model.ModelContext(
    artifacts={
        'config': 'local_model_dir/config.json'
    },
    models={
        'm1': model1,
        'm2': model2,
        'feature_preproc': preproc
    }
)

O registro do modelo serializa essas referências de modelo ao registrar o modelo e as reidrata em tempo de execução. Seu modelo pode recuperar referências a esses modelos subordinados usando, por exemplo, self.context.model_ref('m1'). Se o modelo expuser um método predict, seu código poderá chamá-lo diretamente da referência do modelo recuperado, por exemplo com self.context.model_ref('m1').predict().

Em resumo, o contexto de um modelo personalizado declara que os objetos Python serão serializados junto com os artefatos do modelo, que são arquivos locais usados pelo modelo e que devem ser armazenados no Snowflake junto com o código. Sua classe de modelo usa o contexto para localizar o código e os artefatos; isso funciona quer seu modelo esteja sendo executado localmente ou no Snowflake.

Como escrever a classe de modelo personalizado¶

Para informar ao registro do modelo como registrar e implantar seu modelo personalizado, herde de snowflake.ml.model.custom_model.CustomModel.

Os modelos podem expor vários métodos de inferência (por exemplo, os modelos scikit-learn expõem métodos predict e predict_proba). Para declarar funções de inferência em seu modelo personalizado, defina-as como métodos públicos de sua subclasse e decore-as com @custom_model.inference_api. Este decorador indica que um método faz parte da API pública do modelo, permitindo que seja chamado de Python ou SQL através do registro do modelo. Métodos decorados com inference_api deve aceitar e devolver um DataFrame pandas. Qualquer número de métodos pode ser decorado com inference_api.

Uma classe de modelo personalizado de esqueleto com uma API pública é mostrada abaixo.

Observe o uso de context.path para abrir o arquivo de polarização e self.context.model_ref para obter referências às classes de modelos subordinadas para que seus métodos predict possam ser chamados.

from snowflake.ml.model import custom_model
import pandas as pd

class ExamplePipelineModel(custom_model.CustomModel):

    @custom_model.inference_api
    def predict(self, input: pd.DataFrame) -> pd.DataFrame:
        ...
        return pd.DataFrame(...)

Juntando todas as peças, o que se segue é um modelo personalizado totalmente funcional.

class ExamplePipelineModel(custom_model.CustomModel):
    def __init__(self, context: custom_model.ModelContext) -> None:
        super().__init__(context)
        v = open(context.path('config')).read()
        self.bias = json.loads(v)['bias']

    @custom_model.inference_api
    def predict(self, input: pd.DataFrame) -> pd.DataFrame:
        features = self.context.model_ref('feature_preproc').transform(input)
        model_output = self.context.model_ref('m2').predict(
            self.context.model_ref('m1').predict(features)
        )
        return pd.DataFrame({
            'output': model_output + self.bias})

Uso do modelo personalizado¶

É possível testar seu novo modelo personalizado (pipeline) executando-o localmente da seguinte forma:

my_model_pipeline = ExamplePipelineModel(mc)
output_df = my_model_pipeline.predict(input_df)

Ou registre-o no registro e implante-o no Snowflake. Conforme mostrado no próximo exemplo de código, forneça conda_dependencies (ou pip_requirements) para especificar as bibliotecas que a classe de modelo precisa.

Forneça sample_input_data (um DataFrame pandas) para inferir a assinatura de entrada para o modelo. Como alternativa, forneça uma assinatura do modelo.

reg = Registry(session=sp_session, database_name="ML", schema_name="REGISTRY")
mv = reg.log_model(my_model_pipeline,
            model_name="my_custom_model_pipeline",
            version_name="v1",
            conda_dependencies=["scikit-learn"],
            comment="My Custom ML Model Pipeline",
            sample_input_data=train_features)
output_df = mv.run(input_df)

Inferência de função de tabela¶

A partir do Snowpark ML 1.5.4, é possível registrar modelos com métodos de inferência que retornam várias colunas. Para fazer isso, registre seu modelo com a opção {"function_type": "TABLE_FUNCTION"} e use o decorador @inference_api como acima. No exemplo a seguir, o método decorado retorna um DataFrame pandas que inclui duas colunas de saída.

class ExampleTableFunctionModel(custom_model.CustomModel):

  @custom_model.inference_api
  def predict(self, input: pd.DataFrame) -> pd.DataFrame:
      output_df = pandas.DataFrame([{"OUTPUT1": input["INPUT1"] + 1, input["INPUT2"] + 1}])
      return output_df

my_model = ExampleTableFunctionModel()
reg = Registry(session=sp_session, database_name="ML", schema_name="REGISTRY")
mv = reg.log_model(my_model,
            model_name="my_custom_table_function_model",
            version_name="v1",
            options={"function_type": "TABLE_FUNCTION"},
            sample_input_data=train_features
            )
output_df = mv.run(input_df)

Se o modelo incluir vários métodos de inferência, use a opção method_options para registrar o modelo, indicando quais são FUNCTION e quais são TABLE_FUNCTION:

reg = Registry(session=sp_session, database_name="ML", schema_name="REGISTRY")
mv = reg.log_model(my_model,
            model_name="my_custom_table_function_model",
            version_name="v1",
            options={
              "method_options": {                                 ###
                "METHOD1": {"function_type": "TABLE_FUNCTION"},   ###
                "METHOD2": {"function_type": "FUNCTION"}          ###
              }
            },
            sample_input_data=train_features
            )

O método de inferência da função de tabela do modelo registrado também pode ser invocado via SQL da seguinte forma:

SELECT OUTPUT1, OUTPUT2
  FROM input_table,
      table(
          MY_MODEL!PREDICT(input_table.INPUT1, input_table.INPUT2)
      )

Para saber mais sobre métodos de inferência de modelo particionados, em que o método de inferência usa uma partição de dados como entrada e gera várias linhas e colunas, consulte Modelos personalizados particionados.