__init__.py

"""
Miscovery model package
"""
# Import configuration and model classes
from configuration_miscovery import CustomTransformerConfig
from modeling_miscovery import CustomTransformerModel

# Import pipeline utilities
from pipeline_utils import standard_pipeline, create_language_pipeline

# Register with Auto classes
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
)

# Register configuration
if "miscovery" not in CONFIG_MAPPING:
    CONFIG_MAPPING.register("miscovery", CustomTransformerConfig)

# Register model classes
MODEL_MAPPING.register(CustomTransformerConfig, CustomTransformerModel)
MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING.register(CustomTransformerConfig, CustomTransformerModel)

# Try to register with pipeline
try:
    from transformers.pipelines import PIPELINE_REGISTRY
    
    PIPELINE_REGISTRY.register_model(
        model_type="miscovery",
        task="translation",
        model_class=CustomTransformerModel
    )
except Exception as e:
    print(f"Warning: Could not register with pipeline registry: {e}")
    print("This is not critical; the model should still work with standard_pipeline")

# Export all relevant classes and functions
__all__ = [
    "CustomTransformerConfig",
    "CustomTransformerModel",
    "standard_pipeline",
    "create_language_pipeline"
]