Hugging Face SageMaker Inference

deploy_hub.py

from sagemaker.huggingface import HuggingFaceModel

# Hub Model configuration. https://huggingface.co/models
hub = {
  'MODEL_ID':'distilbert-base-uncased-distilled-squad',
  'TASK':'question-answering'
}

# create Hugging Face Model Class
huggingface_model = HuggingFaceModel(
    transformers_version='4.4',
    pytorch_version='1.6',
    env=hub,
    role=role,
    name=hub['MODEL_ID'], 
)

# deploy model to SageMaker Inference
huggingface_model.deploy(initial_instance_count=1,instance_type="ml.m5.xlarge")

train_and_deploy.py

from sagemaker.huggingface import HuggingFace

# hyperparameters, which are passed into the training job
hyperparameters={'epochs': 1,
                 'train_batch_size': 32,
                 'model_name':'distilbert-base-uncased'
                 }

# creates Hugging Face estimator
huggingface_estimator = HuggingFace(entry_point='train.py',
                            instance_type='ml.p3.2xlarge',
                            instance_count=1,
                            role=role,
                            transformers_version='4.4',
                            pytorch_version='1.6',
                            py_version='py36',
                            hyperparameters = hyperparameters)

# starting the train job with our uploaded datasets as input
huggingface_estimator.fit({'train': training_input_path, 'test': test_input_path})

# 🏋️ Training of the Model.........

# create model instance
huggingface_model = huggingface_estimator.create_model(env={'TASK': 'text-classification'})

# deploy model to SageMaker Inference
huggingface_model.deploy(initial_instance_count=1,instance_type="ml.m5.xlarge")