anthracite 1.0.0__py3-none-any.whl

anthracite/__init__.py

@@ -0,0 +1,222 @@
+"""Anthracite – universal AI training & fine-tuning framework.
+
+Train a model in one call::
+
+    from anthracite import train
+
+    train(
+        model_name="Nutral-GPT-20M",
+        model_type="text_gen",
+        dataset="my_dataset.jsonl",
+        tokens=100_000_000,
+        params=20_000_000,
+        context_length=512,
+        device="auto",
+        batch_size="auto",
+        precision="auto",
+    )
+
+Then use it::
+
+    from anthracite import generate
+    print(generate(model="./models/Nutral-GPT-20M", prompt="Hello, my name is"))
+
+Embedding models work the same way::
+
+    from anthracite import train, embed, similarity
+
+    train(model_name="MyEmbedder", model_type="embedding", dataset="corpus.txt",
+          tokens=20_000_000, params="30M", context_length=256)
+
+    vectors = embed("./models/MyEmbedder", ["a first sentence", "a second one"])
+"""
+
+from __future__ import annotations
+
+from .core.config import ANTHRACITE_VERSION as __version__
+from .core.config import ModelConfig, TrainingConfig
+from .core.exceptions import (
+    AnthraciteError,
+    ArchitectureError,
+    CheckpointError,
+    DatasetNotFoundError,
+    DeviceError,
+    GenerationError,
+    InsufficientMemoryError,
+    InvalidConfigurationError,
+    TPUNotAvailableError,
+    TokenizerError,
+    UnsupportedDatasetError,
+    UnsupportedModelTypeError,
+)
+from .tokenizer.tokenizer import AnthraciteTokenizer
+
+
+def train(
+    model_name: str = "anthracite-model",
+    model_type: str = "text_gen",
+    dataset=None,
+    tokens=10_000_000,
+    params="20M",
+    context_length: int = 512,
+    vocab_size: int = 32768,
+    batch_size="auto",
+    device: str = "auto",
+    precision: str = "auto",
+    **kwargs,
+) -> dict:
+    """Train a new model from scratch and save it as a SafeTensors package.
+
+    Parameters
+    ----------
+    model_name : name of the model (used for the output directory and config).
+    model_type : ``"text_gen"`` or ``"embedding"``.
+    dataset    : path, directory, Hugging Face id, HF Dataset object or list of strings.
+    tokens     : training token budget, e.g. ``100_000_000`` or ``"100M"``.
+    params     : target parameter count, e.g. ``20_000_000`` or ``"20M"``.
+    context_length : sequence length for text models.
+    batch_size : ``"auto"`` or an integer micro-batch size.
+    device     : ``"auto" | "cpu" | "cuda" | "cuda:N" | "tpu"``.
+    precision  : ``"auto" | "fp32" | "fp16" | "bf16"``.
+
+    Extra keyword arguments are forwarded to :class:`~anthracite.core.config.TrainingConfig`
+    (``learning_rate``, ``checkpoint_interval``, ``resume``, ``seed``, ``output_dir``,
+    ``objective``, ``pooling``, ``embedding_dim``, ``val_split`` …), and unknown keys
+    become architecture overrides
+    (``d_model``, ``n_layers``, ``n_heads`` …).
+
+    Returns the training metadata dictionary.
+    """
+    from .core.trainer import train as _train
+
+    return _train(
+        model_name=model_name,
+        model_type=model_type,
+        dataset=dataset,
+        tokens=tokens,
+        params=params,
+        context_length=context_length,
+        vocab_size=vocab_size,
+        batch_size=batch_size,
+        device=device,
+        precision=precision,
+        **kwargs,
+    )
+
+
+def finetune(
+    model=None,
+    dataset=None,
+    tokens=1_000_000,
+    device: str = "auto",
+    batch_size="auto",
+    precision: str = "auto",
+    **kwargs,
+) -> dict:
+    """Continue training an existing Anthracite model on new data.
+
+    ``model`` may be a local model directory, an Anthracite checkpoint, or a
+    Hugging Face repo id.  The base model is never overwritten: the result is
+    written to ``<name>-Finetuned`` unless ``output_dir`` says otherwise.
+    """
+    from .core.finetuner import finetune as _finetune
+
+    return _finetune(
+        model=model,
+        dataset=dataset,
+        tokens=tokens,
+        device=device,
+        batch_size=batch_size,
+        precision=precision,
+        **kwargs,
+    )
+
+
+def load_model(model, device: str = "auto", precision: str | None = None, verbose: bool = False):
+    """Load a trained model (config + architecture + weights + tokenizer) for inference."""
+    from .inference.loader import load_model as _load
+
+    return _load(model, device=device, precision=precision, verbose=verbose)
+
+
+def generate(model, prompt: str = "", **kwargs):
+    """Generate text from a trained ``text_gen`` model."""
+    from .inference.loader import LoadedModel
+    from .inference.loader import load_model as _load
+    from .inference.text import generate_text
+
+    bundle = model if isinstance(model, LoadedModel) else _load(model, device=kwargs.pop("device", "auto"))
+    if bundle.config.model_type == "embedding":
+        raise GenerationError(
+            f"{bundle.config.model_name} is an embedding model, so it has nothing to generate.",
+            "Use embed(), similarity() or search() instead.",
+        )
+    return generate_text(bundle, prompt, **kwargs)
+
+
+def embed(model, texts, batch_size: int = 16, normalize: bool = True, **kwargs):
+    """Encode text into vectors with a trained ``embedding`` model.
+
+    ``texts`` may be one string or a list; returns a numpy array of shape
+    ``(n, dimension)`` (or ``(dimension,)`` for a single string).
+    """
+    from .inference.embedding import embed as _embed
+
+    return _embed(model, texts, batch_size=batch_size, normalize=normalize, **kwargs)
+
+
+def similarity(model, a, b, **kwargs):
+    """Cosine similarity between two strings (or two lists of strings)."""
+    from .inference.embedding import similarity as _similarity
+
+    return _similarity(model, a, b, **kwargs)
+
+
+def search(model, query: str, documents, top_k: int = 5, **kwargs) -> list:
+    """Rank documents against a query; returns ``[(index, document, score)]``."""
+    from .inference.embedding import search as _search
+
+    return _search(model, query, documents, top_k=top_k, **kwargs)
+
+
+def create_interface(model, model_type: str | None = None, **kwargs):
+    """Generate and launch an interface for a trained model (gradio or terminal)."""
+    from .interface.generator import create_interface as _create
+
+    return _create(model, model_type=model_type, **kwargs)
+
+
+def available_model_types() -> list:
+    """Model types this build can train."""
+    from .core import registry
+
+    return registry.available()
+
+
+__all__ = [
+    "train",
+    "finetune",
+    "load_model",
+    "generate",
+    "embed",
+    "similarity",
+    "search",
+    "create_interface",
+    "available_model_types",
+    "AnthraciteTokenizer",
+    "ModelConfig",
+    "TrainingConfig",
+    "AnthraciteError",
+    "DatasetNotFoundError",
+    "UnsupportedDatasetError",
+    "UnsupportedModelTypeError",
+    "TokenizerError",
+    "ArchitectureError",
+    "DeviceError",
+    "TPUNotAvailableError",
+    "InsufficientMemoryError",
+    "InvalidConfigurationError",
+    "CheckpointError",
+    "GenerationError",
+    "__version__",
+]

anthracite/architectures/__init__.py

@@ -0,0 +1,9 @@
+from .anthracite1_text import Anthracite1TextModel, build_text_model
+from .anthracite1_embedding import Anthracite1EmbeddingModel, build_embedding_model
+
+__all__ = [
+    "Anthracite1TextModel",
+    "Anthracite1EmbeddingModel",
+    "build_text_model",
+    "build_embedding_model",
+]

anthracite/architectures/anthracite1_embedding.py

@@ -0,0 +1,231 @@
+"""Anthracite-1 (embedding) – bidirectional encoder for sentence/document vectors.
+
+Shares the Anthracite-1 building blocks with the text generator (RMSNorm,
+rotary positions, SwiGLU) but attention is **bidirectional**, so every token
+sees the whole sequence.  Two training objectives are supported and both live
+in this module:
+
+* ``mlm``         – masked language modelling, for pretraining from raw text
+* ``contrastive`` – in-batch-negative InfoNCE on (anchor, positive) pairs,
+                    for turning a pretrained encoder into a retrieval model
+
+A typical two-stage recipe is: ``train(..., objective="mlm")`` on a corpus,
+then ``finetune(..., objective="contrastive")`` on pairs.
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from ..core.config import ModelConfig
+from .anthracite1_text import RMSNorm, SwiGLU, apply_rope, build_rope_cache
+
+
+class BidirectionalAttention(nn.Module):
+    """Grouped-query attention without the causal mask."""
+
+    def __init__(self, config: ModelConfig):
+        super().__init__()
+        self.n_heads = config.n_heads
+        self.n_kv_heads = config.n_kv_heads
+        self.head_dim = config.d_model // config.n_heads
+        self.repeats = self.n_heads // self.n_kv_heads
+        self.dropout_p = config.dropout
+
+        self.wq = nn.Linear(config.d_model, self.n_heads * self.head_dim, bias=False)
+        self.wk = nn.Linear(config.d_model, self.n_kv_heads * self.head_dim, bias=False)
+        self.wv = nn.Linear(config.d_model, self.n_kv_heads * self.head_dim, bias=False)
+        self.wo = nn.Linear(self.n_heads * self.head_dim, config.d_model, bias=False)
+
+    def forward(self, x, cos, sin, key_padding_mask=None):
+        b, t, _ = x.shape
+        q = self.wq(x).view(b, t, self.n_heads, self.head_dim).transpose(1, 2)
+        k = self.wk(x).view(b, t, self.n_kv_heads, self.head_dim).transpose(1, 2)
+        v = self.wv(x).view(b, t, self.n_kv_heads, self.head_dim).transpose(1, 2)
+
+        q = apply_rope(q, cos, sin)
+        k = apply_rope(k, cos, sin)
+        if self.repeats > 1:
+            k = k.repeat_interleave(self.repeats, dim=1)
+            v = v.repeat_interleave(self.repeats, dim=1)
+
+        attn_mask = None
+        if key_padding_mask is not None:
+            # (B, T) True = real token → (B, 1, 1, T) broadcastable boolean mask
+            attn_mask = key_padding_mask[:, None, None, :].to(torch.bool)
+
+        if hasattr(F, "scaled_dot_product_attention"):
+            out = F.scaled_dot_product_attention(
+                q, k, v, attn_mask=attn_mask, dropout_p=self.dropout_p if self.training else 0.0
+            )
+        else:  # pragma: no cover - very old torch
+            scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
+            if attn_mask is not None:
+                scores = scores.masked_fill(~attn_mask, float("-inf"))
+            out = torch.softmax(scores.float(), dim=-1).to(q.dtype) @ v
+
+        out = out.transpose(1, 2).contiguous().view(b, t, -1)
+        return self.wo(out)
+
+
+class EncoderBlock(nn.Module):
+    def __init__(self, config: ModelConfig):
+        super().__init__()
+        self.attn_norm = RMSNorm(config.d_model, config.norm_eps)
+        self.attn = BidirectionalAttention(config)
+        self.ffn_norm = RMSNorm(config.d_model, config.norm_eps)
+        self.ffn = SwiGLU(config)
+        self.dropout = nn.Dropout(config.dropout) if config.dropout > 0 else nn.Identity()
+
+    def forward(self, x, cos, sin, key_padding_mask=None):
+        x = x + self.dropout(self.attn(self.attn_norm(x), cos, sin, key_padding_mask))
+        x = x + self.dropout(self.ffn(self.ffn_norm(x)))
+        return x
+
+
+def masked_mean(hidden: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+    weights = mask.unsqueeze(-1).to(hidden.dtype)
+    total = (hidden * weights).sum(dim=1)
+    count = weights.sum(dim=1).clamp(min=1e-6)
+    return total / count
+
+
+class Anthracite1EmbeddingModel(nn.Module):
+    """Bidirectional encoder producing fixed-size vectors."""
+
+    model_type = "embedding"
+    architecture = "anthracite-1"
+
+    def __init__(self, config: ModelConfig):
+        super().__init__()
+        config.validate()
+        self.config = config
+        self.token_embedding = nn.Embedding(config.vocab_size, config.d_model)
+        self.blocks = nn.ModuleList([EncoderBlock(config) for _ in range(config.n_layers)])
+        self.final_norm = RMSNorm(config.d_model, config.norm_eps)
+
+        self.projection = (
+            nn.Linear(config.d_model, config.embedding_dim, bias=False)
+            if config.embedding_dim and config.embedding_dim != config.d_model
+            else None
+        )
+        # MLM head: a small transform, output weights tied to the input embeddings
+        self.mlm_transform = nn.Linear(config.d_model, config.d_model, bias=False)
+        self.mlm_norm = RMSNorm(config.d_model, config.norm_eps)
+
+        self.head_dim = config.d_model // config.n_heads
+        self._rope_cache: tuple | None = None
+        self.apply(self._init_weights)
+        scale = (2 * config.n_layers) ** -0.5
+        for block in self.blocks:
+            nn.init.normal_(block.attn.wo.weight, mean=0.0, std=0.02 * scale)
+            nn.init.normal_(block.ffn.w2.weight, mean=0.0, std=0.02 * scale)
+
+    @staticmethod
+    def _init_weights(module: nn.Module) -> None:
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+    # ------------------------------------------------------------------ #
+    @property
+    def dimension(self) -> int:
+        return self.config.embedding_dim or self.config.d_model
+
+    def _rope(self, seq_len: int, device, dtype):
+        if self._rope_cache is None or self._rope_cache[0].size(0) < seq_len or self._rope_cache[0].device != device:
+            cos, sin = build_rope_cache(
+                max(seq_len, self.config.context_length), self.head_dim, self.config.rope_theta,
+                device, torch.float32,
+            )
+            self._rope_cache = (cos, sin)
+        cos, sin = self._rope_cache
+        return cos[:seq_len].to(dtype), sin[:seq_len].to(dtype)
+
+    def encode_hidden(self, input_ids, attention_mask=None):
+        b, t = input_ids.shape
+        if attention_mask is None:
+            attention_mask = input_ids.ne(self.config.pad_token_id)
+        x = self.token_embedding(input_ids)
+        cos, sin = self._rope(t, x.device, x.dtype)
+        for block in self.blocks:
+            x = block(x, cos, sin, attention_mask)
+        return self.final_norm(x), attention_mask
+
+    # -- forward variants ------------------------------------------------ #
+    def forward(self, input_ids, attention_mask=None, normalize: bool = True):
+        """Return pooled sentence vectors."""
+        hidden, mask = self.encode_hidden(input_ids, attention_mask)
+        if self.config.pooling == "cls":
+            pooled = hidden[:, 0]
+        elif self.config.pooling == "max":
+            pooled = hidden.masked_fill(~mask.unsqueeze(-1), float("-inf")).max(dim=1).values
+        else:
+            pooled = masked_mean(hidden, mask)
+        if self.projection is not None:
+            pooled = self.projection(pooled)
+        if normalize:
+            pooled = F.normalize(pooled.float(), p=2, dim=-1).to(pooled.dtype)
+        return pooled
+
+    def mlm_logits(self, hidden: torch.Tensor) -> torch.Tensor:
+        h = self.mlm_norm(F.gelu(self.mlm_transform(hidden)))
+        return F.linear(h, self.token_embedding.weight)
+
+    # -- objectives ------------------------------------------------------ #
+    def mlm_loss(self, input_ids, labels, attention_mask=None):
+        hidden, _ = self.encode_hidden(input_ids, attention_mask)
+        logits = self.mlm_logits(hidden)
+        return F.cross_entropy(
+            logits.reshape(-1, logits.size(-1)).float(), labels.reshape(-1), ignore_index=-100
+        )
+
+    def contrastive_loss(self, anchor_ids, positive_ids, anchor_mask=None, positive_mask=None,
+                         negative_ids=None, negative_mask=None):
+        """Symmetric InfoNCE with in-batch negatives (plus optional hard negatives)."""
+        a = self.forward(anchor_ids, anchor_mask, normalize=True).float()
+        p = self.forward(positive_ids, positive_mask, normalize=True).float()
+        bank = p
+        if negative_ids is not None:
+            n = self.forward(negative_ids, negative_mask, normalize=True).float()
+            bank = torch.cat([p, n], dim=0)
+
+        temperature = max(float(self.config.temperature), 1e-3)
+        logits = (a @ bank.t()) / temperature
+        target = torch.arange(a.size(0), device=a.device)
+        loss = F.cross_entropy(logits, target)
+        if negative_ids is None:  # symmetric direction only makes sense without extra rows
+            loss = 0.5 * (loss + F.cross_entropy(logits.t(), target))
+        return loss
+
+    # -- convenience ----------------------------------------------------- #
+    @torch.no_grad()
+    def embed(self, input_ids, attention_mask=None, normalize: bool = True):
+        self.eval()
+        return self.forward(input_ids, attention_mask, normalize=normalize)
+
+    def num_parameters(self, trainable_only: bool = False) -> int:
+        return sum(p.numel() for p in self.parameters() if p.requires_grad or not trainable_only)
+
+    def resize_token_embeddings(self, new_vocab_size: int) -> None:
+        old = self.token_embedding
+        if new_vocab_size == old.num_embeddings:
+            return
+        new = nn.Embedding(new_vocab_size, self.config.d_model).to(old.weight.device, old.weight.dtype)
+        nn.init.normal_(new.weight, mean=0.0, std=0.02)
+        keep = min(new_vocab_size, old.num_embeddings)
+        with torch.no_grad():
+            new.weight[:keep] = old.weight[:keep]
+        self.token_embedding = new
+        self.config.vocab_size = new_vocab_size
+
+
+def build_embedding_model(config: ModelConfig) -> Anthracite1EmbeddingModel:
+    return Anthracite1EmbeddingModel(config)