| 12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788 |
- from typing import Any
- from ..utils import add_end_docstrings
- from .base import GenericTensor, Pipeline, build_pipeline_init_args
- @add_end_docstrings(
- build_pipeline_init_args(has_tokenizer=True, supports_binary_output=False),
- r"""
- tokenize_kwargs (`dict`, *optional*):
- Additional dictionary of keyword arguments passed along to the tokenizer.
- return_tensors (`bool`, *optional*):
- If `True`, returns a tensor according to the specified framework, otherwise returns a list.""",
- )
- class FeatureExtractionPipeline(Pipeline):
- """
- Feature extraction pipeline uses 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="google-bert/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_length, 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).
- """
- _load_processor = False
- _load_image_processor = False
- _load_feature_extractor = False
- _load_tokenizer = True
- 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]:
- model_inputs = self.tokenizer(inputs, return_tensors="pt", **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):
- # [0] is the first available tensor, logits or last_hidden_state.
- if return_tensors:
- return model_outputs[0]
- return model_outputs[0].tolist()
- def __call__(self, *args: str | list[str], **kwargs: Any) -> Any | list[Any]:
- """
- Extract the features of the input(s) text.
- 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)
|
