jaj

modelLM.py

@@ -22,10 +22,15 @@ class OBILanguageModel(PreTrainedModel):
         self.ln2 = nn.LayerNorm(config.hidden_size)
         self.lm_head = nn.Linear(config.hidden_size, config.vocab_size)  # Use length of SentencePiece vocab
     
-    def forward(self, idx, targets=None):
+    def forward(self, idx, attention_mask=None, targets=None):
         tok_emb = self.token_embedding_table(idx)
         pos_emb = self.position_embedding_table(torch.arange(idx.size(1), device='cpu'))
         x = tok_emb + pos_emb
+
+        # Assuming you need to add attention_mask here
+        if attention_mask is not None:
+            x *= attention_mask.unsqueeze(-1)
+
         x = self.transformer(x, x)
         x = self.ln1(x)
         x = self.ln2(x)