Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Adding support for hidden_states and attentions in unbatching support. #14420

Merged
merged 1 commit into from
Nov 19, 2021
+22 −3
Merged

Adding support for hidden_states and attentions in unbatching support. #14420

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 8 additions & 3 deletions src/transformers/pipelines/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -747,9 +747,14 @@ def loader_batch_item(self):
else:
loader_batched = {}
for k, element in self._loader_batch_data.items():
if k == "past_key_values":
continue
if isinstance(element[self._loader_batch_index], torch.Tensor):
if k in {"hidden_states", "past_key_values", "attentions"} and isinstance(element, tuple):
if isinstance(element[0], torch.Tensor):
loader_batched[k] = tuple(el[self._loader_batch_index].unsqueeze(0) for el in element)
elif isinstance(element[0], np.ndarray):
loader_batched[k] = tuple(
np.expand_dims(el[self._loader_batch_index], 0) for el in element
)
elif isinstance(element[self._loader_batch_index], torch.Tensor):
loader_batched[k] = element[self._loader_batch_index].unsqueeze(0)
elif isinstance(element[self._loader_batch_index], np.ndarray):
loader_batched[k] = np.expand_dims(element[self._loader_batch_index], 0)
Expand Down
14 changes: 14 additions & 0 deletions tests/test_pipelines_common.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
TOKENIZER_MAPPING,
AutoFeatureExtractor,
AutoTokenizer,
DistilBertForSequenceClassification,
IBertConfig,
RobertaConfig,
TextClassificationPipeline,
Expand Down Expand Up @@ -322,6 +323,19 @@ def data(n: int):
results.append(out)
self.assertEqual(len(results), 10)

@require_torch
def test_unbatch_attentions_hidden_states(self):
model = DistilBertForSequenceClassification.from_pretrained(
"Narsil/tiny-distilbert-sequence-classification", output_hidden_states=True, output_attentions=True
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can this leverage the one from hf-internal-testing, or is there a reason you favor this one?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There is not tiny-distilbert in hf-internal-testing, only the random one.
Happy to upload it if you want (as well as all the other Narsil specific tests that probably all belong in hf-internal-testing.

For instance this one was just copied from the text-classification tests.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Right, but the random has random weights but actually has the weights so it wouldn't be initialized randomly every time you load it in memory.

It shouldn't be an issue here given that you're only checking the length of outputs, right?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Well I am using these also in other tests, are you ok if I merge this PR, then open another PR to remove Narsil/ models from a bit everywhere (some check actual output values)

)
tokenizer = AutoTokenizer.from_pretrained("Narsil/tiny-distilbert-sequence-classification")
text_classifier = TextClassificationPipeline(model=model, tokenizer=tokenizer)

# Used to throw an error because `hidden_states` are a tuple of tensors
# instead of the expected tensor.
outputs = text_classifier(["This is great !"] * 20, batch_size=32)
self.assertEqual(len(outputs), 20)


@is_pipeline_test
class PipelinePadTest(unittest.TestCase):
Expand Down