chore(typing): Add type checking to `src/transformers/generation` by tarekziade · Pull Request #44233 · huggingface/transformers

tarekziade · 2026-02-23T16:23:24Z

Extends ty coverage to src/transformers/generation

Added a dedicated type-check wrapper script: utils/check_types.py.
Updated Makefile to run ty checks through the wrapper in both style and check-repo.
merged all typing in a single module src/transformers/_typing.py
- TransformersLogger
- GenerativePreTrainedModel
- WhisperGenerationConfigLike
Aligned model overrides with updated generation method signatures in clvp, musicgen, musicgen_melody

tarekziade · 2026-02-23T16:26:28Z

        self.loop = asyncio.get_running_loop()
-        self.has_asyncio_timeout = hasattr(asyncio, "timeout")
+        timeout_context = getattr(asyncio, "timeout", None)
+        self.has_asyncio_timeout = sys.version_info >= (3, 11) and callable(timeout_context)


we can tighten this with the python version (function was added in 3.11)

HuggingFaceDocBuilderDev · 2026-02-23T16:33:10Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

zucchini-nlp

Apart from a few nits, my only question is about audio-model specific attributes on generation config. We indeed have several special attributes, and I think whisper has the largest amount

Is there any way to not add them all as class attributes in the general GenerationConfig, because they don't need to be accessible when loading a text only model?

tarekziade · 2026-02-24T11:39:28Z

Apart from a few nits, my only question is about audio-model specific attributes on generation config. We indeed have several special attributes, and I think whisper has the largest amount

Is there any way to not add them all as class attributes in the general GenerationConfig, because they don't need to be accessible when loading a text only model?

Yeah we don't want to pollute GenerationConfig for sure,

The right pattern imho:

Keep GenerationConfig generic (no Whisper/audio field declarations).
Define small model-specific Protocols (e.g. WhisperGenerationConfigLike).
Cast locally at Whisper/audio-only use sites.

That avoids polluting text-only configs while keeping static typing strict where those attrs are actually used.
This is what I added for whisper in _typing after your review

Cyrilvallez

Nice in general! Just left a few comments to try to make it a bit simpler!

Cyrilvallez · 2026-03-02T10:51:10Z

-            cache_to_check = self._cache.self_attention_cache if self.config.is_encoder_decoder else self._cache
+            if isinstance(self._cache, EncoderDecoderCache):
+                cache_to_check = self._cache.self_attention_cache
+            elif isinstance(self._cache, StaticCache):
+                cache_to_check = self._cache


I believe both checks are equivalent indeed, but curious to know why we need it to change?

I changed it to use isinstance-based narrowing for two reasons:

static typing => config flags don’t narrow self._cache’s runtime type, so this avoids attribute-type issues;

robustness => if self._cache is ever a StaticCache while config.is_encoder_decoder is True (stale/custom cache state), the old ternary can hit .self_attention_cache on the wrong type. The isinstance version is defensive and keeps behavior the same in expected cases.

tarekziade · 2026-03-02T13:59:14Z

run-slow: clvp, musicgen, musicgen_melody

github-actions · 2026-03-02T14:00:37Z

Workflow Run ⚙️

This comment contains run-slow, running the specified jobs:

models: ["models/clvp", "models/musicgen", "models/musicgen_melody"]
quantizations: []

github-actions · 2026-03-02T14:29:42Z

CI Results

Workflow Run ⚙️

Commit Info

Context	Commit	Description
RUN	2679192e	workflow commit (merge commit)
PR	51ab3d06	branch commit (from PR)
main	744639ab	base commit (on `main`)

✅ No failing test specific to this PR 🎉 👏 !

tarekziade · 2026-03-03T08:18:30Z

run-slow: clvp, musicgen, musicgen_melody

github-actions · 2026-03-03T08:19:49Z

Workflow Run ⚙️

This comment contains run-slow, running the specified jobs:

models: ["models/clvp", "models/musicgen", "models/musicgen_melody"]
quantizations: []

github-actions · 2026-03-03T08:55:31Z

CI Results

Workflow Run ⚙️

Commit Info

Context	Commit	Description
RUN	b35c6e35	workflow commit (merge commit)
PR	c10883fe	branch commit (from PR)
main	28d02a31	base commit (on `main`)

✅ No failing test specific to this PR 🎉 👏 !

Cyrilvallez

Nice! Would love to be able to remove the conditional inheritance though, I believe we cam by annotating a few more self?

Cyrilvallez

Alright, thanks for updating! Let's just avoid redefining a new GenerativePreTrainedModel class in generation/utils.py, everything could and should be inside _typing.py! Feel free to merge afterwards!

Add type declarations for mixin host-class attributes on GenerationMixin, class-level annotations for dynamically-set attributes on GenerationConfig, and fix minor typing issues in candidate_generator, watermarking, and stopping_criteria. Create _typing.py Protocol for documentation/reuse.

…with if guards - Replace assert narrowing with if guards (candidate_generator.py) - Add if-guard for model_kwargs None in _prepare_model_inputs and _maybe_initialize_input_ids_for_generation - Add if-guard for encoder_input_ids None checks - Fix save_directory.split() for PathLike by wrapping with str() - Fix tuple[Tensor] reassignment in compute_transition_scores - Add type annotation for self.model in WhisperTimeStampLogitsProcessor - Add type annotation for self.sequence_bias in SequenceBiasLogitsProcessor - Add no-assert rule to AGENTS.md typing strategy

tarekziade · 2026-03-04T15:35:09Z

run-slow: clvp, musicgen, musicgen_melody

github-actions · 2026-03-04T15:36:28Z

Workflow Run ⚙️

This comment contains run-slow, running the specified jobs:

models: ["models/clvp", "models/musicgen", "models/musicgen_melody"]
quantizations: []

github-actions · 2026-03-04T16:26:32Z

CI Results

Workflow Run ⚙️

Commit Info

Context	Commit	Description
RUN	3c3727ce	workflow commit (merge commit)
PR	881e25d7	branch commit (from PR)
main	fd6bc380	base commit (on `main`)

✅ No failing test specific to this PR 🎉 👏 !

github-actions · 2026-03-04T16:44:12Z

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

run-slow: clvp, musicgen, musicgen_melody

Improve LLM factory typing so generated model classes expose generation APIs for static analysis. - Add GenerationMixin import in the LLM registry module. - Introduce PreTrainedModelWithGenerationMixin as a typed base. - Update ModelAndTokenizer to return this mixed model type. Note: - Recent upstream PR introduced GenerativePreTrainedModel Protocol: huggingface/transformers#44233 - This appears to trigger analyzer warnings from protocol checks. - The warning behavior is suspected upstream typing friction and is temporarily ignored.

tarekziade requested review from Cyrilvallez and zucchini-nlp February 23, 2026 16:23

tarekziade self-assigned this Feb 23, 2026

tarekziade commented Feb 23, 2026

View reviewed changes

tarekziade force-pushed the tarekziade-typing-generation branch from 8c59847 to 1dfddc1 Compare February 23, 2026 19:16

zucchini-nlp reviewed Feb 24, 2026

View reviewed changes

tarekziade force-pushed the tarekziade-typing-generation branch 2 times, most recently from a0a0445 to a212761 Compare February 24, 2026 11:18

redpanda1995 approved these changes Feb 25, 2026

View reviewed changes

Cyrilvallez reviewed Mar 2, 2026

View reviewed changes

tarekziade force-pushed the tarekziade-typing-generation branch from 388547f to b3e6799 Compare March 2, 2026 12:44

tarekziade changed the title ~~chore(qa): Add type checking to src/transformers/generation~~ chore(typing): Add type checking to src/transformers/generation Mar 2, 2026

tarekziade mentioned this pull request Mar 3, 2026

chore(typing): Add type checking to src/transformers/quantizers #44412

Merged

Cyrilvallez reviewed Mar 4, 2026

View reviewed changes

Comment thread src/transformers/generation/utils.py Outdated

Comment thread src/transformers/generation/utils.py

Comment thread src/transformers/generation/utils.py Outdated

tarekziade force-pushed the tarekziade-typing-generation branch from 9c65644 to db703e7 Compare March 4, 2026 13:21

Cyrilvallez approved these changes Mar 4, 2026

View reviewed changes

Comment thread src/transformers/generation/utils.py Outdated

tarekziade added 6 commits March 4, 2026 16:20

extend make file

3f6fc2c

more tweaking

cefc9be

fix get word size

a36c0d4

date

d6ded3e

tarekziade added 13 commits March 4, 2026 16:20

cast value

c420924

@zucchini-nlp reviews

571dd6c

model_kwargs cannot be None, fix type annotations

b234885

removed redundancy

fdec9d8

bypass type checking to reduce code complexity

f55c86c

de-vectorize

fc056a2

unified _typing

c45d22d

explain why we pollute inheritance with a protocol

2756698

do not import torch for type checking, it breaks on CI logging usage

cbfd96c

Add check_types.py to wrap ty type checking

7917083

improved ty usage

d76c95f

get rid of complex inheritance and rely more on self annotations

d503f1a

combine in a single class

ba7ae36

tarekziade force-pushed the tarekziade-typing-generation branch from e21d1d5 to ba7ae36 Compare March 4, 2026 15:27

formatting

881e25d

Merge branch 'main' into tarekziade-typing-generation

050cd1d

Merge branch 'main' into tarekziade-typing-generation

20d5749

tarekziade enabled auto-merge (squash) March 4, 2026 16:47

Cyrilvallez disabled auto-merge March 4, 2026 17:24

Cyrilvallez merged commit a2da4b1 into main Mar 4, 2026
23 of 27 checks passed

Cyrilvallez deleted the tarekziade-typing-generation branch March 4, 2026 17:24

Conversation

tarekziade commented Feb 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

tarekziade Feb 23, 2026

Choose a reason for hiding this comment

Uh oh!

HuggingFaceDocBuilderDev commented Feb 23, 2026

Uh oh!

zucchini-nlp left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

tarekziade commented Feb 24, 2026

Uh oh!

Cyrilvallez left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Cyrilvallez Mar 2, 2026

Choose a reason for hiding this comment

Uh oh!

tarekziade Mar 2, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

tarekziade commented Mar 2, 2026

Uh oh!

github-actions Bot commented Mar 2, 2026

Uh oh!

github-actions Bot commented Mar 2, 2026

CI Results

Commit Info

Uh oh!

tarekziade commented Mar 3, 2026

Uh oh!

github-actions Bot commented Mar 3, 2026

Uh oh!

github-actions Bot commented Mar 3, 2026

CI Results

Commit Info

Uh oh!

Cyrilvallez left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Cyrilvallez left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

tarekziade commented Mar 4, 2026

Uh oh!

github-actions Bot commented Mar 4, 2026

Uh oh!

github-actions Bot commented Mar 4, 2026

CI Results

Commit Info

Uh oh!

github-actions Bot commented Mar 4, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

tarekziade commented Feb 23, 2026 •

edited

Loading