Fix TransformersPredictor.predict when HF_TASK=text-classification

alvarobartt · Mar 6, 2024 · 983c557 · 983c557
1 parent 69d73e2
commit 983c557
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/src/vertex_ai_huggingface_inference_toolkit/predictors/transformers.py b/src/vertex_ai_huggingface_inference_toolkit/predictors/transformers.py
@@ -91,4 +91,8 @@ def load(self, artifacts_uri: Optional[str] = None) -> None:
         )
 
     def predict(self, instances: Dict[str, Any]) -> Dict[str, Any]:
+        # NOTE: temporary patch for `text-classification` until the following PR is merged (if so):
+        # https://github.com/huggingface/transformers/pull/29495
+        if "args" in instances:
+            return self._pipeline(instances.pop("args"), **instances)  # type: ignore
         return self._pipeline(**instances)  # type: ignore