File size: 1,453 Bytes
2db32a1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 | # 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() |