@@ -31,8 +31,6 @@ class FeatureExtractionPipeline(Pipeline):
31
31
If no framework is specified, will default to the one currently installed. If no framework is specified and
32
32
both frameworks are installed, will default to the framework of the `model`, or to PyTorch if no model is
33
33
provided.
34
- return_tensor (`bool`, *optional*):
35
- If `True`, returns a tensor according to the specified framework, otherwise returns a list.
36
34
task (`str`, defaults to `""`):
37
35
A task-identifier for the pipeline.
38
36
args_parser ([`~pipelines.ArgumentHandler`], *optional*):
@@ -42,7 +40,7 @@ class FeatureExtractionPipeline(Pipeline):
42
40
the associated CUDA device id.
43
41
"""
44
42
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 ):
46
44
if tokenize_kwargs is None :
47
45
tokenize_kwargs = {}
48
46
@@ -55,11 +53,7 @@ def _sanitize_parameters(self, truncation=None, tokenize_kwargs=None, return_ten
55
53
56
54
preprocess_params = tokenize_kwargs
57
55
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 , {}, {}
63
57
64
58
def preprocess (self , inputs , ** tokenize_kwargs ) -> Dict [str , GenericTensor ]:
65
59
return_tensors = self .framework
@@ -70,10 +64,8 @@ def _forward(self, model_inputs):
70
64
model_outputs = self .model (** model_inputs )
71
65
return model_outputs
72
66
73
- def postprocess (self , model_outputs , return_tensors = False ):
67
+ def postprocess (self , model_outputs ):
74
68
# [0] is the first available tensor, logits or last_hidden_state.
75
- if return_tensors :
76
- return model_outputs [0 ]
77
69
if self .framework == "pt" :
78
70
return model_outputs [0 ].tolist ()
79
71
elif self .framework == "tf" :
0 commit comments