from transformers import AutoModelForCausalLM

from .llama_modeling import LlamaForCausalLM, LlamaDraftForCausalLM, LlamaDraftConfig, LlamaConfig

AutoModelForCausalLM.register(LlamaConfig, LlamaForCausalLM, exist_ok=True)
AutoModelForCausalLM.register(LlamaDraftConfig, LlamaDraftForCausalLM)
