clyxbox 0.2.0__tar.gz

clyxbox-0.2.0/PKG-INFO

@@ -0,0 +1,12 @@
+Metadata-Version: 2.4
+Name: clyxbox
+Version: 0.2.0
+Summary: Inference library for Clyx language models
+License: Apache-2.0
+Project-URL: Homepage, https://huggingface.co/syntropic-clx
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+Requires-Dist: torch>=2.0
+Requires-Dist: tokenizers>=0.19
+Requires-Dist: huggingface_hub>=0.23
+Requires-Dist: safetensors>=0.4

clyxbox-0.2.0/clyxbox/__init__.py

@@ -0,0 +1,5 @@
+from .model import ClyxModel, ModelConfig
+from .tokenizer import ClyxTokenizer
+
+__version__ = "0.2.0"
+__all__ = ["ClyxModel", "ClyxTokenizer", "ModelConfig"]

clyxbox-0.2.0/clyxbox/model.py

@@ -0,0 +1,311 @@
+import json
+import math
+from dataclasses import dataclass
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint as checkpoint_utils
+
+
+# ── Config ────────────────────────────────────────────────────────────────────
+
+@dataclass
+class ModelConfig:
+    vocab_size: int = 32000
+    hidden_size: int = 768
+    num_layers: int = 12
+    num_heads: int = 12
+    max_position_embeddings: int = 2048
+    intermediate_size: int = 2048
+    norm_eps: float = 1e-6
+    attn_dropout: float = 0.0
+    resid_dropout: float = 0.0
+    tie_word_embeddings: bool = True
+    gradient_checkpointing: bool = False
+
+
+# ── Architecture ──────────────────────────────────────────────────────────────
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight
+
+
+class RotaryEmbedding(nn.Module):
+    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
+        super().__init__()
+        self.dim = dim
+        self.base = base
+        self.max_position_embeddings = max_position_embeddings
+        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float().to(device) / dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self._set_cos_sin_cache(max_position_embeddings, device=device, dtype=torch.get_default_dtype())
+
+    def _set_cos_sin_cache(self, seq_len, device, dtype):
+        self.max_seq_len_cached = seq_len
+        t = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype)
+        freqs = torch.outer(t, self.inv_freq)
+        emb = torch.cat((freqs, freqs), dim=-1)
+        self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
+        self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
+
+    def forward(self, x, seq_len):
+        if seq_len > self.max_seq_len_cached or self.cos_cached.device != x.device:
+            self._set_cos_sin_cache(seq_len, device=x.device, dtype=x.dtype)
+        return self.cos_cached[:seq_len].to(dtype=x.dtype), self.sin_cached[:seq_len].to(dtype=x.dtype)
+
+
+def rotate_half(x):
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids):
+    cos = cos[position_ids].unsqueeze(1)
+    sin = sin[position_ids].unsqueeze(1)
+    return (q * cos) + (rotate_half(q) * sin), (k * cos) + (rotate_half(k) * sin)
+
+
+class CausalSelfAttention(nn.Module):
+    def __init__(self, config: ModelConfig):
+        super().__init__()
+        self.num_heads = config.num_heads
+        self.head_dim = config.hidden_size // config.num_heads
+        self.q_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
+        self.k_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
+        self.v_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
+        self.out_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
+        self.attn_dropout = config.attn_dropout
+        self.resid_dropout = nn.Dropout(config.resid_dropout)
+
+    def forward(self, x, rotary_emb, position_ids, attn_mask=None, past_key_value=None, use_cache=False):
+        bsz, seq_len, hidden = x.size()
+        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_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+
+        max_pos = int(position_ids.max().item()) + 1
+        cos, sin = rotary_emb(q, seq_len=max_pos)
+        q, k = apply_rotary_pos_emb(q, k, cos, sin, position_ids)
+
+        if past_key_value is not None:
+            prev_k, prev_v = past_key_value
+            k = torch.cat([prev_k, k], dim=2)
+            v = torch.cat([prev_v, v], dim=2)
+
+        new_kv = (k, v) if use_cache else None
+        is_causal = attn_mask is None and seq_len > 1
+        y = F.scaled_dot_product_attention(
+            q, k, v,
+            attn_mask=attn_mask,
+            dropout_p=self.attn_dropout if self.training else 0.0,
+            is_causal=is_causal,
+        )
+        y = y.transpose(1, 2).contiguous().view(bsz, seq_len, hidden)
+        return self.resid_dropout(self.out_proj(y)), new_kv
+
+
+class SwiGLUMLP(nn.Module):
+    def __init__(self, config: ModelConfig):
+        super().__init__()
+        self.w1 = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.w2 = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.w3 = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
+        self.dropout = nn.Dropout(config.resid_dropout)
+
+    def forward(self, x):
+        return self.dropout(self.w3(F.silu(self.w1(x)) * self.w2(x)))
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, config: ModelConfig):
+        super().__init__()
+        self.attn_norm = RMSNorm(config.hidden_size, eps=config.norm_eps)
+        self.attn = CausalSelfAttention(config)
+        self.mlp_norm = RMSNorm(config.hidden_size, eps=config.norm_eps)
+        self.mlp = SwiGLUMLP(config)
+
+    def forward(self, x, rotary_emb, position_ids, attn_mask=None, past_key_value=None, use_cache=False):
+        residual = x
+        attn_out, new_kv = self.attn(
+            self.attn_norm(x), rotary_emb, position_ids,
+            attn_mask=attn_mask, past_key_value=past_key_value, use_cache=use_cache,
+        )
+        x = residual + attn_out
+        x = x + self.mlp(self.mlp_norm(x))
+        return x, new_kv
+
+
+class _CoreTransformer(nn.Module):
+    """Internal transformer — matches the checkpoint keys exactly."""
+    def __init__(self, config: ModelConfig):
+        super().__init__()
+        self.config = config
+        self.embed = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.rotary_emb = RotaryEmbedding(
+            dim=config.hidden_size // config.num_heads,
+            max_position_embeddings=config.max_position_embeddings,
+        )
+        self.layers = nn.ModuleList([TransformerBlock(config) for _ in range(config.num_layers)])
+        self.norm = RMSNorm(config.hidden_size, eps=config.norm_eps)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        if config.tie_word_embeddings:
+            self.lm_head.weight = self.embed.weight
+
+    def forward(self, input_ids, position_ids=None, past_key_values=None, use_cache=False):
+        bsz, seq_len = input_ids.size()
+        if position_ids is None:
+            past_len = past_key_values[0][0].size(2) if past_key_values else 0
+            position_ids = torch.arange(past_len, past_len + seq_len, dtype=torch.long, device=input_ids.device)
+            position_ids = position_ids.unsqueeze(0).expand(bsz, -1)
+
+        x = self.embed(input_ids)
+        new_kvs = []
+        for idx, block in enumerate(self.layers):
+            past_kv = past_key_values[idx] if past_key_values else None
+            x, new_kv = block(x, self.rotary_emb, position_ids, past_key_value=past_kv, use_cache=use_cache)
+            new_kvs.append(new_kv)
+
+        x = self.norm(x)
+        logits = self.lm_head(x)
+        return logits, new_kvs
+
+
+# ── Public API ────────────────────────────────────────────────────────────────
+
+class ClyxModel(nn.Module):
+    """
+    Clyx causal language model — ready for inference.
+
+    Usage:
+        model = ClyxModel.from_pretrained("syntropic-clx/Clyx_0.2-115.67M-BASE")
+        model.eval()
+    """
+
+    def __init__(self, config: ModelConfig):
+        super().__init__()
+        self.config = config
+        self._model = _CoreTransformer(config)
+
+    # ── from_pretrained ──────────────────────────────────────────────────────
+
+    @classmethod
+    def from_pretrained(cls, repo_id: str, device: Optional[str] = None, token: Optional[str] = None):
+        """Load a Clyx model from a HuggingFace Hub repository."""
+        try:
+            from huggingface_hub import hf_hub_download
+        except ImportError:
+            raise ImportError("pip install huggingface_hub")
+        try:
+            from safetensors.torch import load_file
+        except ImportError:
+            raise ImportError("pip install safetensors")
+
+        if device is None:
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+
+        # Load config
+        cfg_path = hf_hub_download(repo_id=repo_id, filename="config.json", token=token)
+        with open(cfg_path) as f:
+            raw = json.load(f)
+
+        # Map HF config keys → ModelConfig fields
+        field_map = {
+            "hidden_size": "hidden_size",
+            "num_hidden_layers": "num_layers",
+            "num_attention_heads": "num_heads",
+            "max_position_embeddings": "max_position_embeddings",
+            "intermediate_size": "intermediate_size",
+            "rms_norm_eps": "norm_eps",
+            "vocab_size": "vocab_size",
+            "tie_word_embeddings": "tie_word_embeddings",
+        }
+        cfg_kwargs = {}
+        for hf_key, our_key in field_map.items():
+            if hf_key in raw:
+                cfg_kwargs[our_key] = raw[hf_key]
+            elif our_key in raw:
+                cfg_kwargs[our_key] = raw[our_key]
+
+        config = ModelConfig(**cfg_kwargs)
+        model = cls(config)
+
+        # Load weights
+        weights_path = hf_hub_download(repo_id=repo_id, filename="model.safetensors", token=token)
+        state = load_file(weights_path, device=device)
+        model._model.load_state_dict(state, strict=True)
+        model = model.to(device)
+        return model
+
+    # ── forward ──────────────────────────────────────────────────────────────
+
+    def forward(self, input_ids, **kwargs):
+        return self._model(input_ids, **kwargs)
+
+    # ── generate ─────────────────────────────────────────────────────────────
+
+    @torch.no_grad()
+    def generate(
+        self,
+        input_ids: torch.Tensor,
+        max_new_tokens: int = 200,
+        temperature: float = 1.0,
+        top_p: float = 1.0,
+        top_k: int = 0,
+        repetition_penalty: float = 1.0,
+        eos_token_id: Optional[int] = None,
+    ) -> torch.Tensor:
+        """Auto-regressive generation with temperature / top-p / top-k / repetition penalty."""
+        device = next(self.parameters()).device
+        ids = input_ids.to(device)
+        past_key_values = None
+
+        for _ in range(max_new_tokens):
+            # On the first step pass the full prompt; after that pass only the new token
+            cur_input = ids if past_key_values is None else ids[:, -1:]
+            logits, past_key_values = self._model(cur_input, past_key_values=past_key_values, use_cache=True)
+            next_logits = logits[:, -1, :]  # (B, vocab)
+
+            # Repetition penalty
+            if repetition_penalty != 1.0:
+                for b in range(ids.size(0)):
+                    unique = ids[b].unique()
+                    next_logits[b, unique] = torch.where(
+                        next_logits[b, unique] < 0,
+                        next_logits[b, unique] * repetition_penalty,
+                        next_logits[b, unique] / repetition_penalty,
+                    )
+
+            # Temperature
+            if temperature != 1.0:
+                next_logits = next_logits / temperature
+
+            # Top-k
+            if top_k > 0:
+                topk_vals = torch.topk(next_logits, top_k).values[:, -1, None]
+                next_logits = next_logits.masked_fill(next_logits < topk_vals, float("-inf"))
+
+            # Top-p (nucleus)
+            if top_p < 1.0:
+                sorted_logits, sorted_idx = torch.sort(next_logits, descending=True)
+                cum_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+                remove = cum_probs - F.softmax(sorted_logits, dim=-1) > top_p
+                sorted_logits[remove] = float("-inf")
+                next_logits = torch.zeros_like(next_logits).scatter(1, sorted_idx, sorted_logits)
+
+            probs = F.softmax(next_logits, dim=-1)
+            next_token = torch.multinomial(probs, num_samples=1)
+            ids = torch.cat([ids, next_token], dim=1)
+
+            if eos_token_id is not None and (next_token == eos_token_id).all():
+                break
+
+        return ids

clyxbox-0.2.0/clyxbox/tokenizer.py

@@ -0,0 +1,77 @@
+from typing import List, Optional, Union
+
+import torch
+
+
+SPECIAL_TOKENS = {
+    "PAD": "<PAD>", "UNK": "<UNK>", "BOS": "<BOS>", "EOS": "<EOS>",
+    "USER": "<USER>", "USER_END": "</USER>", "MODEL": "<MODEL>", "MODEL_END": "</MODEL>",
+    "STOP": "<STOP>", "SYSTEM": "<SYSTEM>", "SYSTEM_END": "</SYSTEM>",
+    "MEMORY": "<MEMORY>", "MEMORY_END": "</MEMORY>",
+    "TOOL_CALL": "<TOOL_CALL>", "TOOL_CALL_END": "</TOOL_CALL>",
+    "TOOL_RESULT": "<TOOL_RESULT>", "TOOL_RESULT_END": "</TOOL_RESULT>",
+}
+
+
+class ClyxTokenizer:
+    """
+    Clyx tokenizer backed by a ByteLevel BPE tokenizer.json.
+
+    Usage:
+        tok = ClyxTokenizer.from_pretrained("syntropic-clx/Clyx_0.2-115.67M-BASE")
+        ids = tok.encode("Hello world", return_tensors="pt")
+        text = tok.decode(ids[0])
+    """
+
+    def __init__(self, tokenizer):
+        self._tok = tokenizer
+        for name, token in SPECIAL_TOKENS.items():
+            setattr(self, f"{name.lower()}_id", tokenizer.token_to_id(token))
+
+    # ── from_pretrained ──────────────────────────────────────────────────────
+
+    @classmethod
+    def from_pretrained(cls, repo_id: str, token: Optional[str] = None):
+        """Load tokenizer from a HuggingFace Hub repository."""
+        try:
+            from huggingface_hub import hf_hub_download
+        except ImportError:
+            raise ImportError("pip install huggingface_hub")
+        try:
+            from tokenizers import Tokenizer
+        except ImportError:
+            raise ImportError("pip install tokenizers")
+
+        path = hf_hub_download(repo_id=repo_id, filename="tokenizer.json", token=token)
+        return cls(Tokenizer.from_file(path))
+
+    # ── encode / decode ──────────────────────────────────────────────────────
+
+    def encode(
+        self,
+        text: str,
+        add_special_tokens: bool = False,
+        return_tensors: Optional[str] = None,
+    ) -> Union[List[int], torch.Tensor]:
+        ids = self._tok.encode(text).ids
+        if add_special_tokens:
+            ids = [self.bos_id] + ids + [self.eos_id]
+        if return_tensors == "pt":
+            return torch.tensor([ids], dtype=torch.long)
+        return ids
+
+    def decode(
+        self,
+        ids: Union[List[int], torch.Tensor],
+        skip_special_tokens: bool = True,
+    ) -> str:
+        if isinstance(ids, torch.Tensor):
+            ids = ids.detach().cpu().tolist()
+        return self._tok.decode([int(x) for x in ids], skip_special_tokens=skip_special_tokens)
+
+    def __call__(self, text: str, **kwargs):
+        return self.encode(text, **kwargs)
+
+    @property
+    def vocab_size(self) -> int:
+        return self._tok.get_vocab_size()

clyxbox-0.2.0/clyxbox.egg-info/PKG-INFO

@@ -0,0 +1,12 @@
+Metadata-Version: 2.4
+Name: clyxbox
+Version: 0.2.0
+Summary: Inference library for Clyx language models
+License: Apache-2.0
+Project-URL: Homepage, https://huggingface.co/syntropic-clx
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+Requires-Dist: torch>=2.0
+Requires-Dist: tokenizers>=0.19
+Requires-Dist: huggingface_hub>=0.23
+Requires-Dist: safetensors>=0.4

clyxbox-0.2.0/clyxbox.egg-info/SOURCES.txt

@@ -0,0 +1,9 @@
+pyproject.toml
+clyxbox/__init__.py
+clyxbox/model.py
+clyxbox/tokenizer.py
+clyxbox.egg-info/PKG-INFO
+clyxbox.egg-info/SOURCES.txt
+clyxbox.egg-info/dependency_links.txt
+clyxbox.egg-info/requires.txt
+clyxbox.egg-info/top_level.txt

clyxbox-0.2.0/clyxbox.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

clyxbox-0.2.0/clyxbox.egg-info/requires.txt

@@ -0,0 +1,4 @@
+torch>=2.0
+tokenizers>=0.19
+huggingface_hub>=0.23
+safetensors>=0.4

clyxbox-0.2.0/clyxbox.egg-info/top_level.txt

@@ -0,0 +1 @@
+clyxbox

clyxbox-0.2.0/pyproject.toml

@@ -0,0 +1,24 @@
+[build-system]
+requires = ["setuptools>=68", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "clyxbox"
+version = "0.2.0"
+description = "Inference library for Clyx language models"
+readme = "README.md"
+license = { text = "Apache-2.0" }
+requires-python = ">=3.9"
+dependencies = [
+    "torch>=2.0",
+    "tokenizers>=0.19",
+    "huggingface_hub>=0.23",
+    "safetensors>=0.4",
+]
+
+[project.urls]
+Homepage = "https://huggingface.co/syntropic-clx"
+
+[tool.setuptools.packages.find]
+where = ["."]
+include = ["clyxbox*"]

clyxbox-0.2.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+