| """ | |
| custom_models/tweaktron/__init__.py | |
| Registers Tweaktron with HF Auto classes so flame can load it via | |
| AutoConfig / AutoModelForCausalLM with model_type "tweaktron_text". | |
| """ | |
| from transformers import AutoConfig, AutoModel, AutoModelForCausalLM | |
| from .configuration_tweaktron import TweaktronConfig | |
| from .modeling_tweaktron import TweaktronModel, TweaktronForCausalLM | |
| AutoConfig.register("tweaktron_text", TweaktronConfig) | |
| AutoModel.register(TweaktronConfig, TweaktronModel) | |
| AutoModelForCausalLM.register(TweaktronConfig, TweaktronForCausalLM) | |
| __all__ = ["TweaktronConfig", "TweaktronModel", "TweaktronForCausalLM"] | |