Upload LUAR

config.json

@@ -0,0 +1,13 @@
+{
+  "architectures": [
+    "LUAR"
+  ],
+  "auto_map": {
+    "AutoConfig": "config.LUARConfig",
+    "AutoModel": "model.LUAR"
+  },
+  "embedding_size": 512,
+  "model_type": "LUAR",
+  "torch_dtype": "float32",
+  "transformers_version": "4.33.2"
+}

config.py

@@ -0,0 +1,12 @@
+
+from transformers import PretrainedConfig
+
+class LUARConfig(PretrainedConfig):
+    model_type = "LUAR"
+    
+    def __init__(self,
+        embedding_size: int = 512,
+        **kwargs,
+    ):
+        self.embedding_size = embedding_size
+        super().__init__(**kwargs)

model.py

@@ -0,0 +1,85 @@
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange, reduce, repeat
+from transformers import AutoModel, PreTrainedModel
+
+from .config import LUARConfig
+
+class SelfAttention(nn.Module):
+    """Implements Dot-Product Self-Attention as used in "Attention is all You Need".
+    """
+    def __init__(self):
+        super(SelfAttention, self).__init__()
+
+    def forward(self, k, q, v):
+        d_k = q.size(-1)
+        scores = torch.matmul(k, q.transpose(-2, -1)) / math.sqrt(d_k)
+        p_attn = F.softmax(scores, dim=-1)
+
+        return torch.matmul(p_attn, v)
+
+class LUAR(PreTrainedModel):
+    """Defines the LUAR model.
+    """
+    config_class = LUARConfig
+    
+    def __init__(self, config):
+        super().__init__(config)
+        self.create_transformer()
+        self.attn_fn = SelfAttention()
+        self.linear = nn.Linear(self.hidden_size, config.embedding_size)
+ 
+    def create_transformer(self):
+        """Creates the Transformer backbone.
+        """
+        self.transformer = AutoModel.from_pretrained("sentence-transformers/paraphrase-distilroberta-base-v1")
+        self.hidden_size = self.transformer.config.hidden_size
+        self.num_attention_heads = self.transformer.config.num_attention_heads
+        self.dim_head = self.hidden_size // self.num_attention_heads
+        
+    def mean_pooling(self, token_embeddings, attention_mask):
+        """Mean Pooling as described in the SBERT paper.
+        """
+        input_mask_expanded = repeat(attention_mask, 'b l -> b l d', d=self.hidden_size).float()
+        sum_embeddings = reduce(token_embeddings * input_mask_expanded, 'b l d -> b d', 'sum')
+        sum_mask = torch.clamp(reduce(input_mask_expanded, 'b l d -> b d', 'sum'), min=1e-9)
+        return sum_embeddings / sum_mask
+    
+    def get_episode_embeddings(self, data):
+        """Computes the Author Embedding. 
+        """
+        input, attention_mask = data[0], data[1]
+        B, N, E, _ = attention_mask.shape
+        attention_mask = rearrange(attention_mask, 'b n e l -> (b n e) l')
+
+        input = rearrange(input, 'b n e l -> (b n e) l')
+
+        outputs = self.transformer(
+            input_ids=input,
+            attention_mask=attention_mask,
+            return_dict=True,
+            output_hidden_states=True
+        )
+
+        # at this point, we're embedding individual "comments"
+        comment_embeddings = self.mean_pooling(outputs['last_hidden_state'], attention_mask)
+        comment_embeddings = rearrange(comment_embeddings, '(b n e) l -> (b n) e l', b=B, n=N, e=E)
+
+        # aggregate individual comments embeddings into episode embeddings
+        episode_embeddings = self.attn_fn(comment_embeddings, comment_embeddings, comment_embeddings)
+        episode_embeddings = reduce(episode_embeddings, 'b e l -> b l', 'max')
+        
+        episode_embeddings = self.linear(episode_embeddings)
+        
+        return episode_embeddings, comment_embeddings
+    
+    def forward(self, data):
+        """Calculates a fixed-length feature vector for a batch of episode samples.
+        """
+        output = self.get_episode_embeddings(data)
+
+        return output

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:76a4ca30b8ef7c5d3fb806e7030c8375cdc3f60e2e2a607a2156917ac78e74b4
+size 330083185