[model] feat: add AutoModelForSequenceClassification support in HuggingfaceLoader #256
Conversation
![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 adds support for AutoModelForSequenceClassification in HuggingfaceLoader. The changes are straightforward, involving an import and a new condition to select the correct model class. However, I've identified a significant issue with the logic for selecting the model class. The current implementation, which your change follows, relies on checking only the first architecture from the model's config. This can lead to incorrect model class selection if the desired architecture (e.g., for sequence classification) is not the first in the list. I've left a detailed comment with high severity on this issue. Addressing this would make the model loader much more robust.
veomni/models/loader.py
Outdated
| elif ( | ||
| "ForSequenceClassification" in architecture | ||
| and type(model_config) in AutoModelForSequenceClassification._model_mapping.keys() | ||
| ): | ||
| load_class = AutoModelForSequenceClassification |
There was a problem hiding this comment.
While this change correctly adds support for AutoModelForSequenceClassification, it inherits a potential issue from the existing logic. The architecture variable is derived from _get_model_arch_from_config, which only considers the first element if model_config.architectures is a list. This can cause this check to fail if the ...ForSequenceClassification architecture is not the first in the list. For example, if model_config.architectures is ['MistralModel', 'MistralForSequenceClassification'], architecture will be 'MistralModel' and this condition will fail.
To make this more robust, you should check against all architectures in model_config.architectures. This would likely require refactoring how architecture is obtained and used throughout this if/elif chain, including for the ForCausalLM case. A robust implementation would handle model_config.architectures as a list (or ensure it is one) and check if any of the items contain the target substring. This is a high-severity issue as it can lead to incorrect model loading.
Luosuu
changed the title
veomni/models/loader.py
Outdated
| elif type(model_config) in AutoModelForVision2Seq._model_mapping.keys(): # assume built-in models | ||
| load_class = AutoModelForVision2Seq | ||
| elif ( | ||
| "ForSequenceClassification" in architecture |
There was a problem hiding this comment.
would this break ForCausalLM from loading? or the config would properly point to that?
anyway, is it possible to write a simple unit test for this function?
2 participants
I added AutoModelForSequenceClassification support in our HuggingfaceLoader so that whenever a config’s architectures indicate a ForSequenceClassification class, we now instantiate the corresponding HF classification model instead of falling back to AutoModel.