__init__.py

"""Miscovery model package."""

from .configuration_miscovery import CustomTransformerConfig
from .modeling_miscovery import CustomTransformerModel

# Import necessary mappings for registration
from transformers.models.auto.configuration_auto import CONFIG_MAPPING
from transformers.models.auto.modeling_auto import (
    MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,
    MODEL_MAPPING,
    MODEL_FOR_CAUSAL_LM_MAPPING
)

# Register the model with Hugging Face Auto classes
CONFIG_MAPPING.register("miscovery", CustomTransformerConfig)
MODEL_MAPPING.register(CustomTransformerConfig, CustomTransformerModel)
MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING.register(CustomTransformerConfig, CustomTransformerModel)

# You might need these additional registrations depending on your use case
from transformers.pipelines import SUPPORTED_TASKS, TEXT_GENERATION_MODELS, SUMMARIZATION_MODELS, TRANSLATION_MODELS

# Register model for specific tasks if needed
# This is critical for the pipeline to work properly
if "miscovery" not in TEXT_GENERATION_MODELS:
    TEXT_GENERATION_MODELS.append("miscovery")
    
if "miscovery" not in SUMMARIZATION_MODELS:
    SUMMARIZATION_MODELS.append("miscovery")
    
if "miscovery" not in TRANSLATION_MODELS:
    TRANSLATION_MODELS.append("miscovery")

__all__ = [
    "CustomTransformerConfig",
    "CustomTransformerModel",
]