modeling_ordinal.py

"""HuggingFace-compatible modeling file for Ordinal LLM.

This file enables:
    AutoModelForCausalLM.from_pretrained("KaztoRay/ordinal-5b", trust_remote_code=True)
"""

from __future__ import annotations

import math
import json
from typing import Optional, Tuple
from dataclasses import dataclass


@dataclass
class OrdinalConfig:
    """Configuration for Ordinal model (HuggingFace compatible)."""
    model_type: str = "ordinal"
    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
    vocab_size: int = 50304
    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
    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
    use_source_embeddings: bool = True
    num_source_types: int = 16
    bos_token_id: int = 1
    eos_token_id: int = 2
    pad_token_id: int = 0
    torch_dtype: str = "bfloat16"

    def __post_init__(self):
        if self.verification_layers is None:
            n = self.num_hidden_layers
            self.verification_layers = [n // 3, 2 * n // 3, n - 1]

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs):
        import os
        config_path = os.path.join(pretrained_model_name_or_path, "config.json")
        if os.path.exists(config_path):
            with open(config_path) as f:
                config_dict = json.load(f)
            return cls(**{k: v for k, v in config_dict.items()
                         if k in cls.__dataclass_fields__})
        return cls(**kwargs)


# Placeholder for actual model implementation (requires torch)
# The full implementation is in ordinal_llm/model/architecture/model.py
# This file provides the HuggingFace interface

try:
    import torch
    import torch.nn as nn
    import torch.nn.functional as F
    TORCH_AVAILABLE = True
except ImportError:
    TORCH_AVAILABLE = False


if TORCH_AVAILABLE:
    class OrdinalRMSNorm(nn.Module):
        """Root Mean Square Layer Normalization."""
        def __init__(self, hidden_size: int, eps: float = 1e-5):
            super().__init__()
            self.weight = nn.Parameter(torch.ones(hidden_size))
            self.eps = eps

        def forward(self, x):
            variance = x.float().pow(2).mean(-1, keepdim=True)
            x = x * torch.rsqrt(variance + self.eps)
            return (self.weight * x).to(x.dtype)

    class OrdinalRotaryEmbedding(nn.Module):
        """Rotary Position Embedding (RoPE)."""
        def __init__(self, dim: int, max_seq_len: int = 8192, theta: float = 500000.0):
            super().__init__()
            inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
            self.register_buffer("inv_freq", inv_freq)
            t = torch.arange(max_seq_len).float()
            freqs = torch.outer(t, inv_freq)
            self.register_buffer("cos_cached", freqs.cos())
            self.register_buffer("sin_cached", freqs.sin())

        def forward(self, x, seq_len: int):
            return self.cos_cached[:seq_len], self.sin_cached[:seq_len]

    class OrdinalMLP(nn.Module):
        """SwiGLU MLP."""
        def __init__(self, config: OrdinalConfig):
            super().__init__()
            self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
            self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
            self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)

        def forward(self, x):
            return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))

    class OrdinalAttention(nn.Module):
        """Grouped Query Attention."""
        def __init__(self, config: OrdinalConfig, layer_idx: int = 0):
            super().__init__()
            self.num_heads = config.num_attention_heads
            self.num_kv_heads = config.num_key_value_heads
            self.head_dim = config.head_dim
            self.num_kv_groups = self.num_heads // self.num_kv_heads

            self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=False)
            self.k_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
            self.v_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
            self.o_proj = nn.Linear(self.num_heads * self.head_dim, config.hidden_size, bias=False)
            self.rotary = OrdinalRotaryEmbedding(self.head_dim, config.max_position_embeddings, config.rope_theta)

        def forward(self, x, attention_mask=None, position_ids=None, past_key_value=None):
            bsz, seq_len, _ = x.shape
            q = self.q_proj(x).view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
            k = self.k_proj(x).view(bsz, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
            v = self.v_proj(x).view(bsz, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)

            # Apply RoPE
            cos, sin = self.rotary(x, seq_len)
            # Simplified RoPE application
            q_embed = q * cos.unsqueeze(0).unsqueeze(0) + self._rotate_half(q) * sin.unsqueeze(0).unsqueeze(0)
            k_embed = k * cos.unsqueeze(0).unsqueeze(0) + self._rotate_half(k) * sin.unsqueeze(0).unsqueeze(0)

            # GQA: repeat KV heads
            if self.num_kv_groups > 1:
                k_embed = k_embed.repeat_interleave(self.num_kv_groups, dim=1)
                v = v.repeat_interleave(self.num_kv_groups, dim=1)

            # Attention
            attn_weights = torch.matmul(q_embed, k_embed.transpose(-2, -1)) / math.sqrt(self.head_dim)
            if attention_mask is not None:
                attn_weights = attn_weights + attention_mask
            attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(q.dtype)
            attn_output = torch.matmul(attn_weights, v)

            attn_output = attn_output.transpose(1, 2).reshape(bsz, seq_len, -1)
            return self.o_proj(attn_output)

        @staticmethod
        def _rotate_half(x):
            x1, x2 = x.chunk(2, dim=-1)
            return torch.cat((-x2, x1), dim=-1)

    class OrdinalDecoderLayer(nn.Module):
        """Single transformer decoder layer."""
        def __init__(self, config: OrdinalConfig, layer_idx: int = 0):
            super().__init__()
            self.self_attn = OrdinalAttention(config, layer_idx)
            self.mlp = OrdinalMLP(config)
            self.input_layernorm = OrdinalRMSNorm(config.hidden_size, config.rms_norm_eps)
            self.post_attention_layernorm = OrdinalRMSNorm(config.hidden_size, config.rms_norm_eps)

        def forward(self, hidden_states, attention_mask=None, position_ids=None, past_key_value=None):
            residual = hidden_states
            hidden_states = self.input_layernorm(hidden_states)
            hidden_states = self.self_attn(hidden_states, attention_mask, position_ids, past_key_value)
            hidden_states = residual + hidden_states

            residual = hidden_states
            hidden_states = self.post_attention_layernorm(hidden_states)
            hidden_states = self.mlp(hidden_states)
            hidden_states = residual + hidden_states
            return hidden_states

    class OrdinalConfidenceHead(nn.Module):
        """Per-token confidence scoring (anti-hallucination)."""
        def __init__(self, hidden_size: int):
            super().__init__()
            self.linear1 = nn.Linear(hidden_size, hidden_size // 4)
            self.linear2 = nn.Linear(hidden_size // 4, 1)

        def forward(self, hidden_states):
            x = F.gelu(self.linear1(hidden_states))
            return torch.sigmoid(self.linear2(x))

    class OrdinalModel(nn.Module):
        """Ordinal base model (transformer decoder)."""
        def __init__(self, config: OrdinalConfig):
            super().__init__()
            self.config = config
            self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
            self.layers = nn.ModuleList([
                OrdinalDecoderLayer(config, i) for i in range(config.num_hidden_layers)
            ])
            self.norm = OrdinalRMSNorm(config.hidden_size, config.rms_norm_eps)

        def forward(self, input_ids, attention_mask=None, position_ids=None):
            hidden_states = self.embed_tokens(input_ids)
            for layer in self.layers:
                hidden_states = layer(hidden_states, attention_mask, position_ids)
            return self.norm(hidden_states)

    class OrdinalForCausalLM(nn.Module):
        """Ordinal model for causal language modeling (HF compatible)."""
        def __init__(self, config: OrdinalConfig):
            super().__init__()
            self.config = config
            self.model = OrdinalModel(config)
            self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
            if config.use_confidence_head:
                self.confidence_head = OrdinalConfidenceHead(config.hidden_size)
            if config.tie_word_embeddings:
                self.lm_head.weight = self.model.embed_tokens.weight

        def forward(self, input_ids, attention_mask=None, labels=None, **kwargs):
            hidden_states = self.model(input_ids, attention_mask)
            logits = self.lm_head(hidden_states)

            loss = None
            if labels is not None:
                shift_logits = logits[..., :-1, :].contiguous()
                shift_labels = labels[..., 1:].contiguous()
                loss = F.cross_entropy(
                    shift_logits.view(-1, self.config.vocab_size),
                    shift_labels.view(-1),
                    ignore_index=-100,
                )

            confidence = None
            if hasattr(self, 'confidence_head'):
                confidence = self.confidence_head(hidden_states)

            return {"loss": loss, "logits": logits, "confidence": confidence}

        def generate(self, input_ids, max_new_tokens=512, temperature=0.7, top_p=0.9, **kwargs):
            """Simple autoregressive generation."""
            for _ in range(max_new_tokens):
                outputs = self.forward(input_ids)
                logits = outputs["logits"][:, -1, :] / temperature

                # Top-p sampling
                sorted_logits, sorted_indices = torch.sort(logits, descending=True)
                cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
                sorted_indices_to_remove = cumulative_probs > top_p
                sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
                sorted_indices_to_remove[..., 0] = 0
                indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
                logits[indices_to_remove] = float('-inf')

                probs = F.softmax(logits, dim=-1)
                next_token = torch.multinomial(probs, num_samples=1)
                input_ids = torch.cat([input_ids, next_token], dim=-1)

                if next_token.item() == self.config.eos_token_id:
                    break

                # Confidence-aware: reduce temperature if uncertain
                if hasattr(self, 'confidence_head'):
                    conf = outputs["confidence"][:, -1, 0]
                    if conf.item() < self.config.confidence_threshold:
                        temperature = max(0.3, temperature * 0.9)

            return input_ids

        @classmethod
        def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
            config = OrdinalConfig.from_pretrained(pretrained_model_name_or_path)
            model = cls(config)
            # Load weights if available
            import os
            for weight_file in ["model.safetensors", "pytorch_model.bin"]:
                path = os.path.join(pretrained_model_name_or_path, weight_file)
                if os.path.exists(path):
                    if weight_file.endswith(".safetensors"):
                        from safetensors.torch import load_file
                        state_dict = load_file(path)
                    else:
                        state_dict = torch.load(path, map_location="cpu")
                    model.load_state_dict(state_dict, strict=False)
                    break
            return model

        def num_parameters(self, only_trainable: bool = False):
            return sum(p.numel() for p in self.parameters() if p.requires_grad or not only_trainable)