| from typing import Dict |
|
|
| from .base import GenericTensor, Pipeline |
|
|
|
|
| |
| class FeatureExtractionPipeline(Pipeline): |
| """ |
| Feature extraction pipeline using no model head. This pipeline extracts the hidden states from the base |
| transformer, which can be used as features in downstream tasks. |
| |
| Example: |
| |
| ```python |
| >>> from transformers import pipeline |
| |
| >>> extractor = pipeline(model="bert-base-uncased", task="feature-extraction") |
| >>> result = extractor("This is a simple test.", return_tensors=True) |
| >>> result.shape # This is a tensor of shape [1, sequence_lenth, hidden_dimension] representing the input string. |
| torch.Size([1, 8, 768]) |
| ``` |
| |
| Learn more about the basics of using a pipeline in the [pipeline tutorial](../pipeline_tutorial) |
| |
| This feature extraction pipeline can currently be loaded from [`pipeline`] using the task identifier: |
| `"feature-extraction"`. |
| |
| All models may be used for this pipeline. See a list of all models, including community-contributed models on |
| [huggingface.co/models](https://huggingface.co/models). |
| |
| Arguments: |
| model ([`PreTrainedModel`] or [`TFPreTrainedModel`]): |
| The model that will be used by the pipeline to make predictions. This needs to be a model inheriting from |
| [`PreTrainedModel`] for PyTorch and [`TFPreTrainedModel`] for TensorFlow. |
| tokenizer ([`PreTrainedTokenizer`]): |
| The tokenizer that will be used by the pipeline to encode data for the model. This object inherits from |
| [`PreTrainedTokenizer`]. |
| modelcard (`str` or [`ModelCard`], *optional*): |
| Model card attributed to the model for this pipeline. |
| framework (`str`, *optional*): |
| The framework to use, either `"pt"` for PyTorch or `"tf"` for TensorFlow. The specified framework must be |
| installed. |
| |
| 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_tensors (`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*): |
| Reference to the object in charge of parsing supplied pipeline parameters. |
| device (`int`, *optional*, defaults to -1): |
| Device ordinal for CPU/GPU supports. Setting this to -1 will leverage CPU, a positive will run the model on |
| the associated CUDA device id. |
| tokenize_kwargs (`dict`, *optional*): |
| Additional dictionary of keyword arguments passed along to the tokenizer. |
| """ |
|
|
| def _sanitize_parameters(self, truncation=None, tokenize_kwargs=None, return_tensors=None, **kwargs): |
| if tokenize_kwargs is None: |
| tokenize_kwargs = {} |
|
|
| if truncation is not None: |
| if "truncation" in tokenize_kwargs: |
| raise ValueError( |
| "truncation parameter defined twice (given as keyword argument as well as in tokenize_kwargs)" |
| ) |
| tokenize_kwargs["truncation"] = truncation |
|
|
| preprocess_params = tokenize_kwargs |
|
|
| postprocess_params = {} |
| if return_tensors is not None: |
| postprocess_params["return_tensors"] = return_tensors |
|
|
| return preprocess_params, {}, postprocess_params |
|
|
| def preprocess(self, inputs, **tokenize_kwargs) -> Dict[str, GenericTensor]: |
| return_tensors = self.framework |
| model_inputs = self.tokenizer(inputs, return_tensors=return_tensors, **tokenize_kwargs) |
| return model_inputs |
|
|
| def _forward(self, model_inputs): |
| model_outputs = self.model(**model_inputs) |
| return model_outputs |
|
|
| def postprocess(self, model_outputs, return_tensors=False): |
| |
| if return_tensors: |
| return model_outputs[0] |
| if self.framework == "pt": |
| return model_outputs[0].tolist() |
| elif self.framework == "tf": |
| return model_outputs[0].numpy().tolist() |
|
|
| def __call__(self, *args, **kwargs): |
| """ |
| Extract the features of the input(s). |
| |
| Args: |
| args (`str` or `List[str]`): One or several texts (or one list of texts) to get the features of. |
| |
| Return: |
| A nested list of `float`: The features computed by the model. |
| """ |
| return super().__call__(*args, **kwargs) |
|
|
