Upload folder using huggingface_hub

__init__.py

@@ -14,6 +14,7 @@ from transformers import (
 
 from .configuration_echo import EchoConfig
 from .modeling_echo import EchoForCausalLM, EchoForSequenceClassification, EchoModel
+from .pipelines import ChatTextClassificationPipeline, pipeline
 
 # Register with HuggingFace so AutoClass routing works
 AutoConfig.register("echo", EchoConfig)
@@ -25,4 +26,6 @@ __all__ = [
     "EchoModel",
     "EchoForCausalLM",
     "EchoForSequenceClassification",
+    "ChatTextClassificationPipeline",
+    "pipeline",
 ]

configuration_echo.py

@@ -27,6 +27,7 @@ class EchoConfig(PretrainedConfig):
         id2label: Optional[dict] = None,
         label2id: Optional[dict] = None,
         classifier_dropout: float = 0.0,
+        classification_use_chat_template: bool = True,
         **kwargs,
     ):
         # Synchronize hidden_size / embed_dim (HF synonym pair).

modeling_echo.py

@@ -924,15 +924,22 @@ class EchoForCausalLM(EchoPreTrainedModel, GenerationMixin):
     # tie_word_embeddings=True and get_input/output_embeddings() are both defined.
     _tied_weights_keys = {"lm_head.weight": "model.embedding.weight"}
 
+    # _keys_to_ignore_on_load_missing: some transformers releases (≥5.7)
+    # declare this as a read-only @property, which conflicts with
+    # register_buffer / __setattr__.  We provide a trivial setter
+    # to avoid AttributeError during model.__init__.
+    _keys = []
+
     @property
     def _keys_to_ignore_on_load_missing(self):
-        # When mlp_bias=False (the default, and the setting for all v0.1.2 checkpoints),
-        # bias tensors are not present in the checkpoint and should not trigger warnings.
-        # When mlp_bias=True, these keys WILL exist in the checkpoint — do not silence them.
         if not getattr(self.config, "mlp_bias", False):
             return [r"model\.blocks\.\d+\.mlp_(up|down)\.bias"]
         return []
 
+    @_keys_to_ignore_on_load_missing.setter
+    def _keys_to_ignore_on_load_missing(self, value):
+        self._keys = value
+
     @classmethod
     def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
         model = super().from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
@@ -998,6 +1005,7 @@ class EchoForCausalLM(EchoPreTrainedModel, GenerationMixin):
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
         output_dsrn_telemetry: Optional[bool] = False,
+        skip_logits: Optional[bool] = False,
         **kwargs,
     ) -> Union[Tuple, CausalLMOutputWithPast]:
 
@@ -1063,27 +1071,54 @@ class EchoForCausalLM(EchoPreTrainedModel, GenerationMixin):
             self._latest_c_states = model_out[2]
             self._latest_gate_stats = model_out[3]
 
-        # 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)
+        # ── Surprise-gate temperature scale (computed once for both paths) ──
+        _surprise_scale = None  # (B, T) or None
         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))
+                _surprise_scale = 1.0 + alpha * torch.stack(gate_stats).mean(dim=0)  # (B, T)
 
-        loss = None
-        if labels is not None:
-            # Shift so that tokens < n predict n
-            shift_logits = logits[..., :-1, :].contiguous()
+        # ── skip_logits: avoid materialising the full [B, T, V] logits ──
+        if skip_logits and labels is not None:
+            logits = None
+            shift_hidden = hidden_states[..., :-1, :]
             shift_labels = labels[..., 1:].contiguous()
-            loss_fct = nn.CrossEntropyLoss()
-            loss = loss_fct(shift_logits.view(-1, self.config.vocab_size), shift_labels.view(-1))
+            _CHUNK = 512
+            _flat_hidden = shift_hidden.reshape(-1, shift_hidden.size(-1))
+            _flat_labels = shift_labels.reshape(-1)
+            # Shift surprise scale to align with shifted logits
+            _flat_scale = None
+            if _surprise_scale is not None:
+                _flat_scale = _surprise_scale[..., :-1].reshape(-1)  # (B*(T-1),)
+            _loss_fct = nn.CrossEntropyLoss(ignore_index=-100, reduction="sum")
+            _total_loss = torch.zeros((), dtype=torch.float32, device=_flat_hidden.device)
+            _total_tokens = _flat_hidden.new_zeros((), dtype=torch.long)
+            for _i in range(0, _flat_hidden.size(0), _CHUNK):
+                _ch = _flat_hidden[_i : _i + _CHUNK]
+                _cl = self.lm_head(_ch).float()
+                if _flat_scale is not None:
+                    _cl = _cl / _flat_scale[_i : _i + _CHUNK].unsqueeze(-1)
+                _ll = _flat_labels[_i : _i + _CHUNK]
+                _total_loss = _total_loss + _loss_fct(_cl, _ll)
+                _total_tokens = _total_tokens + (_ll != -100).sum()
+            loss = _total_loss / _total_tokens.clamp(min=1)
+        else:
+            # Project using Causal LM head
+            logits = self.lm_head(hidden_states)
+
+            # ── Surprise-gate temperature modulation ─────────────────────────
+            if _surprise_scale is not None:
+                logits = logits / _surprise_scale.unsqueeze(-1)
+
+            loss = None
+            if labels is not None:
+                # Shift so that tokens < n predict n
+                shift_logits = logits[..., :-1, :].contiguous()
+                shift_labels = labels[..., 1:].contiguous()
+                loss_fct = nn.CrossEntropyLoss()
+                loss = loss_fct(
+                    shift_logits.view(-1, self.config.vocab_size), shift_labels.view(-1)
+                )
 
         if not return_dict:
             output = (logits, new_states)
@@ -1203,12 +1238,27 @@ class EchoForSequenceClassification(EchoPreTrainedModel):
 
     def __init__(self, config: EchoConfig):
         super().__init__(config)
-        self.num_labels = getattr(config, "num_labels", 2)
+        # PretrainedConfig.to_dict() strips num_labels — infer from id2label
+        if config.id2label is not None and len(config.id2label) > 0:
+            self.num_labels = len(config.id2label)
+        else:
+            self.num_labels = getattr(config, "num_labels", 2)
         self.model = EchoModel(config)
 
         classifier_dropout = getattr(config, "classifier_dropout", 0.0)
         self.dropout = nn.Dropout(classifier_dropout) if classifier_dropout > 0.0 else nn.Identity()
-        self.classifier = EchoClassifier(config.embed_dim, self.num_labels, bias=True)
+
+        # Classifier head dimension depends on pooling mode
+        pooling_mode = getattr(config, "pooling_mode", None)
+        if pooling_mode == "mean_c_all":
+            head_dim = config.hidden_size * config.num_heads  # c_all: 2048
+        else:
+            head_dim = config.embed_dim  # hidden state: 512
+
+        self.classifier = EchoClassifier(head_dim, self.num_labels, bias=True)
+
+        # Persist pooling mode from config (survives save/load roundtrip)
+        self._pooling_mode = getattr(config, "pooling_mode", None)
 
         self.post_init()
 
@@ -1269,6 +1319,11 @@ class EchoForSequenceClassification(EchoPreTrainedModel):
         if alpha > 0.0:
             kwargs["output_dsrn_telemetry"] = True
 
+        # mean_c_all pooling needs the recurrent slow state per token
+        pooling_mode = getattr(self, "_pooling_mode", None)
+        if pooling_mode in ("mean_c_all", "hybrid"):
+            kwargs["output_all_states"] = True
+
         model_out = self.model(
             input_ids=input_ids,
             past_key_values=past_key_values,
@@ -1288,32 +1343,41 @@ class EchoForSequenceClassification(EchoPreTrainedModel):
         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
-            seq_lengths = attention_mask.sum(dim=1) - 1  # (B,)
-            seq_lengths = seq_lengths.clamp(min=0)
-        else:
-            # No mask: use the true last token
-            if input_ids is not None:
-                seq_lengths = torch.full(
-                    (hidden_states.size(0),),
-                    hidden_states.size(1) - 1,
-                    dtype=torch.long,
-                    device=hidden_states.device,
+        # --- Pooling ---
+        if pooling_mode == "mean_c_all":
+            # Mean pool the recurrent slow states c_all from the last layer
+            c_all = model_out.all_c_all[-1]  # (B, T, hidden_size * num_heads)
+            if attention_mask is not None:
+                mask_expanded = attention_mask.unsqueeze(-1).expand(c_all.size()).float()
+                pooled = (c_all * mask_expanded).sum(dim=1) / mask_expanded.sum(dim=1).clamp(
+                    min=1e-9
                 )
             else:
-                seq_lengths = torch.full(
-                    (hidden_states.size(0),),
-                    hidden_states.size(1) - 1,
-                    dtype=torch.long,
-                    device=hidden_states.device,
-                )
+                pooled = c_all.mean(dim=1)
+        else:
+            # Default: last non-padding token
+            if attention_mask is not None:
+                seq_lengths = attention_mask.sum(dim=1) - 1  # (B,)
+                seq_lengths = seq_lengths.clamp(min=0)
+            else:
+                if input_ids is not None:
+                    seq_lengths = torch.full(
+                        (hidden_states.size(0),),
+                        hidden_states.size(1) - 1,
+                        dtype=torch.long,
+                        device=hidden_states.device,
+                    )
+                else:
+                    seq_lengths = torch.full(
+                        (hidden_states.size(0),),
+                        hidden_states.size(1) - 1,
+                        dtype=torch.long,
+                        device=hidden_states.device,
+                    )
+            pooled = hidden_states[
+                torch.arange(hidden_states.size(0), device=hidden_states.device), seq_lengths
+            ]  # (B, D)
 
-        # Gather last-token hidden states: (B, D)
-        pooled = hidden_states[
-            torch.arange(hidden_states.size(0), device=hidden_states.device), seq_lengths
-        ]
         pooled = self.dropout(pooled)
         logits = self.classifier(pooled)  # (B, num_labels)
 
@@ -1402,11 +1466,12 @@ class EchoForSequenceClassification(EchoPreTrainedModel):
 
         self.eval()
 
-        # Format text if baked-in templates exist
+        # Format text if baked-in templates exist AND not disabled
         sys_prompt = getattr(self.config, "system_prompt", None)
         usr_template = getattr(self.config, "user_template", None)
+        use_chat = getattr(self.config, "classification_use_chat_template", True)
 
-        if sys_prompt and usr_template:
+        if use_chat and sys_prompt and usr_template:
             messages = [{"role": "system", "content": sys_prompt}]
             messages.append({"role": "user", "content": usr_template.format(text=text)})
             # Format using the tokenizer's chat template
@@ -1575,3 +1640,90 @@ class EchoForSequenceClassification(EchoPreTrainedModel):
             config.dtype = str(src_dtype).replace("torch.", "")
 
         return clf_model
+
+    # ──────────────────────────────────────────────────────────────────
+    # from_embedding — convert an embedding model to a classifier
+    # ──────────────────────────────────────────────────────────────────
+    @classmethod
+    def from_embedding(
+        cls,
+        embed_model,
+        num_labels: int = 60,
+        id2label: Optional[dict] = None,
+        label2id: Optional[dict] = None,
+        classifier_dropout: float = 0.0,
+    ) -> "EchoForSequenceClassification":
+        """
+        Construct an :class:`EchoForSequenceClassification` from an
+        :class:`~echo_embedding.modeling_embedding.EchoModelForSentenceEmbedding`
+        instance (or HF path).
+
+        The backbone weights are copied; the pooling mode (``mean_c_all``)
+        is inherited from the embedding model.  The classifier head is
+        **randomly initialised** — this factory is dataset-agnostic.
+        Fine-tune on your target dataset afterward.
+
+        Parameters
+        ----------
+        embed_model:
+            An ``EchoModelForSentenceEmbedding`` instance or a HuggingFace
+            model path / hub ID.
+        num_labels:
+            Number of output classes.
+        id2label:
+            Optional mapping ``{int -> str}`` for label names.
+        label2id:
+            Optional reverse mapping ``{str -> int}``.
+        classifier_dropout:
+            Dropout probability before the classification head.
+
+        Returns
+        -------
+        EchoForSequenceClassification
+        """
+        # ── 1. Resolve the embedding model ─────────────────────────
+        if isinstance(embed_model, str):
+            from echo_embedding.modeling_embedding import EchoModelForSentenceEmbedding
+
+            embed_model = EchoModelForSentenceEmbedding.from_pretrained(
+                embed_model, trust_remote_code=True
+            )
+
+        if id2label is None:
+            id2label = {i: str(i) for i in range(num_labels)}
+        if label2id is None:
+            label2id = {v: k for k, v in id2label.items()}
+
+        # ── 2. Clone config and inject classification fields ──────
+        config = embed_model.config
+        config.num_labels = num_labels
+        config.id2label = id2label
+        config.label2id = label2id
+        config.classifier_dropout = classifier_dropout
+        config.pooling_mode = getattr(config, "pooling_mode", "c_T")
+        config.classification_use_chat_template = False  # embed→CLF: no chat template
+
+        # Update auto_map so Hub pipeline + AutoModel load correctly
+        config.auto_map = {
+            "AutoConfig": "configuration_echo.EchoConfig",
+            "AutoModel": "modeling_echo.EchoForSequenceClassification",
+            "AutoModelForSequenceClassification": "modeling_echo.EchoForSequenceClassification",
+            "DSRNScan": "triton_scan.DSRNScanTriton",
+        }
+
+        # ── 3. Build classifier (random init) ─────────────────────
+        clf_model = cls(config)
+
+        # ── 4. Copy backbone weights ──────────────────────────────
+        clf_model.model.load_state_dict(embed_model.model.state_dict(), strict=True)
+
+        # ── 5. Cast to source dtype ───────────────────────────────
+        src_dtype = embed_model.dtype
+        if src_dtype != torch.float32:
+            clf_model = clf_model.to(src_dtype)
+            config.dtype = str(src_dtype).replace("torch.", "")
+
+        # ── 6. Store pooling mode for forward pass ─────────────────
+        clf_model._pooling_mode = getattr(config, "pooling_mode", "c_T")
+
+        return clf_model