Skip to content

docs: add doctest for SqueezeBERT#42572

Open
arrdel wants to merge 1 commit into
huggingface:mainfrom
arrdel:docs-add-squeezebert-doctest
Open

docs: add doctest for SqueezeBERT#42572
arrdel wants to merge 1 commit into
huggingface:mainfrom
arrdel:docs-add-squeezebert-doctest

Conversation

@arrdel

@arrdel arrdel commented Dec 2, 2025

Copy link
Copy Markdown

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

What does this PR do?

Fixes #16292 (partial - adds SqueezeBERT)

Adds documentation test for SqueezeBertForMaskedLM as part of the Doc Tests Sprint (huggingface#16292).

- Added example to forward() method showing masked language modeling
- Used squeezebert-uncased checkpoint for testing
- Added file to slow_documentation_tests.txt
@github-actions

github-actions Bot commented Dec 2, 2025

Copy link
Copy Markdown
Contributor

[For maintainers] Suggested jobs to run (before merge)

run-slow: squeezebert

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Pull request overview

Adds a doctest example to SqueezeBertForMaskedLM’s forward() docstring and includes the file in slow documentation tests. This helps validate masked language modeling usage in the docs.

  • Adds a masked LM usage example with squeezebert-uncased.
  • Registers the file in slow_documentation_tests.txt for doctest execution.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File Description
utils/slow_documentation_tests.txt Adds modeling_squeezebert.py to slow documentation tests to run doctest.
src/transformers/models/squeezebert/modeling_squeezebert.py Adds a masked LM doctest example in the forward() docstring.

>>> # 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)

Copilot AI Dec 2, 2025

Copy link

Choose a reason for hiding this comment

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

In the doctest example, logits is a PyTorch tensor, but argmax(axis=-1) uses the NumPy keyword. In PyTorch, the correct keyword is dim. Using axis will raise an error. Suggest updating to argmax(dim=-1) to get the vocabulary index of the predicted token(s).

Suggested change
>>> predicted_token_id = logits[0, mask_token_index].argmax(axis=-1)
>>> predicted_token_id = logits[0, mask_token_index].argmax(dim=-1)

Copilot uses AI. Check for mistakes.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Community Event] Doc Tests Sprint

2 participants