Upload configuration_logos.py

configuration_logos.py

@@ -0,0 +1,114 @@
+"""HuggingFace configuration for Logos."""
+
+from __future__ import annotations
+
+import sys
+import importlib
+from pathlib import Path
+from typing import Any, Dict
+
+from transformers import PretrainedConfig
+
+_MODULE_DIR = Path(__file__).resolve().parent
+if str(_MODULE_DIR) not in sys.path:
+    sys.path.insert(0, str(_MODULE_DIR))
+
+
+_NATIVE_DEFAULTS: Dict[str, Any] = {
+    "vocab_size": 32000,
+    "d_model": 512,
+    "max_seq_len": 2048,
+    "num_layers": 8,
+    "num_heads": 8,
+    "norm_eps": 1e-6,
+    "d_ff": 1364,
+    "use_moe": True,
+    "num_shared_experts": 2,
+    "num_sparse_experts": 64,
+    "top_k": 6,
+    "expert_d_ff": 256,
+    "bias_update_rate": 0.01,
+    "capacity_factor": 2.0,
+    "router_logit_noise_std": 0.1,
+    "router_logit_noise_decay_steps": 2000,
+    "router_init_std": 0.002,
+    "router_bias_error_clip": 1.0,
+    "router_bias_clip": 1.0,
+    "moe_aux_loss_weight": 1e-3,
+    "moe_aux_loss_decay_steps": 2000,
+    "lm_head_chunk_size": 0,
+    "moe_diversity_factor": 0.0,
+    "rope_base": 10000.0,
+    "qk_norm": True,
+    "partial_rope_dim": None,
+    "attention_sink": True,
+    "block_residual_isolate_softmax": False,
+    "head_dim": 64,
+    "conv_size": 4,
+    "chunk_size": 64,
+    "A_init_range": (1, 16),
+    "expand": 2,
+    "swa_window": 256,
+    "swa_every": 4,
+    "swa_offset": 3,
+    "csa_compression": 4,
+    "csa_top_k": 1024,
+    "csa_indexer_heads": 4,
+    "csa_indexer_dim": 32,
+    "csa_indexer_loss_weight": 1.0,
+    "hca_compression": 128,
+    "compressed_query_dim": None,
+    "compressed_head_dim": None,
+    "compressed_rope": False,
+    "indexer_rope": False,
+    "attn_pattern": None,
+    "num_entry_layers": 2,
+    "num_body_layers": 4,
+    "num_exit_layers": 2,
+    "num_loops": 4,
+    "entry_attn_pattern": None,
+    "body_attn_pattern": None,
+    "exit_attn_pattern": None,
+    "entry_top_k": None,
+    "exit_top_k": None,
+    "gradient_checkpointing": False,
+    "ckpt_granularity": "per-block",
+}
+
+
+class LogosConfig(PretrainedConfig):
+    model_type = "logos"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(self, **kwargs: Any):
+        native_values: Dict[str, Any] = {}
+        for name, default in _NATIVE_DEFAULTS.items():
+            native_values[name] = kwargs.pop(name, default)
+
+        tokenizer_encoding = kwargs.pop("tokenizer_encoding", "cl100k_base")
+        kwargs.setdefault("tie_word_embeddings", True)
+        super().__init__(**kwargs)
+
+        for name, value in native_values.items():
+            setattr(self, name, value)
+        self.tokenizer_encoding = tokenizer_encoding
+        self.architectures = ["LogosForCausalLM"]
+        self.auto_map = {
+            "AutoConfig": "configuration_logos.LogosConfig",
+            "AutoModelForCausalLM": "modeling_logos.LogosForCausalLM",
+            "AutoTokenizer": ["tokenization_logos.LogosTokenizer", None],
+        }
+
+    def to_native_config(self):
+        model_dir = getattr(self, "_name_or_path", None) or getattr(self, "name_or_path", None)
+        if model_dir:
+            model_path = Path(str(model_dir)).resolve()
+            if model_path.exists() and str(model_path) not in sys.path:
+                sys.path.insert(0, str(model_path))
+        NativeLogosConfig = importlib.import_module("models.logos").LogosConfig
+
+        data = {name: getattr(self, name) for name in _NATIVE_DEFAULTS}
+        return NativeLogosConfig(**data)
+
+
+__all__ = ["LogosConfig"]