"""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)