ablanx 0.1.0__py3-none-any.whl

ablang_jax/__init__.py

@@ -0,0 +1,32 @@
+"""ablang_jax - a JAX/Flax port of AbLang2 (OPIG).
+
+This package re-implements AbLang2's AbRep encoder in Flax and provides the state_dict key
+mapping needed to load the original BSD-3-Clause PyTorch weights unchanged. The AbLang2 model
+and weights are the work of the AbLang2 authors (OPIG); this is a port. See ATTRIBUTION.md and
+https://github.com/oxpig/AbLang2.
+"""
+from .model import (
+    HEAD_DIM,
+    HID,
+    LN_EPS,
+    N_BLOCK,
+    N_HEAD,
+    VOCAB,
+    Ablanx,
+    Block,
+    load_ablanx_params,
+)
+
+__all__ = [
+    "Ablanx",
+    "Block",
+    "load_ablanx_params",
+    "VOCAB",
+    "HID",
+    "N_HEAD",
+    "N_BLOCK",
+    "HEAD_DIM",
+    "LN_EPS",
+]
+
+__version__ = "0.1.0"

ablang_jax/model.py

@@ -0,0 +1,109 @@
+"""Flax reimplementation of AbLang2's AbRep encoder.
+
+A 12-block pre-norm transformer (hidden 480, 20 heads, SwiGLU feed-forward, rotary positions,
+vocab 26) that reproduces AbLang2's AbRep so the original BSD-3-Clause PyTorch weights load
+unchanged. A single forward returns per-residue hidden states and per-block attention maps.
+
+PyTorch Linear weights [out, in] map to Flax Dense kernels [in, out] by transpose; see
+load_ablanx_params for the full state_dict key mapping. Outputs match reference PyTorch AbLang2 to under
+4e-6 across a 30-Fv panel (test_agreement.py).
+
+AbLang2 is the work of the Oxford Protein Informatics Group; this is a port. See ATTRIBUTION.md and
+https://github.com/oxpig/AbLang2.
+"""
+from __future__ import annotations
+
+import flax.linen as nn
+import jax
+import jax.numpy as jnp
+
+VOCAB = 26
+HID = 480
+N_HEAD = 20
+N_BLOCK = 12
+HEAD_DIM = HID // N_HEAD          # 24
+LN_EPS = 1e-12
+
+
+def _rope_freqs(head_dim, base=10000.0):
+    # inverse frequencies over half the head dim (matches AbLang2's rotary_emb.freqs)
+    half = head_dim // 2
+    return 1.0 / (base ** (jnp.arange(0, half, dtype=jnp.float32) * 2.0 / head_dim))
+
+
+def _apply_rope(x, freqs):
+    """Rotary embedding on x [B, L, H, D], interleaved convention (as AbLang2 uses)."""
+    L = x.shape[1]
+    pos = jnp.arange(L, dtype=jnp.float32)
+    ang = pos[:, None] * freqs[None, :]
+    ang = jnp.repeat(ang, 2, axis=-1)
+    cos = jnp.cos(ang)[None, :, None, :]
+    sin = jnp.sin(ang)[None, :, None, :]
+    x_even = x[..., 0::2]; x_odd = x[..., 1::2]
+    rot = jnp.stack([-x_odd, x_even], axis=-1).reshape(x.shape)
+    return x * cos + rot * sin
+
+
+class Block(nn.Module):
+    @nn.compact
+    def __call__(self, x, mask, freqs, return_attn=False):
+        B, L, _ = x.shape
+        h = nn.LayerNorm(epsilon=LN_EPS, name="pre_attn_layer_norm")(x)
+        q = nn.Dense(HID, name="q_proj")(h).reshape(B, L, N_HEAD, HEAD_DIM)
+        k = nn.Dense(HID, name="k_proj")(h).reshape(B, L, N_HEAD, HEAD_DIM)
+        v = nn.Dense(HID, name="v_proj")(h).reshape(B, L, N_HEAD, HEAD_DIM)
+        q = _apply_rope(q, freqs); k = _apply_rope(k, freqs)
+        # AbLang2's net attention scale is 1/head_dim (q scaled by head_dim^-0.5, then divided by sqrt(head_dim))
+        logits = jnp.einsum("blhd,bmhd->bhlm", q, k) / HEAD_DIM
+        logits = logits + (1.0 - mask[:, None, None, :]) * -1e9
+        attn = jax.nn.softmax(logits, axis=-1)
+        o = jnp.einsum("bhlm,bmhd->blhd", attn, v).reshape(B, L, HID)
+        o = nn.Dense(HID, name="out_proj")(o)
+        x = x + o
+        # SwiGLU: split the feed-forward projection into value and gate, gate with silu
+        h2 = nn.LayerNorm(epsilon=LN_EPS, name="final_layer_norm")(x)
+        up = nn.Dense(3840, name="ffn_in")(h2)
+        val, gate = jnp.split(up, 2, axis=-1)
+        x = x + nn.Dense(HID, name="ffn_out")(jax.nn.silu(gate) * val)
+        return (x, attn) if return_attn else (x, None)
+
+
+class Ablanx(nn.Module):
+    """AbLang2 AbRep encoder. Returns hidden [B, L, 480] and, optionally, attention [N_BLOCK, B, H, L, L]."""
+    @nn.compact
+    def __call__(self, tokens, mask, return_attn=True):
+        freqs = _rope_freqs(HEAD_DIM)
+        x = nn.Embed(VOCAB, HID, name="aa_embed_layer")(tokens.astype("int32"))
+        attns = []
+        for i in range(N_BLOCK):
+            x, a = Block(name=f"encoder_blocks_{i}")(x, mask, freqs, return_attn=return_attn)
+            if return_attn:
+                attns.append(a)
+        # final layer norm after all blocks; this is what rescoding returns
+        x = nn.LayerNorm(epsilon=LN_EPS, name="layer_norm_after_encoder_blocks")(x)
+        return x, (jnp.stack(attns) if return_attn else None)
+
+
+def load_ablanx_params(w):
+    """Map AbLang2 AbRep PyTorch state_dict (dict of numpy arrays, keys 'AbRep....') -> flax param tree.
+    PyTorch Linear weight [out,in] -> flax Dense kernel [in,out] (transpose). LayerNorm weight->scale."""
+    import numpy as np
+    def T(k): return np.asarray(w[k]).T
+    def V(k): return np.asarray(w[k])
+    p = {"aa_embed_layer": {"embedding": V("AbRep.aa_embed_layer.weight")},
+         "layer_norm_after_encoder_blocks": {"scale": V("AbRep.layer_norm_after_encoder_blocks.weight"),
+                                             "bias": V("AbRep.layer_norm_after_encoder_blocks.bias")}}
+    for i in range(N_BLOCK):
+        b = f"AbRep.encoder_blocks.{i}"
+        p[f"encoder_blocks_{i}"] = {
+            "pre_attn_layer_norm": {"scale": V(f"{b}.pre_attn_layer_norm.weight"),
+                                    "bias": V(f"{b}.pre_attn_layer_norm.bias")},
+            "q_proj": {"kernel": T(f"{b}.multihead_attention.q_proj.weight"), "bias": V(f"{b}.multihead_attention.q_proj.bias")},
+            "k_proj": {"kernel": T(f"{b}.multihead_attention.k_proj.weight"), "bias": V(f"{b}.multihead_attention.k_proj.bias")},
+            "v_proj": {"kernel": T(f"{b}.multihead_attention.v_proj.weight"), "bias": V(f"{b}.multihead_attention.v_proj.bias")},
+            "out_proj": {"kernel": T(f"{b}.multihead_attention.out_proj.weight"), "bias": V(f"{b}.multihead_attention.out_proj.bias")},
+            "ffn_in": {"kernel": T(f"{b}.intermediate_layer.0.weight"), "bias": V(f"{b}.intermediate_layer.0.bias")},
+            "ffn_out": {"kernel": T(f"{b}.intermediate_layer.2.weight"), "bias": V(f"{b}.intermediate_layer.2.bias")},
+            "final_layer_norm": {"scale": V(f"{b}.final_layer_norm.weight"), "bias": V(f"{b}.final_layer_norm.bias")},
+        }
+    return p

ablang_jax/precompute.py

@@ -0,0 +1,80 @@
+#!/usr/bin/env python3
+"""ablang_jax.precompute - precompute AbLang2 per-residue embeddings for a set of antibody records.
+
+Runs the reference PyTorch AbLang2 (`pip install ablang2`) over antibody Fv sequences stored as
+sharded .npz files and writes a per-record `seq_emb` field ([L, 480], VH then VL, residue positions
+only). These embeddings can be used as a fixed antibody sequence prior.
+
+This uses the reference PyTorch AbLang2, not the JAX port: it materialises per-residue embeddings
+with the original model for use as a fixed antibody sequence prior.
+
+Expected input: one or more npz shards named `{split}_*.npz` (split in {train, test}) under the data
+directory, each an object-array store with per-record fields:
+    pdb_id      [N]        record id
+    aa          [N][Li]    integer amino-acid indices in the order "ARNDCQEGHILKMFPSTWYV"
+    chain_id    [N][Li]    0 = heavy (VH), 1 = light (VL)
+    res_present [N][Li]    bool mask of resolved residues
+Each shard is rewritten in place with an added `seq_emb` object field.
+
+    python -m ablang_jax.precompute --data /path/to/records
+    # or set ABLANG_JAX_DATA (default: ./out/records)
+"""
+from __future__ import annotations
+
+import argparse
+import glob
+import os
+
+import ablang2
+import numpy as np
+
+_AA_ORDER = "ARNDCQEGHILKMFPSTWYV"   # AlphaFold-style restype index order
+
+_DEFAULT_OUT = os.environ.get("ABLANG_JAX_OUT", "./out")
+_DEFAULT_DATA = os.environ.get("ABLANG_JAX_DATA", os.path.join(_DEFAULT_OUT, "records"))
+
+
+def embed_fv(model, vh_seq, vl_seq):
+    """[L,480] per-residue AbLang2 embedding, residue positions only (ids 1..20), VH then VL."""
+    pair = [[vh_seq, vl_seq]] if vl_seq else [[vh_seq, ""]]
+    emb = np.asarray(model(pair, mode="rescoding")[0], np.float32)          # [T_tok, 480]
+    ids = np.asarray(model.tokenizer(pair, pad=True))[0]                    # [T_tok]
+    keep = (ids >= 1) & (ids <= 20)                                         # 20 standard AAs
+    return emb[: len(ids)][keep[: emb.shape[0]]]                            # [n_residues, 480]
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--data", default=_DEFAULT_DATA,
+                    help="directory of {train,test}_*.npz shards (env: ABLANG_JAX_DATA)")
+    ap.add_argument("--limit-shards", type=int, default=0)
+    args = ap.parse_args()
+    model = ablang2.pretrained()
+
+    for split in ("train", "test"):
+        shards = sorted(glob.glob(os.path.join(args.data, f"{split}_*.npz")))
+        if args.limit_shards:
+            shards = shards[: args.limit_shards]
+        for si, shard in enumerate(shards):
+            z = np.load(shard, allow_pickle=True); arrs = {k: z[k] for k in z.files}
+            n = len(arrs["pdb_id"])
+            embs = []
+            for i in range(n):
+                aa = np.asarray(arrs["aa"][i]); cid = np.asarray(arrs["chain_id"][i])
+                pres = np.asarray(arrs["res_present"][i], bool)
+                vh = "".join(_AA_ORDER[a] for a in aa[(cid == 0) & pres])
+                vl = "".join(_AA_ORDER[a] for a in aa[(cid == 1) & pres])
+                e = embed_fv(model, vh, vl)                                 # [len(vh)+len(vl), 480]
+                exp = len(vh) + len(vl)
+                if e.shape[0] != exp:                                        # alignment guard
+                    e = e[:exp] if e.shape[0] > exp else np.pad(e, [(0, exp - e.shape[0]), (0, 0)])
+                embs.append(e.astype(np.float32))
+            arrs["seq_emb"] = np.array(embs, dtype=object)
+            np.savez_compressed(shard, **arrs)
+            if si % 5 == 0:
+                print(f"{split} shard {si+1}/{len(shards)} emb[0]={embs[0].shape}", flush=True)
+        print(f"{split}: DONE ({len(shards)} shards)", flush=True)
+
+
+if __name__ == "__main__":
+    main()

ablanx-0.1.0.dist-info/METADATA

@@ -0,0 +1,159 @@
+Metadata-Version: 2.4
+Name: ablanx
+Version: 0.1.0
+Summary: JAX/Flax port of AbLang2's AbRep encoder: antibody language-model embeddings and per-block attention.
+Author: Fabricagen
+License: BSD-3-Clause
+Project-URL: Homepage, https://github.com/fabricagen/ablanx
+Project-URL: Repository, https://github.com/fabricagen/ablanx
+Keywords: jax,flax,antibody,protein-language-model,ablang,ablang2
+Classifier: License :: OSI Approved :: BSD License
+Classifier: Programming Language :: Python :: 3
+Classifier: Intended Audience :: Science/Research
+Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: jax>=0.4.20
+Requires-Dist: jaxlib>=0.4.20
+Requires-Dist: flax>=0.7.5
+Requires-Dist: numpy>=1.24
+Provides-Extra: reference
+Requires-Dist: ablang2; extra == "reference"
+Requires-Dist: torch; extra == "reference"
+Dynamic: license-file
+
+# ablanx
+
+ablanx is a faithful JAX/Flax port of the AbLang2 AbRep encoder. It loads the original AbLang2
+weights unchanged and returns per-residue embeddings and per-block attention in a single JAX forward
+pass, so an antibody sequence prior can sit inside a larger JAX/Flax model.
+
+AbLang2 is the work of the Oxford Protein Informatics Group (OPIG); this repository is a port, not a
+new model. Provenance and citation are in [ATTRIBUTION.md](ATTRIBUTION.md) and
+[CITATION.cff](CITATION.cff).
+
+## Why use it
+
+- Drop-in AbLang2 embeddings inside JAX/Flax workflows.
+- No PyTorch dependency at inference (weights load from a plain npz).
+- Pure JAX forward: `jit` / `vmap` / `grad`-compatible and composable in a larger model.
+- Per-block attention returned directly from the same call.
+- Numerically verified against the reference implementation (max abs 3.9e-6 across a 30-Fv panel).
+
+## Quickstart
+
+Install:
+
+    pip install git+https://github.com/fabricagen/ablanx
+    # from a clone: pip install .
+    # imports as `ablang_jax`; PyPI (pip install ablanx) is planned.
+
+Embeddings:
+
+    import numpy as np, jax.numpy as jnp
+    from ablang_jax import Ablanx, load_ablanx_params
+
+    model = Ablanx()
+    params = {"params": load_ablanx_params(dict(np.load("ablang2_weights.npz")))}  # AbRep weights
+    tokens = jnp.array([[...]], dtype=jnp.int32)         # AbLang2 token ids, shape [B, L]
+    mask = jnp.ones_like(tokens, dtype=jnp.float32)      # 1 = keep, 0 = pad
+    hidden, attentions = model.apply(params, tokens, mask, return_attn=True)
+    # hidden:     [B, L, 480]         per-residue embeddings
+    # attentions: [12, B, 20, L, L]   per-block attention maps
+
+See [Weights](#weights) for the npz.
+
+## What it is
+
+AbLang2 is an antibody-specific protein language model. Its encoder (AbRep) is a 12-block pre-norm
+transformer: hidden 480, 20 heads, SwiGLU feed-forward, rotary position embeddings, vocab 26,
+trained on paired antibody sequences. `ablanx` re-expresses AbRep in Flax so the original
+BSD-3-Clause weights load without change, returning per-residue hidden states and per-block attention
+maps in one forward pass.
+
+## Differentiability
+
+The forward is a pure JAX function, so it composes inside a larger model and runs under `jit`,
+`vmap`, and `grad`. Useful gradients flow with respect to:
+
+- the model parameters (`params`), and
+- any continuous input substituted for the token lookup, for example a soft or one-hot embedding
+  passed in place of integer ids.
+
+The public interface takes integer `tokens`, which are discrete and do not themselves admit
+gradients; to optimize a sequence through the encoder, feed a continuous representation rather than
+token ids. In the seam folder, ablanx runs frozen as a fixed sequence prior.
+
+## What works today
+
+- The Flax AbRep encoder, `ablang_jax.model.Ablanx`.
+- The PyTorch-to-Flax weight key mapping, `ablang_jax.model.load_ablanx_params`, which maps an
+  AbLang2 AbRep `state_dict` (as numpy arrays) into the Flax parameter tree.
+- A weight exporter, `export_weights.py`, that writes the AbRep weight npz from the reference model.
+- A precompute script, `ablang_jax.precompute`, that runs the reference PyTorch AbLang2 over a set of
+  antibody records and stores per-residue embeddings.
+
+## Validation
+
+The JAX forward reproduces reference PyTorch AbLang2 to float32 precision. Across a 30-Fv panel (25
+paired VH+VL and 5 VH-only, drawn from public PDB structures), per-residue embeddings match the
+reference with maximum absolute difference 3.9e-6 (median 2.4e-6 per Fv) and cosine above 0.999999,
+using the same weights. The weights are byte-identical to the reference AbRep state_dict (207 tensors,
+exact match). Reproduce with `test_agreement.py`.
+
+![Per-residue embedding agreement across 30 antibody Fvs: identity scatter (cosine 1.000000) and absolute-difference distribution (max 3.9e-6).](agreement.png)
+
+## Weights
+
+The weights are the original AbLang2 AbRep weights, unmodified (207 tensors).
+
+- Export from the reference model (works today, no release needed):
+  `pip install ablang2 torch && python export_weights.py` writes `ablang2_weights.npz`, verified
+  byte-identical to the reference.
+- A torch-free `ablang2_weights.npz` is also attached to the GitHub release once one is tagged.
+
+Point `ABLANG_WEIGHTS` at the npz, or pass `dict(np.load(...))` to `load_ablanx_params`.
+
+## Not included
+
+- Only the AbRep encoder (embeddings and attention) is ported. The amino-acid likelihood head used
+  for the pseudo-log-likelihood (sequence naturalness) is not part of this encoder port; use the
+  reference AbLang2 for likelihoods.
+
+## Precompute
+
+Precompute embeddings for a set of records with the reference PyTorch `ablang2`:
+
+    export ABLANG_JAX_DATA=/path/to/records
+    python -m ablang_jax.precompute --data $ABLANG_JAX_DATA
+
+Input records are npz shards named `{train,test}_*.npz`; see `ablang_jax/precompute.py` for the
+expected per-record fields. Outputs default to `./out/` when paths are not set.
+
+## Tests
+
+    python test_shapes.py                                          # forward, weight-key mapping, masking
+    ABLANG_WEIGHTS=ablang2_weights.npz python test_agreement.py    # agreement vs reference (needs weights)
+
+`test_shapes.py` needs no weights. `test_agreement.py` loads the committed golden panel
+(`golden_ablang2_panel.npz`, reference embeddings for 30 Fvs) and checks the JAX forward matches every
+one; it skips under pytest if `ABLANG_WEIGHTS` is unset.
+
+## Attribution
+
+- Original model, training, and weights: AbLang2, Oxford Protein Informatics Group.
+  https://github.com/oxpig/AbLang2 (BSD-3-Clause, Copyright (c) 2021, Tobias Hegelund Olsen).
+- Paper: Olsen, Moal, Deane, "Addressing the antibody germline bias and its effect on language models
+  for improved antibody design", bioRxiv 2024, doi:10.1101/2024.02.02.578678.
+- This port: Fabricagen.
+
+ablanx is the language-model component of the **seam** bundle, a JAX antibody Fv structure predictor
+that couples jaxfvld with ablanx. For antibody developability screening and repair, see **sift**:
+https://sift.fabricagen.ai (coming soon, ca 08/26)
+
+## License
+
+BSD-3-Clause. This port preserves the original AbLang2 copyright (Copyright (c) 2021, Tobias Hegelund
+Olsen) and adds Copyright (c) 2026, Fabricagen for the port. The upstream AbLang2 license was
+confirmed BSD-3-Clause. See `ATTRIBUTION.md`.

ablanx-0.1.0.dist-info/RECORD

@@ -0,0 +1,8 @@
+ablang_jax/__init__.py,sha256=rqYVJ6ndHNvxWvjfoQRMqfBP-cvw2hFCXJF0vkob-mk,685
+ablang_jax/model.py,sha256=vAxIkeUhpWIgciYke5-pk8geR3mqL6SjT1xNKCUlPPU,5471
+ablang_jax/precompute.py,sha256=M_njZWBc64QhH_qYkXW_gKbAN7PMAq3xRz4lCecOB0s,3795
+ablanx-0.1.0.dist-info/licenses/LICENSE,sha256=p3uEdiNWxk8j27b7COMXbRE8FrqcNc30BELSyZKF38Q,1763
+ablanx-0.1.0.dist-info/METADATA,sha256=7CuM54ElBcSKq9hfeOa6rH15kacXSxYjMUriFkULnxo,7304
+ablanx-0.1.0.dist-info/WHEEL,sha256=K260EYznzXsJYBQGqmI8VTxEdiZYNvDZwW9cBh9-_MA,91
+ablanx-0.1.0.dist-info/top_level.txt,sha256=WASCZ06W0VkKEMRHzov2LbvmuJAJP1E9y_HZ1InHP-I,11
+ablanx-0.1.0.dist-info/RECORD,,

ablanx-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (83.0.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

ablanx-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,33 @@
+BSD 3-Clause License
+
+Copyright (c) 2021, Tobias Hegelund Olsen
+Copyright (c) 2026, Fabricagen
+
+This software is a JAX/Flax port of AbLang2 (Oxford Protein Informatics Group). The
+original AbLang2 model, weights, and copyright are the work of the AbLang2 authors and
+are retained above. See ATTRIBUTION.md for details.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from
+   this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

ablanx-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+ablang_jax