"""Configuration for the HS6 product classifier.

Subclasses XLMRobertaConfig, so every encoder field (hidden_size, num_hidden_layers,
...) keeps its usual name and meaning. Only the classification head and the label
spaces for the coarser HS levels are added on top.
"""

from transformers.models.xlm_roberta.configuration_xlm_roberta import XLMRobertaConfig


class HS6ClassifierConfig(XLMRobertaConfig):
    model_type = "hs6_classifier"

    def __init__(
        self,
        n2: int = 97,
        n4: int = 1266,
        n6: int = 6750,
        head_dropout: float = 0.15,
        pooling: str = "cls",
        recommended_max_length: int = 1024,
        id2hs4=None,
        id2hs2=None,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.n2 = n2
        self.n4 = n4
        self.n6 = n6
        self.head_dropout = head_dropout
        self.pooling = pooling
        # Training used 512 tokens; inference at 1024 is free (the encoder has 8194
        # positions) and worth +2.91 points on texts longer than 1500 characters.
        self.recommended_max_length = recommended_max_length
        # HS4 / HS2 code strings, indexed by class id. Used to name the marginal logits.
        self.id2hs4 = id2hs4 or []
        self.id2hs2 = id2hs2 or []