asas

modelLM.py

@@ -29,20 +29,16 @@ class OBILanguageModel(PreTrainedModel):
 
        
     
-    def forward(self, idx, targets=None):
+    def forward(self, idx, attention_mask=None, targets=None):
         tok_emb = self.token_embedding_table(idx)
-        pos_emb = None  # Initialize pos_emb to None
-        try:
-            pos_emb = self.position_embedding_table(torch.arange(idx.size(1), device='cpu'))
-        except IndexError as e:
-            # Handle the IndexError by initializing pos_emb with zeros
-            print(f"IndexError: {e}")
-            print(f"idx.size(1): {idx.size(1)}")
-            print(f"Positional embedding table shape: {self.position_embedding_table.weight.shape}")
-            pos_emb = torch.zeros((idx.size(1), self.config.hidden_size), device=device)
-
+        pos_emb = self.position_embedding_table(torch.arange(idx.size(1), device='cpu'))
         x = tok_emb + pos_emb
-        x = self.transformer(x, x)
+
+        # Create an attention mask for padding tokens
+        if attention_mask is not None:
+            attention_mask = attention_mask.to(x.device)
+
+        x = self.transformer(x, attn_mask=attention_mask)  # Pass attention_mask to the transformer
         x = self.ln1(x)
         x = self.ln2(x)
         logits = self.lm_head(x)
@@ -52,8 +48,7 @@ class OBILanguageModel(PreTrainedModel):
         else:
             loss = F.cross_entropy(logits.view(-1, self.config.vocab_size), targets.view(-1))
 
-        return (logits, loss)  # Return as a tuple
-
+        return logits, loss
 
 
     def generate(self, idx, max_new_tokens):