File size: 1,037 Bytes
af0603b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 |
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)
|