Update raven_modeling_minimal.py

raven_modeling_minimal.py

@@ -242,7 +242,6 @@ class CausalSelfAttention(torch.nn.Module):
         if past_key_values is not None:
             k, v = past_key_values.update(k, v, step_idx)
 
-        return_attn = False  # hardcode for now
         if return_attn:
             y, attention_map = self.compute_eager_sdpa(q, k, v, attn_mask=mask)
         else:
@@ -369,7 +368,7 @@ class RavenForCausalLM(RavenPreTrainedModel, GenerationMixin):
             "return_latents": True,
             "return_attention": False,
             "return_head": False,
-            "return_stats": True,
+            "return_stats": False,
         },
         use_cache: bool = False,
         cache_position: Optional[torch.Tensor] = None,
@@ -397,7 +396,7 @@ class RavenForCausalLM(RavenPreTrainedModel, GenerationMixin):
         # Non-recurrent prelude
         for block_idx, block in enumerate(self.transformer.prelude):
             input_embeds, attn_map = block(
-                input_embeds, freqs_cis, block_idx, attention_mask, past_key_values, return_attn
+                input_embeds, freqs_cis, block_idx, attention_mask, past_key_values, return_attn=return_attn
             )
             attn_maps[block_idx] = attn_map
 
@@ -411,12 +410,13 @@ class RavenForCausalLM(RavenPreTrainedModel, GenerationMixin):
             past_key_values,
             num_steps,
             attn_maps,
+            return_attn=return_attn,
         )
         latent_states = x.clone().detach()
 
         # Coda layers
         for block_idx, block in enumerate(self.transformer.coda, start=1):
-            x, attn_map = block(x, freqs_cis, -block_idx, attention_mask, past_key_values, return_attn)
+            x, attn_map = block(x, freqs_cis, -block_idx, attention_mask, past_key_values, return_attn=return_attn)
             attn_maps[-block_idx] = attn_map
         x = self.transformer.ln_f(x)
 
@@ -453,6 +453,7 @@ class RavenForCausalLM(RavenPreTrainedModel, GenerationMixin):
         past_key_values: Optional[Cache] = None,
         num_steps: Optional[torch.Tensor] = None,
         attn_maps: dict = {},
+        return_attn: bool = False,
     ):
         x = xk = self.initialize_state(input_embeds) if input_states is None else input_states.clone()
         if num_steps is None:
@@ -470,13 +471,13 @@ class RavenForCausalLM(RavenPreTrainedModel, GenerationMixin):
             for step in range(num_steps_no_grad):
                 xk = x
                 x, block_idx, attn_maps = self.core_block_forward(
-                    xk, input_embeds, freqs_cis, mask, past_key_values, block_idx, attn_maps
+                    xk, input_embeds, freqs_cis, mask, past_key_values, block_idx, attn_maps, return_attn
                 )
 
         for step in range(num_steps_with_grad):
             xk = x
             x, block_idx, attn_maps = self.core_block_forward(
-                xk, input_embeds, freqs_cis, mask, past_key_values, block_idx, attn_maps
+                xk, input_embeds, freqs_cis, mask, past_key_values, block_idx, attn_maps, return_attn
             )
         return self.transformer.ln_f(x), num_steps_no_grad, num_steps_with_grad, xk.detach(), block_idx, attn_maps
 
@@ -489,10 +490,11 @@ class RavenForCausalLM(RavenPreTrainedModel, GenerationMixin):
         past_key_values,
         block_idx: Union[torch.Tensor, int],
         attn_maps: dict = {},
+        return_attn: bool = False,
     ):
         x = self.transformer.adapter(torch.cat([x, input_embeds], dim=-1))
         for idx, block in enumerate(self.transformer.core_block, start=1):
-            x, attn_map = block(x, freqs_cis, block_idx + idx, mask, past_key_values, return_attn=len(attn_maps) > 0)
+            x, attn_map = block(x, freqs_cis, block_idx + idx, mask, past_key_values, return_attn=return_attn)
             attn_maps[block_idx + idx] = attn_map
         return x, block_idx + idx, attn_maps
 
@@ -625,7 +627,7 @@ class RavenForCausalLM(RavenPreTrainedModel, GenerationMixin):
         model_inputs["cache_position"] = cache_position
         current_input_length = input_ids.shape[1]
         if past_key_values is not None:
-            if type(past_key_values) == DynamicCache:
+            if type(past_key_values) != HuginnDynamicCache:
                 # Need to use custom cache, detect and replace HF dynamic cache if generate injects it
                 assert past_key_values.get_seq_length() == 0
                 past_key_values = HuginnDynamicCache()
@@ -645,6 +647,18 @@ class RavenForCausalLM(RavenPreTrainedModel, GenerationMixin):
                 model_inputs[key] = value
         return model_inputs
 
+    @torch.no_grad()
+    def generate(self, *args, **kwargs):
+        """Dispatcher - use HF generate in all normal cases."""
+        if any(
+            k in kwargs
+            for k in ("continuous_compute", "latent_dampening", "criterion", "exit_threshold", "cache_kwargs")
+        ):
+            print("Dispatching to custom generate function call")
+            return self.generate_with_adaptive_compute(*args, **kwargs)
+        else:
+            return super().generate(*args, **kwargs)
+
     @torch.no_grad()
     def generate_minimal(
         self,