Upload pooling_coverage.py

pooling_coverage.py

@@ -0,0 +1,160 @@
+import torch
+
+
+class GaussianCoveragePooling(torch.nn.Module):
+    def __init__(self, coverage_chunks, sigma, alpha):
+        """
+        Custom pooling layer that computes weighted mean pooling using Gaussian-based weights.
+
+        Args:
+            coverage_chunks (int): Number of weighted pooling operations (N).
+            sigma (float): Standard deviation for Gaussian weighting.
+            alpha (float): Weighting factor for merging with standard mean pooling.
+        """
+        super().__init__()
+        self.coverage_chunks = coverage_chunks
+        self.sigma = sigma  # Controls width of Gaussians
+        self.alpha = alpha  # Blends standard mean with weighted mean
+
+    def forward(self, features, chunk_indicators=None):
+        """
+        Computes weighted mean pooling using Gaussian-based weights.
+
+        Args:
+            self (SentenceTransformer): The model.
+            features (dict): The token embeddings and attention mask.
+            chunk_indicators (tensor[bz, 1]): Index indicators to return a specific chunk,
+            leave as None to return embeddings for all chunks. Mainly useful for training,
+            not inference. Leave as None for inference.
+        """
+
+        # Get token embeddings and attention mask
+        token_embeddings = features[
+            "token_embeddings"
+        ]  # (batch_size, seq_len, hidden_dim)
+        attention_mask = (
+            features["attention_mask"].float().unsqueeze(-1)
+        )  # (batch_size, seq_len, 1)
+
+        # Get shapes and devices
+        batch_size, seq_len, hidden_dim = token_embeddings.shape
+        device = token_embeddings.device
+
+        # Compute actual sequence lengths (ignoring padding)
+        # (batch_size, 1)
+        seq_lengths = attention_mask.squeeze(-1).sum(dim=1, keepdim=True)
+        max_seq_length = int(torch.max(seq_lengths).item())
+
+        # Standard mean pooling
+        sum_embeddings = torch.sum(token_embeddings * attention_mask, dim=1)
+        sum_mask = torch.sum(attention_mask, dim=1).clamp(min=1e-9)
+        standard_mean = sum_embeddings / sum_mask  # (batch_size, hidden_dim)
+
+        # Compute chunk centers dynamically based on sequence length
+        chunk_positions = torch.linspace(0, 1, self.coverage_chunks + 2, device=device)[
+            1:-1
+        ]  # Excludes 0 and 1
+        chunk_centers = chunk_positions * seq_lengths  # (batch_size, N)
+
+        # Token positions per sequence (batch_size, seq_len)
+        token_positions = (
+            torch.arange(seq_len, device=device).float().unsqueeze(0)
+        )  # (1, seq_len)
+
+        # Compute Gaussian weights (batch_size, N, seq_len)
+        seq_lengths = seq_lengths.view(seq_lengths.shape[0], 1, 1).repeat(
+            1, self.coverage_chunks, max_seq_length
+        )
+        gaussians = torch.exp(
+            -0.5
+            * (
+                (token_positions.unsqueeze(1) - chunk_centers.unsqueeze(2))
+                / (self.sigma * seq_lengths)
+            )
+            ** 2
+        )
+
+        # Mask out padding and normalize Gaussian weights per sequence
+        # (batch_size, N, seq_len)
+        gaussians = gaussians * attention_mask.squeeze(-1).unsqueeze(1)
+
+        # Normalize against gaussian weights
+        gaussians /= gaussians.sum(dim=2, keepdim=True).clamp(min=1e-9)
+
+        # Compute weighted mean for each chunk (batch_size, N, hidden_dim)
+        weighted_means = torch.einsum(
+            "bns,bsh->bnh", gaussians.to(token_embeddings.dtype), token_embeddings
+        )
+
+        # Blend with standard mean pooling
+        # (batch_size, N, hidden_dim)
+        combined_embeddings = (1 - self.alpha) * standard_mean.unsqueeze(
+            1
+        ) + self.alpha * weighted_means
+
+        # Add an embedding for the entire document at index 0
+        # (batch_size, N+1, hidden_dim)
+        combined_embeddings = torch.cat(
+            [torch.zeros_like(combined_embeddings[:, :1]), combined_embeddings], 1
+        )
+        combined_embeddings[:, 0:1, :] = standard_mean.unsqueeze(1)
+
+        # Select the indicator if provided
+        if chunk_indicators is not None:
+            combined_embeddings = combined_embeddings[
+                torch.arange(combined_embeddings.size(0)), chunk_indicators
+            ]
+
+        # Normalize all the embeddings
+        combined_embeddings = torch.nn.functional.normalize(
+            combined_embeddings, p=2, dim=-1
+        )
+
+        # Flatten final embeddings (batch_size, hidden_dim * (N+1))
+        if chunk_indicators is None:
+            sentence_embedding = combined_embeddings.reshape(
+                batch_size, hidden_dim * (self.coverage_chunks + 1)
+            )
+        else:
+            sentence_embedding = combined_embeddings
+
+        # Return the final flattened entence embedding
+        features["sentence_embedding"] = sentence_embedding
+        return features
+
+
+def use_gaussian_coverage_pooling(m, coverage_chunks=10, sigma=0.05, alpha=1.0):
+    """
+    Add custom pooling layer that computes weighted mean pooling using Gaussian-based weights.
+
+    Args:
+        m (SentenceTransformer): The model to add pooling layer to.
+        coverage_chunks (int): Number of weighted pooling operations (N).
+        sigma (float): Standard deviation for Gaussian weighting.
+        alpha (float): Weighting factor for merging with standard mean pooling.
+    """
+    if isinstance(m[1], GaussianCoveragePooling):
+        m = unuse_gaussian_coverage_pooling(m)
+    word_embedding_model = m[0]
+    custom_pooling = GaussianCoveragePooling(
+        coverage_chunks=coverage_chunks, sigma=sigma, alpha=alpha
+    )
+    old_pooling = m[1]
+    new_m = m.__class__(modules=[word_embedding_model, custom_pooling])
+    new_m.old_pooling = {"old_pooling": old_pooling}
+    return new_m
+
+
+def unuse_gaussian_coverage_pooling(m):
+    """
+    Removes the custom pooling layer.
+
+    Args:
+        m (SentenceTransformer): The model to remove the pooling layer from.
+    """
+
+    if isinstance(m[1], GaussianCoveragePooling):
+        new_m = m.__class__(modules=[m[0], m.old_pooling["old_pooling"]])
+        return new_m
+    else:
+        return m