Changed to bidirectional

llama_diffusion_model.py

@@ -77,7 +77,7 @@ class BidirectionalLlamaAttention(LlamaAttention):
 
 class CustomTransformerConfig(PretrainedConfig):
     def __init__(self, vocab_size=128256, hidden_size=4096, num_layers=32, num_heads=32, prediction_chunk=256, dropout=0,
-                 max_position_embeddings=4096, masking_type="bidirectional_masked", **kwargs):
+                 max_position_embeddings=4096, masking_type="bidirectional", **kwargs):
         super().__init__(**kwargs)
         self.vocab_size = vocab_size
         self.hidden_size = hidden_size
@@ -122,7 +122,7 @@ class CustomTransformerModel(PreTrainedModel):
         # Build attention mask
         device = input_ids.device
 
-        masking_type = getattr(self.config, "masking_type", "bidirectional_masked")
+        masking_type = getattr(self.config, "masking_type", "bidirectional")
         if masking_type == 'bidirectional':
             base_mask = torch.ones(seq_len, seq_len, dtype=torch.bool, device=device)
         elif masking_type == 'bidirectional_masked':