File size: 391 Bytes
cdbec7c | 1 2 3 4 5 6 7 8 9 10 11 12 13 | from transformers import AutoConfig
from transformers.models.mistral import MistralConfig
BIDIR_MISTRAL_TYPE = "bidir_mistral"
class BidirectionalMistralConfig(MistralConfig):
model_type = BIDIR_MISTRAL_TYPE
keys_to_ignore_at_inference = ["past_key_values"]
AutoConfig.register(BIDIR_MISTRAL_TYPE, BidirectionalMistralConfig)
BidirectionalMistralConfig.register_for_auto_class() |