alphagenome-pytorch 0.0.9__tar.gz → 0.0.11__tar.gz

{alphagenome_pytorch-0.0.9 → alphagenome_pytorch-0.0.11}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: alphagenome-pytorch
-Version: 0.0.9
+Version: 0.0.11
 Summary: AlphaGenome
 Project-URL: Homepage, https://pypi.org/project/alphagenome-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/alphagenome
@@ -58,13 +58,13 @@ $ pip install alphagenome-pytorch
 
 ```python
 import torch
-from alphagenome_pytorch import TransformerTower
+from alphagenome_pytorch import AlphaGenome
 
-transformer = TransformerTower(dim = 768, dim_pairwise = 128)
+model = AlphaGenome()
 
-single = torch.randn(2, 512, 768)
+dna = torch.randint(0, 5, (2, 8192))
 
-attended_single, attended_pairwise = transformer(single)
+pred_nucleotide, single, pairwise = model(dna) # (2, 8192, 5), (2, 64, 1536), (2, 4, 4, 1536)
 ```
 
 ## Citations

{alphagenome_pytorch-0.0.9 → alphagenome_pytorch-0.0.11}/README.md

@@ -14,13 +14,13 @@ $ pip install alphagenome-pytorch
 
 ```python
 import torch
-from alphagenome_pytorch import TransformerTower
+from alphagenome_pytorch import AlphaGenome
 
-transformer = TransformerTower(dim = 768, dim_pairwise = 128)
+model = AlphaGenome()
 
-single = torch.randn(2, 512, 768)
+dna = torch.randint(0, 5, (2, 8192))
 
-attended_single, attended_pairwise = transformer(single)
+pred_nucleotide, single, pairwise = model(dna) # (2, 8192, 5), (2, 64, 1536), (2, 4, 4, 1536)
 ```
 
 ## Citations

alphagenome_pytorch-0.0.11/alphagenome_pytorch/__init__.py

@@ -0,0 +1,11 @@
+from alphagenome_pytorch.alphagenome import (
+    AlphaGenome,
+    Attention,
+    PairwiseRowAttention,
+    RelativePosFeatures,
+    RotaryEmbedding,
+    FeedForward,
+    TransformerTower,
+    UpresBlock,
+    DownresBlock,
+)

{alphagenome_pytorch-0.0.9 → alphagenome_pytorch-0.0.11}/alphagenome_pytorch/alphagenome.py

@@ -32,6 +32,9 @@ def exists(v):
 def divisible_by(num, den):
     return (num % den) == 0
 
+def last(arr):
+    return arr[-1]
+
 def is_odd(num):
     return not divisible_by(num, 2)
 
@@ -128,15 +131,15 @@ class UpresBlock(Module):
     def forward(
         self,
         x,
-        unet_skip = None
+        skip = None
     ):
 
         residual = x[:, :-self.pad]
         out = self.conv(x) + residual
 
-        if exists(unet_skip):
+        if exists(skip):
             out = repeat(out, 'b c n -> b c (n upsample)', upsample = 2) * self.residual_scale
-            out = out + self.unet_conv(unet_skip)
+            out = out + self.unet_conv(skip)
 
         return self.conv_out(out) + out
 
@@ -240,6 +243,7 @@ class Attention(Module):
         dim,
         dim_head = 64,
         heads = 8,
+        kv_heads = 1,
         dim_head_qk = 128,
         dim_head_v = 192,
         dim_pairwise = None,
@@ -254,8 +258,13 @@ class Attention(Module):
 
         # splitting and merging of attention heads
 
-        self.split_q_heads = Rearrange('b n (h d) -> b h n d', h = heads)
-        self.merge_heads = Rearrange('b h n d -> b n (h d)')
+        assert divisible_by(heads, kv_heads)
+        groups = heads // kv_heads
+
+        self.split_q_heads = Rearrange('b n (g h d) -> b g h n d', h = kv_heads, g = groups)
+        self.split_kv_heads = Rearrange('b n (h d) -> b h n d', h = kv_heads)
+
+        self.merge_heads = Rearrange('b g h n d -> b n (g h d)')
 
         # projections
 
@@ -274,7 +283,7 @@ class Attention(Module):
             nn.RMSNorm(dim_pairwise), # replace with BatchRMSNorm once crafted
             nn.GELU(),
             LinearNoBias(dim_pairwise, heads),
-            Rearrange('b i j h -> b h i j')
+            Rearrange('b i j (g h) -> b g h i j', g = groups)
         )
         # variables
 
@@ -293,6 +302,7 @@ class Attention(Module):
         # they use multi-query attention, with only 1 key / value head - pretty unconventional, but maybe enough for genomic modeling
 
         q = self.split_q_heads(q)
+        k, v = tuple(self.split_kv_heads(t) for t in (k, v))
 
         q, k, v = self.q_norm(q), self.k_norm(k), self.v_norm(v)
 
@@ -305,7 +315,7 @@ class Attention(Module):
 
         # similarities
 
-        sim = einsum(q, k, 'b h i d, b j d -> b h i j')
+        sim = einsum(q, k, 'b g h i d, b h j d -> b g h i j')
 
         # add attention bias + softclamping
 
@@ -315,7 +325,7 @@ class Attention(Module):
             assert divisible_by(sim.shape[-1], attn_bias.shape[-1])
             expand_factor = sim.shape[-1] // attn_bias.shape[-1]
 
-            attn_bias = repeat(attn_bias, 'b h i j -> b h (i r1) (j r2)', r1 = expand_factor, r2 = expand_factor)
+            attn_bias = repeat(attn_bias, 'b g h i j -> b g h (i r1) (j r2)', r1 = expand_factor, r2 = expand_factor)
 
             sim = softclamp(sim + attn_bias, value = self.softclamp_value)
 
@@ -325,7 +335,7 @@ class Attention(Module):
 
         # aggregate
 
-        out = einsum(attn, v, 'b h i j, b j d -> b h i d')
+        out = einsum(attn, v, 'b g h i j, b h j d -> b g h i d')
 
         out = self.merge_heads(out)
         return self.to_out(out)
@@ -384,7 +394,7 @@ class SingleToPairwise(Module):
             rel_q = relative_shift(einsum(q + q_rel_bias, rel_pos_encoding, 'b n h d, p h d -> b h n p'))
             rel_k = relative_shift(einsum(k + k_rel_bias, rel_pos_encoding, 'b n h d, p h d -> b h n p'))
 
-            rel_sim = add('b h i j, b h j i -> b h i j', rel_q, rel_k) * 0.5
+            rel_sim = add('b h i j, b h j i -> b i j h', rel_q, rel_k) * 0.5
 
             sim = sim + rel_sim
 
@@ -460,7 +470,7 @@ class TransformerTower(Module):
         dropout = 0.,
         ff_expansion_factor = 2.,
         max_positions = 8192,
-        dim_pairwise = None,
+        dim_pairwise = 128,
         pairwise_every_num_single_blocks = 2,   # how often to do a pairwise block
         single_to_pairwise_heads = 32,          # they did 32
         pool_size = 16,
@@ -539,7 +549,7 @@ class TransformerTower(Module):
                 pairwise = maybe_pairwise_attn(pairwise) + pairwise
                 pairwise = maybe_pairwise_ff(pairwise) + pairwise
 
-            single = attn(single, rotary_emb = rotary_emb, pairwise = None) + single
+            single = attn(single, rotary_emb = rotary_emb, pairwise = pairwise) + single
             single = ff(single) + single
 
         return single, pairwise
@@ -556,8 +566,10 @@ class DNAEmbed(Module):
         super().__init__()
         assert is_odd(width)
         self.dim_input = dim_input
-        self.conv = nn.Conv1d(dim_input, dim, width, padding = width // 2)
-        self.pointwise = nn.Conv1d(dim, dim, 1)
+        self.conv = Conv1d(dim_input, dim, width, padding = width // 2)
+        self.pointwise = Conv1d(dim, dim, 1)
+
+        self.pool = Reduce('b d (n pool) -> b d n', 'max', pool = 2)
 
     def forward(
         self,
@@ -568,38 +580,96 @@ class DNAEmbed(Module):
 
         out = self.conv(x)
         out = out + self.pointwise(out)
-        return rearrange(out, 'b d n -> b n d')
+        pooled = self.pool(out) # think they downsample for dna embed block
+
+        return pooled, x
 
 # classes
 
 class AlphaGenome(Module):
     def __init__(
         self,
-        dim = 768,
+        dims: tuple[int, ...] = (
+            768,
+            896,
+            1024,
+            1152,
+            1280,
+            1408,
+            1536
+        ),
         basepairs = 5,
         dna_embed_width = 15,
-        dim_pairwise = None,
         transformer_kwargs: dict = dict()
     ):
         super().__init__()
+
         assert is_odd(dna_embed_width)
 
-        self.to_dna_embed = DNAEmbed(dim, dim_input = basepairs, width = dna_embed_width)
+        assert len(dims) >= 2
+        first_dim, *_, last_dim = dims
+
+        self.dna_embed = DNAEmbed(first_dim, dim_input = basepairs, width = dna_embed_width)
+
+        dim_with_input = (basepairs, *dims)
+        dim_pairs = zip(dim_with_input[:-1], dim_with_input[1:])
+
+        downs = []
+        ups = []
+
+        for layer_num, (dim_in, dim_out) in enumerate(dim_pairs, start = 1):
+            is_first = layer_num == 1
+            channel_diff = dim_out - dim_in
+
+            assert channel_diff > 0
+
+            if not is_first:
+                down = DownresBlock(dim_in, channels_to_add = channel_diff)
+                downs.append(down)
 
-        self.transformer = Transformer(
-            dim = dim,
-            dim_pairwise = dim_pairwise,
+            up = UpresBlock(dim_out, channels_to_remove = channel_diff)
+            ups.insert(0, up)
+
+
+        self.downs = ModuleList(downs)
+        self.ups = ModuleList(ups)
+
+        self.transformer = TransformerTower(
+            dim = last_dim,
             **transformer_kwargs
         )
 
     def forward(
         self,
-        seq,
-        pairwise
+        seq # Int['b n']
     ):
 
-        dna_embed = self.to_dna_embed(seq)
+        skips = []
+
+        # embed with one hot and add skip
+
+        x, skip = self.dna_embed(seq)
+        skips.append(skip)
+
+        # downs
+
+        for down in self.downs:
+            skips.append(x)
+            x = down(x)
+
+        x = rearrange(x, 'b d n -> b n d')
+
+        # attention
+
+        single, pairwise = self.transformer(x)
+
+        # ups with skips from down
+
+        x = rearrange(x, 'b n d -> b d n')
+
+        for up in self.ups:
+            x = up(x, skip = skips.pop())
 
-        attended = self.transformer(dna_embed)
+        pred = rearrange(x, 'b l n -> b n l') # 1bp resolution
 
-        return attended
+        return pred, single, pairwise

{alphagenome_pytorch-0.0.9 → alphagenome_pytorch-0.0.11}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "alphagenome-pytorch"
-version = "0.0.9"
+version = "0.0.11"
 description = "AlphaGenome"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

{alphagenome_pytorch-0.0.9 → alphagenome_pytorch-0.0.11}/tests/test_alphagenome.py

@@ -21,3 +21,16 @@ def test_down_up():
 
     x = torch.randn(1, 64, 8)
     assert up(down(x), x).shape == x.shape
+
+def test_alphagenome():
+    from alphagenome_pytorch import AlphaGenome
+
+    model = AlphaGenome()
+
+    dna = torch.randint(0, 5, (2, 8192))
+
+    pred_nucleotide_logits, single, pairwise = model(dna)
+
+    pred = pred_nucleotide_logits.argmax(dim = -1)
+
+    assert pred.shape == dna.shape

alphagenome_pytorch-0.0.9/alphagenome_pytorch/__init__.py

@@ -1,5 +0,0 @@
-from alphagenome_pytorch.alphagenome import (
-    AlphaGenome,
-    Attention,
-    TransformerTower
-)