Repository: huggingface/transformers
Pull request: Expand model-structure lint rules with a fast AST-based, ruff-like framework

This draft expands `utils/check_modeling_structure.py` into a rule-driven linter for model code, with new checks and tests, while keeping runtime very low. 

Key features:

- The checker is intentionally AST-only (no heavy imports/execution), so it stays fast (<3 seconds against all models)
- It's built in a CI/PR-friendly way, so when devs proposes new models it's easy to provide feedback automatically
- It provides a nice CLI UI to explain every rules, like how Ruff does with `ruff rule XXX`
- it's agentic-friendly

See `utils/check_modeling_structure_rules.toml` for the set of rules

Changed files:
filename: Makefile
patch: @@ -1,7 +1,7 @@
 # make sure to test the local checkout in scripts and not the pre-installed one (don't use quotes!)
 export PYTHONPATH = src
 
-.PHONY: style check-repo fix-repo test test-examples benchmark
+.PHONY: style check-repo check-model-rules check-model-rules-pr check-model-rules-all fix-repo test test-examples benchmark
 
 check_dirs := examples tests src utils scripts benchmark benchmark_v2
 exclude_folders :=  ""
@@ -49,6 +49,18 @@ check-repo:
 	}
 
 
+check-model-rules:
+	python utils/check_modeling_structure.py --changed-only --base-ref origin/main
+
+
+check-model-rules-pr:
+	python utils/check_modeling_structure.py --changed-only --base-ref origin/main --github-annotations
+
+
+check-model-rules-all:
+	python utils/check_modeling_structure.py
+
+
 # Run all repo checks for which there is an automatic fix, most notably modular conversions
 # Note: each line is run in its own shell, and doing `-` before the command ignores the errors if any, continuing with next command
 fix-repo: style

filename: tests/repo_utils/test_check_modeling_structure.py
patch: @@ -0,0 +1,319 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import subprocess
+import sys
+import unittest
+from pathlib import Path
+from unittest.mock import patch
+
+
+git_repo_path = os.path.abspath(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
+sys.path.append(os.path.join(git_repo_path, "utils"))
+
+import check_modeling_structure as cms  # noqa: E402
+
+
+class CheckModelingStructureTest(unittest.TestCase):
+    # --- TRF001: config_class naming consistency (old TRF003) ---
+
+    def test_trf001_valid_config_class(self):
+        source = """
+class FooPreTrainedModel(PreTrainedModel):
+    config_class = FooConfig
+"""
+        file_path = Path("src/transformers/models/foo/modeling_foo.py")
+        violations = cms.analyze_file(file_path, source, enabled_rules={cms.TRF001})
+        trf001 = [v for v in violations if v.rule_id == cms.TRF001]
+        self.assertEqual(trf001, [])
+
+    def test_trf001_invalid_config_class(self):
+        source = """
+class FooPreTrainedModel(PreTrainedModel):
+    config_class = BarConfig
+"""
+        file_path = Path("src/transformers/models/foo/modeling_foo.py")
+        violations = cms.analyze_file(file_path, source, enabled_rules={cms.TRF001})
+        trf001 = [v for v in violations if v.rule_id == cms.TRF001]
+        self.assertEqual(len(trf001), 1)
+        self.assertIn("config_class is BarConfig, expected FooConfig", trf001[0].message)
+
+    # --- TRF002: base_model_prefix (old TRF004) ---
+
+    def test_trf002_valid_prefix(self):
+        source = """
+class FooPreTrainedModel(PreTrainedModel):
+    base_model_prefix = "model"
+"""
+        file_path = Path("src/transformers/models/foo/modeling_foo.py")
+        violations = cms.analyze_file(file_path, source, enabled_rules={cms.TRF002})
+        trf002 = [v for v in violations if v.rule_id == cms.TRF002]
+        self.assertEqual(trf002, [])
+
+    def test_trf002_invalid_empty_prefix(self):
+        source = """
+class FooPreTrainedModel(PreTrainedModel):
+    base_model_prefix = ""
+"""
+        file_path = Path("src/transformers/models/foo/modeling_foo.py")
+        violations = cms.analyze_file(file_path, source, enabled_rules={cms.TRF002})
+        trf002 = [v for v in violations if v.rule_id == cms.TRF002]
+        self.assertEqual(len(trf002), 1)
+        self.assertIn("non-empty canonical token", trf002[0].message)
+
+    # --- TRF003: capture_output enforcement (reworked old TRF005) ---
+
+    def test_trf003_flags_old_return_dict_branching(self):
+        source = """
+class FooPreTrainedModel:
+    pass
+
+class FooModel(FooPreTrainedModel):
+    def forward(self, x, return_dict=None):
+        if not return_dict:
+            return (x,)
+        return x
+"""
+        file_path = Path("src/transformers/models/foo/modeling_foo.py")
+        violations = cms.analyze_file(file_path, source, enabled_rules={cms.TRF003})
+        trf003 = [v for v in violations if v.rule_id == cms.TRF003]
+        self.assertEqual(len(trf003), 1)
+        self.assertIn("old return_dict branching pattern", trf003[0].message)
+
+    def test_trf003_allows_no_return_dict_arg(self):
+        source = """
+class FooPreTrainedModel:
+    pass
+
+class FooModel(FooPreTrainedModel):
+    def forward(self, x):
+        return x
+"""
+        file_path = Path("src/transformers/models/foo/modeling_foo.py")
+        violations = cms.analyze_file(file_path, source, enabled_rules={cms.TRF003})
+        trf003 = [v for v in violations if v.rule_id == cms.TRF003]
+        self.assertEqual(trf003, [])
+
+    def test_trf003_allows_return_dict_without_branching(self):
+        source = """
+class FooPreTrainedModel:
+    pass
+
+class FooModel(FooPreTrainedModel):
+    def forward(self, x, return_dict=None):
+        return x
+"""
+        file_path = Path("src/transformers/models/foo/modeling_foo.py")
+        violations = cms.analyze_file(file_path, source, enabled_rules={cms.TRF003})
+        trf003 = [v for v in violations if v.rule_id == cms.TRF003]
+        self.assertEqual(trf003, [])
+
+    # --- TRF004: tie_weights hard ban (reworked old TRF007) ---
+
+    def test_trf004_flags_any_tie_weights_override(self):
+        source = """
+class FooModel:
+    def tie_weights(self):
+        super().tie_weights()
+"""
+        file_path = Path("src/transformers/models/foo/modeling_foo.py")
+        violations = cms.analyze_file(file_path, source, enabled_rules={cms.TRF004})
+        trf004 = [v for v in violations if v.rule_id == cms.TRF004]
+        self.assertEqual(len(trf004), 1)
+        self.assertIn("overrides tie_weights", trf004[0].message)
+
+    def test_trf004_allows_no_tie_weights(self):
+        source = """
+class FooModel:
+    _tied_weights_keys = ["lm_head.weight"]
+"""
+        file_path = Path("src/transformers/models/foo/modeling_foo.py")
+        violations = cms.analyze_file(file_path, source, enabled_rules={cms.TRF004})
+        trf004 = [v for v in violations if v.rule_id == cms.TRF004]
+        self.assertEqual(trf004, [])
+
+    # --- TRF005: _no_split_modules (old TRF008) ---
+
+    def test_trf005_valid_no_split_modules(self):
+        source = """
+class FooModel:
+    _no_split_modules = ["FooDecoderLayer"]
+"""
+        file_path = Path("src/transformers/models/foo/modeling_foo.py")
+        violations = cms.analyze_file(file_path, source, enabled_rules={cms.TRF005})
+        trf005 = [v for v in violations if v.rule_id == cms.TRF005]
+        self.assertEqual(trf005, [])
+
+    def test_trf005_invalid_empty_string(self):
+        source = """
+class FooModel:
+    _no_split_modules = [""]
+"""
+        file_path = Path("src/transformers/models/foo/modeling_foo.py")
+        violations = cms.analyze_file(file_path, source, enabled_rules={cms.TRF005})
+        trf005 = [v for v in violations if v.rule_id == cms.TRF005]
+        self.assertEqual(len(trf005), 1)
+
+    # --- TRF006: cache args usage (old TRF010) ---
+
+    def test_trf006_catches_unused_cache_args(self):
+        source = """
+class FooPreTrainedModel:
+    pass
+
+class FooModel(FooPreTrainedModel):
+    def forward(self, hidden_states, past_key_value=None, use_cache=False):
+        return hidden_states
+"""
+        file_path = Path("src/transformers/models/foo/modeling_foo.py")
+        violations = cms.analyze_file(file_path, source, enabled_rules={cms.TRF006})
+        trf006 = [v for v in violations if v.rule_id == cms.TRF006]
+        self.assertEqual(len(trf006), 1)
+        self.assertIn("past_key_values/use_cache", trf006[0].message)
+
+    # --- TRF007: post_init order (old TRF011) ---
+
+    def test_trf007_flags_assignment_after_post_init(self):
+        source = """
+class FooPreTrainedModel:
+    pass
+
+class FooModel(FooPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.post_init()
+        self.proj = None
+"""
+        file_path = Path("src/transformers/models/foo/modeling_foo.py")
+        violations = cms.analyze_file(file_path, source, enabled_rules={cms.TRF007})
+        trf007 = [v for v in violations if v.rule_id == cms.TRF007]
+        self.assertEqual(len(trf007), 1)
+        self.assertIn("assigns self.* after self.post_init()", trf007[0].message)
+
+    def test_trf007_allows_post_init_at_end(self):
+        source = """
+class FooPreTrainedModel:
+    pass
+
+class FooModel(FooPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.proj = None
+        self.post_init()
+"""
+        file_path = Path("src/transformers/models/foo/modeling_foo.py")
+        violations = cms.analyze_file(file_path, source, enabled_rules={cms.TRF007})
+        trf007 = [v for v in violations if v.rule_id == cms.TRF007]
+        self.assertEqual(trf007, [])
+
+    # --- TRF009: cross-model imports (old TRF013) ---
+
+    def test_trf009_flags_cross_model_import_in_modeling_file(self):
+        source = """
+from transformers.models.llama.modeling_llama import LlamaAttention
+"""
+        file_path = Path("src/transformers/models/foo/modeling_foo.py")
+        violations = cms.analyze_file(file_path, source, enabled_rules={cms.TRF009})
+        trf009 = [v for v in violations if v.rule_id == cms.TRF009]
+        self.assertEqual(len(trf009), 1)
+        self.assertIn("imports implementation code from `llama`", trf009[0].message)
+
+    def test_trf009_allows_same_model_import_in_modeling_file(self):
+        source = """
+from .configuration_foo import FooConfig
+from transformers.models.foo.configuration_foo import FooConfig as FooConfigAlias
+"""
+        file_path = Path("src/transformers/models/foo/modeling_foo.py")
+        violations = cms.analyze_file(file_path, source, enabled_rules={cms.TRF009})
+        trf009 = [v for v in violations if v.rule_id == cms.TRF009]
+        self.assertEqual(trf009, [])
+
+    def test_trf009_ignores_modular_files(self):
+        source = """
+from transformers.models.llama.modeling_llama import LlamaAttention
+"""
+        file_path = Path("src/transformers/models/foo/modular_foo.py")
+        violations = cms.analyze_file(file_path, source, enabled_rules={cms.TRF009})
+        trf009 = [v for v in violations if v.rule_id == cms.TRF009]
+        self.assertEqual(trf009, [])
+
+    # --- Utility tests ---
+
+    @patch("check_modeling_structure.subprocess.run")
+    def test_get_changed_modeling_files_filters_non_model_files(self, mock_run):
+        mock_run.side_effect = [
+            subprocess.CompletedProcess(
+                args=["git", "diff"],
+                returncode=0,
+                stdout=(
+                    "src/transformers/models/foo/modeling_foo.py\n"
+                    "src/transformers/models/foo/modular_foo.py\n"
+                    "src/transformers/models/foo/configuration_foo.py\n"
+                    "docs/source/en/index.md\n"
+                ),
+                stderr="",
+            ),
+            subprocess.CompletedProcess(args=["git", "diff"], returncode=0, stdout="", stderr=""),
+            subprocess.CompletedProcess(args=["git", "diff", "--cached"], returncode=0, stdout="", stderr=""),
+            subprocess.CompletedProcess(args=["git", "ls-files"], returncode=0, stdout="", stderr=""),
+        ]
+        changed_files = cms.get_changed_modeling_files("origin/main")
+        self.assertEqual(
+            changed_files,
+            {
+                Path("src/transformers/models/foo/modeling_foo.py"),
+                Path("src/transformers/models/foo/modular_foo.py"),
+            },
+        )
+
+    @patch("check_modeling_structure.subprocess.run")
+    def test_get_changed_modeling_files_includes_uncommitted_worktree_changes(self, mock_run):
+        mock_run.side_effect = [
+            subprocess.CompletedProcess(args=["git", "diff"], returncode=0, stdout="", stderr=""),
+            subprocess.CompletedProcess(args=["git", "diff"], returncode=0, stdout="", stderr=""),
+            subprocess.CompletedProcess(
+                args=["git", "diff"],
+                returncode=0,
+                stdout="src/transformers/models/helium/modeling_helium.py\n",
+                stderr="",
+            ),
+            subprocess.CompletedProcess(
+                args=["git", "diff", "--cached"],
+                returncode=0,
+                stdout="src/transformers/models/foo/modular_foo.py\n",
+                stderr="",
+            ),
+            subprocess.CompletedProcess(
+                args=["git", "ls-files"],
+                returncode=0,
+                stdout=("src/transformers/models/bar/modeling_bar.py\ndocs/source/en/index.md\n"),
+                stderr="",
+            ),
+        ]
+
+        changed_files = cms.get_changed_modeling_files("origin/main")
+
+        self.assertEqual(
+            changed_files,
+            {
+                Path("src/transformers/models/helium/modeling_helium.py"),
+                Path("src/transformers/models/foo/modular_foo.py"),
+                Path("src/transformers/models/bar/modeling_bar.py"),
+            },
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

filename: utils/check_modeling_structure.py
patch: @@ -13,23 +13,146 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """
-Utility that ensures that modeling (and modular) files respect some important conventions we have in Transformers.
+Lint modeling and modular files under ``src/transformers/models`` for structural conventions.
+
+How rule registration works
+---------------------------
+- Rule metadata lives in ``utils/check_modeling_structure_rules.toml``.
+- Executable TRF rules are auto-discovered from functions named ``trfXXX_*``.
+- The ``trfXXX`` prefix becomes the rule id (for example ``trf003_check_...`` -> ``TRF003``).
+- Every discovered rule must have a matching entry in the TOML file, and every TOML rule must have a matching
+  ``trfXXX_*`` function. Import-time validation fails if either side is missing.
+- Rule functions receive ``rule_id`` explicitly and should use it for suppression checks and violation messages.
+- Suppressions use ``# trf-ignore: TRFXXX`` on the same line or the line immediately above the flagged construct.
+
+How to add a new TRF rule
+-------------------------
+1. Add a ``[rules.TRFXXX]`` entry to ``utils/check_modeling_structure_rules.toml``.
+2. Fill in ``description``, ``default_enabled``, ``explanation.what_it_does``, ``explanation.why_bad``,
+   ``explanation.bad_example``, and ``explanation.good_example``. Optional model-level exceptions go in
+   ``allowlist_models``.
+3. Implement a new function named ``trfXXX_<descriptive_name>`` with signature
+   ``(tree, file_path, source_lines, violations, rule_id) -> list[Violation]``.
+4. Use ``rule_id`` instead of hardcoding ``"TRFXXX"`` inside the check.
+5. Add or update focused tests in ``tests/repo_utils/test_check_modeling_structure.py``.
+
+CLI usage
+---------
+- ``python utils/check_modeling_structure.py``: check all modeling and modular files.
+- ``python utils/check_modeling_structure.py --changed-only --base-ref origin/main``: only check files changed
+  against a git base ref, plus local staged, unstaged, and untracked modeling files.
+- ``python utils/check_modeling_structure.py --list-rules``: print all available TRF rules and their default state.
+- ``python utils/check_modeling_structure.py --rule TRF001``: show the detailed documentation for one rule from the
+  TOML file.
+- ``python utils/check_modeling_structure.py --enable-rules TRF003``: enable additional rules on top of the defaults.
+- ``python utils/check_modeling_structure.py --enable-all-trf-rules``: enable every TRF rule, including ones disabled
+  by default.
+- ``python utils/check_modeling_structure.py --github-annotations``: emit GitHub Actions error annotations.
 """
 
+import argparse
 import ast
+import subprocess
 import sys
+from collections.abc import Callable
+from contextlib import nullcontext
+from dataclasses import dataclass
 from pathlib import Path
 
 from rich import print
+from rich.console import Console
+
+
+try:
+    import tomllib  # Python >= 3.11
+except ModuleNotFoundError:
+    import tomli as tomllib  # Python 3.10 fallback
 
 
 MODELS_ROOT = Path("src/transformers/models")
 MODELING_PATTERNS = ("modeling_*.py", "modular_*.py")
+RULE_SPECS_PATH = Path(__file__).with_name("check_modeling_structure_rules.toml")
+
+
+def _load_rule_specs() -> dict[str, dict]:
+    data = tomllib.loads(RULE_SPECS_PATH.read_text(encoding="utf-8"))
+    rules = data.get("rules")
+    if not isinstance(rules, dict):
+        raise ValueError(f"Invalid rule spec file: missing [rules] table in {RULE_SPECS_PATH}")
+
+    required_explanation_keys = {"what_it_does", "why_bad", "bad_example", "good_example"}
+    specs: dict[str, dict] = {}
+    for rule_id, spec in rules.items():
+        if not isinstance(spec, dict):
+            raise ValueError(f"Invalid rule spec for {rule_id}: expected table")
+        description = spec.get("description")
+        default_enabled = spec.get("default_enabled")
+        explanation = spec.get("explanation")
+        if not isinstance(description, str) or not description.strip():
+            raise ValueError(f"Invalid rule spec for {rule_id}: missing non-empty description")
+        if not isinstance(default_enabled, bool):
+            raise ValueError(f"Invalid rule spec for {rule_id}: default_enabled must be bool")
+        if not isinstance(explanation, dict) or not required_explanation_keys.issubset(explanation):
+            raise ValueError(f"Invalid rule spec for {rule_id}: incomplete explanation block")
+        if any(not isinstance(explanation[key], str) for key in required_explanation_keys):
+            raise ValueError(f"Invalid rule spec for {rule_id}: explanation values must be strings")
+
+        allowlist_models = spec.get("allowlist_models", [])
+        if not isinstance(allowlist_models, list) or any(not isinstance(item, str) for item in allowlist_models):
+            raise ValueError(f"Invalid rule spec for {rule_id}: allowlist_models must be list[str]")
+
+        specs[rule_id] = {
+            "description": description,
+            "default_enabled": default_enabled,
+            "explanation": explanation,
+            "allowlist_models": set(allowlist_models),
+        }
+
+    return specs
+
+
+TRF_RULE_SPECS = _load_rule_specs()
+TRF_RULES = {rule_id: spec["description"] for rule_id, spec in TRF_RULE_SPECS.items()}
+DEFAULT_ENABLED_TRF_RULES = {rule_id for rule_id, spec in TRF_RULE_SPECS.items() if spec["default_enabled"]}
+# Model-directory baseline allowlists for existing legacy exceptions.
+# Keep each set as small as possible and remove entries when those models are migrated.
+TRF_MODEL_DIR_ALLOWLISTS = {
+    rule_id: spec["allowlist_models"] for rule_id, spec in TRF_RULE_SPECS.items() if spec["allowlist_models"]
+}
+CONSOLE = Console(stderr=True)
+
 
+@dataclass(frozen=True)
+class Violation:
+    file_path: Path
+    line_number: int
+    message: str
+    rule_id: str | None = None
 
-def iter_modeling_files():
-    for pattern in MODELING_PATTERNS:
-        yield from MODELS_ROOT.rglob(pattern)
+
+def _validate_rule_ids(rule_ids: set[str]) -> set[str]:
+    unknown = sorted(rule_id for rule_id in rule_ids if rule_id not in TRF_RULES)
+    if unknown:
+        raise ValueError(f"Unknown rule id(s): {', '.join(unknown)}. Valid rules: {', '.join(sorted(TRF_RULES))}")
+    return rule_ids
+
+
+def _rule_id_from_check_name(name: str) -> str | None:
+    prefix, _, _ = name.partition("_")
+    if len(prefix) != 6 or not prefix.startswith("trf") or not prefix[3:].isdigit():
+        return None
+    return prefix.upper()
+
+
+def iter_modeling_files(paths: set[Path] | None = None):
+    if paths is None:
+        for pattern in MODELING_PATTERNS:
+            yield from MODELS_ROOT.rglob(pattern)
+        return
+
+    for path in sorted(paths):
+        if path.exists():
+            yield path
 
 
 def colored_error_message(file_path: str, line_number: int, message: str) -> str:
@@ -48,7 +171,47 @@ def full_name(node: ast.AST):
         raise ValueError("Not a Name or Attribute node")
 
 
-def check_init_weights(node: ast.AST, violations: list[str], file_path: str) -> list[str]:
+def _simple_name(name: str) -> str:
+    return name.split(".")[-1]
+
+
+def _model_dir_name(file_path: Path) -> str | None:
+    try:
+        relative = file_path.resolve().relative_to(MODELS_ROOT.resolve())
+    except ValueError:
+        try:
+            relative = file_path.relative_to(MODELS_ROOT)
+        except ValueError:
+            return None
+    if len(relative.parts) < 2:
+        return None
+    return relative.parts[0]
+
+
+def _is_rule_allowlisted_for_file(rule_id: str, file_path: Path) -> bool:
+    model_name = _model_dir_name(file_path)
+    if model_name is None:
+        return False
+    return model_name in TRF_MODEL_DIR_ALLOWLISTS.get(rule_id, set())
+
+
+def _known_model_dirs() -> set[str]:
+    return {path.name for path in MODELS_ROOT.iterdir() if path.is_dir()}
+
+
+def _has_rule_suppression(lines: list[str], rule_id: str, line_number: int) -> bool:
+    if line_number <= 0:
+        return False
+
+    token = f"trf-ignore: {rule_id}".lower()
+    candidate_indexes = (line_number - 1, line_number - 2)
+    for idx in candidate_indexes:
+        if 0 <= idx < len(lines) and token in lines[idx].lower():
+            return True
+    return False
+
+
+def check_init_weights(node: ast.AST, violations: list[Violation], file_path: Path) -> list[Violation]:
     """
     Check that `_init_weights` correctly use `init.(...)` patterns to init the weights in-place. This is very important,
     as we rely on the internal flag set on the parameters themselves to check if they need to be re-init or not.
@@ -70,7 +233,7 @@ def check_init_weights(node: ast.AST, violations: list[str], file_path: str) ->
                         "`_init_weights(self, module)` uses an in-place operation on a module's weight. Please use the "
                         "`init` functions primitives instead, usually imported as `from ... import initialization as init`"
                     )
-                    violations.append(colored_error_message(file_path, sub_node.lineno, error_msg))
+                    violations.append(Violation(file_path=file_path, line_number=sub_node.lineno, message=error_msg))
 
     return violations
 
@@ -98,7 +261,7 @@ def is_super_method_call(node: ast.AST, method: str) -> bool:
     )
 
 
-def check_post_init(node: ast.AST, violations: list[str], file_path: str) -> list[str]:
+def check_post_init(node: ast.AST, violations: list[Violation], file_path: Path) -> list[Violation]:
     """
     Check that `self.post_init()` is correctly called at the end of `__init__` for all `PreTrainedModel`s. This is
     very important as we need to do some processing there.
@@ -118,32 +281,722 @@ def check_post_init(node: ast.AST, violations: list[str], file_path: str) -> lis
                 # If we did not break, `post_init` was never called
                 else:
                     error_msg = f"`__init__` of {node.name} does not call `self.post_init`"
-                    violations.append(colored_error_message(file_path, sub_node.lineno, error_msg))
+                    violations.append(Violation(file_path=file_path, line_number=sub_node.lineno, message=error_msg))
                 break
 
     return violations
 
 
-def main():
-    violations: list[str] = []
+def _collect_class_bases(tree: ast.Module) -> dict[str, list[str]]:
+    class_to_bases: dict[str, list[str]] = {}
+    for node in tree.body:
+        if not isinstance(node, ast.ClassDef):
+            continue
+        base_names = []
+        for base in node.bases:
+            try:
+                base_names.append(full_name(base))
+            except ValueError:
+                continue
+        class_to_bases[node.name] = base_names
+    return class_to_bases
 
-    for file_path in iter_modeling_files():
-        try:
-            text = file_path.read_text(encoding="utf-8")
-            tree = ast.parse(text, filename=str(file_path))
-        except Exception as exc:
-            violations.append(f"{file_path}: failed to parse ({exc}).")
+
+def _inherits_pretrained_model(
+    class_name: str, class_to_bases: dict[str, list[str]], visiting: set[str] | None = None
+) -> bool:
+    if visiting is None:
+        visiting = set()
+    if class_name in visiting:
+        return False
+    visiting.add(class_name)
+
+    for base_name in class_to_bases.get(class_name, []):
+        simple_base_name = _simple_name(base_name)
+        if simple_base_name.endswith("PreTrainedModel"):
+            return True
+        if simple_base_name in class_to_bases and _inherits_pretrained_model(
+            simple_base_name, class_to_bases, visiting
+        ):
+            return True
+    return False
+
+
+def _get_class_assignments(class_node: ast.ClassDef) -> dict[str, ast.AST]:
+    assignments: dict[str, ast.AST] = {}
+    for item in class_node.body:
+        if isinstance(item, ast.Assign) and len(item.targets) == 1 and isinstance(item.targets[0], ast.Name):
+            assignments[item.targets[0].id] = item.value
+        elif isinstance(item, ast.AnnAssign) and isinstance(item.target, ast.Name) and item.value is not None:
+            assignments[item.target.id] = item.value
+    return assignments
+
+
+def _class_methods(class_node: ast.ClassDef) -> dict[str, ast.FunctionDef]:
+    return {item.name: item for item in class_node.body if isinstance(item, ast.FunctionDef)}
+
+
+def trf001_check_config_class_consistency(
+    tree: ast.Module, file_path: Path, source_lines: list[str], violations: list[Violation], rule_id: str
+) -> list[Violation]:
+    class_to_bases = _collect_class_bases(tree)
+    for node in tree.body:
+        if not isinstance(node, ast.ClassDef):
+            continue
+        if not node.name.endswith("PreTrainedModel"):
+            continue
+        if not _inherits_pretrained_model(node.name, class_to_bases):
+            continue
+        if _has_rule_suppression(source_lines, rule_id, node.lineno):
+            continue
+
+        assignments = _get_class_assignments(node)
+        config_value = assignments.get("config_class")
+        if config_value is None:
+            continue
+        if not isinstance(config_value, (ast.Name, ast.Attribute)):
             continue
 
-        for node in ast.walk(tree):
-            violations = check_init_weights(node, violations, file_path)
-            violations = check_post_init(node, violations, file_path)
+        config_name = _simple_name(full_name(config_value))
+        expected = f"{node.name.removesuffix('PreTrainedModel')}Config"
+        if config_name != expected:
+            violations.append(
+                Violation(
+                    file_path=file_path,
+                    line_number=getattr(config_value, "lineno", node.lineno),
+                    rule_id=rule_id,
+                    message=(f"{rule_id}: {node.name}.config_class is {config_name}, expected {expected}."),
+                )
+            )
+
+    return violations
+
+
+def trf002_check_base_model_prefix(
+    tree: ast.Module, file_path: Path, source_lines: list[str], violations: list[Violation], rule_id: str
+) -> list[Violation]:
+    class_to_bases = _collect_class_bases(tree)
+    for node in tree.body:
+        if not isinstance(node, ast.ClassDef):
+            continue
+        if not _inherits_pretrained_model(node.name, class_to_bases):
+            continue
+        if _has_rule_suppression(source_lines, rule_id, node.lineno):
+            continue
+
+        assignments = _get_class_assignments(node)
+        prefix_value = assignments.get("base_model_prefix")
+        if prefix_value is None:
+            continue
+        if not (isinstance(prefix_value, ast.Constant) and isinstance(prefix_value.value, str)):
+            violations.append(
+                Violation(
+                    file_path=file_path,
+                    line_number=getattr(prefix_value, "lineno", node.lineno),
+                    rule_id=rule_id,
+                    message=f"{rule_id}: {node.name}.base_model_prefix should be a string literal.",
+                )
+            )
+            continue
+        if not prefix_value.value.strip() or any(char.isspace() for char in prefix_value.value):
+            violations.append(
+                Violation(
+                    file_path=file_path,
+                    line_number=getattr(prefix_value, "lineno", node.lineno),
+                    rule_id=rule_id,
+                    message=f"{rule_id}: {node.name}.base_model_prefix should be a non-empty canonical token.",
+                )
+            )
+
+    return violations
+
+
+def _function_argument_names(function_node: ast.FunctionDef) -> set[str]:
+    names = {arg.arg for arg in function_node.args.args}
+    names.update(arg.arg for arg in function_node.args.kwonlyargs)
+    if function_node.args.vararg is not None:
+        names.add(function_node.args.vararg.arg)
+    if function_node.args.kwarg is not None:
+        names.add(function_node.args.kwarg.arg)
+    return names
+
+
+def _function_uses_name(function_node: ast.FunctionDef, name: str) -> bool:
+    return any(
+        isinstance(node, ast.Name) and node.id == name and isinstance(node.ctx, ast.Load)
+        for node in ast.walk(function_node)
+    )
+
+
+def _has_return_dict_branching(function_node: ast.FunctionDef) -> bool:
+    """Detect the old 'if not return_dict: return (tuple,)' pattern."""
+    for node in ast.walk(function_node):
+        if not isinstance(node, ast.If):
+            continue
+        # Match: `if not return_dict:`
+        test = node.test
+        if isinstance(test, ast.UnaryOp) and isinstance(test.op, ast.Not):
+            operand = test.operand
+            if isinstance(operand, ast.Name) and operand.id == "return_dict":
+                return True
+        # Match: `if return_dict is not None:` or `if return_dict:`
+        if isinstance(test, ast.Name) and test.id == "return_dict":
+            return True
+        if isinstance(test, ast.Compare) and isinstance(test.left, ast.Name) and test.left.id == "return_dict":
+            return True
+    return False
+
+
+def trf003_check_return_dict_usage(
+    tree: ast.Module, file_path: Path, source_lines: list[str], violations: list[Violation], rule_id: str
+) -> list[Violation]:
+    class_to_bases = _collect_class_bases(tree)
+    for node in tree.body:
+        if not isinstance(node, ast.ClassDef):
+            continue
+        if not _inherits_pretrained_model(node.name, class_to_bases):
+            continue
+        if _has_rule_suppression(source_lines, rule_id, node.lineno):
+            continue
+
+        forward_method = _class_methods(node).get("forward")
+        if forward_method is None:
+            continue
+        if "return_dict" not in _function_argument_names(forward_method):
+            continue
+        if not _has_return_dict_branching(forward_method):
+            continue
+
+        violations.append(
+            Violation(
+                file_path=file_path,
+                line_number=forward_method.lineno,
+                rule_id=rule_id,
+                message=(
+                    f"{rule_id}: {node.name}.forward uses old return_dict branching pattern. "
+                    "Use @can_return_tuple or @capture_output decorator instead."
+                ),
+            )
+        )
+
+    return violations
+
+
+def trf004_check_tie_weights_ban(
+    tree: ast.Module, file_path: Path, source_lines: list[str], violations: list[Violation], rule_id: str
+) -> list[Violation]:
+    for node in tree.body:
+        if not isinstance(node, ast.ClassDef):
+            continue
+        if _has_rule_suppression(source_lines, rule_id, node.lineno):
+            continue
+
+        tie_weights = _class_methods(node).get("tie_weights")
+        if tie_weights is None:
+            continue
+
+        violations.append(
+            Violation(
+                file_path=file_path,
+                line_number=tie_weights.lineno,
+                rule_id=rule_id,
+                message=(
+                    f"{rule_id}: {node.name} overrides tie_weights. "
+                    "Use _tied_weights_keys class attribute to declare tied weights instead."
+                ),
+            )
+        )
+
+    return violations
+
+
+def trf005_check_no_split_modules_shape(
+    tree: ast.Module, file_path: Path, source_lines: list[str], violations: list[Violation], rule_id: str
+) -> list[Violation]:
+    for node in tree.body:
+        if not isinstance(node, ast.ClassDef):
+            continue
+        if _has_rule_suppression(source_lines, rule_id, node.lineno):
+            continue
+
+        assignments = _get_class_assignments(node)
+        value = assignments.get("_no_split_modules")
+        if value is None:
+            continue
+        # _no_split_modules = None is valid (means no modules to keep unsplit)
+        if isinstance(value, ast.Constant) and value.value is None:
+            continue
+        if not isinstance(value, (ast.List, ast.Tuple)):
+            violations.append(
+                Violation(
+                    file_path=file_path,
+                    line_number=getattr(value, "lineno", node.lineno),
+                    rule_id=rule_id,
+                    message=f"{rule_id}: {node.name}._no_split_modules should be a list or tuple of strings.",
+                )
+            )
+            continue
+
+        if any(
+            not isinstance(element, ast.Constant) or not isinstance(element.value, str) or not element.value
+            for element in value.elts
+        ):
+            violations.append(
+                Violation(
+                    file_path=file_path,
+                    line_number=getattr(value, "lineno", node.lineno),
+                    rule_id=rule_id,
+                    message=f"{rule_id}: {node.name}._no_split_modules should contain non-empty strings only.",
+                )
+            )
+
+    return violations
+
+
+def trf006_check_cache_argument_usage(
+    tree: ast.Module, file_path: Path, source_lines: list[str], violations: list[Violation], rule_id: str
+) -> list[Violation]:
+    for node in tree.body:
+        if not isinstance(node, ast.ClassDef):
+            continue
+        if _has_rule_suppression(source_lines, rule_id, node.lineno):
+            continue
+
+        forward_method = _class_methods(node).get("forward")
+        if forward_method is None:
+            continue
+
+        arg_names = _function_argument_names(forward_method)
+        cache_state_args = {"past_key_values", "past_key_value"}
+        has_cache_state_arg = bool(arg_names.intersection(cache_state_args))
+        if not has_cache_state_arg:
+            continue
+
+        if "use_cache" in arg_names and _function_uses_name(forward_method, "use_cache"):
+            continue
+        if any(_function_uses_name(forward_method, arg_name) for arg_name in cache_state_args):
+            continue
+
+        violations.append(
+            Violation(
+                file_path=file_path,
+                line_number=forward_method.lineno,
+                rule_id=rule_id,
+                message=(
+                    f"{rule_id}: {node.name}.forward exposes past_key_values/use_cache but does not reference them."
+                ),
+            )
+        )
+
+    return violations
+
+
+def trf009_check_single_file_policy_imports(
+    tree: ast.Module, file_path: Path, source_lines: list[str], violations: list[Violation], rule_id: str
+) -> list[Violation]:
+    if not file_path.name.startswith("modeling_"):
+        return violations
+
+    current_model = _model_dir_name(file_path)
+    if current_model is None:
+        return violations
+
+    known_models = _known_model_dirs()
+
+    for node in ast.walk(tree):
+        if isinstance(node, ast.ImportFrom):
+            if node.module is None:
+                continue
+            if _has_rule_suppression(source_lines, rule_id, node.lineno):
+                continue
+
+            imported_model = None
+            if node.level == 0 and node.module.startswith("transformers.models."):
+                remaining = node.module.split("transformers.models.", 1)[1]
+                imported_model = remaining.split(".", 1)[0]
+            elif node.level >= 2:
+                root_name = node.module.split(".", 1)[0]
+                if root_name in known_models:
+                    imported_model = root_name
+
+            if imported_model is None or imported_model in {current_model, "auto"}:
+                continue
+
+            violations.append(
+                Violation(
+                    file_path=file_path,
+                    line_number=node.lineno,
+                    rule_id=rule_id,
+                    message=(
+                        f"{rule_id}: {file_path.name} imports implementation code from "
+                        f"`{imported_model}`. Keep model logic local to a single modeling file."
+                    ),
+                )
+            )
+            continue
+
+        if isinstance(node, ast.Import):
+            if _has_rule_suppression(source_lines, rule_id, node.lineno):
+                continue
+
+            for alias in node.names:
+                if not alias.name.startswith("transformers.models."):
+                    continue
+                remaining = alias.name.split("transformers.models.", 1)[1]
+                imported_model = remaining.split(".", 1)[0]
+                if imported_model in {current_model, "auto"}:
+                    continue
+                violations.append(
+                    Violation(
+                        file_path=file_path,
+                        line_number=node.lineno,
+                        rule_id=rule_id,
+                        message=(
+                            f"{rule_id}: {file_path.name} imports implementation code from "
+                            f"`{imported_model}`. Keep model logic local to a single modeling file."
+                        ),
+                    )
+                )
+
+    return violations
+
+
+def trf007_check_post_init_order(
+    tree: ast.Module, file_path: Path, source_lines: list[str], violations: list[Violation], rule_id: str
+) -> list[Violation]:
+    class_to_bases = _collect_class_bases(tree)
+    for node in tree.body:
+        if not isinstance(node, ast.ClassDef):
+            continue
+        if not _inherits_pretrained_model(node.name, class_to_bases):
+            continue
+        if _has_rule_suppression(source_lines, rule_id, node.lineno):
+            continue
+
+        init_method = _class_methods(node).get("__init__")
+        if init_method is None:
+            continue
+
+        post_init_index = None
+        for index, statement in enumerate(init_method.body):
+            if isinstance(statement, ast.Expr) and is_self_method_call(statement.value, "post_init"):
+                post_init_index = index
+                break
+        if post_init_index is None:
+            continue
+
+        trailing_statements = init_method.body[post_init_index + 1 :]
+        has_trailing_self_assignments = any(
+            isinstance(statement, (ast.Assign, ast.AnnAssign))
+            and any(
+                isinstance(target, ast.Attribute) and isinstance(target.value, ast.Name) and target.value.id == "self"
+                for target in (statement.targets if isinstance(statement, ast.Assign) else [statement.target])
+            )
+            for statement in trailing_statements
+        )
+        if not has_trailing_self_assignments:
+            continue
+
+        violations.append(
+            Violation(
+                file_path=file_path,
+                line_number=init_method.lineno,
+                rule_id=rule_id,
+                message=f"{rule_id}: {node.name} assigns self.* after self.post_init() in __init__.",
+            )
+        )
+
+    return violations
+
+
+def trf008_check_doc_decorator_usage(
+    tree: ast.Module, file_path: Path, source_lines: list[str], violations: list[Violation], rule_id: str
+) -> list[Violation]:
+    class_to_bases = _collect_class_bases(tree)
+    for node in tree.body:
+        if not isinstance(node, ast.ClassDef):
+            continue
+        if not _inherits_pretrained_model(node.name, class_to_bases):
+            continue
+        if _has_rule_suppression(source_lines, rule_id, node.lineno):
+            continue
+
+        for decorator in node.decorator_list:
+            if not (
+                isinstance(decorator, ast.Call)
+                and isinstance(decorator.func, (ast.Name, ast.Attribute))
+                and _simple_name(full_name(decorator.func)) == "add_start_docstrings"
+            ):
+                continue
+            has_non_empty_string_arg = any(
+                isinstance(arg, ast.Constant) and isinstance(arg.value, str) and arg.value.strip()
+                for arg in decorator.args
+            )
+            if has_non_empty_string_arg:
+                continue
+
+            violations.append(
+                Violation(
+                    file_path=file_path,
+                    line_number=getattr(decorator, "lineno", node.lineno),
+                    rule_id=rule_id,
+                    message=f"{rule_id}: {node.name} uses add_start_docstrings without non-empty docstring arguments.",
+                )
+            )
+            break
+
+    return violations
+
+
+def _is_modeling_candidate(path: Path) -> bool:
+    return path.suffix == ".py" and path.name.startswith(("modeling_", "modular_")) and MODELS_ROOT in path.parents
+
+
+def _git_name_only(command: list[str]) -> list[str]:
+    result = subprocess.run(command, capture_output=True, text=True, check=False)
+    if result.returncode != 0:
+        return []
+    return [line for line in result.stdout.splitlines() if line.strip()]
+
+
+def _git_diff(base_ref: str, triple_dot: bool) -> list[str]:
+    diff_operator = "..." if triple_dot else ".."
+    range_ref = f"{base_ref}{diff_operator}HEAD"
+    return _git_name_only(["git", "diff", "--name-only", "--diff-filter=ACMR", range_ref])
+
+
+def _git_worktree_changes() -> set[Path]:
+    changed_paths = set(_git_name_only(["git", "diff", "--name-only", "--diff-filter=ACMR"]))
+    changed_paths.update(_git_name_only(["git", "diff", "--cached", "--name-only", "--diff-filter=ACMR"]))
+    changed_paths.update(_git_name_only(["git", "ls-files", "--others", "--exclude-standard"]))
+    return {Path(path_str) for path_str in changed_paths}
+
+
+def get_changed_modeling_files(base_ref: str) -> set[Path]:
+    changed_paths = _git_diff(base_ref, triple_dot=True)
+    if not changed_paths:
+        changed_paths = _git_diff(base_ref, triple_dot=False)
+
+    filtered_paths: set[Path] = set()
+    for path in {Path(path_str) for path_str in changed_paths}.union(_git_worktree_changes()):
+        if _is_modeling_candidate(path):
+            filtered_paths.add(path)
+    return filtered_paths
+
+
+# Auto-discover check functions by convention: any function named `trfXXX_*` is
+# registered as the checker for rule TRFXXX. To add a new rule, just define a
+# function following this naming convention with the standard signature:
+#   (tree: ast.Module, file_path: Path, source_lines: list[str], violations: list[Violation], rule_id: str)
+#   -> list[Violation]
+def _build_rule_checks() -> dict[str, Callable[[ast.Module, Path, list[str], list[Violation], str], list[Violation]]]:
+    checks: dict[str, Callable[[ast.Module, Path, list[str], list[Violation], str], list[Violation]]] = {}
+    for name, obj in globals().items():
+        if not callable(obj):
+            continue
+        rule_id = _rule_id_from_check_name(name)
+        if rule_id is None:
+            continue
+        if rule_id not in TRF_RULE_SPECS:
+            raise ValueError(f"Missing rule spec for discovered check function {name} ({rule_id}).")
+        checks[rule_id] = obj
+
+    missing_checks = sorted(set(TRF_RULE_SPECS) - set(checks))
+    if missing_checks:
+        raise ValueError(f"Missing check function(s) for rule id(s): {', '.join(missing_checks)}")
+    return dict(sorted(checks.items()))
+
+
+TRF_RULE_CHECKS = _build_rule_checks()
+globals().update({rule_id: rule_id for rule_id in TRF_RULE_CHECKS})
+
+
+def analyze_file(file_path: Path, text: str, enabled_rules: set[str] | None = None) -> list[Violation]:
+    if enabled_rules is None:
+        enabled_rules = DEFAULT_ENABLED_TRF_RULES
+
+    violations: list[Violation] = []
+    source_lines = text.splitlines()
+    tree = ast.parse(text, filename=str(file_path))
+
+    for node in ast.walk(tree):
+        violations = check_init_weights(node, violations, file_path)
+        violations = check_post_init(node, violations, file_path)
+
+    for rule_id, check_fn in TRF_RULE_CHECKS.items():
+        if rule_id in enabled_rules:
+            violations = check_fn(tree, file_path, source_lines, violations, rule_id)
+
+    return [
+        violation
+        for violation in violations
+        if not (
+            violation.rule_id is not None and _is_rule_allowlisted_for_file(violation.rule_id, violation.file_path)
+        )
+    ]
+
+
+def format_violation(violation: Violation) -> str:
+    return colored_error_message(str(violation.file_path), violation.line_number, violation.message)
+
+
+def emit_violation(violation: Violation, github_annotations: bool):
+    if github_annotations:
+        print(
+            f"::error file={violation.file_path},line={violation.line_number}::{violation.message}",
+            file=sys.stderr,
+        )
+        return
+
+    print(format_violation(violation), file=sys.stderr)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--changed-only",
+        action="store_true",
+        help="Only check changed modeling/modular files compared to --base-ref, plus local worktree changes.",
+    )
+    parser.add_argument(
+        "--base-ref",
+        default="origin/main",
+        help="Base git ref used with --changed-only (default: origin/main).",
+    )
+    parser.add_argument(
+        "--github-annotations",
+        action="store_true",
+        help="Emit GitHub Actions annotation format output.",
+    )
+    parser.add_argument(
+        "--no-progress",
+        action="store_true",
+        help="Disable interactive progress animation.",
+    )
+    parser.add_argument(
+        "--enable-all-trf-rules",
+        action="store_true",
+        help="Enable all TRF rules (defaults already enable most).",
+    )
+    parser.add_argument(
+        "--enable-rules",
+        default="",
+        help="Comma-separated TRF rule ids to enable in addition to defaults (e.g. TRF001,TRF002).",
+    )
+    parser.add_argument(
+        "--list-rules",
+        action="store_true",
+        help="List available TRF rules and exit.",
+    )
+    parser.add_argument(
+        "--rule",
+        default="",
+        help="Show detailed docs for one rule id (e.g. TRF001) and exit.",
+    )
+    return parser.parse_args()
+
+
+def should_show_progress(args: argparse.Namespace) -> bool:
+    return (not args.no_progress) and (not args.github_annotations) and sys.stderr.isatty()
+
+
+def resolve_enabled_rules(args: argparse.Namespace) -> set[str]:
+    if args.enable_all_trf_rules:
+        return _validate_rule_ids(set(TRF_RULES))
+
+    enabled_rules = set(DEFAULT_ENABLED_TRF_RULES)
+    if args.enable_rules.strip():
+        enabled_rules.update(rule_id.strip() for rule_id in args.enable_rules.split(",") if rule_id.strip())
+    return _validate_rule_ids(enabled_rules)
+
+
+def format_rule_summary(rule_id: str) -> str:
+    spec = TRF_RULE_SPECS[rule_id]
+    default_label = "enabled" if spec["default_enabled"] else "disabled"
+    return f"{rule_id}: {spec['description']} (default: {default_label})"
+
+
+def format_rule_details(rule_id: str) -> str:
+    spec = TRF_RULE_SPECS[rule_id]
+    explanation = spec["explanation"]
+    default_label = "yes" if spec["default_enabled"] else "no"
+    return "\n".join(
+        [
+            rule_id,
+            "",
+            f"Summary: {spec['description']}",
+            f"Default enabled: {default_label}",
+            "",
+            "What it does",
+            "",
+            explanation["what_it_does"],
+            "",
+            "Why is this bad?",
+            "",
+            explanation["why_bad"],
+            "",
+            "Example",
+            "",
+            explanation["bad_example"],
+            "",
+            "Use instead:",
+            "",
+            explanation["good_example"],
+        ]
+    )
+
+
+def maybe_handle_rule_docs_cli(args: argparse.Namespace) -> bool:
+    if args.list_rules:
+        for rule_id in sorted(TRF_RULE_SPECS):
+            print(format_rule_summary(rule_id))
+        return True
+
+    if args.rule:
+        rule_id = args.rule.strip().upper()
+        _validate_rule_ids({rule_id})
+        print(format_rule_details(rule_id))
+        return True
+
+    return False
+
+
+def main() -> int:
+    args = parse_args()
+    if maybe_handle_rule_docs_cli(args):
+        return 0
+
+    violations: list[Violation] = []
+    enabled_rules = resolve_enabled_rules(args)
+    selected_paths = get_changed_modeling_files(args.base_ref) if args.changed_only else None
+
+    modeling_files = list(iter_modeling_files(selected_paths))
+
+    show_progress = should_show_progress(args)
+    status_ctx = (
+        CONSOLE.status(f"[bold blue]Checking modeling structure ({len(modeling_files)} files)...[/bold blue]")
+        if show_progress
+        else nullcontext()
+    )
+
+    with status_ctx:
+        for file_path in modeling_files:
+            try:
+                text = file_path.read_text(encoding="utf-8")
+                violations.extend(analyze_file(file_path, text, enabled_rules=enabled_rules))
+            except Exception as exc:
+                violations.append(Violation(file_path=file_path, line_number=1, message=f"failed to parse ({exc})."))
 
     if len(violations) > 0:
-        violations = sorted(violations)
-        print("\n".join(violations), file=sys.stderr)
-        raise ValueError("Some errors in modelings. Check the above message")
+        violations = sorted(violations, key=lambda v: (str(v.file_path), v.line_number, v.message))
+        for violation in violations:
+            emit_violation(violation, github_annotations=args.github_annotations)
+        print(f"Found {len(violations)} modeling structure violation(s).", file=sys.stderr)
+        return 1
+
+    print("OK")
+    return 0
 
 
 if __name__ == "__main__":
-    main()
+    raise SystemExit(main())

filename: utils/check_modeling_structure_rules.toml
patch: @@ -0,0 +1,159 @@
+[rules.TRF001]
+description = "Class-level config_class on <Model>PreTrainedModel should match <Model>Config naming."
+default_enabled = true
+allowlist_models = ["qwen3_omni_moe"]
+
+[rules.TRF001.explanation]
+what_it_does = "Checks naming consistency between <Model>PreTrainedModel and config_class."
+why_bad = "Mismatched config_class can break loading, auto classes, and developer expectations."
+bad_example = '''
+class FooPreTrainedModel(PreTrainedModel):
+    config_class = BarConfig
+'''
+good_example = '''
+class FooPreTrainedModel(PreTrainedModel):
+    config_class = FooConfig
+'''
+
+[rules.TRF002]
+description = "base_model_prefix should be a non-empty canonical string when defined on PreTrainedModel classes."
+default_enabled = true
+allowlist_models = ["lighton_ocr"]
+
+[rules.TRF002.explanation]
+what_it_does = "Checks that base_model_prefix, when set, is a non-empty, whitespace-free string literal."
+why_bad = "Invalid prefixes can break weight loading key mapping and base model access patterns."
+bad_example = '''
+class FooPreTrainedModel(PreTrainedModel):
+    base_model_prefix = ""
+'''
+good_example = '''
+class FooPreTrainedModel(PreTrainedModel):
+    base_model_prefix = "model"
+'''
+
+[rules.TRF003]
+description = "forward() should use capture_output/can_return_tuple decorators instead of manual return_dict branching."
+default_enabled = false
+allowlist_models = []
+
+[rules.TRF003.explanation]
+what_it_does = "Detects forward methods that use the old 'if not return_dict: return (x,)' pattern."
+why_bad = "The old return_dict branching pattern is error-prone and verbose. Use the capture_output or can_return_tuple decorators instead."
+bad_example = '''
+def forward(self, x, return_dict=None):
+    if not return_dict:
+        return (x,)
+    return FooOutput(last_hidden_state=x)
+'''
+good_example = '''
+@can_return_tuple
+def forward(self, x):
+    return FooOutput(last_hidden_state=x)
+'''
+
+[rules.TRF004]
+description = "Models must never override tie_weights. Use _tied_weights_keys instead."
+default_enabled = true
+allowlist_models = ["data2vec", "hubert", "sew", "sew_d", "unispeech", "unispeech_sat", "wav2vec2", "wav2vec2_conformer", "wavlm"]
+
+[rules.TRF004.explanation]
+what_it_does = "Checks that no model class defines a tie_weights method."
+why_bad = "Overriding tie_weights leads to bad consequences for loading, device_map computation, and saving. Use _tied_weights_keys class attribute to declare tied weights instead."
+bad_example = '''
+def tie_weights(self):
+    self.lm_head.weight = self.emb.weight
+'''
+good_example = '''
+class FooForCausalLM(FooPreTrainedModel):
+    _tied_weights_keys = ["lm_head.weight"]
+'''
+
+[rules.TRF005]
+description = "_no_split_modules, when defined, should be a list/tuple of non-empty strings."
+default_enabled = true
+allowlist_models = ["d_fine", "deformable_detr", "glm46v", "lw_detr", "pp_doclayout_v3", "rt_detr", "rt_detr_v2", "voxtral", "voxtral_realtime"]
+
+[rules.TRF005.explanation]
+what_it_does = "Checks the shape of _no_split_modules when present."
+why_bad = "Malformed values can break device-map partitioning and sharding behavior."
+bad_example = '''
+_no_split_modules = [SomeLayerClass, ""]
+'''
+good_example = '''
+_no_split_modules = ["FooDecoderLayer", "FooAttention"]
+'''
+
+[rules.TRF006]
+description = "forward with cache arguments should reference cache control/state variables consistently."
+default_enabled = true
+allowlist_models = ["chinese_clip", "evolla", "idefics2", "llama4"]
+
+[rules.TRF006.explanation]
+what_it_does = "Checks forward signatures that expose cache arguments for usage of those arguments in method body."
+why_bad = "Unused cache arguments can indicate incomplete caching support and inconsistent API behavior."
+bad_example = '''
+def forward(self, x, past_key_values=None, use_cache=False):
+    return x
+'''
+good_example = '''
+def forward(self, x, past_key_values=None, use_cache=False):
+    if use_cache:
+        ...
+    return x
+'''
+
+[rules.TRF007]
+description = "self.post_init() in __init__ should remain at the end of initialization for PreTrainedModel classes."
+default_enabled = true
+allowlist_models = ["distilbert", "lxmert", "mt5", "pix2struct", "pop2piano", "switch_transformers", "t5"]
+
+[rules.TRF007.explanation]
+what_it_does = "Checks for self attribute assignments after self.post_init() in __init__."
+why_bad = "Mutating model structure after post_init can bypass intended initialization/finalization logic."
+bad_example = '''
+def __init__(self, config):
+    ...
+    self.post_init()
+    self.proj = nn.Linear(...)
+'''
+good_example = '''
+def __init__(self, config):
+    ...
+    self.proj = nn.Linear(...)
+    self.post_init()
+'''
+
+[rules.TRF008]
+description = "Doc decorators on PreTrainedModel classes should avoid empty add_start_docstrings usage."
+default_enabled = true
+
+[rules.TRF008.explanation]
+what_it_does = "Checks add_start_docstrings usage on model classes for non-empty docstring arguments."
+why_bad = "Empty decorator usage produces unclear docs and weakens generated API documentation quality."
+bad_example = '''
+@add_start_docstrings("")
+class FooModel(FooPreTrainedModel):
+    ...
+'''
+good_example = '''
+@add_start_docstrings("The Foo model.")
+class FooModel(FooPreTrainedModel):
+    ...
+'''
+
+[rules.TRF009]
+description = "modeling_<name>.py should avoid importing implementation code from another model package."
+default_enabled = true
+allowlist_models = ["dpr", "maskformer", "sam3_video", "vision_text_dual_encoder"]
+
+[rules.TRF009.explanation]
+what_it_does = "Checks modeling files for cross-model imports such as transformers.models.other_model.* or from ..other_model.* imports."
+why_bad = "Cross-model implementation imports violate the single-file policy and make model behavior harder to inspect and maintain."
+bad_example = '''
+from transformers.models.llama.modeling_llama import LlamaAttention
+'''
+good_example = '''
+# Keep implementation local in this modeling file.
+# If code is reused, copy it with an appropriate # Copied from ... statement.
+'''