Skip to content

Commit e88b95f

Browse files
authored
Revert "add return_tensor parameter for feature extraction (#19257)"
This reverts commit 35bd089.
1 parent bf0addc commit e88b95f

File tree

2 files changed

+3
-29
lines changed

2 files changed

+3
-29
lines changed

src/transformers/pipelines/feature_extraction.py

Lines changed: 3 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -31,8 +31,6 @@ class FeatureExtractionPipeline(Pipeline):
3131
If no framework is specified, will default to the one currently installed. If no framework is specified and
3232
both frameworks are installed, will default to the framework of the `model`, or to PyTorch if no model is
3333
provided.
34-
return_tensor (`bool`, *optional*):
35-
If `True`, returns a tensor according to the specified framework, otherwise returns a list.
3634
task (`str`, defaults to `""`):
3735
A task-identifier for the pipeline.
3836
args_parser ([`~pipelines.ArgumentHandler`], *optional*):
@@ -42,7 +40,7 @@ class FeatureExtractionPipeline(Pipeline):
4240
the associated CUDA device id.
4341
"""
4442

45-
def _sanitize_parameters(self, truncation=None, tokenize_kwargs=None, return_tensors=None, **kwargs):
43+
def _sanitize_parameters(self, truncation=None, tokenize_kwargs=None, **kwargs):
4644
if tokenize_kwargs is None:
4745
tokenize_kwargs = {}
4846

@@ -55,11 +53,7 @@ def _sanitize_parameters(self, truncation=None, tokenize_kwargs=None, return_ten
5553

5654
preprocess_params = tokenize_kwargs
5755

58-
postprocess_params = {}
59-
if return_tensors is not None:
60-
postprocess_params["return_tensors"] = return_tensors
61-
62-
return preprocess_params, {}, postprocess_params
56+
return preprocess_params, {}, {}
6357

6458
def preprocess(self, inputs, **tokenize_kwargs) -> Dict[str, GenericTensor]:
6559
return_tensors = self.framework
@@ -70,10 +64,8 @@ def _forward(self, model_inputs):
7064
model_outputs = self.model(**model_inputs)
7165
return model_outputs
7266

73-
def postprocess(self, model_outputs, return_tensors=False):
67+
def postprocess(self, model_outputs):
7468
# [0] is the first available tensor, logits or last_hidden_state.
75-
if return_tensors:
76-
return model_outputs[0]
7769
if self.framework == "pt":
7870
return model_outputs[0].tolist()
7971
elif self.framework == "tf":

tests/pipelines/test_pipelines_feature_extraction.py

Lines changed: 0 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -15,8 +15,6 @@
1515
import unittest
1616

1717
import numpy as np
18-
import tensorflow as tf
19-
import torch
2018

2119
from transformers import (
2220
FEATURE_EXTRACTOR_MAPPING,
@@ -135,22 +133,6 @@ def test_tokenization_small_model_tf(self):
135133
tokenize_kwargs=tokenize_kwargs,
136134
)
137135

138-
@require_torch
139-
def test_return_tensors_pt(self):
140-
feature_extractor = pipeline(
141-
task="feature-extraction", model="hf-internal-testing/tiny-random-distilbert", framework="pt"
142-
)
143-
outputs = feature_extractor("This is a test" * 100, return_tensors=True)
144-
self.assertTrue(torch.is_tensor(outputs))
145-
146-
@require_tf
147-
def test_return_tensors_tf(self):
148-
feature_extractor = pipeline(
149-
task="feature-extraction", model="hf-internal-testing/tiny-random-distilbert", framework="tf"
150-
)
151-
outputs = feature_extractor("This is a test" * 100, return_tensors=True)
152-
self.assertTrue(tf.is_tensor(outputs))
153-
154136
def get_shape(self, input_, shape=None):
155137
if shape is None:
156138
shape = []

0 commit comments

Comments
 (0)