Skip to content

Feature-extraction pipeline to return Tensor #10016

@ierezell

Description

@ierezell

🚀 Feature request

Actually, to code of the feature-extraction pipeline
transformers.pipelines.feature-extraction.FeatureExtractionPipeline l.82 return a super().__call__(*args, **kwargs).tolist()

Which gives a list[float] (or list[list[float]] if list[str] in input)

I guess it's to be framework agnostic, but we can specify framework='pt' in the pipeline config so I was expecting a torch.tensor.

Could we add some logic to return tensors ?

Motivation

Features will be used as input of other models, so keeping them as tensors (even better on GPU) would be profitable.

Thanks in advance for the reply,

Have a great day.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions