Upload UtteranceEmbedings

README.md

@@ -1,16 +1,16 @@
 ---
 license: mit
 tags:
-  - masked-language-modeling
-  - dialogue
-  - speaker-aware
-  - transformer
-  - saute
-  - pytorch
+- masked-language-modeling
+- dialogue
+- speaker-aware
+- transformer
+- saute
+- pytorch
 datasets:
-  - SODA
+- SODA
 language:
-  - en
+- en
 pipeline_tag: fill-mask
 model_type: saute
 library_name: transformers

config.json

@@ -15,9 +15,9 @@
   "max_position_embeddings": 512,
   "max_speakers": 200,
   "model_type": "saute",
-  "num_attention_heads": 1,
+  "num_attention_heads": 8,
   "num_edu_layers": 2,
-  "num_hidden_layers": 1,
+  "num_hidden_layers": 3,
   "num_speaker_embeddings": 512,
   "num_token_layers": 2,
   "speaker_embeddings_size": 768,

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:9406a034ce4cc90e25074e183198a7068a67ba1b3b465e94975252138ac19656
-size 560983656
+oid sha256:4e2ee7cabbb652ec8f13c95a48b0336362ec8b7d4698ca6fefb515229d39a898
+size 605098400

saute_model.py

@@ -2,7 +2,7 @@ import torch
 import torch.nn as nn
 from transformers import PreTrainedModel, BertModel, BertTokenizerFast
 from transformers.modeling_outputs import MaskedLMOutput
-from .saute_config import SAUTEConfig
+from sources.saute_config import SAUTEConfig
 
 activation_to_class = {
     "gelu" : nn.GELU,
@@ -23,12 +23,20 @@ class EDUSpeakerAwareMLM(nn.Module):
             param.requires_grad = False  # frozen encoder
 
         self.d_model = config.hidden_size
-        self.key_proj = nn.Linear(config.hidden_size, config.hidden_size, bias = False)
-        self.val_proj = nn.Linear(config.hidden_size, config.hidden_size, bias = False)
+        # self.key_proj = nn.Linear(config.hidden_size, config.hidden_size, bias = False)
+        # self.val_proj = nn.Linear(config.hidden_size, config.hidden_size, bias = False)
+        self.num_heads = config.num_attention_heads
+        self.head_dim = config.hidden_size // self.num_heads
+
+        self.key_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
+        self.val_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
+        self.query_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
+
         self.query_proj = nn.Linear(config.hidden_size, config.hidden_size, bias = False)
 
         encoder_layer = nn.TransformerEncoderLayer(d_model=config.hidden_size, nhead=config.num_attention_heads, batch_first=True)
         self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=config.num_hidden_layers)
+        # self.out_proj = nn.Linear(config.hidden_size, config.hidden_size)
 
         # self.mlp_proj = nn.Sequential(
         #     nn.Linear(config.hidden_size, 2048),
@@ -59,12 +67,18 @@ class EDUSpeakerAwareMLM(nn.Module):
             token_embeddings = outputs.last_hidden_state  # (B*T, L, D)
 
         token_embeddings = token_embeddings.view(B, T, L, self.d_model)
-        edu_embeddings = token_embeddings.mean(dim=2)  # (B, T, D)
-        query_emb = self.query_proj(token_embeddings)
+        # edu_embeddings = token_embeddings.mean(dim=2)  # (B, T, D)
+        edu_embeddings = token_embeddings[:,:,0] # CLS token
+        # query_emb = self.query_proj(token_embeddings)
 
         # Speaker-aware memory
         speaker_memories = [{} for _ in range(B)]
-        speaker_matrices = torch.zeros(B, T, self.d_model, self.d_model, device=edu_embeddings.device)
+        # speaker_matrices = torch.zeros(B, T, self.d_model, self.d_model, device=edu_embeddings.device)
+        H = self.num_heads
+        d = self.head_dim
+
+        speaker_matrices = torch.zeros(B, T, H, d, d, device=edu_embeddings.device)
+
 
         for b in range(B):
             for t in range(T):
@@ -72,15 +86,22 @@ class EDUSpeakerAwareMLM(nn.Module):
                 e_t = edu_embeddings[b, t]  # (D)
 
                 if speaker not in speaker_memories[b]:
+                    # speaker_memories[b][speaker] = {
+                    #     'kv_sum': torch.zeros(self.d_model, self.d_model, device=e_t.device),
+                    #     # 'k_sum': torch.zeros(self.d_model, device=e_t.device),
+                    # }
                     speaker_memories[b][speaker] = {
-                        'kv_sum': torch.zeros(self.d_model, self.d_model, device=e_t.device),
-                        # 'k_sum': torch.zeros(self.d_model, device=e_t.device),
+                        'kv_sum': torch.zeros(self.num_heads, self.head_dim, self.head_dim, device=e_t.device)
                     }
 
                 mem = speaker_memories[b][speaker]
-                k_t = self.key_proj(e_t)
-                v_t = self.val_proj(e_t)
-                kv_t = torch.outer(k_t, v_t)
+                # k_t = self.key_proj(e_t)
+                # v_t = self.val_proj(e_t)
+                # kv_t = torch.outer(k_t, v_t)
+                k_t = self.key_proj(e_t).view(self.num_heads, self.head_dim)  # (H, d_k)
+                v_t = self.val_proj(e_t).view(self.num_heads, self.head_dim)  # (H, d_v)
+                kv_t = torch.einsum("hd,he->hde", k_t, v_t)  # (H, d_k, d_v)
+
 
                 # with torch.no_grad():
                 mem['kv_sum'] = mem['kv_sum'] + kv_t
@@ -93,11 +114,32 @@ class EDUSpeakerAwareMLM(nn.Module):
                 speaker_matrices[b, t] = mem['kv_sum']
 
         # Apply speaker matrix to each token
-        speaker_matrices_exp = speaker_matrices.unsqueeze(2)  # (B, T, 1, D, D)
-        token_embeddings_exp = query_emb.unsqueeze(-1)  # (B, T, L, D, 1)
-        contextual_tokens = token_embeddings + torch.matmul(speaker_matrices_exp, token_embeddings_exp).squeeze(-1)  # (B, T, L, D)
+        # speaker_matrices_exp = speaker_matrices.unsqueeze(2)  # (B, T, 1, D, D)
+        # token_embeddings_exp = query_emb.unsqueeze(-1)  # (B, T, L, D, 1)
+        # contextual_tokens = token_embeddings + torch.matmul(speaker_matrices_exp, token_embeddings_exp).squeeze(-1)  # (B, T, L, D)
         # contextual_tokens = self.ln1(contextual_tokens)
         # contextual_tokens = self.ln2(contextual_tokens + self.mlp_proj(contextual_tokens))
+        # Project queries
+        query_emb = self.query_proj(token_embeddings)  # (B, T, L, D)
+        query = query_emb.view(B, T, L, H, d)  # (B, T, L, H, d)
+
+        # Apply memory matrices
+        contextual = []
+        for b in range(B):
+            head_outputs = []
+            for t in range(T):
+                speaker = speaker_names[b][t]
+                M = speaker_matrices[b, t]  # (H, d, d)
+                q = query[b, t]  # (L, H, d)
+                q = q.transpose(0, 1)  # (H, L, d)
+                a = torch.matmul(q, M)  # (H, L, d)
+                a = a.transpose(0, 1).contiguous().view(L, -1)  # (L, D)
+                contextual_token = token_embeddings[b, t] + a
+                head_outputs.append(contextual_token)
+            contextual.append(torch.stack(head_outputs))
+        contextual_tokens = torch.stack(contextual)
+        # (B, T, L, D)
+        # contextual_tokens = self.out_proj(contextual_tokens)
 
         # === NEW: EDU-level Transformer ===
         edu_tokens = contextual_tokens.view(B * T, L, self.d_model)  # (B*T, L, D)
@@ -144,4 +186,4 @@ class UtteranceEmbedings(PreTrainedModel):
             # loss = loss_fct(logits.view(-1, self.config.vocab_size), labels.view(-1)) + 1e-3 * flop_penalty
             loss = loss_fct(logits.view(-1, self.config.vocab_size), labels.view(-1))
 
-        return MaskedLMOutput(loss=loss, logits=logits)
+        return MaskedLMOutput(loss=loss, logits=logits)