""" echo_dsrn/modeling_generative_clf.py ──────────────────────────────────────────────────────────────────────────── EchoForGenerativeClassification ──────────────────────────────────────────────────────────────────────────── A sequence classifier built on top of EchoForCausalLM using *constrained scoring* instead of a linear head. Why a new class instead of EchoForSequenceClassification ───────────────────────────────────────────────────────── EchoForSequenceClassification uses a single nn.Linear layer seeded from the lm_head rows. This works perfectly for adapters with single-token labels (e.g. NSFW: "0" / "1"). For multi-token labels like "weather_query" or "iot_hue_lightchange" the mean-pooling approximation loses too much signal. EchoForGenerativeClassification instead computes, for each candidate label L: score(L | x) = Σ_t log P(token_t | x, token_1..t-1) i.e. the sum of log-probabilities of each token in L conditioned on the input and the previously generated label tokens. These 60 scores form the logits returned by forward(), making the model a drop-in AutoModelForSequenceClassification. HuggingFace API compatibility ────────────────────────────── • Registered with AutoModelForSequenceClassification via __init__.py • forward() returns SequenceClassifierOutputWithPast (logits, loss, hidden_states) • classify(text, tokenizer) convenience method mirrors EchoForSequenceClassification • from_causal_lm() factory for converting a merged CausalLM checkpoint • config.auto_map updated by the factory to point here • No new weights are added — the model is the base lm with the adapter merged Scoring algorithm ───────────────── For each label string L_i (from config.id2label): 1. Tokenise L_i (no BOS/EOS, no special tokens) 2. Concatenate [input_ids, label_tokens] along seq_len 3. Run a single forward pass on the combined sequence 4. Sum log-softmax values at the positions where label tokens are predicted (i.e. positions [n_input, n_input+1, ..., n_input+len(label)-1]) This is O(n_labels) forward passes — for 60 MASSIVE intents on short utterances (avg 7 tokens) it runs in ~60ms on GPU, which is acceptable for inference. Training is NOT required — the adapter's generative knowledge is used directly. """ from __future__ import annotations import typing from typing import List, Optional, Tuple, Union if typing.TYPE_CHECKING: # Force HF trust_remote_code to bundle nested dependencies from .triton_scan import triton_dsrn_parallel_scan from .utils import rms_norm_fn import torch import torch.nn.functional as F from transformers.modeling_outputs import SequenceClassifierOutputWithPast from .configuration_echo import EchoConfig from .modeling_echo import EchoForCausalLM class EchoForGenerativeClassification(EchoForCausalLM): """ Intent / multi-label classifier using constrained generative scoring. The model is identical to EchoForCausalLM plus: • A cache of tokenised label sequences built from config.id2label • A forward() that scores all labels and returns (B, num_labels) logits • No new trainable parameters Usage ───── from echo_dsrn import EchoForGenerativeClassification from transformers import AutoTokenizer model = EchoForGenerativeClassification.from_pretrained( "ethicalabs/Echo-SmolTools-114M-Intent-CLF-Gen", trust_remote_code=True, ) tok = AutoTokenizer.from_pretrained( "ethicalabs/Echo-SmolTools-114M-Intent-CLF-Gen", trust_remote_code=True, ) label, probs = model.classify("What time is it in Tokyo?", tok) # → ("datetime_query", tensor([...])) """ # Let HF serialisation know this is a classification model _no_split_modules = [] def __init__(self, config: EchoConfig): super().__init__(config) # Label token cache — populated lazily on first call that needs it # (requires a tokenizer, which we don't have at __init__ time) self._label_token_ids: Optional[List[List[int]]] = None # ------------------------------------------------------------------ # Label token cache # ------------------------------------------------------------------ def _build_label_cache(self, tokenizer) -> List[List[int]]: """ Tokenise every label string in config.id2label and cache the result. Called lazily on the first forward/classify call. """ cache: List[List[int]] = [] n = self.config.num_labels for idx in range(n): label_str = self.config.id2label[idx] # Encode without special tokens — the label is a continuation, # not a standalone sentence. tids = tokenizer.encode(label_str, add_special_tokens=False) if not tids: raise ValueError( f"Label '{label_str}' (id={idx}) tokenises to an empty sequence. " "Please check your tokenizer and label strings." ) cache.append(tids) self._label_token_ids = cache return cache # ------------------------------------------------------------------ # Scoring helpers # ------------------------------------------------------------------ @torch.inference_mode() def _score_labels( self, input_ids: torch.Tensor, # (B, S) attention_mask: torch.Tensor, # (B, S) label_cache: List[List[int]], ) -> torch.Tensor: """ Return a (B, num_labels) tensor of log-probability scores. For each sample b and label L_i: score[b, i] = Σ_t log P(L_i[t] | input_ids[b], L_i[:t]) """ device = input_ids.device B, S = input_ids.shape num_labels = len(label_cache) scores = torch.full((B, num_labels), float("-inf"), device=device) # 1. Process the prompt once and cache the state base_out = super().forward( input_ids=input_ids, attention_mask=attention_mask, use_cache=True, ) prompt_pkv = base_out.past_key_values # Logits predicting the first label token (last token of the prompt) prompt_last_logits = base_out.logits[:, -1:, :] # (B, 1, V) # Helper to clone the custom EchoCache safely def clone_pkv(pkv): if pkv is None: return None if hasattr(pkv, "states"): from .modeling_echo import EchoCache new_states = [tuple(t.clone() for t in state_tuple) for state_tuple in pkv.states] return EchoCache(new_states) elif isinstance(pkv, (list, tuple)): return [tuple(t.clone() for t in state_tuple) for state_tuple in pkv] return pkv for i, label_tids in enumerate(label_cache): L = len(label_tids) label_t = ( torch.tensor(label_tids, dtype=torch.long, device=device).unsqueeze(0).expand(B, -1) ) # (B, L) # 2. Forward pass ONLY the label tokens using the cloned cache label_out = super().forward( input_ids=label_t, attention_mask=None, # EchoModel handles this causally for recurrent steps past_key_values=clone_pkv(prompt_pkv), use_cache=False, ) # 3. Splice the logits together. # The first label token is predicted by prompt_last_logits. # The remaining label tokens are predicted by label_out.logits[:, :-1, :] if L == 1: label_logits = prompt_last_logits else: label_logits = torch.cat( [prompt_last_logits, label_out.logits[:, :-1, :]], dim=1 ) # (B, L, V) log_probs = F.log_softmax(label_logits, dim=-1) # (B, L, V) # Gather log-prob of the correct next token at each label position label_t_expanded = label_t.unsqueeze(-1) # (B, L, 1) token_log_probs = log_probs.gather(dim=-1, index=label_t_expanded).squeeze(-1) # (B, L) # Sum log-probs across label tokens → scalar score per sample scores[:, i] = token_log_probs.sum(dim=-1) # (B,) return scores # (B, num_labels) — these are log-prob sums (higher = better) # ------------------------------------------------------------------ # forward() # ------------------------------------------------------------------ def forward( self, input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, labels: Optional[torch.LongTensor] = None, # Tokenizer is needed for label scoring; injected via classify() or # set once via set_tokenizer(). _tokenizer=None, **kwargs, ) -> SequenceClassifierOutputWithPast: """ Returns SequenceClassifierOutputWithPast with: • logits : (B, num_labels) log-prob sums — higher = more likely intent • loss : cross-entropy loss if `labels` is provided, else None Note: kwargs are accepted but ignored (past_key_values, use_cache, etc.) to maintain drop-in compatibility with the HF pipeline. """ tokenizer = _tokenizer or self._tokenizer_ref if tokenizer is None: raise RuntimeError( "EchoForGenerativeClassification.forward() requires a tokenizer for " "label scoring. Either call model.set_tokenizer(tok) once after loading, " "or use the classify() convenience method." ) if self._label_token_ids is None: self._build_label_cache(tokenizer) logits = self._score_labels(input_ids, attention_mask, self._label_token_ids) loss = None if labels is not None: loss = F.cross_entropy(logits, labels) return SequenceClassifierOutputWithPast( loss=loss, logits=logits, ) # ------------------------------------------------------------------ # Tokenizer binding # ------------------------------------------------------------------ _tokenizer_ref = None # class-level default; overridden per-instance def set_tokenizer(self, tokenizer) -> "EchoForGenerativeClassification": """ Bind a tokenizer so forward() can score label strings. Call this once after loading the model: model = EchoForGenerativeClassification.from_pretrained(...) model.set_tokenizer(tokenizer) """ self._tokenizer_ref = tokenizer if self._label_token_ids is None: self._build_label_cache(tokenizer) return self @classmethod def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs): """ Override from_pretrained to automatically load and bind the tokenizer from the same checkpoint, so that forward() works out of the box when loaded via pipeline() or AutoModelForSequenceClassification without a manual set_tokenizer() call. """ model = super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs) try: from transformers import AutoTokenizer tok = AutoTokenizer.from_pretrained( pretrained_model_name_or_path, trust_remote_code=kwargs.get("trust_remote_code", False), ) model.set_tokenizer(tok) except Exception: # Best-effort: if tokenizer loading fails, the user can still call # set_tokenizer() manually before running forward(). pass return model # ------------------------------------------------------------------ # High-level inference API # ------------------------------------------------------------------ # Default user-message template; matches what the Intent PEFT was trained on. _DEFAULT_USER_TEMPLATE = "Classify the intent of the following request: {utt}" def _format_prompts( self, texts: List[str], tokenizer, system_prompt: Optional[str], user_template: str ) -> List[str]: """ Apply the chat template to a list of raw utterances. Falls back to raw text if the tokenizer has no chat template. """ if not hasattr(tokenizer, "apply_chat_template") or tokenizer.chat_template is None: return texts # bare tokenizer — score raw text directly formatted = [] for utt in texts: messages = [] if system_prompt: messages.append({"role": "system", "content": system_prompt}) messages.append({"role": "user", "content": user_template.format(utt=utt)}) prompt = tokenizer.apply_chat_template( messages, tokenize=False, add_generation_prompt=True ) formatted.append(prompt) return formatted @torch.inference_mode() def classify( self, text: Union[str, List[str]], tokenizer, system_prompt: Optional[str] = None, user_template: Optional[str] = None, max_length: int = 256, ) -> Union[ Tuple[str, torch.Tensor], Tuple[List[str], torch.Tensor], ]: """ Classify text (or a list of texts) into one of the intent classes. The input is formatted with the same chat template used during adapter training. The system_prompt and user_template default to the values baked into config at merge time. Args: text : Raw utterance(s) — no prompt wrapping needed. tokenizer : HF tokenizer for the model. system_prompt : Override the system message (optional). user_template : Override the user message template (optional). Must contain ``{utt}`` as the utterance placeholder. max_length : Max tokenised length (default 256 to fit chat template). Returns (single input): (label_str, probs_tensor) — probs shape (num_labels,) Returns (batch input): (label_list, probs_tensor) — probs shape (B, num_labels) """ self.set_tokenizer(tokenizer) single = isinstance(text, str) texts = [text] if single else text # Resolve prompt components: args > config > defaults sys_prompt = system_prompt or getattr(self.config, "classification_system_prompt", None) usr_template = user_template or getattr( self.config, "classification_user_template", self._DEFAULT_USER_TEMPLATE ) # Apply the chat template the adapter was trained with formatted = self._format_prompts(texts, tokenizer, sys_prompt, usr_template) enc = tokenizer( formatted, return_tensors="pt", padding=True, truncation=True, max_length=max_length, ) device = next(self.parameters()).device input_ids = enc["input_ids"].to(device) attention_mask = enc["attention_mask"].to(device) out = self.forward( input_ids=input_ids, attention_mask=attention_mask, _tokenizer=tokenizer, ) probs = torch.softmax(out.logits, dim=-1) # (B, num_labels) if single: pred_idx = probs[0].argmax().item() label_str = self.config.id2label[pred_idx] return label_str, probs[0] else: pred_idxs = probs.argmax(dim=-1).tolist() label_strs = [self.config.id2label[i] for i in pred_idxs] return label_strs, probs # ------------------------------------------------------------------ # Factory: from_causal_lm # ------------------------------------------------------------------ @classmethod def from_causal_lm( cls, causal_lm_model: EchoForCausalLM, num_labels: int, id2label: dict, label2id: Optional[dict] = None, system_prompt: Optional[str] = None, user_template: Optional[str] = None, ) -> "EchoForGenerativeClassification": """ Construct an EchoForGenerativeClassification from a (possibly adapter-merged) EchoForCausalLM instance. No new weights are added — all parameters come from causal_lm_model. The config is updated in-place to record classification metadata. Args: causal_lm_model : A loaded (and optionally merged) EchoForCausalLM. num_labels : Number of intent classes (e.g. 60 for MASSIVE). id2label : Mapping {int_idx: label_str}. label2id : Optional reverse mapping; auto-derived if None. system_prompt : System message to bake into config (used by classify()). user_template : User message template to bake into config. Must contain ``{utt}`` placeholder. Returns: EchoForGenerativeClassification ready for inference. """ if label2id is None: label2id = {v: int(k) for k, v in id2label.items()} config = causal_lm_model.config config.num_labels = num_labels config.id2label = {int(k): v for k, v in id2label.items()} config.label2id = {v: int(k) for k, v in id2label.items()} # Bake prompt components into config so the model is self-contained if system_prompt is not None: config.classification_system_prompt = system_prompt if user_template is not None: config.classification_user_template = user_template # Point auto_map to this module so HF can find it with trust_remote_code config.auto_map = { **getattr(config, "auto_map", {}), "AutoModelForSequenceClassification": ( "modeling_generative_clf.EchoForGenerativeClassification" ), } # Reuse all weights — no state_dict copy needed, just change the class gen_clf = cls(config) gen_clf.load_state_dict(causal_lm_model.state_dict()) # Cast to same dtype as source src_dtype = next(causal_lm_model.parameters()).dtype gen_clf.to(src_dtype) return gen_clf