[Python] Hugging Face pipeline support

sdks/python/apache_beam/examples/inference/huggingface_question_answering.py

@@ -36,6 +36,7 @@
 from apache_beam.ml.inference.base import PredictionResult
 from apache_beam.ml.inference.base import RunInference
 from apache_beam.ml.inference.huggingface_inference import HuggingFacePipelineModelHandler
+from apache_beam.ml.inference.huggingface_inference import PipelineTask
 from apache_beam.options.pipeline_options import PipelineOptions
 from apache_beam.options.pipeline_options import SetupOptions
 from apache_beam.runners.runner import PipelineResult
@@ -102,7 +103,10 @@ def run(
   if not test_pipeline:
     pipeline = beam.Pipeline(options=pipeline_options)
 
-  model_handler = HuggingFacePipelineModelHandler(model=known_args.model_name, )
+  model_handler = HuggingFacePipelineModelHandler(
+      task=PipelineTask.QuestionAnswering,
+      model=known_args.model_name,
+      load_model_args={'framework': 'pt'})
   if not known_args.input:
     text = (
         pipeline | 'CreateSentences' >> beam.Create([

sdks/python/apache_beam/ml/inference/huggingface_inference.py

@@ -20,6 +20,7 @@
 import logging
 import sys
 from collections import defaultdict
+from enum import Enum
 from typing import Any
 from typing import Callable
 from typing import Dict
@@ -71,6 +72,39 @@
     Iterable[PredictionResult]]
 
 
+class PipelineTask(str, Enum):
+  AudioClassification = 'audio-classification'
+  AutomaticSpeechRecognition = 'automatic-speech-recognition'
+  Conversational = 'conversational'
+  DepthEstimation = 'depth-estimation'
+  DocumentQuestionAnswering = 'document-question-answering'
+  FeatureExtraction = 'feature-extraction'
+  FillMask = 'fill-mask'
+  ImageClassification = 'image-classification'
+  ImageSegmentation = 'image-segmentation'
+  ImageToText = 'image-to-text'
+  MaskGeneration = 'mask-generation'
+  NER = 'ner'
+  ObjectDetection = 'object-detection'
+  QuestionAnswering = 'question-answering'
+  SentimentAnalysis = 'sentiment-analysis'
+  Summarization = 'summarization'
+  TableQuestionAnswering = 'table-question-answering'
+  TextClassification = 'text-classification'
+  TextGeneration = 'text-generation'
+  Text2TextGeneration = 'text2text-generation'
+  TokenClassification = 'token-classification'
+  Translation = 'translation'
+  VideoClassification = 'video-classification'
+  VisualQuestionAnswering = 'visual-question-answering'
+  VQA = 'vqa'
+  ZeroShotAudioClassification = 'zero-shot-audio-classification'
+  ZeroShotClassification = 'zero-shot-classification'
+  ZeroShotImageClassification = 'zero-shot-image-classification'
+  ZeroShotObjectDetection = 'zero-shot-object-detection'
+  Translation_XX_to_YY = 'translation_XX_to_YY'
+
+
 def _validate_constructor_args(model_uri, model_class):
   message = (
       "Please provide both model class and model uri to load the model."
@@ -525,7 +559,7 @@ class HuggingFacePipelineModelHandler(ModelHandler[str,
                                                    Pipeline]):
   def __init__(
       self,
-      task: str = "",
+      task: Union[str, PipelineTask] = "",
       model=None,
       *,
       inference_fn: PipelineInferenceFn = _default_pipeline_inference_fn,
@@ -547,7 +581,8 @@ def __init__(
         task="fill-mask"))
 
     Args:
-      task (str): task supported by HuggingFace Pipelines.
+      task (str or Enum): task supported by HuggingFace Pipelines.
+        Accepts a string task or an Enum from PipelineTask.
       model : path to pretrained model on Hugging Face Models Hub to use custom
         model for the chosen task. If the model already defines the task then
         no need to specify the task parameter.

sdks/python/apache_beam/ml/inference/huggingface_inference_it_test.py

@@ -39,8 +39,8 @@
 
 @pytest.mark.uses_transformers
 @pytest.mark.it_postcommit
+@pytest.mark.timeout(1800)
 class HuggingFaceInference(unittest.TestCase):
-  @pytest.mark.timeout(1800)
   def test_hf_language_modeling(self):
     test_pipeline = TestPipeline(is_integration_test=True)
     # Path to text file containing some sentences