File size: 817 Bytes
063be31 |
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 |
"""NanoGPT HuggingFace Integration"""
from transformers import AutoConfig, AutoModel, AutoModelForCausalLM
# Import our classes
try:
from .configuration_nanogpt import NanoGPTConfig
from .modeling_nanogpt import (
NanoGPTModel,
NanoGPTForCausalLM,
NanoGPTPreTrainedModel
)
except ImportError:
from configuration_nanogpt import NanoGPTConfig
from modeling_nanogpt import (
NanoGPTModel,
NanoGPTForCausalLM,
NanoGPTPreTrainedModel
)
# Register the model with Auto* classes
AutoConfig.register("nanogpt", NanoGPTConfig)
AutoModel.register(NanoGPTConfig, NanoGPTModel)
AutoModelForCausalLM.register(NanoGPTConfig, NanoGPTForCausalLM)
__all__ = [
"NanoGPTConfig",
"NanoGPTModel",
"NanoGPTForCausalLM",
"NanoGPTPreTrainedModel"
]
|