From 158cef2caf5fa19400d42eb3a7b5e4ed8d05704e Mon Sep 17 00:00:00 2001 From: arrdel Date: Tue, 2 Dec 2025 12:27:02 -0500 Subject: [PATCH] docs: add doctest for SqueezeBERT Adds documentation test for SqueezeBertForMaskedLM as part of the Doc Tests Sprint (#16292). - Added example to forward() method showing masked language modeling - Used squeezebert-uncased checkpoint for testing - Added file to slow_documentation_tests.txt --- .../squeezebert/modeling_squeezebert.py | 21 +++++++++++++++++++ utils/slow_documentation_tests.txt | 1 + 2 files changed, 22 insertions(+) diff --git a/src/transformers/models/squeezebert/modeling_squeezebert.py b/src/transformers/models/squeezebert/modeling_squeezebert.py index 94d5f36891d8..7b27cc1b4e04 100644 --- a/src/transformers/models/squeezebert/modeling_squeezebert.py +++ b/src/transformers/models/squeezebert/modeling_squeezebert.py @@ -534,6 +534,27 @@ def forward( Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ..., config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]` + + Example: + + ```python + >>> from transformers import AutoTokenizer, SqueezeBertForMaskedLM + >>> import torch + + >>> tokenizer = AutoTokenizer.from_pretrained("squeezebert/squeezebert-uncased") + >>> model = SqueezeBertForMaskedLM.from_pretrained("squeezebert/squeezebert-uncased") + + >>> inputs = tokenizer("The capital of France is [MASK].", return_tensors="pt") + >>> with torch.no_grad(): + ... logits = model(**inputs).logits + + >>> # retrieve index of [MASK] + >>> mask_token_index = (inputs.input_ids == tokenizer.mask_token_id)[0].nonzero(as_tuple=True)[0] + + >>> predicted_token_id = logits[0, mask_token_index].argmax(axis=-1) + >>> tokenizer.decode(predicted_token_id) # doctest: +SKIP + 'paris' + ``` """ return_dict = return_dict if return_dict is not None else self.config.use_return_dict diff --git a/utils/slow_documentation_tests.txt b/utils/slow_documentation_tests.txt index 03ef646f04bc..5a079b0f790b 100644 --- a/utils/slow_documentation_tests.txt +++ b/utils/slow_documentation_tests.txt @@ -13,3 +13,4 @@ src/transformers/models/idefics2/modeling_idefics2.py src/transformers/models/kosmos2/modeling_kosmos2.py src/transformers/models/musicgen_melody/modeling_musicgen_melody.py src/transformers/models/musicgen_melody/processing_musicgen_melody.py +src/transformers/models/squeezebert/modeling_squeezebert.py