from transformers.models.auto.modeling_auto import AutoModelForTokenClassification
from piiceetah.models import BertCRFForTokenClassification
from piiceetah.angular_models import (
    AngularBertForTokenClassification, 
    AngularRobertaForTokenClassification, 
    AngularDebertaV2ForTokenClassification,
    AngularModernBertForTokenClassification
)

# 註冊CRF模型
AutoModelForTokenClassification.register(BertCRFForTokenClassification.config_class, BertCRFForTokenClassification)

# 註冊Angular模型
AutoModelForTokenClassification.register(AngularBertForTokenClassification.config_class, AngularBertForTokenClassification)
AutoModelForTokenClassification.register(AngularRobertaForTokenClassification.config_class, AngularRobertaForTokenClassification)
AutoModelForTokenClassification.register(AngularDebertaV2ForTokenClassification.config_class, AngularDebertaV2ForTokenClassification)
AutoModelForTokenClassification.register(AngularModernBertForTokenClassification.config_class, AngularModernBertForTokenClassification)