Upload UtteranceEmbedings

README.md

@@ -0,0 +1,199 @@
+---
+library_name: transformers
+tags: []
+---
+
+# Model Card for Model ID
+
+<!-- Provide a quick summary of what the model is/does. -->
+
+
+
+## Model Details
+
+### Model Description
+
+<!-- Provide a longer summary of what this model is. -->
+
+This is the model card of a 🤗 transformers model that has been pushed on the Hub. This model card has been automatically generated.
+
+- **Developed by:** [More Information Needed]
+- **Funded by [optional]:** [More Information Needed]
+- **Shared by [optional]:** [More Information Needed]
+- **Model type:** [More Information Needed]
+- **Language(s) (NLP):** [More Information Needed]
+- **License:** [More Information Needed]
+- **Finetuned from model [optional]:** [More Information Needed]
+
+### Model Sources [optional]
+
+<!-- Provide the basic links for the model. -->
+
+- **Repository:** [More Information Needed]
+- **Paper [optional]:** [More Information Needed]
+- **Demo [optional]:** [More Information Needed]
+
+## Uses
+
+<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
+
+### Direct Use
+
+<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
+
+[More Information Needed]
+
+### Downstream Use [optional]
+
+<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
+
+[More Information Needed]
+
+### Out-of-Scope Use
+
+<!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
+
+[More Information Needed]
+
+## Bias, Risks, and Limitations
+
+<!-- This section is meant to convey both technical and sociotechnical limitations. -->
+
+[More Information Needed]
+
+### Recommendations
+
+<!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
+
+Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
+
+## How to Get Started with the Model
+
+Use the code below to get started with the model.
+
+[More Information Needed]
+
+## Training Details
+
+### Training Data
+
+<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
+
+[More Information Needed]
+
+### Training Procedure
+
+<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
+
+#### Preprocessing [optional]
+
+[More Information Needed]
+
+
+#### Training Hyperparameters
+
+- **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
+
+#### Speeds, Sizes, Times [optional]
+
+<!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
+
+[More Information Needed]
+
+## Evaluation
+
+<!-- This section describes the evaluation protocols and provides the results. -->
+
+### Testing Data, Factors & Metrics
+
+#### Testing Data
+
+<!-- This should link to a Dataset Card if possible. -->
+
+[More Information Needed]
+
+#### Factors
+
+<!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
+
+[More Information Needed]
+
+#### Metrics
+
+<!-- These are the evaluation metrics being used, ideally with a description of why. -->
+
+[More Information Needed]
+
+### Results
+
+[More Information Needed]
+
+#### Summary
+
+
+
+## Model Examination [optional]
+
+<!-- Relevant interpretability work for the model goes here -->
+
+[More Information Needed]
+
+## Environmental Impact
+
+<!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
+
+Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
+
+- **Hardware Type:** [More Information Needed]
+- **Hours used:** [More Information Needed]
+- **Cloud Provider:** [More Information Needed]
+- **Compute Region:** [More Information Needed]
+- **Carbon Emitted:** [More Information Needed]
+
+## Technical Specifications [optional]
+
+### Model Architecture and Objective
+
+[More Information Needed]
+
+### Compute Infrastructure
+
+[More Information Needed]
+
+#### Hardware
+
+[More Information Needed]
+
+#### Software
+
+[More Information Needed]
+
+## Citation [optional]
+
+<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
+
+**BibTeX:**
+
+[More Information Needed]
+
+**APA:**
+
+[More Information Needed]
+
+## Glossary [optional]
+
+<!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
+
+[More Information Needed]
+
+## More Information [optional]
+
+[More Information Needed]
+
+## Model Card Authors [optional]
+
+[More Information Needed]
+
+## Model Card Contact
+
+[More Information Needed]

config.json

@@ -2,8 +2,26 @@
   "architectures": [
     "UtteranceEmbedings"
   ],
-  "auto_map" : {
-    "AutoModel" : "saute--model.UtteranceEmbedding",
-    "AutoConfig": "saute--saute_config.SAUTEConfig"
-  }
+  "attention_probs_dropout_prob": 0.1,
+  "auto_map": {
+    "AutoConfig": "saute_config.SAUTEConfig",
+    "AutoModel": "saute_model.UtteranceEmbedings"
+  },
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 768,
+  "intermediate_size": 3072,
+  "max_edu_length": 128,
+  "max_edus_per_dialog": 100,
+  "max_position_embeddings": 512,
+  "max_speakers": 200,
+  "model_type": "saute",
+  "num_attention_heads": 1,
+  "num_edu_layers": 2,
+  "num_hidden_layers": 1,
+  "num_speaker_embeddings": 512,
+  "num_token_layers": 2,
+  "speaker_embeddings_size": 768,
+  "torch_dtype": "float32",
+  "transformers_version": "4.52.4",
+  "vocab_size": 30522
 }

model.safetensors

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

saute_model.py

@@ -0,0 +1,147 @@
+import torch
+import torch.nn as nn
+from transformers import PreTrainedModel, BertModel, BertTokenizerFast
+from transformers.modeling_outputs import MaskedLMOutput
+from sources.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):
+        super().__init__()
+        # model_name="sentence-transformers/all-MiniLM-L6-v2"
+        model_name = "bert-base-uncased"
+
+        self.edu_encoder = AutoModel.from_pretrained(model_name)
+        for param in self.edu_encoder.parameters():
+            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.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):
+        """
+        input_ids: (B, T, L)
+        attention_mask: (B, T, L)
+        speaker_names: list of list of strings, shape (B, T)
+        """
+        B, T, L = input_ids.shape
+
+        # Encode EDUs using frozen encoder
+        with torch.no_grad():
+            input_ids_flat = input_ids.view(B * T, L)
+            attention_mask_flat = attention_mask.view(B * T, L)
+            outputs = self.edu_encoder(input_ids=input_ids_flat, attention_mask=attention_mask_flat)
+            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)
+
+        # 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)
+
+        for b in range(B):
+            for t in range(T):
+                speaker = speaker_names[b][t]
+                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),
+                    }
+
+                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)
+
+                # with torch.no_grad():
+                mem['kv_sum'] = mem['kv_sum'] + kv_t
+                # mem['k_sum'] = mem['k_sum'] + k_t
+
+                # z = torch.clamp(mem['k_sum'] @ k_t, min=1e-6)
+                # M_s = mem['kv_sum'] / z  # (D, D)
+
+                # speaker_matrices[b, t] = M_s
+                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)
+        # 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
+
+
+class UtteranceEmbedings(PreTrainedModel):
+    config_class = SAUTEConfig
+
+    def __init__(self, config : SAUTEConfig):
+        super().__init__(config)
+        
+        self.lm_head    = nn.Linear(config.hidden_size, config.vocab_size)
+        self.saute_unit = EDUSpeakerAwareMLM(config)
+
+        self.config : SAUTEConfig = config
+        
+        self.init_weights()
+
+    def forward(
+        self,
+        input_ids       : torch.Tensor,
+        speaker_names   : list[str],
+        attention_mask  : torch.Tensor  = None,
+        labels          : torch.Tensor  = None
+    ):
+        # print(input_ids.shape)
+        X, flop_penalty = self.saute_unit.forward(
+            input_ids       =   input_ids,
+            speaker_names   =   speaker_names,
+            attention_mask  =   attention_mask,
+            # hidden_state    =   None
+        )
+        # print(X.shape)
+        
+        logits = self.lm_head(X)
+
+        loss = None
+        if labels is not None:
+            loss_fct = nn.CrossEntropyLoss()
+            # 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)