Add support for a new language model

In addition to the no-code language models, Foundry enables users to integrate with popular language model frameworks such as Hugging Face ↗ or spaCy ↗. Any Python framework available through Conda or Pip can be utilized in Foundry. However, since many frameworks download pre-trained models, corpora, or other configurations from the Internet, they often require additional steps to become fully functional within Foundry. This is a result of Foundry's security architecture that, by default, denies access to the public Internet for user-written Python code. The additional steps required depend on the particulars of the language model framework.

Hugging Face

To add support for a Hugging Face language model that is not available as a no-code language model, two options are available:

1. Import the model files as raw datasets.
2. Allowlist the Hugging Face domains.

Import the model files as a dataset

You can import any model from the Hugging Face model hub as a dataset. You can use that dataset in Code Repositories or Code Workspaces.

First, download ↗ the model files from Hugging Face. Then, upload the model as a schema-less dataset to bring the files into Foundry. These files can be uploaded either through a frontend upload (New > Dataset > Import > Select all files) or through a Data Connection sync if model files are stored on a shared drive in your private network.

The import dataset should contain all files from the Files and versions tab of the model details in Hugging Face. However, only one binary model file is required (for example, pytorch_model.bin or tf_model.h5). In most cases, we recommend using the PyTorch model as the binary model file.

Once the model files are stored in a dataset, you can use the dataset as an input in your transform. Depending on the size of your model, you may need to specify a Spark profile like DRIVER_MEMORY_MEDIUM to load the model into memory. The code below uses a utility from foundry-huggingface-adapters:

Copied!
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
from palantir_models.transforms import ModelOutput
from transforms.api import transform, Input
from transformers import AutoTokenizer, AutoModel
from huggingface_adapters.utils import copy_model_to_driver


@transform(
    model_output=ModelOutput("/path/to/output/model_asset"),
    hugging_face_raw=Input("/path/to/input/dataset"),
)
def compute(model_output, hugging_face_raw):
    temp_dir = copy_model_to_driver(hugging_face_raw.filesystem())
    tokenizer = AutoTokenizer.from_pretrained(temp_dir)
    model = AutoModel.from_pretrained(temp_dir)

    # Wrap the model with a model adapter and save as a model
    # model_output.publish(...)

Depending on the use case, you can use one of the language model adapters like EmbeddingAdapter:

Copied!
1
2
3
4
5
6
7
# other imports
from huggingface_adapters.embedding_adapter import EmbeddingAdapter
# ...
    model_output.publish(
        model_adapter=EmbeddingAdapter(tokenizer, model),
        change_type=ModelVersionChangeType.MINOR
    )

Allowlist Hugging Face domains

Alternatively, you can allowlist the relevant Hugging Face domains in your Foundry enrollment’s network egress policies. The relevant domains to be allowlisted are https://huggingface.co ↗ and https://cdn-lfs.huggingface.co ↗. Find details about network egress configuration in the Foundry administration documentation.

In addition, the code repository that loads the model from Hugging Face must have the transforms-external-systems library added and be configured accordingly to use the newly created egress policy. Once the configuration is set up, open-source language models can be loaded in a Python transform.

Copied!
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
from palantir_models.transforms import ModelOutput
from transforms.api import transform, Input
from transforms.external.systems import use_external_systems, EgressPolicy, ExportControl
from transformers import AutoTokenizer, AutoModel
import tempfile

@use_external_systems(
    export_control=ExportControl(markings=['<marking ID>']),
    egress=EgressPolicy(<policy RID>),
)
@transform(
    model_output=ModelOutput('/path/to/output/model_asset'),
    text_input=Input('/path/to/input/dataset'),
)
def compute(export_control, egress, model_output, text_input):
    CACHE_DIR = tempfile.mkdtemp()
    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased", cache_dir=CACHE_DIR)
    model = AutoModel.from_pretrained("bert-base-cased", cache_dir=CACHE_DIR)

    # Wrap the model instance with a model adapter and save it in the Palantir platform
    # model_output.publish(...)

Usage in Foundry

Once you have access to the language model, either through a dataset or through the Hugging Face domains, you can integrate with it as a Palantir model by wrapping the language model with a model adapter as defined in the model training in Code Repositories documentation.