Update model.py

model.py

@@ -1,8 +1,16 @@
 import torch
 import torch.nn as nn
-from transformers import PreTrainedModel, AutoModel
+from transformers import PreTrainedModel, BertModel, BertTokenizerFast
 from transformers.modeling_outputs import MaskedLMOutput
-from sources.saute_config import SAUTEConfig
+from saute_config import SAUTEConfig
+
+activation_to_class = {
+    "gelu" : nn.GELU,
+    "relu" : nn.ReLU,
+    "sigmoid" : nn.Sigmoid
+}
+
+from transformers import AutoModel
 
 class EDUSpeakerAwareMLM(nn.Module):
     def __init__(self, config):
@@ -15,12 +23,25 @@ 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.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.saute = SAUTE(config)
+        # self.mlp_proj = nn.Sequential(
+        #     nn.Linear(config.hidden_size, 2048),
+        #     activation_to_class["gelu"](),
+        #     # nn.Dropout(0.1),
+        #     nn.Linear(2048, config.hidden_size),
+        #     # nn.Dropout(0.1),
+        # )
+        self.ln1 = nn.LayerNorm(config.hidden_size)
+        # self.ln2 = nn.LayerNorm(config.hidden_size)
+
+        # self.speaker_memory = {}  # Will be filled per batch
+        # self.lm_head = nn.Linear(config.hidden_size, self.edu_encoder.config.vocab_size)
 
     def forward(self, input_ids, attention_mask, speaker_names):
         """
@@ -39,41 +60,11 @@ class EDUSpeakerAwareMLM(nn.Module):
 
         token_embeddings = token_embeddings.view(B, T, L, self.d_model)
         edu_embeddings = token_embeddings.mean(dim=2)  # (B, T, D)
-
-        contextual_tokens = self.saute(input_ids, speaker_names, token_embeddings, edu_embeddings)
-
-        # === NEW: EDU-level Transformer ===
-        edu_tokens = contextual_tokens.view(B * T, L, self.d_model)  # (B*T, L, D)
-        encoded_edu = self.transformer(edu_tokens)  # (B*T, L, D)
-        encoded = encoded_edu.view(B, T, L, self.d_model)  # (B, T, L, D)
-
-        return encoded, 0
-
-class SAUTE(nn.Module):
-    def __init__(self,
-            config      : SAUTEConfig
-        ):
-        super().__init__()
-
-        self.d_model = config.hidden_size
-        
-        self.query_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)
-
-    def forward(self,
-            input_ids           : torch.Tensor,
-            speaker_names       : list[str],
-            token_embeddings    : torch.Tensor,
-            edu_embeddings      : torch.Tensor
-        ):
+        query_emb = self.query_proj(token_embeddings)
 
         # Speaker-aware memory
-        B, T, L = input_ids.shape
-
         speaker_memories = [{} for _ in range(B)]
         speaker_matrices = torch.zeros(B, T, self.d_model, self.d_model, device=edu_embeddings.device)
-        query_embeddings = self.query_proj(token_embeddings)
 
         for b in range(B):
             for t in range(T):
@@ -103,10 +94,18 @@ class SAUTE(nn.Module):
 
         # Apply speaker matrix to each token
         speaker_matrices_exp = speaker_matrices.unsqueeze(2)  # (B, T, 1, D, D)
-        token_embeddings_exp = query_embeddings.unsqueeze(-1)  # (B, T, L, D, 1)
+        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))
+
+        # === NEW: EDU-level Transformer ===
+        edu_tokens = contextual_tokens.view(B * T, L, self.d_model)  # (B*T, L, D)
+        encoded_edu = self.transformer(edu_tokens)  # (B*T, L, D)
+        encoded = encoded_edu.view(B, T, L, self.d_model)  # (B, T, L, D)
+
+        return encoded, 0
 
-        return contextual_tokens
 
 class UtteranceEmbedings(PreTrainedModel):
     config_class = SAUTEConfig
@@ -135,6 +134,7 @@ class UtteranceEmbedings(PreTrainedModel):
             attention_mask  =   attention_mask,
             # hidden_state    =   None
         )
+        # print(X.shape)
         
         logits = self.lm_head(X)