Adding doctest for fill-mask pipeline. (#20241)

This commit is contained in:
Nicolas Patry 2022-11-16 09:51:20 +01:00 committed by GitHub
parent 5e080c11bf
commit 529037fda5
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -36,6 +36,22 @@ class FillMaskPipeline(Pipeline):
Masked language modeling prediction pipeline using any `ModelWithLMHead`. See the [masked language modeling
examples](../task_summary#masked-language-modeling) for more information.
Example:
```python
>>> from transformers import pipeline
>>> fill_masker = pipeline(model="bert-base-uncased")
>>> potential_words = fill_masker("This is a simple [MASK].")
>>> from transformers.testing_utils import nested_simplify
>>> nested_simplify(potential_words) # The scores might vary slightly across pytorch/tensorflow versions.
[{'score': 0.042, 'token': 3291, 'token_str': 'problem', 'sequence': 'this is a simple problem.'}, {'score': 0.031, 'token': 3160, 'token_str': 'question', 'sequence': 'this is a simple question.'}, {'score': 0.03, 'token': 8522, 'token_str': 'equation', 'sequence': 'this is a simple equation.'}, {'score': 0.027, 'token': 2028, 'token_str': 'one', 'sequence': 'this is a simple one.'}, {'score': 0.024, 'token': 3627, 'token_str': 'rule', 'sequence': 'this is a simple rule.'}]
```
[Using pipelines in a webserver or with a dataset](../pipeline_tutorial)
This mask filling pipeline can currently be loaded from [`pipeline`] using the following task identifier:
`"fill-mask"`.