Fix custom model registration by adding custom_models/__init__.py
#37
Conversation
WalkthroughA new step was added to the README for custom model integration, instructing users to update Changes
Estimated code review effort🎯 2 (Simple) | ⏱️ ~7 minutes Poem
Note ⚡️ Unit Test Generation is now available in beta!Learn more here, or try it out under "Finishing Touches" below. ✨ Finishing Touches
🧪 Generate unit tests
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
README.md (1)
485-491: Step 3 vs Step 6 duplicate the same instruction – collapse them for clarityStep 3 already tells users to “Register your models in
__init__.py”.
Step 6 repeats the same requirement, which can confuse first-time readers.Consider merging the bullets or re-ordering:
3. Edit `custom_models/__init__.py`: - `from . import my_model` - set `__all__ = ["my_model"]`Keeps the checklist concise and avoids readers missing one of the two steps.
train.sh (1)
69-70: Python one-liner can drop the extra jq dependencySince the sole purpose is to obtain
model_type, do it directly in Python and avoid thejqpipe (makes the script portable to systems without jq):-model=$( - python -c "import fla, sys, custom_models; from transformers import AutoConfig; print(AutoConfig.from_pretrained(sys.argv[1]).to_json_string())" "$config" | jq -r '.model_type' -) +model=$( + python - <<'PY' "$config" +import sys, custom_models # noqa: F401 – import for side-effect registration +from transformers import AutoConfig +print(AutoConfig.from_pretrained(sys.argv[1]).model_type) +PY +)No behavioural change, but one fewer external dependency and slightly simpler.
| from . import sba | ||
|
|
||
| __all__ = ["sba"] |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Import fails hard if ‘sba’ is absent – make registration resilient / configurable
__init__.py eagerly imports sba, raising ModuleNotFoundError the moment a user removes or renames that directory.
Given the README now tells users to edit this file for their own models, a safer pattern is to:
- Gate the import in a
try/exceptand emit a clear error message, or - Provide a minimal lazy-import helper that looks up modules listed in
__all__only when accessed.
This prevents a broken install/run‐time experience when someone forgets to adjust the import while experimenting.
-from . import sba
-
-__all__ = ["sba"]
+import importlib
+import warnings
+
+__all__ = ["sba"] # update this list for your own modules
+
+def __getattr__(name):
+ if name in __all__:
+ try:
+ module = importlib.import_module(f".{name}", __name__)
+ globals()[name] = module
+ return module
+ except ModuleNotFoundError:
+ raise ImportError(
+ f"custom_models: requested sub-module '{name}' not found – "
+ "did you forget to create/rename it?"
+ )
+ raise AttributeError(name)Keeps side-effects minimal and offers a clearer failure mode.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| from . import sba | |
| __all__ = ["sba"] | |
| import importlib | |
| import warnings | |
| __all__ = ["sba"] # update this list for your own modules | |
| def __getattr__(name): | |
| if name in __all__: | |
| try: | |
| module = importlib.import_module(f".{name}", __name__) | |
| globals()[name] = module | |
| return module | |
| except ModuleNotFoundError: | |
| raise ImportError( | |
| f"custom_models: requested sub-module '{name}' not found – " | |
| "did you forget to create/rename it?" | |
| ) | |
| raise AttributeError(name) |
🤖 Prompt for AI Agents
In custom_models/__init__.py at lines 1 to 3, the current eager import of 'sba'
causes a ModuleNotFoundError if 'sba' is missing. To fix this, wrap the import
of 'sba' in a try/except block that catches ImportError and logs a clear,
user-friendly error message explaining the missing module. Alternatively,
implement a lazy-import mechanism that only imports modules listed in __all__
when they are accessed, preventing immediate import failures and improving
resilience when users modify the models.
This PR adds
custom_models/__init__.pyto ensure proper model registration withAutoConfigandAutoModelForCausalLMAlso updates
train.shto explicitly importcustom_modelswhen parsing the model name (used in wandb run name)related issue: #35
Summary by CodeRabbit
Documentation
New Features
Chores