Search code examples
pythonartificial-intelligenceamazon-sagemakerhuggingface-transformershuggingface-tokenizers

Huggingface sagemaker


I am trying to use the text2text (translation) model facebook/m2m100_418M to run on sagemaker.

So if you click on deploy and then sagemaker there is some boilerplate code that works well but I can't seem to find how to pass it the arguments src_lang="en", tgt_lang="fr" just like when using the pipeline or transformers. So right now it translates into random languages.

I'm guessing I should add it in here somehow but it's not documented.

predictor.predict({
    'inputs': "The answer to the universe is"
})

Does anybody have an idea of how to pass arguments to the predict method?

Edit

This is the code that was wrong where you will need to change the HF_TASK:

import sagemaker

role = sagemaker.get_execution_role()
# Hub Model configuration. https://huggingface.co/models
hub = {
    'HF_MODEL_ID':'facebook/m2m100_418M',
    'HF_TASK':'text2text-generation'
}

# create Hugging Face Model Class
huggingface_model = HuggingFaceModel(
    transformers_version='4.6.1',
    pytorch_version='1.7.1',
    py_version='py36',
    env=hub,
    role=role, 
)

# deploy model to SageMaker Inference
predictor = huggingface_model.deploy(
    initial_instance_count=1, # number of instances
    instance_type='ml.m5.xlarge' # ec2 instance type
)```

Solution

  • Ok I figured it out. The task was wrong, the source and target language need to be in the task (HF_TASK)

    So for example: 'HF_TASK': 'translation_en_to_fr'