local models

[FelipeAdachi]

Langkit downloads models from HF Hub, which will hit an error in network-restricted environments. Even though modules such as toxicity,themes, and input_output allow passing the path of local model, the auto-initialization on import will hit an error before we're able to pass the local path in the config.

To deal with this, this change lazy initializes the models so it will fetch them when it's actually needed, when the udf is called. When using local model, one can do as below:

from langkit import LangKitConfig

local_config = LangKitConfig(toxicity_model_path="local-toxicity-model",
              transformer_name="local-sentence-transformers")

toxicity.init(config=local_config)
themes.init(config=local_config)
input_output.init(config=local_config)

provided the supported models were already downloaded and stored in local-toxicity-model and local-sentence-transformers. For reference, these are the HF models that are currently being used in Langkit:

martin-ha/toxic-comment-model

sentence-transformers/all-MiniLM-L6-v2

[jamie256]

Left some preliminary comments.

[jamie256]

Curios if this is noticeably slower or not?

[FelipeAdachi]

I did a very simple test, iterating on 5k samples (single row extraction), and init_model() takes, on average 0.0048 ms (4.8e-6 seconds) - ignoring the first call, which will actually load the model artifacts

update: on recent changes, the lazy initialization is now done for each row, even if it's a batch - it was changed to concile with the detoxify addition in another PR. If the added latency is prohibitive, we need to rethink the design proposed in this PR

[jamie256]

LGTM!

Let's file an issue for local model support on other toxicity models.

Also consider a post_init() or other way to trigger initialization outside of an actual request. We could document that you can call predict on these models as a first time call that should trigger any downloads or initialization.