portugalfuturista/replica-omnisciente
2
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 70
replica-omnisciente/dirac/evals/opencode/opencode_refactor_latency

782 lines
33 KiB
Text
Raw Permalink Blame History

diff --git a/src/transformers/pipelines/any_to_any.py b/src/transformers/pipelines/any_to_any.py
index 4ae91d5a17..c5d26777cd 100644
--- a/src/transformers/pipelines/any_to_any.py
+++ b/src/transformers/pipelines/any_to_any.py
@@ -14,6 +14,7 @@
import enum
import re
+import time
from typing import Any, Union, overload
import numpy as np
@@ -422,7 +423,13 @@ class AnyToAnyPipeline(Pipeline):
if "generation_config" not in generate_kwargs:
generate_kwargs["generation_config"] = self.generation_config
- generated_sequence = self.model.generate(**model_inputs, **generate_kwargs)
+ if self._record_latency:
+ start_time = time.perf_counter()
+ generated_sequence = self.model.generate(**model_inputs, **generate_kwargs)
+ latency = (time.perf_counter() - start_time) * 1000
+ logger.info(f"[{self.__class__.__name__}] Inference latency: {latency:.2f} ms")
+ else:
+ generated_sequence = self.model.generate(**model_inputs, **generate_kwargs)
return {"generated_sequence": generated_sequence, "prompt_text": prompt_text, "input_ids": input_ids}
def postprocess(
diff --git a/src/transformers/pipelines/audio_classification.py b/src/transformers/pipelines/audio_classification.py
index 6e173111aa..36c0d49084 100644
--- a/src/transformers/pipelines/audio_classification.py
+++ b/src/transformers/pipelines/audio_classification.py
@@ -12,6 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import subprocess
+import time
from typing import Any
import httpx
@@ -240,6 +241,12 @@ class AudioClassificationPipeline(Pipeline):
return processed
def _forward(self, model_inputs):
+ if self._record_latency:
+ start_time = time.perf_counter()
+ model_outputs = self.model(**model_inputs)
+ latency = (time.perf_counter() - start_time) * 1000
+ logger.info(f"[{self.__class__.__name__}] Inference latency: {latency:.2f} ms")
+ return model_outputs
model_outputs = self.model(**model_inputs)
return model_outputs
diff --git a/src/transformers/pipelines/automatic_speech_recognition.py b/src/transformers/pipelines/automatic_speech_recognition.py
index 58349d0b10..af4c6f259f 100644
--- a/src/transformers/pipelines/automatic_speech_recognition.py
+++ b/src/transformers/pipelines/automatic_speech_recognition.py
@@ -11,6 +11,7 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
+import time
from collections import defaultdict
from typing import TYPE_CHECKING, Any, Union
@@ -516,7 +517,13 @@ class AutomaticSpeechRecognitionPipeline(ChunkPipeline):
"attention_mask": attention_mask,
**generate_kwargs,
}
- tokens = self.model.generate(**generate_kwargs)
+ if self._record_latency:
+ start_time = time.perf_counter()
+ tokens = self.model.generate(**generate_kwargs)
+ latency = (time.perf_counter() - start_time) * 1000
+ logger.info(f"[{self.__class__.__name__}] Inference latency: {latency:.2f} ms")
+ else:
+ tokens = self.model.generate(**generate_kwargs)
# whisper longform generation stores timestamps in "segments"
if return_timestamps == "word" and self.type == "seq2seq_whisper":
@@ -539,7 +546,13 @@ class AutomaticSpeechRecognitionPipeline(ChunkPipeline):
self.model.main_input_name: model_inputs.pop(self.model.main_input_name),
"attention_mask": attention_mask,
}
- outputs = self.model(**inputs)
+ if self._record_latency:
+ start_time = time.perf_counter()
+ outputs = self.model(**inputs)
+ latency = (time.perf_counter() - start_time) * 1000
+ logger.info(f"[{self.__class__.__name__}] Inference latency: {latency:.2f} ms")
+ else:
+ outputs = self.model(**inputs)
logits = outputs.logits
if self.type == "ctc_with_lm":
diff --git a/src/transformers/pipelines/base.py b/src/transformers/pipelines/base.py
index 916576c4c6..8c5d026ea0 100644
--- a/src/transformers/pipelines/base.py
+++ b/src/transformers/pipelines/base.py
@@ -696,6 +696,9 @@ def build_pipeline_init_args(
binary_output (`bool`, *optional*, defaults to `False`):
Flag indicating if the output the pipeline should happen in a serialized format (i.e., pickle) or as
the raw output data e.g. text."""
+ docstring += r"""
+ record_latency (`bool`, *optional*, defaults to `False`):
+ Whether or not to record the latency of the model inference."""
return docstring
@@ -776,6 +779,7 @@ class Pipeline(_ScikitCompat, PushToHubMixin):
task: str = "",
device: int | torch.device | None = None,
binary_output: bool = False,
+ record_latency: bool = False,
**kwargs,
):
# We need to pop them for _sanitize_parameters call later
@@ -844,6 +848,7 @@ class Pipeline(_ScikitCompat, PushToHubMixin):
logger.debug(f"Device set to use {self.device}")
self.binary_output = binary_output
+ self._record_latency = record_latency
# We shouldn't call `model.to()` for models loaded with accelerate as well as the case that model is already on device
if (
diff --git a/src/transformers/pipelines/depth_estimation.py b/src/transformers/pipelines/depth_estimation.py
index 03ee70673d..ed63e818b2 100644
--- a/src/transformers/pipelines/depth_estimation.py
+++ b/src/transformers/pipelines/depth_estimation.py
@@ -1,3 +1,4 @@
+import time
from typing import Any, Union, overload
from ..utils import (
@@ -122,7 +123,13 @@ class DepthEstimationPipeline(Pipeline):
def _forward(self, model_inputs):
target_size = model_inputs.pop("target_size")
- model_outputs = self.model(**model_inputs)
+ if self._record_latency:
+ start_time = time.perf_counter()
+ model_outputs = self.model(**model_inputs)
+ latency = (time.perf_counter() - start_time) * 1000
+ logger.info(f"[{self.__class__.__name__}] Inference latency: {latency:.2f} ms")
+ else:
+ model_outputs = self.model(**model_inputs)
model_outputs["target_size"] = target_size
return model_outputs
diff --git a/src/transformers/pipelines/document_question_answering.py b/src/transformers/pipelines/document_question_answering.py
index de976f9d87..468e4fbe50 100644
--- a/src/transformers/pipelines/document_question_answering.py
+++ b/src/transformers/pipelines/document_question_answering.py
@@ -13,6 +13,7 @@
# limitations under the License.
import re
+import time
from typing import Any, Union, overload
import numpy as np
@@ -568,9 +569,21 @@ class DocumentQuestionAnsweringPipeline(ChunkPipeline):
if "generation_config" not in generate_kwargs:
generate_kwargs["generation_config"] = self.generation_config
- model_outputs = self.model.generate(**model_inputs, **generate_kwargs)
+ if self._record_latency:
+ start_time = time.perf_counter()
+ model_outputs = self.model.generate(**model_inputs, **generate_kwargs)
+ latency = (time.perf_counter() - start_time) * 1000
+ logger.info(f"[{self.__class__.__name__}] Inference latency: {latency:.2f} ms")
+ else:
+ model_outputs = self.model.generate(**model_inputs, **generate_kwargs)
else:
- model_outputs = self.model(**model_inputs)
+ if self._record_latency:
+ start_time = time.perf_counter()
+ model_outputs = self.model(**model_inputs)
+ latency = (time.perf_counter() - start_time) * 1000
+ logger.info(f"[{self.__class__.__name__}] Inference latency: {latency:.2f} ms")
+ else:
+ model_outputs = self.model(**model_inputs)
model_outputs = dict(model_outputs.items())
model_outputs["p_mask"] = p_mask
diff --git a/src/transformers/pipelines/feature_extraction.py b/src/transformers/pipelines/feature_extraction.py
index a37f147605..4cd12d5522 100644
--- a/src/transformers/pipelines/feature_extraction.py
+++ b/src/transformers/pipelines/feature_extraction.py
@@ -1,9 +1,13 @@
+import time
from typing import Any
-from ..utils import add_end_docstrings
+from ..utils import add_end_docstrings, logging
from .base import GenericTensor, Pipeline, build_pipeline_init_args
+logger = logging.get_logger(__name__)
+
+
@add_end_docstrings(
build_pipeline_init_args(has_tokenizer=True, supports_binary_output=False),
r"""
@@ -66,6 +70,12 @@ class FeatureExtractionPipeline(Pipeline):
return model_inputs
def _forward(self, model_inputs):
+ if self._record_latency:
+ start_time = time.perf_counter()
+ model_outputs = self.model(**model_inputs)
+ latency = (time.perf_counter() - start_time) * 1000
+ logger.info(f"[{self.__class__.__name__}] Inference latency: {latency:.2f} ms")
+ return model_outputs
model_outputs = self.model(**model_inputs)
return model_outputs
diff --git a/src/transformers/pipelines/fill_mask.py b/src/transformers/pipelines/fill_mask.py
index 1ea7c487be..b333269e72 100644
--- a/src/transformers/pipelines/fill_mask.py
+++ b/src/transformers/pipelines/fill_mask.py
@@ -1,3 +1,4 @@
+import time
from typing import Any, overload
import numpy as np
@@ -118,7 +119,13 @@ class FillMaskPipeline(Pipeline):
return model_inputs
def _forward(self, model_inputs):
- model_outputs = self.model(**model_inputs)
+ if self._record_latency:
+ start_time = time.perf_counter()
+ model_outputs = self.model(**model_inputs)
+ latency = (time.perf_counter() - start_time) * 1000
+ logger.info(f"[{self.__class__.__name__}] Inference latency: {latency:.2f} ms")
+ else:
+ model_outputs = self.model(**model_inputs)
model_outputs["input_ids"] = model_inputs["input_ids"]
return model_outputs
diff --git a/src/transformers/pipelines/image_classification.py b/src/transformers/pipelines/image_classification.py
index 18a570df6e..33ce255fab 100644
--- a/src/transformers/pipelines/image_classification.py
+++ b/src/transformers/pipelines/image_classification.py
@@ -11,6 +11,7 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
+import time
from typing import Any, Union, overload
import numpy as np
@@ -187,6 +188,12 @@ class ImageClassificationPipeline(Pipeline):
return model_inputs
def _forward(self, model_inputs):
+ if self._record_latency:
+ start_time = time.perf_counter()
+ model_outputs = self.model(**model_inputs)
+ latency = (time.perf_counter() - start_time) * 1000
+ logger.info(f"[{self.__class__.__name__}] Inference latency: {latency:.2f} ms")
+ return model_outputs
model_outputs = self.model(**model_inputs)
return model_outputs
diff --git a/src/transformers/pipelines/image_feature_extraction.py b/src/transformers/pipelines/image_feature_extraction.py
index d049957a41..52833346a2 100644
--- a/src/transformers/pipelines/image_feature_extraction.py
+++ b/src/transformers/pipelines/image_feature_extraction.py
@@ -1,9 +1,13 @@
+import time
from typing import Any, Union
-from ..utils import add_end_docstrings, is_vision_available
+from ..utils import add_end_docstrings, is_vision_available, logging
from .base import GenericTensor, Pipeline, build_pipeline_init_args
+logger = logging.get_logger(__name__)
+
+
if is_vision_available():
from PIL import Image
@@ -71,6 +75,12 @@ class ImageFeatureExtractionPipeline(Pipeline):
return model_inputs
def _forward(self, model_inputs):
+ if self._record_latency:
+ start_time = time.perf_counter()
+ model_outputs = self.model(**model_inputs)
+ latency = (time.perf_counter() - start_time) * 1000
+ logger.info(f"[{self.__class__.__name__}] Inference latency: {latency:.2f} ms")
+ return model_outputs
model_outputs = self.model(**model_inputs)
return model_outputs
diff --git a/src/transformers/pipelines/image_segmentation.py b/src/transformers/pipelines/image_segmentation.py
index 49854beb5a..2cf975e455 100644
--- a/src/transformers/pipelines/image_segmentation.py
+++ b/src/transformers/pipelines/image_segmentation.py
@@ -1,3 +1,4 @@
+import time
from typing import Any, Union, overload
import numpy as np
@@ -172,7 +173,13 @@ class ImageSegmentationPipeline(Pipeline):
def _forward(self, model_inputs):
target_size = model_inputs.pop("target_size")
- model_outputs = self.model(**model_inputs)
+ if self._record_latency:
+ start_time = time.perf_counter()
+ model_outputs = self.model(**model_inputs)
+ latency = (time.perf_counter() - start_time) * 1000
+ logger.info(f"[{self.__class__.__name__}] Inference latency: {latency:.2f} ms")
+ else:
+ model_outputs = self.model(**model_inputs)
model_outputs["target_size"] = target_size
return model_outputs
diff --git a/src/transformers/pipelines/image_text_to_text.py b/src/transformers/pipelines/image_text_to_text.py
index 7d28b91ab2..c86df235b3 100644
--- a/src/transformers/pipelines/image_text_to_text.py
+++ b/src/transformers/pipelines/image_text_to_text.py
@@ -13,6 +13,7 @@
# limitations under the License.
import enum
+import time
from typing import Any, Union, overload
from ..generation import GenerationConfig
@@ -391,7 +392,13 @@ class ImageTextToTextPipeline(Pipeline):
if "generation_config" not in generate_kwargs:
generate_kwargs["generation_config"] = self.generation_config
- generated_sequence = self.model.generate(**model_inputs, **generate_kwargs)
+ if self._record_latency:
+ start_time = time.perf_counter()
+ generated_sequence = self.model.generate(**model_inputs, **generate_kwargs)
+ latency = (time.perf_counter() - start_time) * 1000
+ logger.info(f"[{self.__class__.__name__}] Inference latency: {latency:.2f} ms")
+ else:
+ generated_sequence = self.model.generate(**model_inputs, **generate_kwargs)
return {"generated_sequence": generated_sequence, "prompt_text": prompt_text, "input_ids": input_ids}
diff --git a/src/transformers/pipelines/keypoint_matching.py b/src/transformers/pipelines/keypoint_matching.py
index d75656a7db..07a14a70af 100644
--- a/src/transformers/pipelines/keypoint_matching.py
+++ b/src/transformers/pipelines/keypoint_matching.py
@@ -12,16 +12,20 @@
# See the License for the specific language governing permissions and
# limitations under the License.
+import time
from collections.abc import Sequence
from typing import Any, TypeAlias, TypedDict, Union
from typing_extensions import overload
from ..image_utils import is_pil_image
-from ..utils import is_vision_available, requires_backends
+from ..utils import is_vision_available, logging, requires_backends
from .base import Pipeline
+logger = logging.get_logger(__name__)
+
+
if is_vision_available():
from PIL import Image
@@ -152,7 +156,13 @@ class KeypointMatchingPipeline(Pipeline):
def _forward(self, preprocess_outputs):
model_inputs = preprocess_outputs["model_inputs"]
- model_outputs = self.model(**model_inputs)
+ if self._record_latency:
+ start_time = time.perf_counter()
+ model_outputs = self.model(**model_inputs)
+ latency = (time.perf_counter() - start_time) * 1000
+ logger.info(f"[{self.__class__.__name__}] Inference latency: {latency:.2f} ms")
+ else:
+ model_outputs = self.model(**model_inputs)
forward_outputs = {"model_outputs": model_outputs, "target_sizes": [preprocess_outputs["target_sizes"]]}
return forward_outputs
diff --git a/src/transformers/pipelines/mask_generation.py b/src/transformers/pipelines/mask_generation.py
index 920ce9d0a3..39b201b616 100644
--- a/src/transformers/pipelines/mask_generation.py
+++ b/src/transformers/pipelines/mask_generation.py
@@ -1,3 +1,4 @@
+import time
from collections import defaultdict
from typing import TYPE_CHECKING, Any, Union, overload
@@ -256,7 +257,13 @@ class MaskGenerationPipeline(ChunkPipeline):
reshaped_input_sizes = model_inputs.pop("reshaped_input_sizes", None)
reshaped_input_sizes = reshaped_input_sizes.tolist() if reshaped_input_sizes is not None else None
- model_outputs = self.model(**model_inputs)
+ if self._record_latency:
+ start_time = time.perf_counter()
+ model_outputs = self.model(**model_inputs)
+ latency = (time.perf_counter() - start_time) * 1000
+ logger.info(f"[{self.__class__.__name__}] Inference latency: {latency:.2f} ms")
+ else:
+ model_outputs = self.model(**model_inputs)
# post processing happens here in order to avoid CPU GPU copies of ALL the masks
low_resolution_masks = model_outputs["pred_masks"]
diff --git a/src/transformers/pipelines/object_detection.py b/src/transformers/pipelines/object_detection.py
index 0a4fba996d..90dddf8df1 100644
--- a/src/transformers/pipelines/object_detection.py
+++ b/src/transformers/pipelines/object_detection.py
@@ -1,3 +1,4 @@
+import time
from typing import TYPE_CHECKING, Any, Union, overload
from ..utils import add_end_docstrings, is_torch_available, is_vision_available, logging, requires_backends
@@ -126,7 +127,13 @@ class ObjectDetectionPipeline(Pipeline):
def _forward(self, model_inputs):
target_size = model_inputs.pop("target_size")
- outputs = self.model(**model_inputs)
+ if self._record_latency:
+ start_time = time.perf_counter()
+ outputs = self.model(**model_inputs)
+ latency = (time.perf_counter() - start_time) * 1000
+ logger.info(f"[{self.__class__.__name__}] Inference latency: {latency:.2f} ms")
+ else:
+ outputs = self.model(**model_inputs)
model_outputs = outputs.__class__({"target_size": target_size, **outputs})
if self.tokenizer is not None:
model_outputs["bbox"] = model_inputs["bbox"]
diff --git a/src/transformers/pipelines/table_question_answering.py b/src/transformers/pipelines/table_question_answering.py
index 96bcc863cb..507cae4738 100644
--- a/src/transformers/pipelines/table_question_answering.py
+++ b/src/transformers/pipelines/table_question_answering.py
@@ -1,4 +1,5 @@
import collections
+import time
import types
import numpy as np
@@ -7,11 +8,15 @@ from ..generation import GenerationConfig
from ..utils import (
add_end_docstrings,
is_torch_available,
+ logging,
requires_backends,
)
from .base import ArgumentHandler, Dataset, Pipeline, PipelineException, build_pipeline_init_args
+logger = logging.get_logger(__name__)
+
+
if is_torch_available():
import torch
@@ -325,6 +330,9 @@ class TableQuestionAnsweringPipeline(Pipeline):
def _forward(self, model_inputs, sequential=False, **generate_kwargs):
table = model_inputs.pop("table")
+ if self._record_latency:
+ start_time = time.perf_counter()
+
if self.type == "tapas":
if sequential:
outputs = self.sequential_inference(**model_inputs)
@@ -336,6 +344,11 @@ class TableQuestionAnsweringPipeline(Pipeline):
generate_kwargs["generation_config"] = self.generation_config
outputs = self.model.generate(**model_inputs, **generate_kwargs)
+
+ if self._record_latency:
+ latency = (time.perf_counter() - start_time) * 1000
+ logger.info(f"[{self.__class__.__name__}] Inference latency: {latency:.2f} ms")
+
model_outputs = {"model_inputs": model_inputs, "table": table, "outputs": outputs}
return model_outputs
diff --git a/src/transformers/pipelines/text_classification.py b/src/transformers/pipelines/text_classification.py
index ab9a5d8efb..dd08e06723 100644
--- a/src/transformers/pipelines/text_classification.py
+++ b/src/transformers/pipelines/text_classification.py
@@ -1,9 +1,10 @@
import inspect
+import time
from typing import Any
import numpy as np
-from ..utils import ExplicitEnum, add_end_docstrings, is_torch_available
+from ..utils import ExplicitEnum, add_end_docstrings, is_torch_available, logging
from .base import GenericTensor, Pipeline, build_pipeline_init_args
@@ -11,6 +12,9 @@ if is_torch_available():
from ..models.auto.modeling_auto import MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING_NAMES
+logger = logging.get_logger(__name__)
+
+
def sigmoid(_outputs):
return 1.0 / (1.0 + np.exp(-_outputs))
@@ -173,6 +177,13 @@ class TextClassificationPipeline(Pipeline):
model_forward = self.model.forward
if "use_cache" in inspect.signature(model_forward).parameters:
model_inputs["use_cache"] = False
+
+ if self._record_latency:
+ start_time = time.perf_counter()
+ outputs = self.model(**model_inputs)
+ latency = (time.perf_counter() - start_time) * 1000
+ logger.info(f"[{self.__class__.__name__}] Inference latency: {latency:.2f} ms")
+ return outputs
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/text_generation.py b/src/transformers/pipelines/text_generation.py
index 6a0b2966d0..070c7ca392 100644
--- a/src/transformers/pipelines/text_generation.py
+++ b/src/transformers/pipelines/text_generation.py
@@ -1,12 +1,16 @@
import enum
+import time
from typing import Any, overload
from ..generation import GenerationConfig
-from ..utils import ModelOutput, add_end_docstrings, is_torch_available
+from ..utils import ModelOutput, add_end_docstrings, is_torch_available, logging
from ..utils.chat_template_utils import Chat, ChatType
from .base import Pipeline, build_pipeline_init_args
+logger = logging.get_logger(__name__)
+
+
if is_torch_available():
import torch
@@ -400,7 +404,13 @@ class TextGenerationPipeline(Pipeline):
if "generation_config" not in generate_kwargs:
generate_kwargs["generation_config"] = self.generation_config
- output = self.model.generate(input_ids=input_ids, attention_mask=attention_mask, **generate_kwargs)
+ if self._record_latency:
+ start_time = time.perf_counter()
+ output = self.model.generate(input_ids=input_ids, attention_mask=attention_mask, **generate_kwargs)
+ latency = (time.perf_counter() - start_time) * 1000
+ logger.info(f"[{self.__class__.__name__}] Inference latency: {latency:.2f} ms")
+ else:
+ output = self.model.generate(input_ids=input_ids, attention_mask=attention_mask, **generate_kwargs)
if isinstance(output, ModelOutput):
generated_sequence = output.sequences
diff --git a/src/transformers/pipelines/text_to_audio.py b/src/transformers/pipelines/text_to_audio.py
index 81c6e34f95..50e34397a7 100644
--- a/src/transformers/pipelines/text_to_audio.py
+++ b/src/transformers/pipelines/text_to_audio.py
@@ -12,15 +12,19 @@
# See the License for the specific language governing permissions and
# limitations under the License.from typing import List, Union
+import time
from typing import Any, TypedDict, overload
from ..audio_utils import AudioInput
from ..generation import GenerationConfig
-from ..utils import is_torch_available
+from ..utils import is_torch_available, logging
from ..utils.chat_template_utils import Chat, ChatType
from .base import Pipeline
+logger = logging.get_logger(__name__)
+
+
if is_torch_available():
import torch
@@ -211,7 +215,13 @@ class TextToAudioPipeline(Pipeline):
if "output_audio" not in forward_params:
forward_params["output_audio"] = True
- output = self.model.generate(**model_inputs, **forward_params)
+ if self._record_latency:
+ start_time = time.perf_counter()
+ output = self.model.generate(**model_inputs, **forward_params)
+ latency = (time.perf_counter() - start_time) * 1000
+ logger.info(f"[{self.__class__.__name__}] Inference latency: {latency:.2f} ms")
+ else:
+ output = self.model.generate(**model_inputs, **forward_params)
else:
if len(generate_kwargs):
raise ValueError(
@@ -219,7 +229,13 @@ class TextToAudioPipeline(Pipeline):
"empty. For forward-only TTA models, please use `forward_params` instead of `generate_kwargs`. "
f"For reference, the `generate_kwargs` used here are: {generate_kwargs.keys()}"
)
- output = self.model(**model_inputs, **forward_params)[0]
+ if self._record_latency:
+ start_time = time.perf_counter()
+ output = self.model(**model_inputs, **forward_params)[0]
+ latency = (time.perf_counter() - start_time) * 1000
+ logger.info(f"[{self.__class__.__name__}] Inference latency: {latency:.2f} ms")
+ else:
+ output = self.model(**model_inputs, **forward_params)[0]
if self.vocoder is not None:
# in that case, the output is a spectrogram that needs to be converted into a waveform
diff --git a/src/transformers/pipelines/token_classification.py b/src/transformers/pipelines/token_classification.py
index 7deca9dc90..3082fcec41 100644
--- a/src/transformers/pipelines/token_classification.py
+++ b/src/transformers/pipelines/token_classification.py
@@ -1,3 +1,4 @@
+import time
import types
import warnings
from typing import Any, overload
@@ -9,10 +10,14 @@ from ..utils import (
ExplicitEnum,
add_end_docstrings,
is_torch_available,
+ logging,
)
from .base import ArgumentHandler, ChunkPipeline, Dataset, build_pipeline_init_args
+logger = logging.get_logger(__name__)
+
+
if is_torch_available():
import torch
@@ -308,7 +313,13 @@ class TokenClassificationPipeline(ChunkPipeline):
word_ids = model_inputs.pop("word_ids", None)
word_to_chars_map = model_inputs.pop("word_to_chars_map", None)
- output = self.model(**model_inputs)
+ if self._record_latency:
+ start_time = time.perf_counter()
+ output = self.model(**model_inputs)
+ latency = (time.perf_counter() - start_time) * 1000
+ logger.info(f"[{self.__class__.__name__}] Inference latency: {latency:.2f} ms")
+ else:
+ output = self.model(**model_inputs)
logits = output["logits"] if isinstance(output, dict) else output[0]
return {
diff --git a/src/transformers/pipelines/video_classification.py b/src/transformers/pipelines/video_classification.py
index 4e2ff77cbe..e05a71967b 100644
--- a/src/transformers/pipelines/video_classification.py
+++ b/src/transformers/pipelines/video_classification.py
@@ -11,6 +11,7 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
+import time
from io import BytesIO
from typing import Any, overload
@@ -150,6 +151,12 @@ class VideoClassificationPipeline(Pipeline):
return model_inputs
def _forward(self, model_inputs):
+ if self._record_latency:
+ start_time = time.perf_counter()
+ model_outputs = self.model(**model_inputs)
+ latency = (time.perf_counter() - start_time) * 1000
+ logger.info(f"[{self.__class__.__name__}] Inference latency: {latency:.2f} ms")
+ return model_outputs
model_outputs = self.model(**model_inputs)
return model_outputs
diff --git a/src/transformers/pipelines/zero_shot_audio_classification.py b/src/transformers/pipelines/zero_shot_audio_classification.py
index 03c1a8d1c1..f09a486581 100644
--- a/src/transformers/pipelines/zero_shot_audio_classification.py
+++ b/src/transformers/pipelines/zero_shot_audio_classification.py
@@ -11,6 +11,7 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
+import time
from collections import UserDict
from typing import Any
@@ -138,7 +139,13 @@ class ZeroShotAudioClassificationPipeline(Pipeline):
# Batching case.
text_inputs = text_inputs[0][0]
- outputs = self.model(**text_inputs, **model_inputs)
+ if self._record_latency:
+ start_time = time.perf_counter()
+ outputs = self.model(**text_inputs, **model_inputs)
+ latency = (time.perf_counter() - start_time) * 1000
+ logger.info(f"[{self.__class__.__name__}] Inference latency: {latency:.2f} ms")
+ else:
+ outputs = self.model(**text_inputs, **model_inputs)
model_outputs = {
"candidate_labels": candidate_labels,
diff --git a/src/transformers/pipelines/zero_shot_classification.py b/src/transformers/pipelines/zero_shot_classification.py
index d88772310d..2d15d0f9dd 100644
--- a/src/transformers/pipelines/zero_shot_classification.py
+++ b/src/transformers/pipelines/zero_shot_classification.py
@@ -1,4 +1,5 @@
import inspect
+import time
import numpy as np
@@ -222,7 +223,14 @@ class ZeroShotClassificationPipeline(ChunkPipeline):
model_forward = self.model.forward
if "use_cache" in inspect.signature(model_forward).parameters:
model_inputs["use_cache"] = False
- outputs = self.model(**model_inputs)
+
+ if self._record_latency:
+ start_time = time.perf_counter()
+ outputs = self.model(**model_inputs)
+ latency = (time.perf_counter() - start_time) * 1000
+ logger.info(f"[{self.__class__.__name__}] Inference latency: {latency:.2f} ms")
+ else:
+ outputs = self.model(**model_inputs)
model_outputs = {
"candidate_label": candidate_label,
diff --git a/src/transformers/pipelines/zero_shot_image_classification.py b/src/transformers/pipelines/zero_shot_image_classification.py
index d129c58365..551d9a9218 100644
--- a/src/transformers/pipelines/zero_shot_image_classification.py
+++ b/src/transformers/pipelines/zero_shot_image_classification.py
@@ -1,3 +1,4 @@
+import time
from collections import UserDict
from typing import Any, Union, overload
@@ -166,7 +167,13 @@ class ZeroShotImageClassificationPipeline(Pipeline):
# Batching case.
text_inputs = text_inputs[0][0]
- outputs = self.model(**text_inputs, **model_inputs)
+ if self._record_latency:
+ start_time = time.perf_counter()
+ outputs = self.model(**text_inputs, **model_inputs)
+ latency = (time.perf_counter() - start_time) * 1000
+ logger.info(f"[{self.__class__.__name__}] Inference latency: {latency:.2f} ms")
+ else:
+ outputs = self.model(**text_inputs, **model_inputs)
model_outputs = {
"candidate_labels": candidate_labels,
diff --git a/src/transformers/pipelines/zero_shot_object_detection.py b/src/transformers/pipelines/zero_shot_object_detection.py
index 7f353afd74..28a7dd6418 100644
--- a/src/transformers/pipelines/zero_shot_object_detection.py
+++ b/src/transformers/pipelines/zero_shot_object_detection.py
@@ -1,3 +1,4 @@
+import time
from typing import Any, Union, overload
from ..utils import add_end_docstrings, is_torch_available, is_vision_available, logging, requires_backends
@@ -195,7 +196,13 @@ class ZeroShotObjectDetectionPipeline(ChunkPipeline):
candidate_label = model_inputs.pop("candidate_label")
is_last = model_inputs.pop("is_last")
- outputs = self.model(**model_inputs)
+ if self._record_latency:
+ start_time = time.perf_counter()
+ outputs = self.model(**model_inputs)
+ latency = (time.perf_counter() - start_time) * 1000
+ logger.info(f"[{self.__class__.__name__}] Inference latency: {latency:.2f} ms")
+ else:
+ outputs = self.model(**model_inputs)
model_outputs = {"target_size": target_size, "candidate_label": candidate_label, "is_last": is_last, **outputs}
return model_outputs
Reference in a new issue View git blame Copy permalink