Upload model

config.json

@@ -0,0 +1,26 @@
+{
+  "architectures": [
+    "CompressionModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_compression.CompressionConfig",
+    "AutoModel": "modeling_compression.CompressionModel"
+  },
+  "compression_sizes": [
+    512,
+    256,
+    128,
+    64,
+    32
+  ],
+  "dropout": 0.1,
+  "input_size": 768,
+  "loss_k_vals": [
+    10,
+    100,
+    256
+  ],
+  "model_type": "compression_head",
+  "torch_dtype": "float32",
+  "transformers_version": "4.38.2"
+}

configuration_compression.py

@@ -0,0 +1,21 @@
+from transformers import PretrainedConfig
+from typing import List
+
+class CompressionConfig(PretrainedConfig):
+    model_type = "compression_head"
+    
+    def __init__(self,
+                 input_size: int = 768,
+                 compression_sizes: List[int] = [512, 256, 128, 64, 32],
+                 dropout: float = 0.1,
+                 loss_k_vals: List[int] = [],
+                 **kwargs
+                ):
+        
+        self.input_size = input_size
+        self.compression_sizes = compression_sizes
+        self.dropout = dropout
+        self.loss_k_vals = loss_k_vals
+        
+        super().__init__(**kwargs)
+        

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7d62042eb9cbd70af97e6e2abbcfe3fa25972b969d17c421ad173348fee8b4ba
+size 10557544

modeling_compression.py

@@ -0,0 +1,132 @@
+import torch
+import torch.nn as nn 
+import torch.nn.functional as F 
+
+from typing import Tuple, Optional, List
+from dataclasses import dataclass
+
+from transformers import PreTrainedModel
+from transformers.utils import ModelOutput
+
+from .configuration_compression import CompressionConfig
+
+def cosine_pairwise(embeddings):
+    return F.cosine_similarity(embeddings.unsqueeze(1), embeddings.unsqueeze(0), dim=2)
+
+def cov(tensor, rowvar=True, bias=False):
+    """Estimate a covariance matrix (np.cov)"""
+    tensor = tensor if rowvar else tensor.transpose(-1, -2)
+    tensor = tensor - tensor.mean(dim=-1, keepdim=True)
+    factor = 1 / (tensor.shape[-1] - int(not bool(bias)))
+    return factor * tensor @ tensor.transpose(-1, -2).conj()
+
+def remove_diag(x):
+    n = x.shape[0]
+    return x.masked_select(~torch.eye(n, dtype=bool, device=x.device)).view(n, n - 1)
+
+def corrcoef(tensor, rowvar=True):
+    """Get Pearson product-moment correlation coefficients (np.corrcoef)"""
+    covariance = cov(tensor, rowvar=rowvar)
+    variance = covariance.diagonal(0, -1, -2)
+    if variance.is_complex():
+        variance = variance.real
+    stddev = variance.sqrt()
+    covariance /= stddev.unsqueeze(-1)
+    covariance /= stddev.unsqueeze(-2)
+    if covariance.is_complex():
+        covariance.real.clip_(-1, 1)
+        covariance.imag.clip_(-1, 1)
+    else:
+        covariance.clip_(-1, 1)
+    return covariance
+
+def compute_correlation(base_sims, compressed_sims, rm_diag=True):
+    if rm_diag:
+        base_sims = remove_diag(base_sims)
+        compressed_sims = remove_diag(compressed_sims)
+
+    inputs = torch.stack([base_sims, 
+                          compressed_sims], dim=1)
+    return (1-corrcoef(inputs)[:, 0, 1]).mean()
+
+def loss_function(base_sims, compressed_sims, k_vals):
+    outputs =  [compute_correlation(base_sims, compressed_sims)]
+
+    if k_vals:
+        base_ranks = base_sims.argsort(-1, descending=True)[:, 1:]
+        n = base_ranks.shape[1]
+        for k in k_vals:
+            base_sims_k = torch.gather(base_sims, 1, base_ranks[:, :k])
+            compressed_sims_k = torch.gather(compressed_sims, 1, base_ranks[:, :k])
+            outputs.append(compute_correlation(base_sims_k, compressed_sims_k, rm_diag=False))
+
+    return torch.stack(outputs).unsqueeze(0)
+
+class FeedForward(nn.Module):
+    def __init__(self, d_in, d_out):
+        super().__init__()
+        self.fc1 = nn.Linear(d_in, d_out*2)
+        self.fc2 = nn.Linear(d_out, d_out)
+        
+    def forward(self, x):
+        x = self.fc1(x)
+        x1, x2 = x.chunk(2, dim=-1)
+        x = self.fc2(F.silu(x1) * x2)
+        return x
+
+class CompressionHead(nn.Module):
+    def __init__(self, d_in, d_out, dropout=0.1):
+        super().__init__()
+        self.ff = FeedForward(d_in, d_out)
+        self.skip = nn.Linear(d_in, d_out)
+        self.dropout = nn.Dropout(dropout)
+        
+    def forward(self, x):
+        x = self.dropout(x)
+        x = self.ff(x) + self.skip(x)
+        return x
+
+@dataclass
+class CompressionModelOutput(ModelOutput):
+    loss: Optional[torch.FloatTensor] = None
+    losses: Optional[List[torch.FloatTensor]] = None
+    base_embedding: Optional[torch.FloatTensor] = None
+    compressed_embeddings: Optional[List[torch.FloatTensor]] = None
+
+class CompressionModel(PreTrainedModel):
+    config_class = CompressionConfig
+    def __init__(self, config):
+        super().__init__(config)
+        self.heads = nn.ModuleList([CompressionHead(config.input_size, i, config.dropout)
+                      for i in config.compression_sizes])
+        
+    def forward(self, embedding, compute_loss=True, return_dict=True):
+        outputs = []
+        losses = None
+
+        if compute_loss:
+            losses = []
+            emb_sims = cosine_pairwise(embedding)
+
+        for head in self.heads:
+            compressed_embedding = head(embedding)
+            outputs.append(compressed_embedding)
+
+            if compute_loss:
+                comp_sims = cosine_pairwise(compressed_embedding)
+                loss = loss_function(emb_sims, comp_sims, self.config.loss_k_vals)
+                losses.append(loss)
+
+        loss = torch.cat(losses).sum()
+
+        if not return_dict:
+            return (loss, losses, embedding, outputs)
+        
+        return CompressionModelOutput(loss=loss,
+                                        losses=losses,
+                                        base_embedding=embedding,
+                                        compressed_embeddings=outputs)
+
+
+
+