broccoli-ml 5.1.4__tar.gz → 5.2.0__tar.gz

{broccoli_ml-5.1.4 → broccoli_ml-5.2.0}/PKG-INFO

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

{broccoli_ml-5.1.4 → broccoli_ml-5.2.0}/broccoli/vit.py

@@ -549,13 +549,12 @@ class ViT(nn.Module):
     def attention_logits(self, x):
         return self.encoder.attention_logits(x)
 
-    def head_to_bos_token_attention(self, x):
+    def head_to_bos_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[-1]._bos_tokens
-        just_bos = sequence_averages[:, :, :n_bos_tokens]
-        return F.softmax(just_bos, dim=1)  # (layer, head, bos_token)
+        return sequence_averages[:, :, :n_bos_tokens]  # (layer, head, bos_token)
 
     def reset_parameters(self):
         self.encoder.reset_parameters()

{broccoli_ml-5.1.4 → broccoli_ml-5.2.0}/pyproject.toml

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