Avoid some pipeline tasks to use use_cache=True (#24893)

* fix * fix * fix * fix * fix --------- Co-authored-by: ydshieh <[email protected]>
huggingface · Jul 19, 2023 · 129cb6d · 129cb6d
1 parent 476be08
commit 129cb6d
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 0 deletions.
diff --git a/src/transformers/pipelines/question_answering.py b/src/transformers/pipelines/question_answering.py
@@ -1,3 +1,4 @@
+import inspect
 import types
 import warnings
 from collections.abc import Iterable
@@ -510,6 +511,10 @@ def preprocess(self, example, padding="do_not_pad", doc_stride=None, max_questio
     def _forward(self, inputs):
         example = inputs["example"]
         model_inputs = {k: inputs[k] for k in self.tokenizer.model_input_names}
+        # `XXXForSequenceClassification` models should not use `use_cache=True` even if it's supported
+        model_forward = self.model.forward if self.framework == "pt" else self.model.call
+        if "use_cache" in inspect.signature(model_forward).parameters.keys():
+            model_inputs["use_cache"] = False
         output = self.model(**model_inputs)
         if isinstance(output, dict):
             return {"start": output["start_logits"], "end": output["end_logits"], "example": example, **inputs}

diff --git a/src/transformers/pipelines/text_classification.py b/src/transformers/pipelines/text_classification.py
@@ -1,3 +1,4 @@
+import inspect
 import warnings
 from typing import Dict
 
@@ -179,6 +180,10 @@ def preprocess(self, inputs, **tokenizer_kwargs) -> Dict[str, GenericTensor]:
         return self.tokenizer(inputs, return_tensors=return_tensors, **tokenizer_kwargs)
 
     def _forward(self, model_inputs):
+        # `XXXForSequenceClassification` models should not use `use_cache=True` even if it's supported
+        model_forward = self.model.forward if self.framework == "pt" else self.model.call
+        if "use_cache" in inspect.signature(model_forward).parameters.keys():
+            model_inputs["use_cache"] = False
         return self.model(**model_inputs)
 
     def postprocess(self, model_outputs, function_to_apply=None, top_k=1, _legacy=True):

diff --git a/src/transformers/pipelines/zero_shot_classification.py b/src/transformers/pipelines/zero_shot_classification.py
@@ -1,3 +1,4 @@
+import inspect
 from typing import List, Union
 
 import numpy as np
@@ -221,6 +222,10 @@ def _forward(self, inputs):
         candidate_label = inputs["candidate_label"]
         sequence = inputs["sequence"]
         model_inputs = {k: inputs[k] for k in self.tokenizer.model_input_names}
+        # `XXXForSequenceClassification` models should not use `use_cache=True` even if it's supported
+        model_forward = self.model.forward if self.framework == "pt" else self.model.call
+        if "use_cache" in inspect.signature(model_forward).parameters.keys():
+            model_inputs["use_cache"] = False
         outputs = self.model(**model_inputs)
 
         model_outputs = {