configuration_sentiment.py

"""Configuration class for SentimentClassifier."""

from typing import Optional

from transformers import PretrainedConfig


class SentimentClassifierConfig(PretrainedConfig):
    """
    Configuration class for SentimentClassifier model.

    This class stores the configuration of a :class:`~SentimentClassifier` model.
    It is used to instantiate a SentimentClassifier model according to the specified
    arguments, defining the model architecture.

    Args:
        pretrained_model (:obj:`str`, defaults to :obj:`"xlm-roberta-base"`):
            Name of the pre-trained transformer model to use as encoder.
        num_labels (:obj:`int`, defaults to :obj:`3`):
            Number of sentiment classes (positive/neutral/negative).
        dropout (:obj:`float`, defaults to :obj:`0.1`):
            Dropout probability for the classification head.
        hidden_size (:obj:`int`, optional):
            Hidden size of the encoder model. If None, will be auto-detected from encoder config.
        model_type (:obj:`str`, defaults to :obj:`"sentiment-classifier"`):
            Model type identifier for the Hugging Face Hub.
    """

    model_type = "sentiment-classifier"

    # Enable trust_remote_code by mapping to the custom classes
    # This tells HuggingFace where to find the custom model and config classes
    auto_map = {
        "AutoConfig": "configuration_sentiment.SentimentClassifierConfig",
        "AutoModelForSequenceClassification": "sentiment_classifier.SentimentClassifier",
    }

    def __init__(
        self,
        pretrained_model: str = "xlm-roberta-base",
        num_labels: int = 3,
        dropout: float = 0.1,
        hidden_size: Optional[int] = None,
        **kwargs,
    ):
        """Initialize SentimentClassifierConfig."""
        # Set auto_map in kwargs before calling super().__init__
        # This ensures it gets saved to config.json
        if "auto_map" not in kwargs:
            kwargs["auto_map"] = {
                "AutoConfig": "configuration_sentiment.SentimentClassifierConfig",
                "AutoModelForSequenceClassification": "sentiment_classifier.SentimentClassifier",
            }

        super().__init__(**kwargs)

        self.pretrained_model = pretrained_model
        self.num_labels = num_labels
        self.dropout = dropout
        self.hidden_size = hidden_size