From f2ecb9eec42d9f55bd4190f3755f63dab816a083 Mon Sep 17 00:00:00 2001 From: Sylvain Gugger <35901082+sgugger@users.noreply.github.com> Date: Mon, 17 Oct 2022 11:56:29 -0400 Subject: [PATCH] Revert "add return_tensor parameter for feature extraction (#19257)" (#19680) This reverts commit 35bd089a241788a43a43e27de1ef3f5cede7954b. --- .../pipelines/feature_extraction.py | 14 +++----------- .../test_pipelines_feature_extraction.py | 18 ------------------ 2 files changed, 3 insertions(+), 29 deletions(-) diff --git a/src/transformers/pipelines/feature_extraction.py b/src/transformers/pipelines/feature_extraction.py index 48f7735b6ce0b4..c7d50c971955fd 100644 --- a/src/transformers/pipelines/feature_extraction.py +++ b/src/transformers/pipelines/feature_extraction.py @@ -31,8 +31,6 @@ class FeatureExtractionPipeline(Pipeline): If no framework is specified, will default to the one currently installed. If no framework is specified and both frameworks are installed, will default to the framework of the `model`, or to PyTorch if no model is provided. - return_tensor (`bool`, *optional*): - If `True`, returns a tensor according to the specified framework, otherwise returns a list. task (`str`, defaults to `""`): A task-identifier for the pipeline. args_parser ([`~pipelines.ArgumentHandler`], *optional*): @@ -42,7 +40,7 @@ class FeatureExtractionPipeline(Pipeline): the associated CUDA device id. """ - def _sanitize_parameters(self, truncation=None, tokenize_kwargs=None, return_tensors=None, **kwargs): + def _sanitize_parameters(self, truncation=None, tokenize_kwargs=None, **kwargs): if tokenize_kwargs is None: tokenize_kwargs = {} @@ -55,11 +53,7 @@ def _sanitize_parameters(self, truncation=None, tokenize_kwargs=None, return_ten preprocess_params = tokenize_kwargs - postprocess_params = {} - if return_tensors is not None: - postprocess_params["return_tensors"] = return_tensors - - return preprocess_params, {}, postprocess_params + return preprocess_params, {}, {} def preprocess(self, inputs, **tokenize_kwargs) -> Dict[str, GenericTensor]: return_tensors = self.framework @@ -70,10 +64,8 @@ def _forward(self, model_inputs): model_outputs = self.model(**model_inputs) return model_outputs - def postprocess(self, model_outputs, return_tensors=False): + def postprocess(self, model_outputs): # [0] is the first available tensor, logits or last_hidden_state. - if return_tensors: - return model_outputs[0] if self.framework == "pt": return model_outputs[0].tolist() elif self.framework == "tf": diff --git a/tests/pipelines/test_pipelines_feature_extraction.py b/tests/pipelines/test_pipelines_feature_extraction.py index 2e431fa1d48612..f75af6808bccdf 100644 --- a/tests/pipelines/test_pipelines_feature_extraction.py +++ b/tests/pipelines/test_pipelines_feature_extraction.py @@ -15,8 +15,6 @@ import unittest import numpy as np -import tensorflow as tf -import torch from transformers import ( FEATURE_EXTRACTOR_MAPPING, @@ -135,22 +133,6 @@ def test_tokenization_small_model_tf(self): tokenize_kwargs=tokenize_kwargs, ) - @require_torch - def test_return_tensors_pt(self): - feature_extractor = pipeline( - task="feature-extraction", model="hf-internal-testing/tiny-random-distilbert", framework="pt" - ) - outputs = feature_extractor("This is a test" * 100, return_tensors=True) - self.assertTrue(torch.is_tensor(outputs)) - - @require_tf - def test_return_tensors_tf(self): - feature_extractor = pipeline( - task="feature-extraction", model="hf-internal-testing/tiny-random-distilbert", framework="tf" - ) - outputs = feature_extractor("This is a test" * 100, return_tensors=True) - self.assertTrue(tf.is_tensor(outputs)) - def get_shape(self, input_, shape=None): if shape is None: shape = []