Added Model.from_pretrained method to load a pretrained model - #169
Conversation
Summary of ChangesHello @vfdev-5, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request introduces a standardized Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
7a5f8fc to
5f091d3
Compare
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a from_pretrained method across several models, which is a fantastic enhancement for usability. The implementation is largely consistent and simplifies model loading from external sources like Hugging Face Hub and timm. My review focuses on improving docstrings and code structure for maintainability. I've also identified a few significant issues in the dinov3 model implementation related to loading weights from timm, which will need to be addressed for the feature to work correctly.
I am having trouble creating individual review comments. Click here to see my feedback.
bonsai/models/dinov3/params.py (59-93)
The weight name mapping for timm models appears to be incorrect and incomplete.
- Many regex patterns (e.g., for
k_proj,v_proj,o_proj) seem to be copy-pasted from theoriginalmapping and uselayer...instead of thetimmconventionblocks.... - The mapping for
blocks.([0-9]+).attn.proj.weighttoq_proj.kernelis likely wrong. Intimm,projis the output projection, so it should map too_proj.kernel. - There is no mapping for
timm'sqkv.weight, which combines query, key, and value weights. This tensor would need to be split and mapped toq_proj,k_proj, andv_projin the bonsai model.
This is a critical issue that will prevent timm models from being loaded correctly.
bonsai/models/dinov3/modeling.py (352)
There's an issue with loading timm models like timm/vit_small_patch16_dinov3.lvd1689m. The create_model_from_safe_tensors function is called without specifying the mapping_type, so it defaults to "original". This will use the wrong weight name mapping for a timm model, causing loading to fail.
You should detect the model source (e.g., by checking if model_name starts with "timm/") and pass the appropriate mapping_type to create_model_from_safe_tensors.
bonsai/models/convnext/modeling.py (145-161)
There are a couple of improvements that can be made here:
- Docstring: The docstring could be more descriptive, following a standard format (e.g., Google style) to explain the method's purpose, arguments, and return value.
config_map: This dictionary is redefined on every method call. It would be more efficient to define it as a class-level constant.
Here is an example of how you could refactor this:
class ConvNeXt(nnx.Module):
_PRETRAINED_CONFIGS = {
"facebook/convnext-tiny-224": ModelConfig.convnext_tiny_224,
"facebook/convnext-small-224": ModelConfig.convnext_small_224,
"facebook/convnext-base-224": ModelConfig.convnext_base_224,
"facebook/convnext-large-224": ModelConfig.convnext_large_224,
}
# ... other methods
@classmethod
def from_pretrained(cls, model_name: str, config: ModelConfig | None = None):
"""Loads a pretrained ConvNeXt model from a Hugging Face Hub repository.
Args:
model_name: The *model id* of a pretrained model on huggingface.co.
config: An optional `ModelConfig` to override the default configuration.
Returns:
A `ConvNeXt` model instance with pretrained weights.
"""
from huggingface_hub import snapshot_download
from bonsai.models.convnext import params
if config is None:
if model_name not in cls._PRETRAINED_CONFIGS:
raise ValueError(f"Model name '{model_name}' is unknown, please provide config argument")
config = cls._PRETRAINED_CONFIGS[model_name]()
# ... rest of the methodbonsai/models/densenet121/modeling.py (139-154)
Similar to other models in this PR, the docstring could be more descriptive, and the config_map could be defined as a class-level constant for efficiency and better organization.
Here's an example of how you could apply this pattern:
class DenseNet(nnx.Module):
_PRETRAINED_CONFIGS = {
"keras/densenet_121_imagenet": ModelConfig.densenet_121,
"keras/densenet_169_imagenet": ModelConfig.densenet_169,
"keras/densenet_201_imagenet": ModelConfig.densenet_201,
}
# ... other methods
@classmethod
def from_pretrained(cls, model_name: str, config: ModelConfig | None = None):
"""Loads a pretrained DenseNet model from a Hugging Face Hub repository.
Args:
model_name: The *model id* of a pretrained model on huggingface.co.
config: An optional `ModelConfig` to override the default configuration.
Returns:
A `DenseNet` model instance with pretrained weights.
"""
from huggingface_hub import snapshot_download
from bonsai.models.densenet121 import params
if config is None:
if model_name not in cls._PRETRAINED_CONFIGS:
raise ValueError(f"Model name '{model_name}' is unknown, please provide config argument")
config = cls._PRETRAINED_CONFIGS[model_name]()
# ... rest of the methodbonsai/models/dinov3/modeling.py (329-349)
For consistency and maintainability, consider improving the docstring and moving the config_map to a class-level constant. This avoids redefining the dictionary on each call.
Here's a suggested refactoring:
class Dinov3ViTModel(nnx.Module):
_PRETRAINED_CONFIGS = {
"facebook/dinov3-vits16-pretrain-lvd1689m": ModelConfig.dinov3_vits16,
"timm/vit_small_patch16_dinov3.lvd1689m": ModelConfig.dinov3_vits16,
# ... other models
}
# ... other methods
@classmethod
def from_pretrained(cls, model_name: str, config: ModelConfig | None = None):
"""Loads a pretrained DINOv3 model from a Hugging Face Hub repository.
Args:
model_name: The *model id* of a pretrained model on huggingface.co.
config: An optional `ModelConfig` to override the default configuration.
Returns:
A `Dinov3ViTModel` instance with pretrained weights.
"""
# ... implementation using cls._PRETRAINED_CONFIGSbonsai/models/efficientnet/modeling.py (330-353)
Similar to other models in this PR, the config_map is redefined on every call. It would be more efficient to define it as a class-level constant. The docstring is good, but for consistency with other models, I'm suggesting a similar refactor.
class EfficientNet(nnx.Module):
_PRETRAINED_CONFIGS = {
"efficientnet_b0": ModelConfig.b0,
"efficientnet_b1": ModelConfig.b1,
"efficientnet_b2": ModelConfig.b2,
"efficientnet_b3": ModelConfig.b3,
"efficientnet_b4": ModelConfig.b4,
"efficientnet_b5": ModelConfig.b5,
"efficientnet_b6": ModelConfig.b6,
"efficientnet_b7": ModelConfig.b7,
}
# ... other methods
@classmethod
def from_pretrained(cls, model_name: str, config: ModelConfig | None = None):
# ...
if config is not None:
raise ValueError("config must be None when using EfficientNet.from_pretrained")
if model_name not in cls._PRETRAINED_CONFIGS:
raise ValueError(f"Model name '{model_name}' is unknown, please provide config argument")
config = cls._PRETRAINED_CONFIGS[model_name]()
return params._create_model_from_timm(model_name, config)
f5f2901 to
5ce1a9d
Compare
vfdev-5
changed the title
|
@vfdev-5 Hi Victor, could you fix the merge conflicts above? |
c802c48 to
585ca40
Compare
|
@jenriver Hi Jen, I fixed the conflict, we can move forward with this PR, thanks! |
b0e06cf to
0cdb0e2
Compare
3 participants
Description:
from bonsai.models import Qwen3VLForConditionalGeneration