Serialization for models

To reuse models across workflows, Palantir needs to be able to deserialize stored weights. Since the serialization process can be specific to model types, the author of the model adapter is expected to detail how this should happen. Note that this is not a concern for container models, where serialization is typically a part of the container lifecycle, or external models, which have externally hosted weights.

Auto serialization

To simplify serialization and deserialization for model weights within the platform. Palantir ships a palantir_models_serializers library that provides many default serialization methods for saving and loading models:

• Dill
• Cloudpickle
• JSON
• YAML
• Hugging Face
• PyTorch
• Tensorflow
• XGBoost

How to use a default serializer

A default serializer can be used by annotating the __init__ method on a model adapter with the @auto_serialize annotation. Palantir will automatically save and load your model adapter contents with your chosen serialization method. You can auto_serialize multiple arguments, each with different serializers. Every argument in your __init__ method must have an equivalent argument in your @auto_serialize definition.

Example model definition using auto_serialization

The below is a valid Python transform that trains and publishes a model. The AutoSerializationAdapter uses the DillSerializer to serialize the model.

Copied!
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
from transforms.api import transform
from palantir_models.transforms import ModelOutput

from sklearn.linear_model import LinearRegression
import numpy as np


@transform(
    model_output=ModelOutput("/Foundry/path/to/model_asset"),
)
def compute(model_output):
    x_values = [i for i in range(100)]
    y_values = [2 * i for i in x_values]

    X = np.array(x_values, dtype=np.float32).reshape(-1, 1)
    y = np.array(y_values, dtype=np.float32).reshape(-1, 1)

    model = LinearRegression()
    model.fit(X, y)

    model_output.publish(
        model_adapter=AutoSerializationAdapter(
            model,
            {'prediction_column': 'prediction'}
        )
    )

import palantir_models as pm
from palantir_models_serializers import DillSerializer

class AutoSerializationAdapter(pm.ModelAdapter):

    @pm.auto_serialize(
        model=DillSerializer()
    )
    def __init__(self, model):
        self.model = model

    @classmethod
    def api(cls):
        inputs = {"df_in": pm.Pandas()}
        outputs = {"df_out": pm.Pandas()}
        return inputs, outputs

    def predict(self, df_in):
        df_in['prediction'] = self.model.predict(df_in)
        return df_in

Example model definition using multiple auto_serialization

The above model adapter can be extended to include an optional config dictionary that is serialized with the default JsonSerializer.

Copied!
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
import palantir_models as pm
from palantir_models_serializers import DillSerializer, JsonSerializer

class AutoSerializationAdapter(pm.ModelAdapter):

    @pm.auto_serialize(
        model=DillSerializer(),
        config=JsonSerializer()
    )
    def __init__(self, model, config={}):
        self.model = model
        self.prediction_column = config['prediction_column'] if 'prediction_column' in config else 'prediction'

    @classmethod
    def api(cls):
        inputs = {"df_in": pm.Pandas()}
        outputs = {"df_out": pm.Pandas()}
        return inputs, outputs

    def predict(self, df_in):
        df_in[self.prediction_column] = self.model.predict(df_in)
        return df_in

Default serializers

Dill

The palantir_models_serializers.DillSerializer will serialize a Python object with dill ↗ by calling dill.dump with dill.load to save and load your object to disk.

The DillSerializer class can be used to serialize many Python objects, including scikit-learn and statsmodels.

Cloudpickle

The palantir_models_serializers.CloudPickleSerializer will serialize a Python object with Cloudpickle ↗ by calling cloudpickle.dump and cloudpickle.load to save and load your object to disk.

The CloudPickleSerializer class can be used to serialize many Python objects, including scikit-learn and statsmodels.

JSON

The palantir_models_serializers.JsonSerializer will serialize a Python Dictionary as JSON by calling yaml.safe_dump and json.safe_load on your Python Dictionary.

YAML

The palantir_models_serializers.YamlSerializer will serialize a Python Dictionary with JSON by calling yaml.safe_dump and yaml.safe_load on your Python Dictionary.

Hugging Face

The palantir_models_serializers library currently provides three default serializers for Hugging Face models ↗: HfPipelineSerializer, HfAutoTokenizerSerializer, and HfAutoModelSerializer. All three Hugging Face serializers require that the transformers library be added as a dependency to your Python environment.

HfPipelineSerializer

The palantir_models_serializers.HfPipelineSerializer will serialize a transformers.pipeline ↗ object with save_pretrained and reinstantiate your pipeline object at load.

The HfPipelineSerializer has one required string parameter representing the task ↗ of the pipeline. Any additional kwargs will be used for loading the pipeline.

Copied!
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
import palantir_models as pm
from palantir_models_serializers import HfPipelineSerializer
from transformers import pipeline
import pandas as pd

class HFNerAdapter(pm.ModelAdapter):

    @pm.auto_serialize(
        pipeline=HfPipelineSerializer("ner"),
    )
    def __init__(self, pipeline):
        self.pipeline = pipeline

    @classmethod
    def api(cls):
        inputs = {"df_in" : pm.Pandas([("text", str)])}
        outputs = {"df_out" : pm.Pandas([("text", str), ("generation", str)])}
        return inputs, outputs

    def predict(self, df_in: pd.DataFrame):
        result = self.pipeline(df_in["text"].tolist())
        df_in["generation"] = result
        return df_in

HfAutoModelSerializer and HfAutoTokenizerSerializer

The palantir_models_serializers.HfAutoModelSerializer and palantir_models_serializers.HfAutoTokenizerSerializer will serialize a transformers.AutoModel and transformers.AutoTokenizer model with save_pretrained and from_pretrained.

Note, for some models and tokenizers it is recommended to use the specific model or tokenizer classes rather than the generic ones; in these cases the specific classes can be passed to the serializers as the first model_class argument.

Copied!
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
import palantir_models as pm
from palantir_models_serializers import HfAutoModelSerializer, HfAutoTokenizerSerializer
from transformers import AutoModelForSeq2SeqLM
import pandas as pd
import torch


class HFTextGenerationAdapter(pm.ModelAdapter):

    @pm.auto_serialize(
        model=HfAutoModelSerializer(AutoModelForSeq2SeqLM),
        tokenizer=HfAutoTokenizerSerializer()
    )
    def __init__(self, model, tokenizer):
        self.model = model
        self.tokenizer = tokenizer

    @classmethod
    def api(cls):
        inputs = {"df_in" : pm.Pandas([("text", str)])}
        outputs = {"df_out" : pm.Pandas([("text", str), ("generation", str)])}
        return inputs, outputs

    def predict(self, df_in: pd.DataFrame):
        input_ids = self.tokenizer(
            df_in["text"].tolist(),
            return_tensors="pt",
            padding=True
        ).input_ids

        outputs = self.model.generate(input_ids)
        result = self.tokenizer.batch_decode(outputs, skip_special_tokens=True)

        df_in["generation"] = result
        return df_in

PyTorch

The palantir_models_serializers.PytorchStateSerializer will serialize a torch.nn.Module ↗ with torch.save and torch.load to save and load your object to disk.

TensorFlow

The palantir_models_serializers.TensorflowKerasSerializer will serialize a tensorflow.keras.Model ↗ with obj.save and tensorflow.keras.models.load_model to save and load your model to disk. When your model is being deserialized, Foundry will call obj.compile().

XGBoost

The palantir_models_serializers.XGBoostSerializer will serialize a xgboost.sklearn.XGBModel ↗ with save_model and load_model to save and load your model to disk. The XGBModel class is the base class for many XGBoost models including xgboost.XGBClassifier, xgboost.XGBRegressor, and xgboost.XGBRanker.

Writing your own AutoSerializer

See the full implementation of each of the serializers and documentation on how to add a new default serializer.

If you believe there should be an additional default serializer that is not listed above, contact your Palantir representative.

Custom Serialization

In cases where the default serializers are insufficient, users can implement the load() and save() methods. Refer to the API reference for additional details and examples.