Delete pooling_coverage.py

pooling_coverage.py

@@ -1,160 +0,0 @@
-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