From 06e242933555628803c8e901e12feb41f1f60934 Mon Sep 17 00:00:00 2001 From: Nicolas Patry Date: Tue, 15 Nov 2022 17:51:00 +0100 Subject: [PATCH 1/2] Adding doctest for `feature-extraction`. --- src/transformers/pipelines/feature_extraction.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/src/transformers/pipelines/feature_extraction.py b/src/transformers/pipelines/feature_extraction.py index 48f7735b6ce0..9900cc27374c 100644 --- a/src/transformers/pipelines/feature_extraction.py +++ b/src/transformers/pipelines/feature_extraction.py @@ -9,6 +9,19 @@ 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 + + >>> classifier = pipeline(model="bert-base-uncased", task="feature-extraction") + >>> result = classifier("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]) + ``` + + [Using pipelines in a webserver or with a dataset](../pipeline_tutorial) + This feature extraction pipeline can currently be loaded from [`pipeline`] using the task identifier: `"feature-extraction"`. From cae2d2472258b7205ba9aeec944975753eca40c6 Mon Sep 17 00:00:00 2001 From: Nicolas Patry Date: Tue, 15 Nov 2022 18:04:39 +0100 Subject: [PATCH 2/2] Update feature_extraction.py --- src/transformers/pipelines/feature_extraction.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/transformers/pipelines/feature_extraction.py b/src/transformers/pipelines/feature_extraction.py index 9900cc27374c..f46ccfb51112 100644 --- a/src/transformers/pipelines/feature_extraction.py +++ b/src/transformers/pipelines/feature_extraction.py @@ -14,8 +14,8 @@ class FeatureExtractionPipeline(Pipeline): ```python >>> from transformers import pipeline - >>> classifier = pipeline(model="bert-base-uncased", task="feature-extraction") - >>> result = classifier("This is a simple test.", return_tensors=True) + >>> 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]) ```