Update modeling_quasrav4.py

modeling_quasrav4.py

@@ -9,7 +9,7 @@ from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutpu
 from transformers.generation import GenerationMixin
 from transformers.utils import logging
 
-from .configuration_quasrav4 import InfinityFormerConfig
+from configuration_quasrav4 import InfinityFormerConfig
 
 logger = logging.get_logger(__name__)
 
@@ -136,13 +136,26 @@ class InfinityFormerLayer(nn.Module):
         self.embed_dim = config.hidden_size
         self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
         self.self_attn = LinearAttention(config, layer_idx)
+
+        print(f"DEBUG: Layer {layer_idx}, use_memory_attention={config.use_memory_attention}") # DEBUG
+        if config.use_memory_attention:
+            self.mem_attn_layer_norm = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
+            self.mem_attn = LinearAttention(config, layer_idx)
+        
         self.ffn = GatedFeedForward(config)
 
     def forward(self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, **kwargs) -> Tuple[torch.Tensor, ...]:
         residual = hidden_states
-        hidden_states_ln = self.self_attn_layer_norm(hidden_states)
-        attn_outputs, _ = self.self_attn(hidden_states=hidden_states_ln, attention_mask=attention_mask, **kwargs)
-        hidden_states = residual + attn_outputs
+        hidden_states = self.self_attn_layer_norm(hidden_states)
+        hidden_states = self.self_attn(hidden_states, attention_mask=attention_mask, **kwargs)[0]
+        hidden_states = residual + hidden_states
+
+        if hasattr(self, 'mem_attn'):
+            mem_residual = hidden_states
+            hidden_states = self.mem_attn_layer_norm(hidden_states)
+            hidden_states = self.mem_attn(hidden_states, attention_mask=attention_mask, **kwargs)[0]
+            hidden_states = mem_residual + hidden_states
+
         hidden_states = self.ffn(hidden_states)
         return (hidden_states,)