configuration_roformer_classification.py

"""RoFormer configuration for classification with projection head."""

from transformers import RoFormerConfig


class RoFormerClassificationConfig(RoFormerConfig):
    """Configuration for RoFormer with contrastive projection head.

    Extends RoFormerConfig with additional parameters for the projection head
    used in contrastive learning for file type classification.
    """

    model_type = "roformer-classification"

    def __init__(
        self,
        projection_dim: int = 256,
        num_labels: int = 106,
        **kwargs,
    ):
        """Initialize configuration.

        Args:
            projection_dim: Dimension of the projection head output (for embeddings)
            num_labels: Number of classification labels (MIME types)
            **kwargs: Additional arguments passed to RoFormerConfig
        """
        super().__init__(**kwargs)
        self.projection_dim = projection_dim
        self.num_labels = num_labels