__init__.py

# coding=utf-8
"""Wiola: a small language model with SRPE, GCLA, ATM, DSFF and WiolaRMSNorm.

Importing this package registers Wiola with the HuggingFace Auto* classes so
that ``AutoModelForCausalLM.from_pretrained("oscowlai/wiola-360m")`` works once
weights are published.
"""

from .configuration_wiola import WiolaConfig
from .modeling_wiola import (
    WiolaDecoderLayer,
    WiolaForCausalLM,
    WiolaModel,
    WiolaPreTrainedModel,
)

__version__ = "0.1.1"

__all__ = [
    "WiolaConfig",
    "WiolaModel",
    "WiolaForCausalLM",
    "WiolaPreTrainedModel",
    "WiolaDecoderLayer",
]


def _register_auto_classes():
    try:
        from transformers import AutoConfig, AutoModel, AutoModelForCausalLM
    except Exception:  # transformers not installed
        return
    try:
        AutoConfig.register("wiola", WiolaConfig)
        AutoModel.register(WiolaConfig, WiolaModel)
        AutoModelForCausalLM.register(WiolaConfig, WiolaForCausalLM)
    except Exception:
        # Already registered (e.g. re-import) — safe to ignore.
        pass


# The auto_map in config.json now handles Hub loading; we no longer call
# register_for_auto_class to avoid save_pretrained() packaging source code.
# WiolaConfig.register_for_auto_class("AutoConfig")
# WiolaModel.register_for_auto_class("AutoModel")
# WiolaForCausalLM.register_for_auto_class("AutoModelForCausalLM")

_register_auto_classes()