broccoli-ml 5.1.0__tar.gz → 7.0.0__tar.gz

{broccoli_ml-5.1.0 → broccoli_ml-7.0.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: broccoli-ml
-Version: 5.1.0
+Version: 7.0.0
 Summary: Some useful Pytorch models, circa 2025
 License: MIT
 Author: Nicholas Bailey

{broccoli_ml-5.1.0 → broccoli_ml-7.0.0}/broccoli/transformer.py

@@ -13,6 +13,7 @@ from .rope import RotaryEmbedding, apply_rotary_emb
 try:
     from flash_attn import flash_attn_func
 
+    print("Using flash-attn.")
     FLASH_ATTN = True
 except ImportError:
     pass
@@ -76,7 +77,7 @@ class MHAttention(nn.Module):
         causal=False,
         seq_len=None,
         linear_module: nn.Module = nn.Linear,
-        bos_tokens=0,
+        utility_tokens=0,
         rotary_embedding=None,
         source_size=None,
         scaling="d",
@@ -129,7 +130,7 @@ class MHAttention(nn.Module):
             )
         self.rotary_embedding = rotary_embedding
         self.source_size = source_size
-        self.bos_tokens = bos_tokens
+        self.utility_tokens = utility_tokens
 
         self.reset_parameters()
 
@@ -156,7 +157,7 @@ class MHAttention(nn.Module):
         self, q: torch.Tensor, k: torch.Tensor
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         """
-        Apply Axial RoPE to all tokens except BOS tokens
+        Apply Axial RoPE to all tokens except utility tokens
         """
 
         if len(self.source_size) == 1:
@@ -180,8 +181,8 @@ class MHAttention(nn.Module):
                 "`source_size` must be a tuple of 1, 2 or 3 integers"
             )
 
-        q_bos, q_img = q[:, : self.bos_tokens, :], q[:, self.bos_tokens :, :]
-        k_bos, k_img = k[:, : self.bos_tokens, :], k[:, self.bos_tokens :, :]
+        q_util, q_img = q[:, : self.utility_tokens, :], q[:, self.utility_tokens :, :]
+        k_util, k_img = k[:, : self.utility_tokens, :], k[:, self.utility_tokens :, :]
 
         q_img = rearrange(
             q_img,
@@ -208,9 +209,9 @@ class MHAttention(nn.Module):
             f"b {spatial_dimension_names} d -> b ({spatial_dimension_names}) d",
         )
 
-        # Re-combine the BOS tokens and the RoPE-enhanced image tokens
-        q = torch.cat([q_bos, q_img], dim=1)
-        k = torch.cat([k_bos, k_img], dim=1)
+        # Re-combine the utility tokens and the RoPE-enhanced sequence tokens
+        q = torch.cat([q_util, q_img], dim=1)
+        k = torch.cat([k_util, k_img], dim=1)
 
         return q, k
 
@@ -284,7 +285,7 @@ class MHAttention(nn.Module):
 
             return self.out_proj(output_without_heads)
 
-    def attention_scores(self, q, k, v):
+    def attention_logits(self, q, k, v):
 
         q, k, v = self.project_qkv(q, k, v)
 
@@ -301,8 +302,6 @@ class MHAttention(nn.Module):
         if self.causal:
             qk_scores.masked_fill_(self.mask, float("-inf"))
 
-        qk_scores = F.softmax(qk_scores, dim=-1)
-
         return qk_scores  # (batch, head, seq_len, seq_len)
 
     def reset_parameters(self):
@@ -416,7 +415,7 @@ class TransformerBlock(nn.Module):
         n_heads,
         relative_position_embedding=False,
         source_size=None,
-        bos_tokens=0,
+        utility_tokens=0,
         mlp_ratio=4,
         activation: nn.Module = nn.ReLU,
         activation_kwargs: Optional[dict] = None,
@@ -474,7 +473,7 @@ class TransformerBlock(nn.Module):
             linear_module=linear_module,
             rotary_embedding=self.rotary_embedding,
             source_size=source_size,
-            bos_tokens=bos_tokens,
+            utility_tokens=utility_tokens,
             scaling=msa_scaling,
         )
 
@@ -529,15 +528,15 @@ class TransformerBlock(nn.Module):
 
         return x
 
-    def attention_scores(self, x):
+    def attention_logits(self, x):
         """
         Give back the attention scores used in this layer.
         """
         if self.pre_norm:
             x = self.layer_norm_1(x)
-            return self.attn(x, x, x)
+            return self.attn.attention_logits(x, x, x)
         else:
-            return self.attn(x, x, x)
+            return self.attn.attention_logits(x, x, x)
 
     def reset_parameters(self):
         self.layer_norm_1.reset_parameters()
@@ -573,8 +572,8 @@ class TransformerEncoder(nn.Module):
         stochastic_depth=0.0,
         causal=False,
         linear_module=nn.Linear,
-        bos_tokens=0,
-        return_bos_tokens=False,
+        utility_tokens=0,
+        return_utility_tokens=False,
         pre_norm=True,
         post_norm=False,
         normformer=False,
@@ -592,22 +591,33 @@ class TransformerEncoder(nn.Module):
         if relative_position_embedding and (source_size is None):
             raise ValueError(
                 "`source_size` for TransformerEncoder cannot be None if"
-                " `position_embedding_type` is relative"
+                " `relative_position_embedding` is True"
+            )
+
+        if absolute_position_embedding and (seq_len is None):
+            raise ValueError(
+                "`seq_len` for TransformerEncoder cannot be None if"
+                " `absolute_position_embedding` is True"
             )
 
         super().__init__()
         self.seq_len = seq_len
         self.n_heads = n_heads
-        self._bos_tokens = bos_tokens
-        self.return_bos_tokens = return_bos_tokens
-
-        # Initialise BOS tokens with normal init, like usual Pytorch embeddings
-        if self._bos_tokens:
-            self._bos_embedding = nn.Parameter(torch.empty(self._bos_tokens, d_model))
-            nn.init.normal_(self._bos_embedding, mean=0.0, std=1.0)
-            self.full_sequence_length = self.seq_len + self._bos_tokens
+        self._utility_tokens = utility_tokens
+        self.return_utility_tokens = return_utility_tokens
+
+        # Initialise utility tokens with normal init, like usual Pytorch embeddings
+        if self._utility_tokens:
+            self._utility_token_embedding = nn.Parameter(
+                torch.empty(self._utility_tokens, d_model)
+            )
+            nn.init.normal_(self._utility_token_embedding, mean=0.0, std=1.0)
+        else:
+            self._utility_token_embedding = None
+
+        if self._utility_tokens and (self.seq_len is not None):
+            self.full_sequence_length = self.seq_len + self._utility_tokens
         else:
-            self._bos_embedding = None
             self.full_sequence_length = self.seq_len
 
         self.d_model = d_model
@@ -641,7 +651,7 @@ class TransformerEncoder(nn.Module):
                     n_heads,
                     relative_position_embedding=relative_position_embedding,
                     source_size=source_size,
-                    bos_tokens=bos_tokens,
+                    utility_tokens=utility_tokens,
                     mlp_ratio=mlp_ratio,
                     activation=activation,
                     activation_kwargs=activation_kwargs,
@@ -669,8 +679,10 @@ class TransformerEncoder(nn.Module):
         return ",".join([str(block._kv_distance) for block in self.blocks])
 
     def preprocess(self, x):
-        if self._bos_tokens:
-            x = torch.cat([self._bos_embedding.expand(x.size(0), -1, -1), x], dim=1)
+        if self._utility_tokens:
+            x = torch.cat(
+                [self._utility_token_embedding.expand(x.size(0), -1, -1), x], dim=1
+            )
         else:
             x = x
 
@@ -683,6 +695,8 @@ class TransformerEncoder(nn.Module):
                 )  # to shape (1, seq_len) to broadcast over batch
             )
 
+        return x
+
     def forward(self, x):
 
         x = self.preprocess(x)
@@ -690,12 +704,12 @@ class TransformerEncoder(nn.Module):
         for block in self.blocks:
             x = block(x)
 
-        if self._bos_tokens and not self.return_bos_tokens:
-            return x[:, self._bos_tokens :, :]
+        if self._utility_tokens and not self.return_utility_tokens:
+            return x[:, self._utility_tokens :, :]
         else:
             return x
 
-    def attention_scores(self, x):
+    def attention_logits(self, x):
 
         x = self.preprocess(x)
 
@@ -703,15 +717,15 @@ class TransformerEncoder(nn.Module):
 
         for block in self.blocks:
             # Get attention scores with shape (batch, 1, head, seq_len, seq_len)
-            layer_attention_scores = block.attention_scores(x).unsqueeze(1)
-            layer_scores.append(layer_attention_scores)
+            layer_attention_logits = block.attention_logits(x).unsqueeze(1)
+            layer_scores.append(layer_attention_logits)
             x = block(x)
 
         return torch.cat(layer_scores, dim=1)  # (batch, layer, head, seq_len, seq_len)
 
     def reset_parameters(self):
-        if self._bos_embedding is not None:
-            nn.init.normal_(self._bos_embedding, mean=0.0, std=1.0)
+        if self._utility_token_embedding is not None:
+            nn.init.normal_(self._utility_token_embedding, mean=0.0, std=1.0)
 
         if self.absolute_position_embedding is not None:
             self.absolute_position_embedding.reset_parameters()

{broccoli_ml-5.1.0 → broccoli_ml-7.0.0}/broccoli/vit.py

@@ -11,7 +11,6 @@ from einops.layers.torch import Rearrange
 
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 
 
 class GetCLSToken(nn.Module):
@@ -39,6 +38,9 @@ class SequencePool(nn.Module):
         weights = self.attention(x)
         return einsum(weights, x, "batch seq, batch seq d_model -> batch d_model")
 
+    def attention_scores(self, x):
+        return self.attention(x)
+
     def reset_parameters(self):
         # Iterate over modules in the sequential block
         for module in self.attention:
@@ -169,8 +171,8 @@ class ViTEncoder(nn.Module):
         transformer_layers=7,
         transformer_heads=4,
         transformer_mlp_ratio=2,
-        transformer_bos_tokens=0,
-        transformer_return_bos_tokens=False,
+        transformer_utility_tokens=0,
+        transformer_return_utility_tokens=False,
         transformer_activation: nn.Module = SquaredReLU,
         transformer_activation_kwargs: Optional[dict] = None,
         transformer_ff_linear_module_up=None,
@@ -336,8 +338,8 @@ class ViTEncoder(nn.Module):
                 stochastic_depth=transformer_stochastic_depth,
                 causal=False,
                 linear_module=linear_module,
-                bos_tokens=transformer_bos_tokens,
-                return_bos_tokens=transformer_return_bos_tokens,
+                utility_tokens=transformer_utility_tokens,
+                return_utility_tokens=transformer_return_utility_tokens,
                 pre_norm=transformer_pre_norm,
                 normformer=transformer_normformer,
                 post_norm=transformer_post_norm,
@@ -400,9 +402,9 @@ class ViTEncoder(nn.Module):
     def forward(self, x):
         return self.encoder(x)
 
-    def attention_scores(self, x):
+    def attention_logits(self, x):
         x = self.encoder[:-1](x)
-        return self.encoder[-1].attention_scores(x)
+        return self.encoder[-1].attention_logits(x)
 
     def reset_parameters(self):
         for module in self.encoder:
@@ -449,8 +451,8 @@ class ViT(nn.Module):
         transformer_layers=7,
         transformer_heads=4,
         transformer_mlp_ratio=2,
-        transformer_bos_tokens=0,
-        transformer_return_bos_tokens=False,
+        transformer_utility_tokens=0,
+        transformer_return_utility_tokens=False,
         transformer_activation: nn.Module = SquaredReLU,
         transformer_activation_kwargs: Optional[dict] = None,
         transformer_ff_linear_module_up=None,
@@ -516,8 +518,8 @@ class ViT(nn.Module):
             transformer_layers=transformer_layers,
             transformer_heads=transformer_heads,
             transformer_mlp_ratio=transformer_mlp_ratio,
-            transformer_bos_tokens=transformer_bos_tokens,
-            transformer_return_bos_tokens=transformer_return_bos_tokens,
+            transformer_utility_tokens=transformer_utility_tokens,
+            transformer_return_utility_tokens=transformer_return_utility_tokens,
             transformer_activation=transformer_activation,
             transformer_activation_kwargs=transformer_activation_kwargs,
             transformer_ff_linear_module_up=transformer_ff_linear_module_up,
@@ -546,16 +548,26 @@ class ViT(nn.Module):
     def forward(self, x):
         return self.pool(self.encoder(x))
 
-    def attention_scores(self, x):
-        return self.encoder.attention_scores(x)
+    def attention_logits(self, x):
+        return self.encoder.attention_logits(x)
+
+    def pool_attention(self, x):
+        if hasattr(self.pool.summarize, "attention"):
+            return self.pool.summarize.attention(self.encoder(x))
+        else:
+            raise NotImplementedError(
+                "`pool_attention` is currently only implemented where"
+                " head class is SequencePoolClassificationHead"
+            )
 
-    def head_to_bos_token_attention(self, x):
-        all_attention = self.attention_scores(x)
+    def head_to_utility_token_attention_logits(self, x):
+        all_attention = self.attention_logits(x)
         batch_averages = torch.mean(all_attention, dim=0, keepdim=False)
         sequence_averages = torch.mean(batch_averages, dim=-1, keepdim=False)
-        n_bos_tokens = self.encoder.encoder._bos_tokens
-        just_bos = sequence_averages[:, :, :n_bos_tokens]
-        return F.softmax(just_bos, dim=-1)  # (layer, head, bos_token)
+        n_utility_tokens = self.encoder.encoder[-1]._utility_tokens
+        return sequence_averages[
+            :, :, :n_utility_tokens
+        ]  # (layer, head, utility_tokens)
 
     def reset_parameters(self):
         self.encoder.reset_parameters()

{broccoli_ml-5.1.0 → broccoli_ml-7.0.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "broccoli-ml"
-version = "5.1.0"
+version = "7.0.0"
 description = "Some useful Pytorch models, circa 2025"
 authors = [
     {name = "Nicholas Bailey"}