aws/sagemaker-huggingface-inference-toolkit

Support for image-classification tasks

akuma12 opened this issue · 2 comments

I'm attempting to deploy a Huggingface vision transformer model to Sagemaker, and it works great if I get the model directly from the Huggingface Hub. However, if I try to provide my own model with a custom code/inference.py file, I get an error saying Task couldn't be inferenced from BeitForImageClassification. Inference Toolkit can only inference tasks from architectures ending with <display list of architectures here>.

Is it just a matter of adding "ForImageClassification": "image-classification" to the ARCHITECTURES_2_TASK dict in transformers_utils.py? If so I'll happily add a PR.

Hey @akuma12 we are not yet supporting Image tasks yet since they are a few more things to add, like serializer and deserializer to properly work with image types. But we are working on supporting this soon!
in the meantime, you could create a custom inference.py and add the logic there.

Sounds good. I started work on my own inference.py file, so that should work for now. I'll close this issue since you're already working on it.