Load a gated or private model from Hugging Face
hf_access_token
model.py
use_auth_token
config.yaml
secrets: hf_access_token: null
from transformers import pipeline class Model: def __init__(self, **kwargs) -> None: self._secrets = kwargs["secrets"] self._model = None def load(self): self._model = pipeline( "fill-mask", model="baseten/docs-example-gated-model", use_auth_token=self._secrets["hf_access_token"] ) def predict(self, model_input): return self._model(model_input)
truss init private-bert && cd private-bert
requirements: - torch==2.0.1 - transformers==4.30.2
truss push
truss predict -d '"It is a [MASK] world"'
Was this page helpful?