configuration_ordinal.py

"""HuggingFace-compatible configuration for Ordinal LLM.

This file enables AutoConfig.from_pretrained() with trust_remote_code=True.
Follows HuggingFace PretrainedConfig pattern exactly.
"""

from __future__ import annotations
import json
import os


class OrdinalConfig:
    """Ordinal model configuration (HuggingFace compatible)."""

    model_type = "ordinal"

    def __init__(
        self,
        vocab_size: int = 50304,
        hidden_size: int = 3584,
        intermediate_size: int = 9216,
        num_hidden_layers: int = 36,
        num_attention_heads: int = 28,
        num_key_value_heads: int = 4,
        head_dim: int = 128,
        max_position_embeddings: int = 8192,
        rms_norm_eps: float = 1e-5,
        rope_theta: float = 500000.0,
        hidden_act: str = "silu",
        tie_word_embeddings: bool = False,
        use_cache: bool = True,
        bos_token_id: int = 1,
        eos_token_id: int = 2,
        pad_token_id: int = 0,
        torch_dtype: str = "bfloat16",
        # Anti-hallucination features
        use_confidence_head: bool = True,
        confidence_threshold: float = 0.7,
        use_retrieval_attention: bool = True,
        retrieval_dim: int = 256,
        num_retrieval_heads: int = 4,
        use_fact_verification_layer: bool = True,
        verification_layers: list | None = None,
        use_source_embeddings: bool = True,
        num_source_types: int = 16,
        # Sliding window
        sliding_window: int | None = None,
        **kwargs,
    ):
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.intermediate_size = intermediate_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.num_key_value_heads = num_key_value_heads
        self.head_dim = head_dim
        self.max_position_embeddings = max_position_embeddings
        self.rms_norm_eps = rms_norm_eps
        self.rope_theta = rope_theta
        self.hidden_act = hidden_act
        self.tie_word_embeddings = tie_word_embeddings
        self.use_cache = use_cache
        self.bos_token_id = bos_token_id
        self.eos_token_id = eos_token_id
        self.pad_token_id = pad_token_id
        self.torch_dtype = torch_dtype
        self.use_confidence_head = use_confidence_head
        self.confidence_threshold = confidence_threshold
        self.use_retrieval_attention = use_retrieval_attention
        self.retrieval_dim = retrieval_dim
        self.num_retrieval_heads = num_retrieval_heads
        self.use_fact_verification_layer = use_fact_verification_layer
        self.verification_layers = verification_layers or self._default_verification_layers()
        self.use_source_embeddings = use_source_embeddings
        self.num_source_types = num_source_types
        self.sliding_window = sliding_window

    def _default_verification_layers(self) -> list[int]:
        n = self.num_hidden_layers
        return [n // 3, 2 * n // 3, n - 1]

    def to_dict(self) -> dict:
        d = self.__dict__.copy()
        d["model_type"] = self.model_type
        d["architectures"] = ["OrdinalForCausalLM"]
        d["auto_map"] = {
            "AutoConfig": "configuration_ordinal.OrdinalConfig",
            "AutoModelForCausalLM": "modeling_ordinal.OrdinalForCausalLM",
        }
        d["transformers_version"] = "4.45.0"
        return d

    def save_pretrained(self, save_directory: str) -> None:
        os.makedirs(save_directory, exist_ok=True)
        with open(os.path.join(save_directory, "config.json"), "w", encoding="utf-8") as f:
            json.dump(self.to_dict(), f, indent=2)

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs):
        config_file = os.path.join(pretrained_model_name_or_path, "config.json")
        if os.path.isfile(config_file):
            with open(config_file, encoding="utf-8") as f:
                config_dict = json.load(f)
            # Filter to valid init params
            valid_keys = set(cls.__init__.__code__.co_varnames) - {"self", "kwargs"}
            filtered = {k: v for k, v in config_dict.items() if k in valid_keys}
            return cls(**filtered, **kwargs)
        return cls(**kwargs)

    @classmethod
    def from_dict(cls, config_dict: dict):
        valid_keys = set(cls.__init__.__code__.co_varnames) - {"self", "kwargs"}
        filtered = {k: v for k, v in config_dict.items() if k in valid_keys}
        return cls(**filtered)

    # Preset configurations
    @classmethod
    def ordinal_128m(cls): return cls(hidden_size=768, intermediate_size=2048, num_hidden_layers=12, num_attention_heads=12, num_key_value_heads=4, head_dim=64)
    @classmethod
    def ordinal_256m(cls): return cls(hidden_size=1024, intermediate_size=2816, num_hidden_layers=16, num_attention_heads=16, num_key_value_heads=4, head_dim=64)
    @classmethod
    def ordinal_512m(cls): return cls(hidden_size=1536, intermediate_size=4096, num_hidden_layers=20, num_attention_heads=16, num_key_value_heads=4, head_dim=96)
    @classmethod
    def ordinal_1b(cls): return cls(hidden_size=2048, intermediate_size=5504, num_hidden_layers=24, num_attention_heads=16, num_key_value_heads=4, head_dim=128)
    @classmethod
    def ordinal_2b(cls): return cls(hidden_size=2560, intermediate_size=6912, num_hidden_layers=28, num_attention_heads=20, num_key_value_heads=4, head_dim=128)
    @classmethod
    def ordinal_4b(cls): return cls(hidden_size=3072, intermediate_size=8192, num_hidden_layers=32, num_attention_heads=24, num_key_value_heads=4, head_dim=128)
    @classmethod
    def ordinal_5b(cls): return cls(hidden_size=3584, intermediate_size=9216, num_hidden_layers=36, num_attention_heads=28, num_key_value_heads=4, head_dim=128)
    @classmethod
    def ordinal_7b(cls): return cls(hidden_size=4096, intermediate_size=11008, num_hidden_layers=32, num_attention_heads=32, num_key_value_heads=8, head_dim=128)
    @classmethod
    def ordinal_13b(cls): return cls(hidden_size=5120, intermediate_size=13824, num_hidden_layers=40, num_attention_heads=40, num_key_value_heads=8, head_dim=128)
    @classmethod
    def ordinal_20b(cls): return cls(hidden_size=6144, intermediate_size=16384, num_hidden_layers=52, num_attention_heads=48, num_key_value_heads=8, head_dim=128)
    @classmethod
    def ordinal_33b(cls): return cls(hidden_size=6656, intermediate_size=17920, num_hidden_layers=64, num_attention_heads=56, num_key_value_heads=8, head_dim=128)
    @classmethod
    def ordinal_48b(cls): return cls(hidden_size=8192, intermediate_size=22016, num_hidden_layers=72, num_attention_heads=64, num_key_value_heads=8, head_dim=128)