Upload folder using huggingface_hub

__init__.py

@@ -0,0 +1,28 @@
+"""
+echo_dsrn/__init__.py
+────────────────────────────────────────────────────────────────────────────
+Package init: registers echo_dsrn classes with HuggingFace AutoClass so
+that AutoConfig.from_pretrained() and AutoModelForCausalLM.from_pretrained()
+work transparently without trust_remote_code=True.
+"""
+
+from transformers import (
+    AutoConfig,
+    AutoModelForCausalLM,
+    AutoModelForSequenceClassification,
+)
+
+from .configuration_echo import EchoConfig
+from .modeling_echo import EchoForCausalLM, EchoForSequenceClassification, EchoModel
+
+# Register with HuggingFace so AutoClass routing works
+AutoConfig.register("echo", EchoConfig)
+AutoModelForCausalLM.register(EchoConfig, EchoForCausalLM)
+AutoModelForSequenceClassification.register(EchoConfig, EchoForSequenceClassification)
+
+__all__ = [
+    "EchoConfig",
+    "EchoModel",
+    "EchoForCausalLM",
+    "EchoForSequenceClassification",
+]

chat_template.jinja

@@ -28,7 +28,11 @@
                 {%- if tool_call['function'] is defined %}
                     {%- set tool_call = tool_call['function'] %}
                 {%- endif %}
-                {{- '\n<tool_call>\n{"name": "' + tool_call['name'] + '", "arguments": ' + tool_call['arguments'] | tojson + '}\n</tool_call>' }}
+                {%- if tool_call['arguments'] is string %}
+                    {{- '\n<tool_call>\n{"name": "' + tool_call['name'] + '", "arguments": ' + tool_call['arguments'] + '}\n</tool_call>' }}
+                {%- else %}
+                    {{- '\n<tool_call>\n{"name": "' + tool_call['name'] + '", "arguments": ' + tool_call['arguments'] | tojson + '}\n</tool_call>' }}
+                {%- endif %}
             {%- endfor %}
         {%- endif %}
         {{- '<|end|>\n' }}

configuration_echo.py

@@ -17,6 +17,11 @@ class EchoConfig(PretrainedConfig):
         use_hybrid_attention=True,
         use_rmsnorm=True,
         mlp_bias: bool = False,
+        pooling_mode: str = "c_T",
+        attention_masking: str = "causal",
+        # --- DSpark speculative decoding integration ---
+        output_surprise_gate_logits: bool = False,
+        surprise_temperature_alpha: float = 0.0,
         # --- Classification fields (optional, ignored by CausalLM) ---
         num_labels: int = 2,
         id2label: Optional[dict] = None,
@@ -52,6 +57,10 @@ class EchoConfig(PretrainedConfig):
         self.use_hybrid_attention = use_hybrid_attention
         self.use_rmsnorm = use_rmsnorm
         self.mlp_bias = mlp_bias
+        self.pooling_mode = pooling_mode
+        self.attention_masking = attention_masking
+        self.output_surprise_gate_logits = output_surprise_gate_logits
+        self.surprise_temperature_alpha = surprise_temperature_alpha
         self.classifier_dropout = classifier_dropout
 
         # Standard HF aliases

modeling_echo.py

@@ -13,8 +13,9 @@ from transformers.modeling_outputs import (
 from .configuration_echo import EchoConfig
 
 if TYPE_CHECKING:
-    # Force HF trust_remote_code AST parser to bundle triton_scan.py
-    pass
+    # Force HF trust_remote_code AST parser to bundle triton_scan.py and utils.py
+    from .triton_scan import triton_dsrn_parallel_scan
+    from .utils import rms_norm_fn
 
 try:
     # pyrefly: ignore [missing-import]
@@ -280,7 +281,7 @@ def dsrn_parallel_kernel_legacy(
     # Enabled on Legacy to fix Disconnected Slow State bug while keeping LayerNorm
     x_out = x_out + model_block.linear_read(c_all)
 
-    return x_out, h_new, c_new, gate_stats, h_all, c_all
+    return x_out, h_new, c_new, gate_stats, h_all, c_all, gate_logits
 
 
 def dsrn_parallel_kernel_hybrid(
@@ -359,7 +360,7 @@ def dsrn_parallel_kernel_hybrid(
     if model_block.use_hybrid_attention:
         x_out = x_out + model_block.linear_read(c_all)
 
-    return x_out, h_new, c_new, gate_stats, h_all, c_all
+    return x_out, h_new, c_new, gate_stats, h_all, c_all, gate_logits
 
 
 def dsrn_parallel_kernel(
@@ -632,8 +633,10 @@ class DSRNBlock(nn.Module):
             # Placeholder for Triton
             pass
 
+        output_gate_logits = kwargs.get("output_gate_logits", False)
+
         # Use Parallel Kernel
-        x_out, h_new, c_new, gate_stats, h_all, c_all = dsrn_parallel_kernel(
+        x_out, h_new, c_new, gate_stats, h_all, c_all, gate_logits = dsrn_parallel_kernel(
             self, x, h_prev, c_prev
         )
 
@@ -662,7 +665,11 @@ class DSRNBlock(nn.Module):
             h_new_full = (h_new, c_new)
 
         if kwargs.get("output_all_states", False):
+            if output_gate_logits:
+                return x_out, h_new_full, gate_stats, h_all, c_all, gate_logits
             return x_out, h_new_full, gate_stats, h_all, c_all
+        if output_gate_logits:
+            return x_out, h_new_full, gate_stats, gate_logits
         return x_out, h_new_full, gate_stats
 
 
@@ -804,6 +811,10 @@ class EchoModel(EchoPreTrainedModel):
         all_h_all = [] if output_all_states else None
         all_c_all = [] if output_all_states else None
 
+        # Gate logits for DSpark speculative decoding integration
+        _output_gl = getattr(self.config, "output_surprise_gate_logits", False)
+        all_gate_logits = [] if _output_gl else None
+
         # Layer-Major Execution
         for i, block in enumerate(self.blocks):
 
@@ -834,10 +845,17 @@ class EchoModel(EchoPreTrainedModel):
                     state_i,
                     use_reentrant=False,
                     output_all_states=output_all_states,
+                    output_gate_logits=_output_gl,
                     **kwargs,
                 )
             else:
-                out = block(x, state_i, output_all_states=output_all_states, **kwargs)
+                out = block(
+                    x,
+                    state_i,
+                    output_all_states=output_all_states,
+                    output_gate_logits=_output_gl,
+                    **kwargs,
+                )
 
             x = out[0]
             next_states.append(out[1])
@@ -846,6 +864,11 @@ class EchoModel(EchoPreTrainedModel):
                 all_gate_stats.append(out[2])
                 all_c_states.append(out[1][1])
 
+            if _output_gl:
+                # gate_logits is at index 5 when output_all_states=False,
+                # and at index 5 when output_all_states=True (before h_all, c_all)
+                all_gate_logits.append(out[3] if not output_all_states else out[5])
+
             if output_all_states:
                 all_h_all.append(out[3])
                 all_c_all.append(out[4])
@@ -864,6 +887,10 @@ class EchoModel(EchoPreTrainedModel):
                 if output_all_states:
                     return x, next_states, all_c_states, all_gate_stats, all_h_all, all_c_all
                 return x, next_states, all_c_states, all_gate_stats
+            if _output_gl:
+                if output_all_states:
+                    return x, next_states, all_gate_logits, all_h_all, all_c_all
+                return x, next_states, all_gate_logits
             if output_all_states:
                 return x, next_states, all_h_all, all_c_all
             return x, next_states
@@ -878,6 +905,8 @@ class EchoModel(EchoPreTrainedModel):
         if output_dsrn_telemetry:
             output_obj.all_c_states = all_c_states
             output_obj.all_gate_stats = all_gate_stats
+        if _output_gl:
+            output_obj.all_gate_logits = all_gate_logits
         if output_all_states:
             output_obj.all_h_all = all_h_all
             output_obj.all_c_all = all_c_all
@@ -990,6 +1019,12 @@ class EchoForCausalLM(EchoPreTrainedModel, GenerationMixin):
             else getattr(self.config, "use_return_dict", True)
         )
 
+        _output_gl = getattr(self.config, "output_surprise_gate_logits", False)
+        alpha = getattr(self.config, "surprise_temperature_alpha", 0.0)
+        # Force telemetry when alpha > 0 (needed for gate_stats)
+        if alpha > 0.0:
+            output_dsrn_telemetry = True
+
         '''
         If kwargs is getting overloaded with extra args HF generate passes,
         we safely extract kwargs here.
@@ -1013,9 +1048,12 @@ class EchoForCausalLM(EchoPreTrainedModel, GenerationMixin):
         if hasattr(model_out, "last_hidden_state"):
             hidden_states = model_out.last_hidden_state
             new_states = model_out.past_key_values
+            # Extract gate_logits when available (DSpark speculative decoding)
+            gate_logits = getattr(model_out, "all_gate_logits", None)
         else:
             hidden_states = model_out[0]
             new_states = model_out[1]
+            gate_logits = model_out[2] if len(model_out) > 2 and _output_gl else None
 
         # Extract telemetry if model returned raw tuple (or via custom properties)
         if hasattr(model_out, "all_c_states"):
@@ -1028,6 +1066,17 @@ class EchoForCausalLM(EchoPreTrainedModel, GenerationMixin):
         # Project using Causal LM head
         logits = self.lm_head(hidden_states)
 
+        # ── Surprise-gate temperature modulation ─────────────────────────
+        # When alpha > 0, the surprise gate λ_t modulates the output logits:
+        #   logits = logits / (1 + α · λ_t)
+        # High surprise flattens the distribution; low surprise leaves it alone.
+        alpha = getattr(self.config, "surprise_temperature_alpha", 0.0)
+        if alpha > 0.0:
+            gate_stats = getattr(model_out, "all_gate_stats", None)
+            if gate_stats is not None and len(gate_stats) > 0:
+                gate_mean = torch.stack(gate_stats).mean(dim=0)  # (B, T)
+                logits = logits / (1.0 + alpha * gate_mean.unsqueeze(-1))
+
         loss = None
         if labels is not None:
             # Shift so that tokens < n predict n
@@ -1038,15 +1087,20 @@ class EchoForCausalLM(EchoPreTrainedModel, GenerationMixin):
 
         if not return_dict:
             output = (logits, new_states)
+            if gate_logits is not None:
+                output = output + (gate_logits,)
             return ((loss,) + output) if loss is not None else output
 
-        return CausalLMOutputWithPast(
+        out = CausalLMOutputWithPast(
             loss=loss,
             logits=logits,
             past_key_values=new_states if use_cache else None,
             hidden_states=(hidden_states,) if output_hidden_states else None,
             attentions=None,
         )
+        if gate_logits is not None:
+            out.all_gate_logits = gate_logits
+        return out
 
     def prepare_inputs_for_generation(
         self, input_ids, past_key_values=None, attention_mask=None, **kwargs
@@ -1220,6 +1274,15 @@ class EchoForSequenceClassification(EchoPreTrainedModel):
         hidden_states = model_out[0]  # (B, T, D)
         new_states = model_out[1]
 
+        # Extract gate_logits when available (DSpark speculative decoding)
+        _output_gl = getattr(self.config, "output_surprise_gate_logits", False)
+        if hasattr(model_out, "all_gate_logits"):
+            gate_logits = model_out.all_gate_logits
+        elif isinstance(model_out, tuple) and len(model_out) > 2 and _output_gl:
+            gate_logits = model_out[2]
+        else:
+            gate_logits = None
+
         # --- Pooling: last non-padding token ---
         if attention_mask is not None:
             # Find the index of the last 1 in each row of attention_mask
@@ -1267,15 +1330,20 @@ class EchoForSequenceClassification(EchoPreTrainedModel):
 
         if not return_dict:
             output = (logits, new_states)
+            if gate_logits is not None:
+                output = output + (gate_logits,)
             return ((loss,) + output) if loss is not None else output
 
-        return SequenceClassifierOutputWithPast(
+        out = SequenceClassifierOutputWithPast(
             loss=loss,
             logits=logits,
             past_key_values=new_states if use_cache else None,
             hidden_states=None,
             attentions=None,
         )
+        if gate_logits is not None:
+            out.all_gate_logits = gate_logits
+        return out
 
     # ------------------------------------------------------------------
     # Convenience inference API

utils.py

@@ -0,0 +1,51 @@
+def visualize_masked_samples(dataset, tokenizer, num_samples=2):
+    """
+    Prints tokenized samples and color-codes tokens based on whether they are masked for loss.
+    """
+    # ANSI escape codes
+    RESET = '\033[0m'
+    MASKED = '\033[90m\033[9m'  # Dark gray + strikethrough
+    LOSS = '\033[92m\033[1m'  # Bright green + bold
+
+    print("\n" + "=" * 80)
+    print(f"🔍 PREVIEWING TOKENIZED DATASET MASKS ({min(num_samples, len(dataset))} samples)")
+    print(
+        f"Legend: {MASKED}Gray Strikethrough{RESET} = Masked (prompt, ignored), {LOSS}Bright Green{RESET} = Loss Calculated (completion)"
+    )
+    print("=" * 80 + "\n")
+
+    for i in range(min(num_samples, len(dataset))):
+        row = dataset[i]
+        ids = row['input_ids']
+        labels = row['labels']
+
+        out = []
+        span_ids = []
+        # Fallback if somehow empty
+        if not labels:
+            continue
+
+        span_masked = labels[0] == -100
+
+        for tid, lbl in zip(ids, labels):
+            is_masked = lbl == -100
+            if is_masked == span_masked:
+                span_ids.append(tid)
+            else:
+                text = tokenizer.decode(span_ids, skip_special_tokens=False)
+                color = MASKED if span_masked else LOSS
+                # Ensure newlines don't break the ANSI formatting
+                text = text.replace('\n', f'{RESET}\n{color}')
+                out.append(f"{color}{text}{RESET}")
+                span_ids = [tid]
+                span_masked = is_masked
+
+        if span_ids:
+            text = tokenizer.decode(span_ids, skip_special_tokens=False)
+            color = MASKED if span_masked else LOSS
+            text = text.replace('\n', f'{RESET}\n{color}')
+            out.append(f"{color}{text}{RESET}")
+
+        print(f"--- Sample {i} ({len(ids)} tokens, {labels.count(-100)} masked) ---")
+        print("".join(out))
+        print("\n" + "-" * 80 + "\n")