"""
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"]