From 018c84100e45df0346bf3ea4e501b16216ed7f46 Mon Sep 17 00:00:00 2001
From: Ingvar <igor.stepanov2000@gmail.com>
Date: Fri, 2 Oct 2026 10:27:17 +0300
Subject: [PATCH 01/12] implement recurrence for decoder kv scorer

---
 gliclass/config.py   |  22 +++++++++
 gliclass/model.py    | 111 +++++++++++++++++++++++++++++++++++++++++--
 gliclass/scorers.py  | 106 +++++++++++++++++++++++++++++++++++++++--
 gliclass/training.py |  25 ++++++++--
 4 files changed, 251 insertions(+), 13 deletions(-)

diff --git a/gliclass/config.py b/gliclass/config.py
index dd11728..524344a 100644
--- a/gliclass/config.py
+++ b/gliclass/config.py
@@ -60,6 +60,14 @@ def __init__(
         dropout=0.1,
         use_segment_embeddings=False,
         scorer_encoder_num_layers=2,
+        recurrent_steps=1,
+        recurrent_min_steps=1,
+        recurrent_inference=True,
+        recurrent_inference_max_steps=None,
+        recurrent_halt_threshold=0.01,
+        recurrent_improvement_coef=0.1,
+        recurrent_improvement_margin=0.0,
+        recurrent_bptt_steps=None,
         **kwargs,
     ):
         if isinstance(encoder_config, dict):
@@ -149,4 +157,18 @@ def __init__(
 
         self.scorer_encoder_num_layers = scorer_encoder_num_layers
 
+        # Recurrent hidden reasoning for the decoder-kv scorer (recurrent_steps=1 disables it).
+        # Training depth is sampled uniformly from [recurrent_min_steps, recurrent_steps].
+        self.recurrent_steps = recurrent_steps
+        self.recurrent_min_steps = recurrent_min_steps
+        # Inference: run up to recurrent_inference_max_steps (default recurrent_steps, may exceed it) and
+        # stop an example once no label probability changes by recurrent_halt_threshold (None/0 = never).
+        self.recurrent_inference = recurrent_inference
+        self.recurrent_inference_max_steps = recurrent_inference_max_steps
+        self.recurrent_halt_threshold = recurrent_halt_threshold
+        # Hinge relu(l_t - (1 - margin) * sg(l_{t-1})) for steps t >= 3, on top of a loss at every step.
+        self.recurrent_improvement_coef = recurrent_improvement_coef
+        self.recurrent_improvement_margin = recurrent_improvement_margin
+        self.recurrent_bptt_steps = recurrent_bptt_steps
+
         super().__init__(problem_type=problem_type, **kwargs)
diff --git a/gliclass/model.py b/gliclass/model.py
index 0757381..cab2146 100644
--- a/gliclass/model.py
+++ b/gliclass/model.py
@@ -62,6 +62,9 @@ class GLiClassOutput(SequenceClassifierOutput):
     text_embeddings: torch.Tensor | None = None
     class_embeddings: torch.Tensor | None = None
     past_key_values: Any | None = None
+    recurrent_logits: Tuple[torch.Tensor, ...] | None = None
+    recurrent_losses: torch.Tensor | None = None
+    recurrent_num_steps: torch.Tensor | None = None
 
 
 class GLiClassPreTrainedModel(PreTrainedModel):
@@ -1042,6 +1045,78 @@ def get_loss(self, logits, labels):
                 raise NotImplementedError(f"{self.config.problem_type} is not implemented.")
         return loss
 
+    def _resolve_recurrence(self, use_recurrence=None, max_recurrent_steps=None, halt_threshold=None):
+        """Pick (num_steps, halt_threshold) for this call.
+
+        Training: depth sampled uniformly from [recurrent_min_steps, max], no halting.
+        Inference: up to max steps, halting per example once probabilities stop changing.
+        """
+        if not self.scorer.recurrent:
+            return 1, None
+        if use_recurrence is None:
+            use_recurrence = self.training or self.config.recurrent_inference
+        if not use_recurrence:
+            return 1, None
+
+        if self.training:
+            max_steps = max_recurrent_steps or self.config.recurrent_steps
+            min_steps = min(self.config.recurrent_min_steps or 1, max_steps)
+            return int(torch.randint(min_steps, max_steps + 1, (1,)).item()), None
+
+        max_steps = max_recurrent_steps or self.config.recurrent_inference_max_steps or self.config.recurrent_steps
+        if halt_threshold is None:
+            halt_threshold = self.config.recurrent_halt_threshold
+        return max_steps, halt_threshold
+
+    def _per_sample_loss(self, logits, labels):
+        """Unreduced loss per example, shape (batch,)."""
+        if self.config.problem_type == "multi_label_classification":
+            from .loss_functions import focal_loss_with_logits
+
+            all_losses = focal_loss_with_logits(
+                logits,
+                labels,
+                self.config.focal_loss_alpha,
+                self.config.focal_loss_gamma,
+                "none",
+            )
+            valid = labels.ne(self.config.ignore_index)
+            return all_losses.sum(-1) / valid.sum(-1).clamp_min(1)
+        elif self.config.problem_type == "single_label_classification":
+            return nn.functional.cross_entropy(
+                logits, labels.view(-1), ignore_index=self.config.ignore_index, reduction="none"
+            )
+        raise NotImplementedError(f"{self.config.problem_type} is not implemented.")
+
+    def get_recurrent_loss(self, step_logits, labels):
+        """Loss over all reasoning steps.
+
+        loss = mean_t L_t                                                 (focal loss after every step)
+             + coef * mean_{t>=3} relu(l_t - (1 - margin) * sg(l_{t-1}))  (deeper must beat shallower)
+
+        The improvement term starts at step 3 vs step 2, so step 1 is only ever trained by its own
+        loss. The stop-gradient on l_{t-1} means the hinge can only push the deeper step down; it
+        never rewards making the earlier step worse.
+
+        Returns:
+            loss: scalar
+            step_losses: (num_steps,) detached mean loss per step, for monitoring
+        """
+        num_labels = step_logits[-1].shape[-1]
+        if labels.shape[-1] != num_labels:
+            labels = labels[:, :num_labels]
+
+        per_sample = torch.stack([self._per_sample_loss(logits, labels) for logits in step_logits])  # (T, B)
+        step_losses = per_sample.mean(dim=1)
+        loss = step_losses.mean()
+
+        if len(step_logits) >= 3 and self.config.recurrent_improvement_coef > 0:
+            target = (1.0 - self.config.recurrent_improvement_margin) * per_sample[1:-1].detach()
+            improvement = torch.relu(per_sample[2:] - target).mean()
+            loss = loss + self.config.recurrent_improvement_coef * improvement
+
+        return loss, step_losses.detach()
+
     def _extract_label_section(
         self,
         hidden_states: torch.Tensor,
@@ -1131,6 +1206,9 @@ def classify_from_decoder_cache(
         label_mask: torch.Tensor,
         position_ids: torch.Tensor,
         past_key_values,
+        use_recurrence: bool | None = None,
+        max_recurrent_steps: int | None = None,
+        halt_threshold: float | None = None,
     ) -> GLiClassOutput:
         """Classify label tokens against a text cache without persisting labels."""
         decoder_outputs = self.decoder_model(
@@ -1144,12 +1222,15 @@ def classify_from_decoder_cache(
 
         label_hidden = decoder_outputs.last_hidden_state[:, -input_ids.shape[1] :, :]
         scorer_device = next(self.scorer.parameters()).device
-        logits = self.scorer(
+        num_steps, halt_threshold = self._resolve_recurrence(use_recurrence, max_recurrent_steps, halt_threshold)
+        step_logits, steps_taken = self.scorer.recurrent_forward(
             hidden_states=label_hidden.to(scorer_device),
             input_ids=input_ids.to(scorer_device),
             attention_mask=label_mask.to(scorer_device),
+            num_steps=num_steps,
+            halt_threshold=halt_threshold,
         )
-        return GLiClassOutput(logits=logits)
+        return GLiClassOutput(logits=step_logits[-1], recurrent_num_steps=steps_taken)
 
     def forward(
         self,
@@ -1159,6 +1240,10 @@ def forward(
         labels: torch.Tensor | None = None,
         return_dict: bool | None = None,
         use_cache: bool = False,
+        use_recurrence: bool | None = None,
+        max_recurrent_steps: int | None = None,
+        halt_threshold: float | None = None,
+        output_recurrent_logits: bool = False,
         **kwargs,
     ):
         """Forward pass.
@@ -1170,9 +1255,13 @@ def forward(
             labels: Classification labels
             return_dict: Return dict output
             use_cache: Whether to return updated cache
+            use_recurrence: Enable/disable scorer recurrence for this call (default: config)
+            max_recurrent_steps: Max reasoning steps for this call (training samples depth up to it)
+            halt_threshold: Inference early-stop threshold on label probability change (0 disables)
+            output_recurrent_logits: Return logits of every reasoning step
 
         Returns:
-            GLiClassOutput with logits, loss, and optionally past_key_values
+            GLiClassOutput with logits (of the last step), loss, and optionally past_key_values
         """
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
@@ -1189,16 +1278,25 @@ def forward(
 
         label_hidden, label_ids, label_mask = self._extract_label_section(hidden_states, input_ids, attention_mask)
 
-        logits = self.scorer(
+        num_steps, halt_threshold = self._resolve_recurrence(use_recurrence, max_recurrent_steps, halt_threshold)
+        step_logits, steps_taken = self.scorer.recurrent_forward(
             hidden_states=label_hidden,
             input_ids=label_ids,
             attention_mask=label_mask,
+            num_steps=num_steps,
+            halt_threshold=halt_threshold,
+            return_all_steps=labels is not None or output_recurrent_logits,
         )
+        logits = step_logits[-1]
 
         if labels is not None:
             self.num_labels = logits.shape[-1]
 
-        loss = self.get_loss(logits, labels)
+        recurrent_losses = None
+        if len(step_logits) > 1 and labels is not None:
+            loss, recurrent_losses = self.get_recurrent_loss(step_logits, labels)
+        else:
+            loss = self.get_loss(logits, labels)
 
         if not return_dict:
             output = (logits,)
@@ -1212,6 +1310,9 @@ def forward(
             hidden_states=decoder_outputs.hidden_states if hasattr(decoder_outputs, "hidden_states") else None,
             attentions=decoder_outputs.attentions if hasattr(decoder_outputs, "attentions") else None,
             past_key_values=decoder_outputs.past_key_values if use_cache else None,
+            recurrent_logits=tuple(step_logits) if output_recurrent_logits else None,
+            recurrent_losses=recurrent_losses,
+            recurrent_num_steps=steps_taken if num_steps > 1 else None,
         )
 
 
diff --git a/gliclass/scorers.py b/gliclass/scorers.py
index 584e7c0..82afb8c 100644
--- a/gliclass/scorers.py
+++ b/gliclass/scorers.py
@@ -221,6 +221,13 @@ def __init__(self, config, **kwargs):
 
         self.epsilon = 1e-8
 
+        # GRU-style recurrent reasoning over scorer_encoder. Gates are step-agnostic (no step
+        # embeddings), so inference can run more steps than were seen during training.
+        self.recurrent = getattr(config, "recurrent_steps", 1) > 1
+        if self.recurrent:
+            self.reset_gate = nn.Linear(config.hidden_size * 2, config.hidden_size)
+            self.update_gate = nn.Linear(config.hidden_size * 2, config.hidden_size)
+
     def _extract_representations(self, hidden_states, input_ids, attention_mask):
         """
         Extract text and label representations from hidden states.
@@ -267,14 +274,107 @@ def forward(self, hidden_states, input_ids, attention_mask, **kwargs):
             hidden_states: (batch_size, seq_length, hidden_size) from decoder backbone
             input_ids: (batch_size, seq_length)
             attention_mask: (batch_size, seq_length)
+            **kwargs: recurrence options, see recurrent_forward
 
         Returns:
             logits: (batch_size, num_labels)
         """
-        encoder_outputs = self.scorer_encoder(hidden_states, attention_mask=attention_mask, return_dict=True)
-
-        contextualized_hidden_states = encoder_outputs.last_hidden_state
+        step_logits, _ = self.recurrent_forward(hidden_states, input_ids, attention_mask, **kwargs)
+        return step_logits[-1]
+
+    def _encode(self, hidden_states, attention_mask):
+        return self.scorer_encoder(hidden_states, attention_mask=attention_mask, return_dict=True).last_hidden_state
+
+    def _recurrent_step(self, state, inputs, attention_mask):
+        """h_t = (1 - z) * h_{t-1} + z * Enc(r * h_{t-1} + (1 - r) * x), x = decoder hidden states."""
+        reset = torch.sigmoid(self.reset_gate(torch.cat([state, inputs], dim=-1)))
+        candidate = self._encode(reset * state + (1 - reset) * inputs, attention_mask)
+        update = torch.sigmoid(self.update_gate(torch.cat([state, candidate], dim=-1)))
+        return (1 - update) * state + update * candidate
+
+    def _valid_label_mask(self, input_ids, attention_mask, num_labels):
+        label_counts = (input_ids.eq(self.config.class_token_index) & attention_mask.bool()).sum(dim=1)
+        return torch.arange(num_labels, device=input_ids.device).unsqueeze(0) < label_counts.unsqueeze(1)
+
+    def _probabilities(self, logits, valid_labels):
+        if self.config.problem_type == "single_label_classification":
+            return torch.softmax(logits.masked_fill(~valid_labels, float("-inf")), dim=-1).nan_to_num(0.0)
+        return torch.sigmoid(logits)
+
+    def recurrent_forward(
+        self,
+        hidden_states,
+        input_ids,
+        attention_mask,
+        num_steps=1,
+        halt_threshold=None,
+        return_all_steps=False,
+        **kwargs,
+    ):
+        """Run step 1 (one scorer_encoder pass) and up to num_steps - 1 gated recurrent steps.
+
+        With halt_threshold, an example stops once no valid label probability moves by
+        halt_threshold or more between consecutive steps; its state and logits are frozen.
 
+        Returns:
+            step_logits: list of (batch_size, num_labels) logits for every step if
+                return_all_steps, otherwise a single-element list with the final logits
+            steps_taken: (batch_size,) number of steps each example ran
+        """
+        if num_steps > 1 and not self.recurrent:
+            raise ValueError("num_steps > 1 requires a scorer built with config.recurrent_steps > 1.")
+
+        bptt_steps = getattr(self.config, "recurrent_bptt_steps", None)
+
+        state = self._encode(hidden_states, attention_mask)
+        logits = self._score(state, input_ids, attention_mask)
+        step_logits = [logits]
+
+        batch_size, num_labels = logits.shape
+        steps_taken = torch.ones(batch_size, dtype=torch.long, device=logits.device)
+        active = torch.ones(batch_size, dtype=torch.bool, device=logits.device)
+        if halt_threshold:
+            valid_labels = self._valid_label_mask(input_ids, attention_mask, num_labels)
+
+        for step in range(1, num_steps):
+            if bptt_steps and step % bptt_steps == 0:
+                # truncated BPTT: later losses stop shaping earlier iterations
+                state = state.detach()
+
+            new_state = self._recurrent_step(state, hidden_states, attention_mask)
+            new_logits = self._score(new_state, input_ids, attention_mask)
+
+            if halt_threshold:
+                change = (
+                    (self._probabilities(new_logits, valid_labels) - self._probabilities(logits, valid_labels))
+                    .abs()
+                    .masked_fill(~valid_labels, 0.0)
+                    .amax(dim=-1)
+                )
+                new_state = torch.where(active[:, None, None], new_state, state)
+                new_logits = torch.where(active[:, None], new_logits, logits)
+                steps_taken = steps_taken + active.long()
+                active = active & change.ge(halt_threshold)
+            else:
+                steps_taken = steps_taken + 1
+
+            state, logits = new_state, new_logits
+            if return_all_steps:
+                step_logits.append(logits)
+            else:
+                step_logits[-1] = logits
+
+            if halt_threshold and not active.any():
+                break
+
+        if self.recurrent and self.training:
+            # keep the gates in the graph for DDP when the sampled depth is 1
+            unused = sum(param.sum() for param in (*self.reset_gate.parameters(), *self.update_gate.parameters()))
+            step_logits[-1] = step_logits[-1] + 0.0 * unused
+
+        return step_logits, steps_taken
+
+    def _score(self, contextualized_hidden_states, input_ids, attention_mask):
         text_repr, label_repr = self._extract_representations(contextualized_hidden_states, input_ids, attention_mask)
 
         text_repr = self.text_projector(text_repr)
diff --git a/gliclass/training.py b/gliclass/training.py
index 4a88e02..36e26c1 100644
--- a/gliclass/training.py
+++ b/gliclass/training.py
@@ -319,6 +319,8 @@ def __init__(self, ewc: EWC | None = None, prev_dataset=None, *args, **kwargs):
         self.ewc = ewc
         self.prev_dataset = prev_dataset
         self._ewc_initialized = ewc is not None
+        # running sums of per-step losses of recurrent scorers, flushed in log()
+        self._recurrent_loss_sums = {}
 
     def _maybe_initialize_ewc(self):
         """Initialize EWC if needed and not already initialized."""
@@ -369,11 +371,8 @@ def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
             Loss tensor, or tuple of (loss, outputs) if return_outputs=True
         """
         # Get base loss from parent
-        if return_outputs:
-            loss, outputs = super().compute_loss(model, inputs, return_outputs=True, **kwargs)
-        else:
-            loss = super().compute_loss(model, inputs, return_outputs=False, **kwargs)
-            outputs = None
+        loss, outputs = super().compute_loss(model, inputs, return_outputs=True, **kwargs)
+        self._track_recurrent_losses(model, outputs)
 
         # Add EWC penalty if enabled
         if self.ewc is not None and self.args.use_ewc:
@@ -385,6 +384,22 @@ def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
             return loss, outputs
         return loss
 
+    def _track_recurrent_losses(self, model, outputs):
+        step_losses = outputs.get("recurrent_losses") if isinstance(outputs, dict) else None
+        if step_losses is None or not model.training:
+            return
+        for step, value in enumerate(step_losses.tolist(), start=1):
+            total, count = self._recurrent_loss_sums.get(step, (0.0, 0))
+            self._recurrent_loss_sums[step] = (total + value, count + 1)
+
+    def log(self, logs, *args, **kwargs):
+        """Add mean per-step losses of recurrent scorers to training logs."""
+        if "loss" in logs and self._recurrent_loss_sums:
+            for step, (total, count) in sorted(self._recurrent_loss_sums.items()):
+                logs[f"loss_step_{step}"] = round(total / count, 4)
+            self._recurrent_loss_sums = {}
+        return super().log(logs, *args, **kwargs)
+
     def train(self, *args, **kwargs):
         """Train with EWC initialization."""
         # Initialize EWC before training starts

From ee5894e2fddeb8218181d04d426440e1e97c19ba Mon Sep 17 00:00:00 2001
From: Ingvar <igor.stepanov2000@gmail.com>
Date: Fri, 2 Oct 2026 12:40:45 +0300
Subject: [PATCH 02/12] make recurrenses more robust

---
 gliclass/config.py   | 15 +++++++-
 gliclass/model.py    | 89 ++++++++++++++++++++++++++++++++++++++-----
 gliclass/scorers.py  | 90 +++++++++++++++++++++++++++++++++++++-------
 gliclass/training.py | 37 ++++++++++++++----
 4 files changed, 198 insertions(+), 33 deletions(-)

diff --git a/gliclass/config.py b/gliclass/config.py
index 524344a..e75e4c5 100644
--- a/gliclass/config.py
+++ b/gliclass/config.py
@@ -66,8 +66,12 @@ def __init__(
         recurrent_inference_max_steps=None,
         recurrent_halt_threshold=0.01,
         recurrent_improvement_coef=0.1,
-        recurrent_improvement_margin=0.0,
+        recurrent_improvement_margin=0.05,
         recurrent_bptt_steps=None,
+        recurrent_read_text=False,
+        recurrent_confidence_coef=0.0,
+        recurrent_confidence_mode="relative",
+        recurrent_confidence_margin=0.0,
         **kwargs,
     ):
         if isinstance(encoder_config, dict):
@@ -167,8 +171,17 @@ def __init__(
         self.recurrent_inference_max_steps = recurrent_inference_max_steps
         self.recurrent_halt_threshold = recurrent_halt_threshold
         # Hinge relu(l_t - (1 - margin) * sg(l_{t-1})) for steps t >= 3, on top of a loss at every step.
+        # margin > 0 makes "no change" cost something, so the recurrence cannot settle into a copy of step 2.
         self.recurrent_improvement_coef = recurrent_improvement_coef
         self.recurrent_improvement_margin = recurrent_improvement_margin
         self.recurrent_bptt_steps = recurrent_bptt_steps
+        # False: recurrent steps see only the label section; True: each step also cross-attends to the text
+        self.recurrent_read_text = recurrent_read_text
+        # Confidence loss on recurrent steps (0 = off): push label probabilities away from 0.5.
+        # "relative": each step's mean label entropy must fall below (1 - margin) * the previous step's;
+        # "absolute": minimize the entropy of every recurrent step.
+        self.recurrent_confidence_coef = recurrent_confidence_coef
+        self.recurrent_confidence_mode = recurrent_confidence_mode
+        self.recurrent_confidence_margin = recurrent_confidence_margin
 
         super().__init__(problem_type=problem_type, **kwargs)
diff --git a/gliclass/model.py b/gliclass/model.py
index cab2146..237c7b6 100644
--- a/gliclass/model.py
+++ b/gliclass/model.py
@@ -64,6 +64,7 @@ class GLiClassOutput(SequenceClassifierOutput):
     past_key_values: Any | None = None
     recurrent_logits: Tuple[torch.Tensor, ...] | None = None
     recurrent_losses: torch.Tensor | None = None
+    recurrent_entropies: torch.Tensor | None = None
     recurrent_num_steps: torch.Tensor | None = None
 
 
@@ -1088,23 +1089,44 @@ def _per_sample_loss(self, logits, labels):
             )
         raise NotImplementedError(f"{self.config.problem_type} is not implemented.")
 
-    def get_recurrent_loss(self, step_logits, labels):
+    def _per_sample_entropy(self, logits, valid_labels):
+        """Mean uncertainty per example over its real labels, shape (batch,).
+
+        Multi-label: binary entropy of each sigmoid probability (0 when p is 0 or 1, max at 0.5).
+        Single-label: entropy of the softmax over the example's labels.
+        """
+        if self.config.problem_type == "single_label_classification":
+            log_probs = torch.log_softmax(logits.masked_fill(~valid_labels, float("-inf")), dim=-1)
+            return -(log_probs.exp() * log_probs.masked_fill(~valid_labels, 0.0)).sum(-1)
+        probs = torch.sigmoid(logits)
+        entropy = nn.functional.binary_cross_entropy_with_logits(logits, probs, reduction="none")
+        return (entropy * valid_labels).sum(-1) / valid_labels.sum(-1).clamp_min(1)
+
+    def get_recurrent_loss(self, step_logits, labels, valid_labels=None):
         """Loss over all reasoning steps.
 
         loss = mean_t L_t                                                 (focal loss after every step)
              + coef * mean_{t>=3} relu(l_t - (1 - margin) * sg(l_{t-1}))  (deeper must beat shallower)
+             + conf_coef * confidence term on recurrent steps t >= 2      (deeper must be more certain)
 
         The improvement term starts at step 3 vs step 2, so step 1 is only ever trained by its own
         loss. The stop-gradient on l_{t-1} means the hinge can only push the deeper step down; it
         never rewards making the earlier step worse.
 
+        Confidence term (recurrent_confidence_coef > 0), with H_t the mean label entropy of step t:
+            "relative": relu(H_t - (1 - conf_margin) * sg(H_{t-1}))  each recurrence less uncertain than the last
+            "absolute": H_t                                          every recurrence as certain as possible
+
         Returns:
             loss: scalar
             step_losses: (num_steps,) detached mean loss per step, for monitoring
+            step_entropies: (num_steps,) detached mean label entropy per step, for monitoring
         """
         num_labels = step_logits[-1].shape[-1]
         if labels.shape[-1] != num_labels:
             labels = labels[:, :num_labels]
+        if valid_labels is None:
+            valid_labels = torch.ones_like(step_logits[-1], dtype=torch.bool)
 
         per_sample = torch.stack([self._per_sample_loss(logits, labels) for logits in step_logits])  # (T, B)
         step_losses = per_sample.mean(dim=1)
@@ -1115,14 +1137,28 @@ def get_recurrent_loss(self, step_logits, labels):
             improvement = torch.relu(per_sample[2:] - target).mean()
             loss = loss + self.config.recurrent_improvement_coef * improvement
 
-        return loss, step_losses.detach()
+        entropies = torch.stack([self._per_sample_entropy(logits, valid_labels) for logits in step_logits])  # (T, B)
+        confidence_coef = getattr(self.config, "recurrent_confidence_coef", 0.0)
+        if len(step_logits) >= 2 and confidence_coef > 0:
+            mode = getattr(self.config, "recurrent_confidence_mode", "relative")
+            if mode == "relative":
+                margin = getattr(self.config, "recurrent_confidence_margin", 0.0)
+                confidence = torch.relu(entropies[1:] - (1.0 - margin) * entropies[:-1].detach()).mean()
+            elif mode == "absolute":
+                confidence = entropies[1:].mean()
+            else:
+                raise ValueError(f"Unknown recurrent_confidence_mode: {mode}")
+            loss = loss + confidence_coef * confidence
+
+        return loss, step_losses.detach(), entropies.mean(dim=1).detach()
 
     def _extract_label_section(
         self,
         hidden_states: torch.Tensor,
         input_ids: torch.Tensor,
         attention_mask: torch.Tensor,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        return_text_section: bool = False,
+    ) -> tuple[torch.Tensor, ...]:
         """
         Extract label section hidden states from full-sequence decoder output.
 
@@ -1136,12 +1172,15 @@ def _extract_label_section(
             padded_hidden: (batch, max_label_len, hidden_size)
             padded_ids:    (batch, max_label_len)
             label_mask:    (batch, max_label_len)
+            with return_text_section, also everything before the label section:
+            text_hidden:   (batch, max_text_len, hidden_size)
+            text_mask:     (batch, max_text_len)
         """
         batch_size = hidden_states.shape[0]
         sep_id = self.sep_token_id
         label_id = self.config.class_token_index
 
-        slices_h, slices_ids = [], []
+        slices_h, slices_ids, text_slices = [], [], []
 
         for i in range(batch_size):
             real_len = int(attention_mask[i].sum().item())
@@ -1154,6 +1193,7 @@ def _extract_label_section(
                 # Fallback: use everything up to real_len
                 slices_h.append(hidden_states[i, :real_len])
                 slices_ids.append(ids_i)
+                text_slices.append(hidden_states[i, :real_len])
                 continue
 
             first_label = label_pos[0].item()
@@ -1163,12 +1203,15 @@ def _extract_label_section(
             if seps_before.numel() == 0:
                 slices_h.append(hidden_states[i, :real_len])
                 slices_ids.append(ids_i)
+                text_slices.append(hidden_states[i, :real_len])
                 continue
 
             # start right after that SEP → first token of "label1<<LABEL>>...<<SEP>>"
             start = int(seps_before[-1].item()) + 1
             slices_h.append(hidden_states[i, start:real_len])
             slices_ids.append(ids_i[start:real_len])
+            # text section ends with that SEP, so it is never empty
+            text_slices.append(hidden_states[i, :start])
 
         padded_hidden = pad_sequence(slices_h, batch_first=True)
         padded_ids = pad_sequence(slices_ids, batch_first=True)
@@ -1180,7 +1223,14 @@ def _extract_label_section(
         positions = torch.arange(padded_hidden.shape[1], device=attention_mask.device)
         label_mask = (positions.unsqueeze(0) < section_lengths.unsqueeze(1)).to(attention_mask.dtype)
 
-        return padded_hidden, padded_ids, label_mask
+        if not return_text_section:
+            return padded_hidden, padded_ids, label_mask
+
+        text_hidden = pad_sequence(text_slices, batch_first=True)
+        text_lengths = torch.tensor([section.shape[0] for section in text_slices], device=attention_mask.device)
+        text_positions = torch.arange(text_hidden.shape[1], device=attention_mask.device)
+        text_mask = (text_positions.unsqueeze(0) < text_lengths.unsqueeze(1)).to(attention_mask.dtype)
+        return padded_hidden, padded_ids, label_mask, text_hidden, text_mask
 
     def update_decoder_cache(
         self,
@@ -1209,8 +1259,14 @@ def classify_from_decoder_cache(
         use_recurrence: bool | None = None,
         max_recurrent_steps: int | None = None,
         halt_threshold: float | None = None,
+        text_hidden_states: torch.Tensor | None = None,
+        text_attention_mask: torch.Tensor | None = None,
     ) -> GLiClassOutput:
-        """Classify label tokens against a text cache without persisting labels."""
+        """Classify label tokens against a text cache without persisting labels.
+
+        With config.recurrent_read_text, pass the decoder's last hidden states of the cached text
+        (text_hidden_states / text_attention_mask); the KV cache alone does not contain them.
+        """
         decoder_outputs = self.decoder_model(
             input_ids=input_ids,
             attention_mask=attention_mask,
@@ -1229,6 +1285,8 @@ def classify_from_decoder_cache(
             attention_mask=label_mask.to(scorer_device),
             num_steps=num_steps,
             halt_threshold=halt_threshold,
+            text_hidden_states=None if text_hidden_states is None else text_hidden_states.to(scorer_device),
+            text_attention_mask=None if text_attention_mask is None else text_attention_mask.to(scorer_device),
         )
         return GLiClassOutput(logits=step_logits[-1], recurrent_num_steps=steps_taken)
 
@@ -1276,9 +1334,16 @@ def forward(
 
         hidden_states = decoder_outputs.last_hidden_state
 
-        label_hidden, label_ids, label_mask = self._extract_label_section(hidden_states, input_ids, attention_mask)
-
         num_steps, halt_threshold = self._resolve_recurrence(use_recurrence, max_recurrent_steps, halt_threshold)
+
+        text_hidden = text_mask = None
+        if self.scorer.read_text and num_steps > 1:
+            label_hidden, label_ids, label_mask, text_hidden, text_mask = self._extract_label_section(
+                hidden_states, input_ids, attention_mask, return_text_section=True
+            )
+        else:
+            label_hidden, label_ids, label_mask = self._extract_label_section(hidden_states, input_ids, attention_mask)
+
         step_logits, steps_taken = self.scorer.recurrent_forward(
             hidden_states=label_hidden,
             input_ids=label_ids,
@@ -1286,15 +1351,18 @@ def forward(
             num_steps=num_steps,
             halt_threshold=halt_threshold,
             return_all_steps=labels is not None or output_recurrent_logits,
+            text_hidden_states=text_hidden,
+            text_attention_mask=text_mask,
         )
         logits = step_logits[-1]
 
         if labels is not None:
             self.num_labels = logits.shape[-1]
 
-        recurrent_losses = None
+        recurrent_losses = recurrent_entropies = None
         if len(step_logits) > 1 and labels is not None:
-            loss, recurrent_losses = self.get_recurrent_loss(step_logits, labels)
+            valid_labels = self.scorer._valid_label_mask(label_ids, label_mask, logits.shape[-1])
+            loss, recurrent_losses, recurrent_entropies = self.get_recurrent_loss(step_logits, labels, valid_labels)
         else:
             loss = self.get_loss(logits, labels)
 
@@ -1312,6 +1380,7 @@ def forward(
             past_key_values=decoder_outputs.past_key_values if use_cache else None,
             recurrent_logits=tuple(step_logits) if output_recurrent_logits else None,
             recurrent_losses=recurrent_losses,
+            recurrent_entropies=recurrent_entropies,
             recurrent_num_steps=steps_taken if num_steps > 1 else None,
         )
 
diff --git a/gliclass/scorers.py b/gliclass/scorers.py
index 82afb8c..230d6c2 100644
--- a/gliclass/scorers.py
+++ b/gliclass/scorers.py
@@ -169,6 +169,63 @@ def forward(self, text_rep, label_rep, text_mask=None, **kwargs):
         return self.score_mlp(torch.cat([context, label_rep], dim=-1)).squeeze(-1)
 
 
+class RecurrentReasoningCell(nn.Module):
+    """GRU-style reasoning step around the scorer's (shared) DeBERTa encoder.
+
+        x    = LN(decoder label-section states)               scale-matched to the encoder state
+        read = CrossAttn(q=h, kv=LN(text states))             only with read_text
+        r    = sigmoid(W_r [h; x])
+        h~   = Enc(r * h + (1 - r) * x + read)
+        z    = sigmoid(W_z [h; h~])
+        h'   = (1 - z) * h + z * h~
+
+    No parameter depends on the step index, so any number of steps can be run.
+    """
+
+    def __init__(self, hidden_size, read_text=False):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.reset_gate = nn.Linear(hidden_size * 2, hidden_size)
+        self.update_gate = nn.Linear(hidden_size * 2, hidden_size)
+
+        self.read_text = read_text
+        if read_text:
+            self.num_heads = max(1, hidden_size // 64)
+            self.head_dim = hidden_size // self.num_heads
+            self.read_q = nn.Linear(hidden_size, hidden_size)
+            self.read_k = nn.Linear(hidden_size, hidden_size)
+            self.read_v = nn.Linear(hidden_size, hidden_size)
+            self.read_out = nn.Linear(hidden_size, hidden_size)
+
+    def normalize(self, hidden_states):
+        return nn.functional.layer_norm(hidden_states, (self.hidden_size,))
+
+    def prepare_text(self, text_states):
+        """Project normalized text states to keys/values once per forward pass."""
+        batch_size, text_length, _ = text_states.shape
+        text_states = self.normalize(text_states)
+        keys = self.read_k(text_states).view(batch_size, text_length, self.num_heads, self.head_dim)
+        values = self.read_v(text_states).view(batch_size, text_length, self.num_heads, self.head_dim)
+        return keys, values
+
+    def read(self, state, text_memory, text_mask):
+        batch_size, length, _ = state.shape
+        keys, values = text_memory
+        queries = self.read_q(self.normalize(state)).view(batch_size, length, self.num_heads, self.head_dim)
+        context = attn_padded(queries, keys, values, key_padding_mask=text_mask)
+        return self.read_out(context.reshape(batch_size, length, self.hidden_size))
+
+    def forward(self, state, inputs, attention_mask, encode, text_memory=None, text_mask=None):
+        """One step; `inputs` must already be normalized with self.normalize."""
+        reset = torch.sigmoid(self.reset_gate(torch.cat([state, inputs], dim=-1)))
+        candidate_input = reset * state + (1 - reset) * inputs
+        if self.read_text:
+            candidate_input = candidate_input + self.read(state, text_memory, text_mask)
+        candidate = encode(candidate_input, attention_mask)
+        update = torch.sigmoid(self.update_gate(torch.cat([state, candidate], dim=-1)))
+        return (1 - update) * state + update * candidate
+
+
 class DecoderKVScorer(nn.Module):
     """
     Scorer for decoder-kv architecture with built-in bidirectional encoder and representation extraction.
@@ -221,12 +278,11 @@ def __init__(self, config, **kwargs):
 
         self.epsilon = 1e-8
 
-        # GRU-style recurrent reasoning over scorer_encoder. Gates are step-agnostic (no step
-        # embeddings), so inference can run more steps than were seen during training.
+        # Recurrent reasoning over scorer_encoder (see RecurrentReasoningCell). Step 1 is the plain scorer.
         self.recurrent = getattr(config, "recurrent_steps", 1) > 1
+        self.read_text = self.recurrent and getattr(config, "recurrent_read_text", False)
         if self.recurrent:
-            self.reset_gate = nn.Linear(config.hidden_size * 2, config.hidden_size)
-            self.update_gate = nn.Linear(config.hidden_size * 2, config.hidden_size)
+            self.reasoning_cell = RecurrentReasoningCell(config.hidden_size, read_text=self.read_text)
 
     def _extract_representations(self, hidden_states, input_ids, attention_mask):
         """
@@ -285,13 +341,6 @@ def forward(self, hidden_states, input_ids, attention_mask, **kwargs):
     def _encode(self, hidden_states, attention_mask):
         return self.scorer_encoder(hidden_states, attention_mask=attention_mask, return_dict=True).last_hidden_state
 
-    def _recurrent_step(self, state, inputs, attention_mask):
-        """h_t = (1 - z) * h_{t-1} + z * Enc(r * h_{t-1} + (1 - r) * x), x = decoder hidden states."""
-        reset = torch.sigmoid(self.reset_gate(torch.cat([state, inputs], dim=-1)))
-        candidate = self._encode(reset * state + (1 - reset) * inputs, attention_mask)
-        update = torch.sigmoid(self.update_gate(torch.cat([state, candidate], dim=-1)))
-        return (1 - update) * state + update * candidate
-
     def _valid_label_mask(self, input_ids, attention_mask, num_labels):
         label_counts = (input_ids.eq(self.config.class_token_index) & attention_mask.bool()).sum(dim=1)
         return torch.arange(num_labels, device=input_ids.device).unsqueeze(0) < label_counts.unsqueeze(1)
@@ -309,12 +358,16 @@ def recurrent_forward(
         num_steps=1,
         halt_threshold=None,
         return_all_steps=False,
+        text_hidden_states=None,
+        text_attention_mask=None,
         **kwargs,
     ):
         """Run step 1 (one scorer_encoder pass) and up to num_steps - 1 gated recurrent steps.
 
         With halt_threshold, an example stops once no valid label probability moves by
         halt_threshold or more between consecutive steps; its state and logits are frozen.
+        With config.recurrent_read_text, every recurrent step also attends to text_hidden_states
+        (decoder states of the text before the label section).
 
         Returns:
             step_logits: list of (batch_size, num_labels) logits for every step if
@@ -336,12 +389,21 @@ def recurrent_forward(
         if halt_threshold:
             valid_labels = self._valid_label_mask(input_ids, attention_mask, num_labels)
 
+        if num_steps > 1:
+            inputs = self.reasoning_cell.normalize(hidden_states)
+            text_memory = text_mask = None
+            if self.read_text:
+                if text_hidden_states is None or text_attention_mask is None:
+                    raise ValueError("config.recurrent_read_text requires text_hidden_states and text_attention_mask.")
+                text_memory = self.reasoning_cell.prepare_text(text_hidden_states)
+                text_mask = text_attention_mask.bool()
+
         for step in range(1, num_steps):
             if bptt_steps and step % bptt_steps == 0:
                 # truncated BPTT: later losses stop shaping earlier iterations
                 state = state.detach()
 
-            new_state = self._recurrent_step(state, hidden_states, attention_mask)
+            new_state = self.reasoning_cell(state, inputs, attention_mask, self._encode, text_memory, text_mask)
             new_logits = self._score(new_state, input_ids, attention_mask)
 
             if halt_threshold:
@@ -368,8 +430,8 @@ def recurrent_forward(
                 break
 
         if self.recurrent and self.training:
-            # keep the gates in the graph for DDP when the sampled depth is 1
-            unused = sum(param.sum() for param in (*self.reset_gate.parameters(), *self.update_gate.parameters()))
+            # keep the reasoning cell in the graph for DDP when the sampled depth is 1
+            unused = sum(param.sum() for param in self.reasoning_cell.parameters())
             step_logits[-1] = step_logits[-1] + 0.0 * unused
 
         return step_logits, steps_taken
diff --git a/gliclass/training.py b/gliclass/training.py
index 36e26c1..746181e 100644
--- a/gliclass/training.py
+++ b/gliclass/training.py
@@ -275,6 +275,9 @@ class TrainingArguments(transformers.TrainingArguments):
     optim: str = field(default="adamw_torch")
     others_lr: float | None = None
     others_weight_decay: float | None = 0.0
+    recurrent_lr: float | None = field(
+        default=None, metadata={"help": "Learning rate for the decoder-kv scorer's recurrent reasoning cell."}
+    )
 
     use_ewc: bool = field(
         default=False, metadata={"help": "Whether to use Elastic Weight Consolidation (EWC) for continual learning."}
@@ -385,18 +388,21 @@ def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
         return loss
 
     def _track_recurrent_losses(self, model, outputs):
-        step_losses = outputs.get("recurrent_losses") if isinstance(outputs, dict) else None
-        if step_losses is None or not model.training:
+        if not isinstance(outputs, dict) or not model.training:
             return
-        for step, value in enumerate(step_losses.tolist(), start=1):
-            total, count = self._recurrent_loss_sums.get(step, (0.0, 0))
-            self._recurrent_loss_sums[step] = (total + value, count + 1)
+        for key, prefix in (("recurrent_losses", "loss_step"), ("recurrent_entropies", "entropy_step")):
+            values = outputs.get(key)
+            if values is None:
+                continue
+            for step, value in enumerate(values.tolist(), start=1):
+                total, count = self._recurrent_loss_sums.get((prefix, step), (0.0, 0))
+                self._recurrent_loss_sums[(prefix, step)] = (total + value, count + 1)
 
     def log(self, logs, *args, **kwargs):
-        """Add mean per-step losses of recurrent scorers to training logs."""
+        """Add mean per-step losses and label entropies of recurrent scorers to training logs."""
         if "loss" in logs and self._recurrent_loss_sums:
-            for step, (total, count) in sorted(self._recurrent_loss_sums.items()):
-                logs[f"loss_step_{step}"] = round(total / count, 4)
+            for (prefix, step), (total, count) in sorted(self._recurrent_loss_sums.items()):
+                logs[f"{prefix}_{step}"] = round(total / count, 4)
             self._recurrent_loss_sums = {}
         return super().log(logs, *args, **kwargs)
 
@@ -604,12 +610,27 @@ def create_optimizer(self):
                     },
                 ]
 
+            if getattr(self.args, "recurrent_lr", None) is not None:
+                optimizer_grouped_parameters = self._split_recurrent_groups(opt_model, optimizer_grouped_parameters)
+
             optimizer_cls, optimizer_kwargs = Trainer.get_optimizer_cls_and_kwargs(self.args)
 
             self.optimizer = optimizer_cls(optimizer_grouped_parameters, **optimizer_kwargs)
 
         return self.optimizer
 
+    def _split_recurrent_groups(self, model, groups):
+        """Move reasoning-cell parameters into their own groups with args.recurrent_lr."""
+        recurrent_ids = {id(p) for name, p in model.named_parameters() if "reasoning_cell" in name}
+        split = []
+        for group in groups:
+            kept = [p for p in group["params"] if id(p) not in recurrent_ids]
+            moved = [p for p in group["params"] if id(p) in recurrent_ids]
+            split.append({**group, "params": kept})
+            if moved:
+                split.append({**group, "params": moved, "lr": self.args.recurrent_lr})
+        return split
+
 
 @dataclass
 class RLTrainerConfig(TrainingArguments):

From 495852f2998e51369ef7ff38e2988410d415de4b Mon Sep 17 00:00:00 2001
From: Ingvar <igor.stepanov2000@gmail.com>
Date: Fri, 2 Oct 2026 13:32:56 +0300
Subject: [PATCH 03/12] implement calibrator

---
 calibrate.py                   | 167 ++++++++++++++
 gliclass/calibration.py        | 411 +++++++++++++++++++++++++++++++++
 gliclass/config.py             |  14 ++
 gliclass/model.py              | 101 +++++++-
 gliclass/pipeline.py           |  35 ++-
 gliclass/scorers.py            |  18 +-
 gliclass/streaming/pipeline.py |   2 +
 7 files changed, 725 insertions(+), 23 deletions(-)
 create mode 100644 calibrate.py
 create mode 100644 gliclass/calibration.py

diff --git a/calibrate.py b/calibrate.py
new file mode 100644
index 0000000..c1135b3
--- /dev/null
+++ b/calibrate.py
@@ -0,0 +1,167 @@
+"""Fit a post-hoc calibrator for a trained GLiClass model with its backbone frozen.
+
+The frozen model is run once over the calibration data; the calibrator MLP is then trained on the cached
+(text repr, label repr, logits, gold) tuples by minimizing the binary NLL of sigmoid(beta * logit + bias) plus
+a prior towards a global temperature / Platt scaling fit (see gliclass/calibration.py). The saved model returns
+calibrated logits and keeps the raw ones in `uncalibrated_logits`; the pipeline reports both scores.
+
+Example:
+    python calibrate.py --model_path models/checkpoint-1000 --data_path data/val.json
+        --save_path models/checkpoint-1000-calibrated
+"""
+
+import os
+
+os.environ["TOKENIZERS_PARALLELISM"] = "true"
+import json
+import random
+import argparse
+
+import torch
+from transformers import AutoTokenizer
+from torch.utils.data import DataLoader
+
+from gliclass import GLiClassModel
+from gliclass.calibration import (
+    fit_calibrator,
+    fit_global_scaling,
+    evaluate_calibrator,
+    collect_calibration_examples,
+)
+from gliclass.data_processing import GLiClassDataset, AugmentationConfig, DataCollatorWithPadding
+
+DTYPES = {"float32": torch.float32, "bfloat16": torch.bfloat16, "float16": torch.float16}
+
+
+def build_dataset(data, model, tokenizer, args):
+    config = model.config
+    labels_tokenizer = None
+    if config.label_model_name is not None:
+        labels_tokenizer = AutoTokenizer.from_pretrained(config.label_model_name)
+    return GLiClassDataset(
+        data,
+        tokenizer,
+        AugmentationConfig(enabled=False),
+        max_length=args.max_length,
+        problem_type=config.problem_type,
+        architecture_type=config.architecture_type,
+        prompt_first=config.prompt_first,
+        labels_tokenizer=labels_tokenizer,
+        shuffle_labels=False,
+    )
+
+
+def main(args):
+    random.seed(args.seed)
+    torch.manual_seed(args.seed)
+    device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
+
+    model = GLiClassModel.from_pretrained(args.model_path, torch_dtype=DTYPES[args.dtype]).to(device)
+    tokenizer = AutoTokenizer.from_pretrained(args.model_path)
+    model.eval()
+    model.requires_grad_(False)
+
+    with open(args.data_path) as f:
+        data = json.load(f)
+    random.shuffle(data)
+    if args.max_examples is not None:
+        data = data[: args.max_examples]
+
+    # split by text so labels of one text never land on both sides
+    num_eval = max(1, int(len(data) * args.eval_ratio))
+    train_data, eval_data = data[num_eval:], data[:num_eval]
+    print(f"Calibration texts: {len(train_data)} train, {len(eval_data)} held-out")
+
+    collator = DataCollatorWithPadding(device="cpu", config=model.config)
+
+    def collect(split):
+        loader = DataLoader(
+            build_dataset(split, model, tokenizer, args),
+            batch_size=args.batch_size,
+            shuffle=False,
+            collate_fn=collator,
+            num_workers=args.num_workers,
+        )
+        return collect_calibration_examples(model, loader, device)
+
+    train_examples = collect(train_data)
+    eval_examples = collect(eval_data)
+
+    base_beta, base_bias = fit_global_scaling(train_examples, fit_bias=args.use_bias)
+    method = "Platt" if args.use_bias else "temperature"
+    print(f"Global {method} scaling: beta_0 = {base_beta:.4f}, b_0 = {base_bias:.4f}")
+
+    calibrator = model.add_calibrator(
+        hidden_size=args.hidden_size, beta_max=args.beta_max, dropout=args.dropout, use_bias=args.use_bias
+    )
+    fit_calibrator(
+        calibrator,
+        train_examples,
+        eval_examples,
+        device,
+        base_beta=base_beta,
+        base_bias=base_bias,
+        prior_weight=args.prior_weight,
+        bias_prior_weight=args.bias_prior_weight,
+        num_epochs=args.num_epochs,
+        batch_size=args.calibrator_batch_size,
+        lr=args.lr,
+        weight_decay=args.weight_decay,
+        patience=args.patience,
+        seed=args.seed,
+    )
+
+    report = {
+        "base_beta": base_beta,
+        "base_bias": base_bias,
+        "train": evaluate_calibrator(calibrator, train_examples, device, base_beta, base_bias),
+        "held_out": evaluate_calibrator(calibrator, eval_examples, device, base_beta, base_bias),
+        "args": vars(args),
+    }
+    print(json.dumps({key: report[key] for key in ("base_beta", "base_bias", "held_out")}, indent=2))
+
+    os.makedirs(args.save_path, exist_ok=True)
+    model.save_pretrained(args.save_path)
+    tokenizer.save_pretrained(args.save_path)
+    with open(os.path.join(args.save_path, "calibration_report.json"), "w") as f:
+        json.dump(report, f, indent=2)
+    print(f"Saved calibrated model to {args.save_path}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Fit a GLiClass calibrator with a frozen backbone")
+    parser.add_argument("--model_path", type=str, required=True)
+    parser.add_argument("--data_path", type=str, required=True, help="JSON in the training data format")
+    parser.add_argument("--save_path", type=str, required=True)
+    parser.add_argument("--eval_ratio", type=float, default=0.2, help="Share of texts held out for early stopping")
+    parser.add_argument("--max_examples", type=int, default=None)
+    parser.add_argument("--max_length", type=int, default=1024)
+    parser.add_argument("--batch_size", type=int, default=8, help="Backbone batch size for feature collection")
+    parser.add_argument("--num_workers", type=int, default=4)
+    parser.add_argument("--dtype", type=str, default="float32", choices=list(DTYPES))
+
+    parser.add_argument("--hidden_size", type=int, default=256)
+    parser.add_argument("--beta_max", type=float, default=3.0, help="Upper bound of the inverse temperature")
+    parser.add_argument("--dropout", type=float, default=0.1)
+    parser.add_argument(
+        "--prior_weight",
+        type=float,
+        default=0.01,
+        help="Weight of (log beta - log beta_0)^2; large -> temperature scaling",
+    )
+    parser.add_argument(
+        "--use_bias",
+        action=argparse.BooleanOptionalAction,
+        default=True,
+        help="Also predict a per-pair logit bias (fixes base-rate miscalibration; may flip decisions at 0.5)",
+    )
+    parser.add_argument(
+        "--bias_prior_weight", type=float, default=0.01, help="Weight of (bias - b_0)^2, b_0 from global Platt scaling"
+    )
+    parser.add_argument("--num_epochs", type=int, default=20)
+    parser.add_argument("--calibrator_batch_size", type=int, default=64)
+    parser.add_argument("--lr", type=float, default=1e-3)
+    parser.add_argument("--weight_decay", type=float, default=0.01)
+    parser.add_argument("--patience", type=int, default=3)
+    parser.add_argument("--seed", type=int, default=42)
+    main(parser.parse_args())
diff --git a/gliclass/calibration.py b/gliclass/calibration.py
new file mode 100644
index 0000000..34b3054
--- /dev/null
+++ b/gliclass/calibration.py
@@ -0,0 +1,411 @@
+"""Post-hoc calibration of GLiClass scores with a frozen backbone.
+
+The calibrator predicts an inverse temperature beta in (0, beta_max) for every (text, label) pair and,
+optionally, a bias b; the calibrated probability is sigmoid(beta * logit + b). It is fitted by minimizing the
+binary NLL on held-out data: for a pair the backbone gets wrong the NLL decreases as beta -> 0 (probability
+-> sigmoid(b)), for a correct pair it decreases as beta grows, so the MLP learns where the backbone can be
+trusted. A prior pulls (log beta, b) towards a global fit (log beta_0, b_0), which keeps the per-pair optimum
+from collapsing to {0, beta_max}.
+
+Without the bias, beta > 0 never flips the sign of a logit: decisions at 0.5 are unchanged and a distrusted
+pair goes to 0.5. Scaling alone can only shrink towards 0.5, so it cannot fix a backbone that systematically
+over- or under-predicts positives; the bias shifts the base rate (global Platt scaling sigmoid(a * z + c) is
+then the baseline), at the cost of possibly flipping decisions and pulling distrusted pairs to sigmoid(b)
+instead of 0.5.
+"""
+
+import math
+import random
+from dataclasses import dataclass
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from torch.nn.utils.rnn import pad_sequence
+
+NUM_LOGIT_FEATURES = 6
+
+
+def calibrator_representation_size(config):
+    """Size of the text/label representations the architecture exposes to the calibrator."""
+    if config.architecture_type == "decoder-kv":
+        return config.hidden_size
+    return config.encoder_config.hidden_size
+
+
+def logit_features(logits, label_mask):
+    """Per-pair confidence features: (batch, num_labels, NUM_LOGIT_FEATURES).
+
+    [z, |z|, z - max_valid(z), z - mean_valid(z), mean binary entropy of the text's labels, log(#labels)]
+    """
+    mask = label_mask.bool()
+    maskf = mask.to(logits.dtype)
+    num_labels = maskf.sum(dim=-1, keepdim=True).clamp_min(1.0)
+
+    max_logit = logits.masked_fill(~mask, float("-inf")).amax(dim=-1, keepdim=True)
+    max_logit = torch.where(mask.any(dim=-1, keepdim=True), max_logit, torch.zeros_like(max_logit))
+    mean_logit = (logits * maskf).sum(dim=-1, keepdim=True) / num_labels
+
+    probs = torch.sigmoid(logits)
+    entropy = -(probs * F.logsigmoid(logits) + (1 - probs) * F.logsigmoid(-logits))
+    mean_entropy = (entropy * maskf).sum(dim=-1, keepdim=True) / num_labels
+
+    features = [
+        logits,
+        logits.abs(),
+        logits - max_logit,
+        logits - mean_logit,
+        mean_entropy.expand_as(logits),
+        num_labels.log().expand_as(logits),
+    ]
+    return torch.stack(features, dim=-1) * maskf.unsqueeze(-1)
+
+
+class GLiClassCalibrator(nn.Module):
+    """MLP over (text repr, label repr, logit statistics) -> inverse temperature beta in (0, beta_max)
+    and, with use_bias, an additive logit bias b. Calibrated logit: beta * z + b.
+    """
+
+    def __init__(self, representation_size, hidden_size=256, beta_max=3.0, dropout=0.1, use_bias=False):
+        super().__init__()
+        self.beta_max = beta_max
+        self.use_bias = use_bias
+        self.text_norm = nn.LayerNorm(representation_size)
+        self.label_norm = nn.LayerNorm(representation_size)
+        self.mlp = nn.Sequential(
+            nn.Linear(4 * representation_size + NUM_LOGIT_FEATURES, hidden_size),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_size, hidden_size),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_size, 2 if use_bias else 1),
+        )
+        self.reset_to_global_scaling(1.0)
+
+    def reset_to_global_scaling(self, base_beta, base_bias=0.0):
+        """Zero the output layer so (beta, b) == (base_beta, base_bias) everywhere: training starts from the
+        global temperature / Platt scaling fit.
+        """
+        ratio = min(max(base_beta / self.beta_max, 1e-4), 1 - 1e-4)
+        output = self.mlp[-1]
+        with torch.no_grad():
+            output.weight.zero_()
+            output.bias[0] = math.log(ratio / (1 - ratio))
+            if self.use_bias:
+                output.bias[1] = base_bias
+
+    def forward(self, text_repr, label_repr, logits, label_mask):
+        """Predict the inverse temperature of every (text, label) pair.
+
+        Args:
+            text_repr: (batch, dim)
+            label_repr: (batch, num_labels, dim)
+            logits: (batch, num_labels) uncalibrated logits
+            label_mask: (batch, num_labels), 1 for real labels
+
+        Returns:
+            beta: (batch, num_labels) inverse temperatures (1.0 on padded labels)
+            bias: (batch, num_labels) logit biases (0.0 on padded labels and without use_bias)
+        """
+        dtype = self.mlp[0].weight.dtype
+        num_labels = logits.shape[-1]
+        label_repr = label_repr[:, :num_labels]
+        label_mask = label_mask[:, :num_labels].bool()
+
+        text = self.text_norm(text_repr.to(dtype)).unsqueeze(1).expand_as(label_repr)
+        label = self.label_norm(label_repr.to(dtype))
+        stats = logit_features(logits.float(), label_mask).to(dtype)
+        features = torch.cat([text, label, text * label, (text - label).abs(), stats], dim=-1)
+
+        output = self.mlp(features).float()
+        beta = self.beta_max * torch.sigmoid(output[..., 0])
+        bias = output[..., 1] if self.use_bias else torch.zeros_like(beta)
+        return beta.masked_fill(~label_mask, 1.0), bias.masked_fill(~label_mask, 0.0)
+
+
+def build_calibrator(config):
+    return GLiClassCalibrator(
+        calibrator_representation_size(config),
+        hidden_size=config.calibrator_hidden_size,
+        beta_max=config.calibrator_beta_max,
+        dropout=config.calibrator_dropout,
+        use_bias=getattr(config, "calibrator_use_bias", False),
+    )
+
+
+# ---------------------------------------------------------------------------------------------------------
+# Fitting on cached backbone outputs
+# ---------------------------------------------------------------------------------------------------------
+
+
+@dataclass
+class CalibrationExample:
+    text_repr: torch.Tensor  # (dim,)
+    label_repr: torch.Tensor  # (num_labels, dim)
+    logits: torch.Tensor  # (num_labels,)
+    targets: torch.Tensor  # (num_labels,) in [0, 1]
+
+
+def _targets_from_labels(labels, num_labels, problem_type):
+    if problem_type == "single_label_classification" or labels.dim() == 1:
+        return F.one_hot(labels.long().view(-1), num_classes=num_labels).float()
+    return labels[:, :num_labels].float()
+
+
+@torch.no_grad()
+def collect_calibration_examples(model, dataloader, device, show_progress=True):
+    """Run the frozen model once over a dataloader and cache per-text calibration inputs on CPU.
+
+    Uses the uncalibrated logits, so it also works for re-fitting a model that already has a calibrator.
+    """
+    from tqdm import tqdm
+
+    model.eval()
+    examples = []
+    iterator = tqdm(dataloader, desc="Collecting calibration features") if show_progress else dataloader
+    for batch in iterator:
+        labels = batch["labels"]
+        inputs = {
+            key: value.to(device) if isinstance(value, torch.Tensor) else value
+            for key, value in batch.items()
+            if key not in {"labels", "labels_text", "input_texts"}
+        }
+        outputs = model.model(
+            **inputs,
+            output_text_embeddings=True,
+            output_class_embeddings=True,
+            return_dict=True,
+        )
+        logits = outputs.logits.float()
+        num_labels = min(logits.shape[-1], labels.shape[-1]) if labels.dim() > 1 else logits.shape[-1]
+        targets = _targets_from_labels(labels, logits.shape[-1], model.config.problem_type)[:, :num_labels]
+        label_mask = outputs.class_mask[:, :num_labels].bool().cpu()
+        if labels.dim() > 1 and "labels_text" in batch:
+            counts = torch.tensor([len(item) for item in batch["labels_text"]])
+            label_mask &= torch.arange(num_labels).unsqueeze(0) < counts.unsqueeze(1)
+
+        text_repr = outputs.text_embeddings.float().cpu()
+        label_repr = outputs.class_embeddings.float().cpu()
+        logits = logits.cpu()
+        for i in range(logits.shape[0]):
+            valid = label_mask[i]
+            if not valid.any():
+                continue
+            examples.append(
+                CalibrationExample(
+                    text_repr=text_repr[i],
+                    label_repr=label_repr[i, :num_labels][valid],
+                    logits=logits[i, :num_labels][valid],
+                    targets=targets[i][valid],
+                )
+            )
+    return examples
+
+
+def collate_calibration_examples(examples):
+    label_repr = pad_sequence([ex.label_repr for ex in examples], batch_first=True)
+    logits = pad_sequence([ex.logits for ex in examples], batch_first=True)
+    targets = pad_sequence([ex.targets for ex in examples], batch_first=True)
+    label_mask = pad_sequence([torch.ones_like(ex.logits, dtype=torch.bool) for ex in examples], batch_first=True)
+    return {
+        "text_repr": torch.stack([ex.text_repr for ex in examples]),
+        "label_repr": label_repr,
+        "logits": logits,
+        "targets": targets,
+        "label_mask": label_mask,
+    }
+
+
+def _masked_text_mean(values, label_mask):
+    """Average over labels inside each text, then over texts: every text weighs the same."""
+    maskf = label_mask.float()
+    per_text = (values * maskf).sum(dim=-1) / maskf.sum(dim=-1).clamp_min(1.0)
+    return per_text.mean()
+
+
+def calibration_loss(
+    beta,
+    bias,
+    logits,
+    targets,
+    label_mask,
+    base_beta,
+    base_bias=0.0,
+    prior_weight=0.01,
+    bias_prior_weight=0.01,
+):
+    """Binary NLL of sigmoid(beta * z + b)
+    + prior_weight * (log beta - log base_beta)^2 + bias_prior_weight * (b - base_bias)^2.
+    """
+    nll = F.binary_cross_entropy_with_logits(beta * logits + bias, targets, reduction="none")
+    prior = prior_weight * (beta.clamp_min(1e-6).log() - math.log(base_beta)) ** 2
+    prior = prior + bias_prior_weight * (bias - base_bias) ** 2
+    nll = _masked_text_mean(nll, label_mask)
+    prior = _masked_text_mean(prior, label_mask)
+    return nll + prior, nll
+
+
+def fit_global_scaling(examples, fit_bias=False, max_iter=100):
+    """The single (beta_0, b_0) minimizing the binary NLL of sigmoid(beta_0 * z + b_0).
+
+    fit_bias=False is temperature scaling (b_0 = 0), fit_bias=True is Platt scaling with a positive slope.
+    """
+    logits = torch.cat([ex.logits for ex in examples])
+    targets = torch.cat([ex.targets for ex in examples])
+    log_beta = torch.zeros((), requires_grad=True)
+    bias = torch.zeros((), requires_grad=fit_bias)
+    params = [log_beta, bias] if fit_bias else [log_beta]
+    optimizer = torch.optim.LBFGS(params, lr=0.5, max_iter=max_iter, line_search_fn="strong_wolfe")
+
+    def closure():
+        optimizer.zero_grad()
+        loss = F.binary_cross_entropy_with_logits(log_beta.exp() * logits + bias, targets)
+        loss.backward()
+        return loss
+
+    optimizer.step(closure)
+    return float(log_beta.detach().exp()), float(bias.detach())
+
+
+def calibration_metrics(probs, targets, num_bins=15):
+    """NLL, Brier, ECE and accuracy at 0.5 over flattened pairs."""
+    probs = probs.clamp(1e-7, 1 - 1e-7)
+    nll = F.binary_cross_entropy(probs, targets).item()
+    brier = ((probs - targets) ** 2).mean().item()
+    accuracy = ((probs >= 0.5).float() == (targets >= 0.5).float()).float().mean().item()
+
+    confidence = torch.maximum(probs, 1 - probs)
+    correct = ((probs >= 0.5) == (targets >= 0.5)).float()
+    bins = torch.clamp(((confidence - 0.5) * 2 * num_bins).long(), max=num_bins - 1)
+    ece = 0.0
+    for b in range(num_bins):
+        in_bin = bins == b
+        if in_bin.any():
+            ece += in_bin.float().mean().item() * abs(confidence[in_bin].mean().item() - correct[in_bin].mean().item())
+    return {"nll": nll, "brier": brier, "ece": ece, "accuracy": accuracy}
+
+
+@torch.no_grad()
+def evaluate_calibrator(calibrator, examples, device, base_beta=None, base_bias=0.0, batch_size=256):
+    """Metrics for the raw model, the global scaling fit (if base_beta) and the calibrator."""
+    calibrator.eval()
+    raw, scaled, calibrated, targets, betas, biases = [], [], [], [], [], []
+    for start in range(0, len(examples), batch_size):
+        batch = collate_calibration_examples(examples[start : start + batch_size])
+        batch = {key: value.to(device) for key, value in batch.items()}
+        beta, bias = calibrator(batch["text_repr"], batch["label_repr"], batch["logits"], batch["label_mask"])
+        mask = batch["label_mask"]
+        logits = batch["logits"][mask]
+        raw.append(torch.sigmoid(logits).cpu())
+        if base_beta is not None:
+            scaled.append(torch.sigmoid(base_beta * logits + base_bias).cpu())
+        calibrated.append(torch.sigmoid(beta[mask] * logits + bias[mask]).cpu())
+        targets.append(batch["targets"][mask].cpu())
+        betas.append(beta[mask].cpu())
+        biases.append(bias[mask].cpu())
+
+    targets = torch.cat(targets)
+    results = {"uncalibrated": calibration_metrics(torch.cat(raw), targets)}
+    if base_beta is not None:
+        results["global_scaling"] = calibration_metrics(torch.cat(scaled), targets)
+    results["calibrated"] = calibration_metrics(torch.cat(calibrated), targets)
+    for name, values in (("beta", torch.cat(betas)), ("bias", torch.cat(biases))):
+        results[name] = {"mean": values.mean().item(), "min": values.min().item(), "max": values.max().item()}
+    return results
+
+
+def fit_calibrator(
+    calibrator,
+    train_examples,
+    eval_examples,
+    device,
+    base_beta,
+    base_bias=0.0,
+    prior_weight=0.01,
+    bias_prior_weight=0.01,
+    num_epochs=20,
+    batch_size=64,
+    lr=1e-3,
+    weight_decay=0.01,
+    patience=3,
+    seed=42,
+    min_base_beta=0.05,
+    log_fn=print,
+):
+    """Train the calibrator MLP on cached features; keeps the weights with the best held-out NLL.
+
+    base_beta is floored at min_base_beta: near beta = 0 the sigmoid parameterization has no gradient,
+    so starting from a collapsed global temperature would leave the MLP unable to move.
+    """
+    if base_beta < min_base_beta:
+        log_fn(f"beta_0={base_beta:.3g} is below {min_base_beta}; using {min_base_beta} for the start and the prior")
+        base_beta = min_base_beta
+    rng = random.Random(seed)
+    calibrator.to(device)
+    if not calibrator.use_bias:
+        base_bias = 0.0
+    calibrator.reset_to_global_scaling(base_beta, base_bias)
+    optimizer = torch.optim.AdamW(calibrator.parameters(), lr=lr, weight_decay=weight_decay)
+
+    def eval_nll():
+        calibrator.eval()
+        total, count = 0.0, 0
+        with torch.no_grad():
+            for start in range(0, len(eval_examples), 256):
+                batch = collate_calibration_examples(eval_examples[start : start + 256])
+                batch = {key: value.to(device) for key, value in batch.items()}
+                beta, bias = calibrator(
+                    batch["text_repr"], batch["label_repr"], batch["logits"], batch["label_mask"]
+                )
+                _, nll = calibration_loss(
+                    beta, bias, batch["logits"], batch["targets"], batch["label_mask"], base_beta, base_bias
+                )
+                total += nll.item() * batch["logits"].shape[0]
+                count += batch["logits"].shape[0]
+        return total / max(count, 1)
+
+    best_nll = eval_nll()
+    best_state = {key: value.detach().clone() for key, value in calibrator.state_dict().items()}
+    log_fn(f"epoch 0: held-out NLL {best_nll:.5f} (global scaling, beta_0={base_beta:.4f}, b_0={base_bias:.4f})")
+    bad_epochs = 0
+
+    for epoch in range(1, num_epochs + 1):
+        calibrator.train()
+        order = list(range(len(train_examples)))
+        rng.shuffle(order)
+        train_loss = 0.0
+        for start in range(0, len(order), batch_size):
+            batch = collate_calibration_examples([train_examples[i] for i in order[start : start + batch_size]])
+            batch = {key: value.to(device) for key, value in batch.items()}
+            beta, bias = calibrator(batch["text_repr"], batch["label_repr"], batch["logits"], batch["label_mask"])
+            loss, _ = calibration_loss(
+                beta,
+                bias,
+                batch["logits"],
+                batch["targets"],
+                batch["label_mask"],
+                base_beta,
+                base_bias,
+                prior_weight,
+                bias_prior_weight,
+            )
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+            train_loss += loss.item() * batch["logits"].shape[0]
+
+        held_out_nll = eval_nll()
+        log_fn(f"epoch {epoch}: train loss {train_loss / len(order):.5f}, held-out NLL {held_out_nll:.5f}")
+        if held_out_nll < best_nll - 1e-5:
+            best_nll = held_out_nll
+            best_state = {key: value.detach().clone() for key, value in calibrator.state_dict().items()}
+            bad_epochs = 0
+        else:
+            bad_epochs += 1
+            if bad_epochs >= patience:
+                break
+
+    calibrator.load_state_dict(best_state)
+    calibrator.eval()
+    return calibrator
diff --git a/gliclass/config.py b/gliclass/config.py
index e75e4c5..aad3e2c 100644
--- a/gliclass/config.py
+++ b/gliclass/config.py
@@ -72,6 +72,11 @@ def __init__(
         recurrent_confidence_coef=0.0,
         recurrent_confidence_mode="relative",
         recurrent_confidence_margin=0.0,
+        use_calibrator=False,
+        calibrator_hidden_size=256,
+        calibrator_beta_max=3.0,
+        calibrator_dropout=0.1,
+        calibrator_use_bias=False,
         **kwargs,
     ):
         if isinstance(encoder_config, dict):
@@ -184,4 +189,13 @@ def __init__(
         self.recurrent_confidence_mode = recurrent_confidence_mode
         self.recurrent_confidence_margin = recurrent_confidence_margin
 
+        # Post-hoc calibrator (see gliclass/calibration.py), fitted with the backbone frozen.
+        # It predicts an inverse temperature beta in (0, calibrator_beta_max) per (text, label) pair and,
+        # with calibrator_use_bias, an additive logit bias: calibrated logit = beta * logit + bias.
+        self.use_calibrator = use_calibrator
+        self.calibrator_hidden_size = calibrator_hidden_size
+        self.calibrator_beta_max = calibrator_beta_max
+        self.calibrator_dropout = calibrator_dropout
+        self.calibrator_use_bias = calibrator_use_bias
+
         super().__init__(problem_type=problem_type, **kwargs)
diff --git a/gliclass/model.py b/gliclass/model.py
index 237c7b6..2463df7 100644
--- a/gliclass/model.py
+++ b/gliclass/model.py
@@ -24,6 +24,7 @@
 from .layers import FeaturesProjector, BiEncoderProjector, LayerwiseAttention, LstmSeq2SeqEncoder
 from .scorers import SCORER2OBJECT
 from .poolings import POOLING2OBJECT
+from .calibration import build_calibrator
 from .loss_functions import focal_loss_with_logits, sequence_contrastive_loss
 
 IS_LLM2VEC = is_module_available("llm2vec")
@@ -66,6 +67,13 @@ class GLiClassOutput(SequenceClassifierOutput):
     recurrent_losses: torch.Tensor | None = None
     recurrent_entropies: torch.Tensor | None = None
     recurrent_num_steps: torch.Tensor | None = None
+    # (batch, num_labels), 1 for real labels; set whenever text/class embeddings are requested
+    class_mask: torch.Tensor | None = None
+    # Set when the model has a calibrator: logits are then calibrated,
+    # inverse_temperatures * uncalibrated_logits + calibration_biases
+    uncalibrated_logits: torch.Tensor | None = None
+    inverse_temperatures: torch.Tensor | None = None
+    calibration_biases: torch.Tensor | None = None
 
 
 class GLiClassPreTrainedModel(PreTrainedModel):
@@ -471,7 +479,7 @@ def process_encoder_output(self, input_ids, attention_mask, encoder_layer, label
             logits = logits * self.logit_scale.to(classes_embedding.device)
 
         loss = self.get_loss(logits, labels, classes_embedding, classes_embedding_mask)
-        return (logits, loss, pooled_output, classes_embedding)
+        return (logits, loss, pooled_output, classes_embedding, classes_embedding_mask)
 
     def forward(
         self,
@@ -530,7 +538,7 @@ def forward(
             hidden_states = outputs.hidden_states
             loss = 0
             for encoder_layer in hidden_states:
-                logits, layer_loss, pooled_output, classes_embedding = self.process_encoder_output(
+                logits, layer_loss, pooled_output, classes_embedding, classes_mask = self.process_encoder_output(
                     input_ids, attention_mask, encoder_layer, labels, max_num_classes
                 )
                 loss += layer_loss
@@ -542,7 +550,7 @@ def forward(
                     encoder_layer = outputs[0]
             else:
                 encoder_layer = outputs.hidden_states[self.config.encoder_layer_id]
-            logits, loss, pooled_output, classes_embedding = self.process_encoder_output(
+            logits, loss, pooled_output, classes_embedding, classes_mask = self.process_encoder_output(
                 input_ids, attention_mask, encoder_layer, labels, max_num_classes
             )
 
@@ -557,6 +565,7 @@ def forward(
             attentions=outputs.attentions,
             text_embeddings=pooled_output if output_text_embeddings else None,
             class_embeddings=classes_embedding if output_class_embeddings else None,
+            class_mask=classes_mask,
         )
 
 
@@ -675,6 +684,7 @@ def forward(
             attentions=outputs.decoder_attentions,
             text_embeddings=pooled_output if output_text_embeddings else None,
             class_embeddings=classes_embedding if output_class_embeddings else None,
+            class_mask=classes_embedding_mask,
         )
 
 
@@ -768,6 +778,7 @@ def forward(
             attentions=outputs.decoder_attentions,
             text_embeddings=pooled_output if output_text_embeddings else None,
             class_embeddings=classes_embedding if output_class_embeddings else None,
+            class_mask=classes_embedding_mask,
         )
 
 
@@ -880,6 +891,7 @@ def forward(
             logits=logits,
             text_embeddings=text_embeddings if output_text_embeddings else None,
             class_embeddings=class_embeddings if output_class_embeddings else None,
+            class_mask=labels_mask if labels_mask is not None else torch.ones_like(logits, dtype=torch.long),
         )
 
 
@@ -954,6 +966,7 @@ def forward(
             logits=logits,
             text_embeddings=text_embeddings if output_text_embeddings else None,
             class_embeddings=class_embeddings if output_class_embeddings else None,
+            class_mask=labels_mask if labels_mask is not None else torch.ones_like(logits, dtype=torch.long),
         )
 
 
@@ -1261,6 +1274,8 @@ def classify_from_decoder_cache(
         halt_threshold: float | None = None,
         text_hidden_states: torch.Tensor | None = None,
         text_attention_mask: torch.Tensor | None = None,
+        output_text_embeddings: bool | None = None,
+        output_class_embeddings: bool | None = None,
     ) -> GLiClassOutput:
         """Classify label tokens against a text cache without persisting labels.
 
@@ -1279,16 +1294,26 @@ def classify_from_decoder_cache(
         label_hidden = decoder_outputs.last_hidden_state[:, -input_ids.shape[1] :, :]
         scorer_device = next(self.scorer.parameters()).device
         num_steps, halt_threshold = self._resolve_recurrence(use_recurrence, max_recurrent_steps, halt_threshold)
-        step_logits, steps_taken = self.scorer.recurrent_forward(
+        label_ids = input_ids.to(scorer_device)
+        label_mask = label_mask.to(scorer_device)
+        step_logits, steps_taken, (text_repr, label_repr) = self.scorer.recurrent_forward(
             hidden_states=label_hidden.to(scorer_device),
-            input_ids=input_ids.to(scorer_device),
-            attention_mask=label_mask.to(scorer_device),
+            input_ids=label_ids,
+            attention_mask=label_mask,
             num_steps=num_steps,
             halt_threshold=halt_threshold,
             text_hidden_states=None if text_hidden_states is None else text_hidden_states.to(scorer_device),
             text_attention_mask=None if text_attention_mask is None else text_attention_mask.to(scorer_device),
+            return_representations=True,
+        )
+        logits = step_logits[-1]
+        return GLiClassOutput(
+            logits=logits,
+            recurrent_num_steps=steps_taken,
+            text_embeddings=text_repr if output_text_embeddings else None,
+            class_embeddings=label_repr if output_class_embeddings else None,
+            class_mask=self.scorer._valid_label_mask(label_ids, label_mask, logits.shape[-1]).long(),
         )
-        return GLiClassOutput(logits=step_logits[-1], recurrent_num_steps=steps_taken)
 
     def forward(
         self,
@@ -1302,6 +1327,9 @@ def forward(
         max_recurrent_steps: int | None = None,
         halt_threshold: float | None = None,
         output_recurrent_logits: bool = False,
+        output_text_embeddings: bool | None = None,
+        output_class_embeddings: bool | None = None,
+        max_num_classes: int | None = None,
         **kwargs,
     ):
         """Forward pass.
@@ -1317,6 +1345,9 @@ def forward(
             max_recurrent_steps: Max reasoning steps for this call (training samples depth up to it)
             halt_threshold: Inference early-stop threshold on label probability change (0 disables)
             output_recurrent_logits: Return logits of every reasoning step
+            output_text_embeddings / output_class_embeddings: Return the scorer's text / label
+                representations of the final step
+            max_num_classes: Unused, accepted for pipeline compatibility
 
         Returns:
             GLiClassOutput with logits (of the last step), loss, and optionally past_key_values
@@ -1344,7 +1375,7 @@ def forward(
         else:
             label_hidden, label_ids, label_mask = self._extract_label_section(hidden_states, input_ids, attention_mask)
 
-        step_logits, steps_taken = self.scorer.recurrent_forward(
+        step_logits, steps_taken, (text_repr, label_repr) = self.scorer.recurrent_forward(
             hidden_states=label_hidden,
             input_ids=label_ids,
             attention_mask=label_mask,
@@ -1353,6 +1384,7 @@ def forward(
             return_all_steps=labels is not None or output_recurrent_logits,
             text_hidden_states=text_hidden,
             text_attention_mask=text_mask,
+            return_representations=True,
         )
         logits = step_logits[-1]
 
@@ -1382,6 +1414,9 @@ def forward(
             recurrent_losses=recurrent_losses,
             recurrent_entropies=recurrent_entropies,
             recurrent_num_steps=steps_taken if num_steps > 1 else None,
+            text_embeddings=text_repr if output_text_embeddings else None,
+            class_embeddings=label_repr if output_class_embeddings else None,
+            class_mask=self.scorer._valid_label_mask(label_ids, label_mask, logits.shape[-1]).long(),
         )
 
 
@@ -1400,8 +1435,41 @@ def __init__(self, config, from_pretrained=False):
             self.model = GLiClassEncoderDecoderCLS(config, from_pretrained)
         elif config.architecture_type == "decoder-kv":
             self.model = GLiClassDecoderKV(config, from_pretrained)
+        self.calibrator = build_calibrator(config) if getattr(config, "use_calibrator", False) else None
         self.post_init()
 
+    def add_calibrator(self, hidden_size=None, beta_max=None, dropout=None, use_bias=None):
+        """Attach a fresh calibrator (beta == 1, bias == 0, i.e. identity) and record it in the config."""
+        if hidden_size is not None:
+            self.config.calibrator_hidden_size = hidden_size
+        if beta_max is not None:
+            self.config.calibrator_beta_max = beta_max
+        if dropout is not None:
+            self.config.calibrator_dropout = dropout
+        if use_bias is not None:
+            self.config.calibrator_use_bias = use_bias
+        self.config.use_calibrator = True
+        reference = next(self.model.parameters())
+        self.calibrator = build_calibrator(self.config).to(device=reference.device)
+        return self.calibrator
+
+    def remove_calibrator(self):
+        self.config.use_calibrator = False
+        self.calibrator = None
+
+    def _calibrate(self, outputs, keep_text_embeddings, keep_class_embeddings):
+        logits = outputs.logits
+        beta, bias = self.calibrator(outputs.text_embeddings, outputs.class_embeddings, logits, outputs.class_mask)
+        outputs.uncalibrated_logits = logits
+        outputs.inverse_temperatures = beta
+        outputs.calibration_biases = bias
+        outputs.logits = (beta * logits.float() + bias).to(logits.dtype)
+        if not keep_text_embeddings:
+            outputs.text_embeddings = None
+        if not keep_class_embeddings:
+            outputs.class_embeddings = None
+        return outputs
+
     def get_input_embeddings(self):
         if self.config.architecture_type in {"uni-encoder"}:
             return self.model.encoder_model.get_input_embeddings()
@@ -1496,8 +1564,14 @@ def resize_token_embeddings(self, new_num_tokens: int | None = None, pad_to_mult
     def forward(self, *args, **kwargs):
         if kwargs.get("adapter_ids") is None:
             kwargs.pop("adapter_ids", None)
+        if self.calibrator is None:
+            return self.model(*args, **kwargs)
+
+        keep_text = bool(kwargs.get("output_text_embeddings"))
+        keep_class = bool(kwargs.get("output_class_embeddings"))
+        kwargs.update(output_text_embeddings=True, output_class_embeddings=True, return_dict=True)
         outputs = self.model(*args, **kwargs)
-        return outputs
+        return self._calibrate(outputs, keep_text, keep_class)
 
     def update_decoder_cache(self, **kwargs):
         """Extend a decoder-KV text cache."""
@@ -1509,4 +1583,11 @@ def classify_from_decoder_cache(self, **kwargs) -> GLiClassOutput:
         """Classify labels against a decoder-KV text cache."""
         if self.config.architecture_type != "decoder-kv":
             raise ValueError("Cached classification requires architecture_type='decoder-kv'.")
-        return self.model.classify_from_decoder_cache(**kwargs)
+        if self.calibrator is None:
+            return self.model.classify_from_decoder_cache(**kwargs)
+
+        keep_text = bool(kwargs.get("output_text_embeddings"))
+        keep_class = bool(kwargs.get("output_class_embeddings"))
+        kwargs.update(output_text_embeddings=True, output_class_embeddings=True)
+        outputs = self.model.classify_from_decoder_cache(**kwargs)
+        return self._calibrate(outputs, keep_text, keep_class)
diff --git a/gliclass/pipeline.py b/gliclass/pipeline.py
index 19150f6..3298f2a 100644
--- a/gliclass/pipeline.py
+++ b/gliclass/pipeline.py
@@ -265,18 +265,31 @@ def _postprocess_logits(
         labels: List[str],
         classification_type: str,
         threshold: float,
+        initial_logits: torch.Tensor | None = None,
     ) -> tuple[List[Dict[str, Any]], Dict[str, float]]:
-        """Convert one row of logits into predictions and a complete score map."""
-        logits = logits[: len(labels)]
+        """Convert one row of logits into predictions and a complete score map.
+
+        With initial_logits (uncalibrated logits of a calibrated model), every prediction also carries
+        "initial_score"; "score" and the threshold use the calibrated logits.
+        """
+        activation = (lambda x: torch.softmax(x, dim=-1)) if classification_type == "single-label" else torch.sigmoid
+        scores = activation(logits[: len(labels)].float()).tolist()
+        initial_scores = None
+        if initial_logits is not None:
+            initial_scores = activation(initial_logits[: len(labels)].float()).tolist()
+
+        def prediction(index):
+            item = {"label": labels[index], "score": scores[index]}
+            if initial_scores is not None:
+                item["initial_score"] = initial_scores[index]
+            return item
+
+        all_scores = dict(zip(labels, scores, strict=True))
         if classification_type == "single-label":
-            scores = torch.softmax(logits, dim=-1)
-            all_scores = {label: float(score.detach()) for label, score in zip(labels, scores, strict=True)}
-            best = int(scores.argmax().item())
-            return [{"label": labels[best], "score": float(scores[best].detach())}], all_scores
-
-        scores = torch.sigmoid(logits)
-        all_scores = {label: float(score.detach()) for label, score in zip(labels, scores, strict=True)}
-        predictions = [{"label": label, "score": score} for label, score in all_scores.items() if score >= threshold]
+            best = max(range(len(scores)), key=scores.__getitem__)
+            return [prediction(best)], all_scores
+
+        predictions = [prediction(i) for i, score in enumerate(scores) if score >= threshold]
         return predictions, all_scores
 
     def _process_labels(
@@ -505,6 +518,7 @@ def __call__(
                 adapter_ids=batch_adapter_ids,
             )
             logits = model_output.logits
+            initial_logits = getattr(model_output, "uncalibrated_logits", None)
 
             for i in range(len(batch_texts)):
                 global_idx = idx + i
@@ -521,6 +535,7 @@ def __call__(
                     curr_labels,
                     item_classification_type,
                     item_threshold,
+                    initial_logits=None if initial_logits is None else initial_logits[i],
                 )
                 results.append(predictions)
                 if return_hierarchical:
diff --git a/gliclass/scorers.py b/gliclass/scorers.py
index 230d6c2..cc7de5c 100644
--- a/gliclass/scorers.py
+++ b/gliclass/scorers.py
@@ -360,6 +360,7 @@ def recurrent_forward(
         return_all_steps=False,
         text_hidden_states=None,
         text_attention_mask=None,
+        return_representations=False,
         **kwargs,
     ):
         """Run step 1 (one scorer_encoder pass) and up to num_steps - 1 gated recurrent steps.
@@ -373,6 +374,8 @@ def recurrent_forward(
             step_logits: list of (batch_size, num_labels) logits for every step if
                 return_all_steps, otherwise a single-element list with the final logits
             steps_taken: (batch_size,) number of steps each example ran
+            representations: only with return_representations, the (text_repr, label_repr) the
+                final logits were scored from
         """
         if num_steps > 1 and not self.recurrent:
             raise ValueError("num_steps > 1 requires a scorer built with config.recurrent_steps > 1.")
@@ -380,7 +383,7 @@ def recurrent_forward(
         bptt_steps = getattr(self.config, "recurrent_bptt_steps", None)
 
         state = self._encode(hidden_states, attention_mask)
-        logits = self._score(state, input_ids, attention_mask)
+        logits, text_repr, label_repr = self._score(state, input_ids, attention_mask, return_representations=True)
         step_logits = [logits]
 
         batch_size, num_labels = logits.shape
@@ -404,7 +407,9 @@ def recurrent_forward(
                 state = state.detach()
 
             new_state = self.reasoning_cell(state, inputs, attention_mask, self._encode, text_memory, text_mask)
-            new_logits = self._score(new_state, input_ids, attention_mask)
+            new_logits, new_text_repr, new_label_repr = self._score(
+                new_state, input_ids, attention_mask, return_representations=True
+            )
 
             if halt_threshold:
                 change = (
@@ -415,12 +420,15 @@ def recurrent_forward(
                 )
                 new_state = torch.where(active[:, None, None], new_state, state)
                 new_logits = torch.where(active[:, None], new_logits, logits)
+                new_text_repr = torch.where(active[:, None], new_text_repr, text_repr)
+                new_label_repr = torch.where(active[:, None, None], new_label_repr, label_repr)
                 steps_taken = steps_taken + active.long()
                 active = active & change.ge(halt_threshold)
             else:
                 steps_taken = steps_taken + 1
 
             state, logits = new_state, new_logits
+            text_repr, label_repr = new_text_repr, new_label_repr
             if return_all_steps:
                 step_logits.append(logits)
             else:
@@ -434,9 +442,11 @@ def recurrent_forward(
             unused = sum(param.sum() for param in self.reasoning_cell.parameters())
             step_logits[-1] = step_logits[-1] + 0.0 * unused
 
+        if return_representations:
+            return step_logits, steps_taken, (text_repr, label_repr)
         return step_logits, steps_taken
 
-    def _score(self, contextualized_hidden_states, input_ids, attention_mask):
+    def _score(self, contextualized_hidden_states, input_ids, attention_mask, return_representations=False):
         text_repr, label_repr = self._extract_representations(contextualized_hidden_states, input_ids, attention_mask)
 
         text_repr = self.text_projector(text_repr)
@@ -454,6 +464,8 @@ def _score(self, contextualized_hidden_states, input_ids, attention_mask):
 
         logits = self.mlp(combined_rep).squeeze(-1)
 
+        if return_representations:
+            return logits, text_repr, label_repr
         return logits
 
 
diff --git a/gliclass/streaming/pipeline.py b/gliclass/streaming/pipeline.py
index c18310a..df81b4b 100644
--- a/gliclass/streaming/pipeline.py
+++ b/gliclass/streaming/pipeline.py
@@ -412,11 +412,13 @@ def _classify(
         for local_index, batch_index in enumerate(triggered_indices):
             labels = batch["labels"][batch_index]
             logits = model_output.logits[local_index, : len(labels)]
+            initial_logits = getattr(model_output, "uncalibrated_logits", None)
             predictions, all_scores = self._postprocess_logits(
                 logits,
                 labels,
                 batch["classification_types"][batch_index],
                 batch["thresholds"][batch_index],
+                initial_logits=None if initial_logits is None else initial_logits[local_index, : len(labels)],
             )
             if batch["return_hierarchical"][batch_index]:
                 predictions = build_hierarchical_output(

From fc8bd5683b963c7cda4e1559e85404e99d9d4a81 Mon Sep 17 00:00:00 2001
From: Ingvar <igor.stepanov2000@gmail.com>
Date: Fri, 2 Oct 2026 16:56:55 +0300
Subject: [PATCH 04/12] implement multi-modals support for gliclass

---
 gliclass/config.py          |  16 +++-
 gliclass/data_processing.py |  44 +++++++----
 gliclass/model.py           |  26 +++++--
 gliclass/multimodal.py      | 150 ++++++++++++++++++++++++++++++++++++
 gliclass/pipeline.py        | 112 +++++++++++++++++++++------
 gliclass/training.py        |   3 +
 6 files changed, 306 insertions(+), 45 deletions(-)
 create mode 100644 gliclass/multimodal.py

diff --git a/gliclass/config.py b/gliclass/config.py
index aad3e2c..ecb4b87 100644
--- a/gliclass/config.py
+++ b/gliclass/config.py
@@ -77,6 +77,8 @@ def __init__(
         calibrator_beta_max=3.0,
         calibrator_dropout=0.1,
         calibrator_use_bias=False,
+        multimodal=False,
+        freeze_media_encoders=True,
         **kwargs,
     ):
         if isinstance(encoder_config, dict):
@@ -109,13 +111,16 @@ def __init__(
             self.label_model_config = None
         self.label_model_name = label_model_name
 
+        # Multi-modal backbones have a composite config; text attributes live in its text config
+        text_config = self.encoder_config.get_text_config()
+
         if hidden_size is None:
-            self.hidden_size = self.encoder_config.hidden_size
+            self.hidden_size = text_config.hidden_size
         else:
             self.hidden_size = hidden_size
 
         if vocab_size is None:
-            self.vocab_size = self.encoder_config.vocab_size
+            self.vocab_size = text_config.vocab_size
         else:
             self.vocab_size = vocab_size
 
@@ -160,7 +165,7 @@ def __init__(
         self.layer_wise = layer_wise
         self.encoder_layer_id = encoder_layer_id
         self.embed_class_token = embed_class_token
-        self.pad_token_id = self.encoder_config.pad_token_id
+        self.pad_token_id = text_config.pad_token_id
         self.dropout = dropout
         self.use_segment_embeddings = use_segment_embeddings
 
@@ -198,4 +203,9 @@ def __init__(
         self.calibrator_dropout = calibrator_dropout
         self.calibrator_use_bias = calibrator_use_bias
 
+        # decoder-kv: keep the backbone's vision / audio encoders (Qwen3_5Model, Gemma4Model) so images and
+        # audio can be supplied next to text; freeze_media_encoders keeps those encoders fixed during training.
+        self.multimodal = multimodal
+        self.freeze_media_encoders = freeze_media_encoders
+
         super().__init__(problem_type=problem_type, **kwargs)
diff --git a/gliclass/data_processing.py b/gliclass/data_processing.py
index 12569b4..1da30e1 100644
--- a/gliclass/data_processing.py
+++ b/gliclass/data_processing.py
@@ -6,10 +6,15 @@
 from torch.utils.data import Dataset
 from torch.nn.utils.rnn import pad_sequence
 
+from .multimodal import MEDIA_KEYS, collate_media, is_processor, process_multimodal_with_budget
 
-def format_decoder_kv_context(text: str, prompt: str = "", examples: str = "") -> str:
-    """Format the persistent context portion of a decoder-KV sequence."""
-    return f"{prompt}{examples}{text}"
+
+def format_decoder_kv_context(text: str, prompt: str = "", examples: str = "", media: str = "") -> str:
+    """Format the persistent context portion of a decoder-KV sequence.
+
+    media: image / audio placeholder tokens (see gliclass.multimodal.format_media_prefix), placed before the text.
+    """
+    return f"{prompt}{examples}{media}{text}"
 
 
 def format_decoder_kv_labels(
@@ -34,9 +39,10 @@ def format_decoder_kv_sequence(
     examples: str = "",
     label_token: str = "<<LABEL>>",
     sep_token: str = "<<SEP>>",
+    media: str = "",
 ) -> str:
     """Format a complete decoder-KV sequence for training or classic inference."""
-    return format_decoder_kv_context(text, prompt, examples) + format_decoder_kv_labels(
+    return format_decoder_kv_context(text, prompt, examples, media) + format_decoder_kv_labels(
         labels,
         label_token=label_token,
         sep_token=sep_token,
@@ -341,17 +347,25 @@ def tokenize_and_prepare_labels_for_decoder_kv(self, example):
         examples_text = self.format_examples(example)
         text = str(example["text"])
 
-        input_text = format_decoder_kv_sequence(
-            text,
-            example["all_labels"],
-            prompt=prompt,
-            examples=examples_text,
-            label_token=self.label_token,
-            sep_token=self.sep_token,
-        )
+        def build(text, media=""):
+            return format_decoder_kv_sequence(
+                text,
+                example["all_labels"],
+                prompt=prompt,
+                examples=examples_text,
+                label_token=self.label_token,
+                sep_token=self.sep_token,
+                media=media,
+            )
+
         label2idx = {label: idx for idx, label in enumerate(example["all_labels"])}
 
-        tokenized_inputs = self.tokenize(input_text)
+        if is_processor(self.tokenizer):
+            tokenized_inputs = process_multimodal_with_budget(
+                self.tokenizer, build, text, example.get("images") or [], example.get("audio") or [], self.max_length
+            )
+        else:
+            tokenized_inputs = self.tokenize(build(text))
         tokenized_inputs["labels"] = self.prepare_labels(example, label2idx, self.problem_type)
         tokenized_inputs["labels_text"] = example["all_labels"]
         tokenized_inputs["input_texts"] = example["text"]
@@ -480,7 +494,8 @@ def _resolve_max_num_classes(self, batch):
         return None  # 'fixed': model uses config.max_num_classes
 
     def __call__(self, batch):
-        keys = batch[0].keys()
+        # Media tensors are indexed by image / audio item and only present in samples that have media
+        keys = [key for key in batch[0].keys() if key not in MEDIA_KEYS]
         padded_batch = {key: [] for key in keys}
 
         for key in keys:
@@ -508,5 +523,6 @@ def __call__(self, batch):
             else:
                 raise TypeError(f"Unsupported data type: {type(key_data[0])}")
 
+        padded_batch.update(collate_media(batch))
         padded_batch["max_num_classes"] = self._resolve_max_num_classes(batch)
         return padded_batch
diff --git a/gliclass/model.py b/gliclass/model.py
index 2463df7..688b4df 100644
--- a/gliclass/model.py
+++ b/gliclass/model.py
@@ -996,7 +996,15 @@ def __init__(self, config: GLiClassModelConfig, from_pretrained=False):
 
         config_name = config.encoder_config.__class__.__name__
 
-        if config_name == "Qwen3_5TextConfig":
+        if getattr(config, "multimodal", False):
+            # Base multi-modal model (no LM head): encodes images / audio and merges them into the sequence
+            if config_name == "Qwen3_5Config":
+                from transformers.models.qwen3_5 import Qwen3_5Model as ModelClass
+            elif config_name == "Gemma4Config":
+                from transformers.models.gemma4 import Gemma4Model as ModelClass
+            else:
+                raise ValueError(f"multimodal decoder-kv requires Qwen3.5 or Gemma 4. Got: {config_name}")
+        elif config_name == "Qwen3_5TextConfig":
             from transformers.models.qwen3_5 import Qwen3_5TextModel
 
             ModelClass = Qwen3_5TextModel
@@ -1016,12 +1024,19 @@ def __init__(self, config: GLiClassModelConfig, from_pretrained=False):
             self.decoder_model = ModelClass.from_pretrained(config.encoder_model_name)
         else:
             self.decoder_model = ModelClass(config.encoder_config)
+            # multi-modal backbones build their sub-models in the config's dtype (e.g. bf16); match the scorer
+            self.decoder_model.to(torch.get_default_dtype())
 
         if config.vocab_size is not None and hasattr(self.decoder_model, "resize_token_embeddings"):
-            current_vocab = self.decoder_model.config.vocab_size
+            current_vocab = self.decoder_model.config.get_text_config().vocab_size
             if current_vocab != config.vocab_size:
                 self.decoder_model.resize_token_embeddings(config.vocab_size)
 
+        if getattr(config, "multimodal", False) and getattr(config, "freeze_media_encoders", True):
+            from .multimodal import freeze_media_encoders
+
+            freeze_media_encoders(self.decoder_model)
+
         from .scorers import DecoderKVScorer
 
         self.scorer = DecoderKVScorer(config)
@@ -1037,7 +1052,8 @@ def get_loss(self, logits, labels):
         if labels is not None:
             # Sequence truncation may cut label tokens → align labels to logits
             num_labels = logits.shape[-1]
-            if labels.shape[-1] != num_labels:
+            # single-label targets are class indices (batch,); only per-label targets need aligning
+            if labels.dim() > 1 and labels.shape[-1] != num_labels:
                 labels = labels[:, :num_labels]
 
             if self.config.problem_type == "multi_label_classification":
@@ -1136,7 +1152,7 @@ def get_recurrent_loss(self, step_logits, labels, valid_labels=None):
             step_entropies: (num_steps,) detached mean label entropy per step, for monitoring
         """
         num_labels = step_logits[-1].shape[-1]
-        if labels.shape[-1] != num_labels:
+        if labels.dim() > 1 and labels.shape[-1] != num_labels:
             labels = labels[:, :num_labels]
         if valid_labels is None:
             valid_labels = torch.ones_like(step_logits[-1], dtype=torch.bool)
@@ -1556,7 +1572,7 @@ def resize_token_embeddings(self, new_num_tokens: int | None = None, pad_to_mult
             model_embeds = self.model.decoder_model.resize_token_embeddings(new_num_tokens, pad_to_multiple_of)
         else:
             raise NotImplementedError("Resizing is not implemented for bi-encoder architecture")
-        self.config.encoder_config.vocab_size = model_embeds.num_embeddings
+        self.config.encoder_config.get_text_config().vocab_size = model_embeds.num_embeddings
         self.config.vocab_size = model_embeds.num_embeddings
         self.vocab_size = model_embeds.num_embeddings
         return model_embeds
diff --git a/gliclass/multimodal.py b/gliclass/multimodal.py
new file mode 100644
index 0000000..0e2aefd
--- /dev/null
+++ b/gliclass/multimodal.py
@@ -0,0 +1,150 @@
+"""Multi-modal (image / audio) input helpers for the decoder-kv architecture.
+
+Media enter a decoder-kv sequence as placeholder tokens in the context part:
+
+    [prompt][examples]<image><image><audio>text<<SEP>>label1<<LABEL>>...<<SEP>>
+
+The HF processor expands every placeholder into the item's soft tokens and returns the media tensors;
+the multi-modal backbone (Qwen3_5Model, Gemma4Model) encodes them and writes the features into those
+positions. Placeholders are matched to media in order, so the i-th image placeholder is the i-th image.
+"""
+
+import torch
+from transformers.processing_utils import ProcessorMixin
+
+# Processor outputs indexed by media item rather than by sample: concatenated across the batch.
+MEDIA_KEYS = (
+    "pixel_values",
+    "image_grid_thw",
+    "image_position_ids",
+    "pixel_values_videos",
+    "video_grid_thw",
+    "video_position_ids",
+    "input_features",
+    "input_features_mask",
+)
+# Padding values for trailing dims of media tensors that differ between samples ((-1, -1) = padded patch).
+MEDIA_PAD_VALUES = {"image_position_ids": -1, "video_position_ids": -1}
+
+
+def is_processor(obj) -> bool:
+    return isinstance(obj, ProcessorMixin)
+
+
+def get_tokenizer(tokenizer_or_processor):
+    """Text tokenizer of a processor, or the tokenizer itself."""
+    if is_processor(tokenizer_or_processor):
+        return tokenizer_or_processor.tokenizer
+    return tokenizer_or_processor
+
+
+def media_placeholder(processor, kind: str) -> str:
+    """The placeholder to write in the text for one image or audio item.
+
+    Qwen expects <|vision_start|><|image_pad|><|vision_end|>; Gemma adds its begin/end markers itself.
+    """
+    if kind == "image":
+        token = processor.image_token
+        if getattr(processor, "vision_start_token", None) is not None:
+            return f"{processor.vision_start_token}{token}{processor.vision_end_token}"
+        return token
+    if kind == "audio":
+        token = getattr(processor, "audio_token", None)
+        if token is None:
+            raise ValueError(f"{type(processor).__name__} does not support audio inputs.")
+        return token
+    raise ValueError(f"Unknown media kind: {kind}")
+
+
+def format_media_prefix(processor, num_images: int = 0, num_audio: int = 0) -> str:
+    prefix = media_placeholder(processor, "image") * num_images if num_images else ""
+    if num_audio:
+        prefix += media_placeholder(processor, "audio") * num_audio
+    return prefix
+
+
+def load_images(images):
+    """Load images given as paths, URLs or PIL images."""
+    from transformers.image_utils import load_image
+
+    return [load_image(image) for image in images or []]
+
+
+def load_audio(audio, sampling_rate: int):
+    """Load audio given as paths, URLs or 1D arrays (arrays must already be at `sampling_rate`)."""
+    from transformers.audio_utils import load_audio as _load_audio
+
+    return [_load_audio(item, sampling_rate=sampling_rate) for item in audio or []]
+
+
+def process_multimodal(processor, text: str, images=None, audio=None) -> dict:
+    """Run the processor on one sample; returns 1D input_ids/attention_mask/mm_token_type_ids plus media tensors.
+
+    Media may be paths, URLs, PIL images or arrays. The text must already contain one placeholder per item.
+    """
+    images = load_images(images)
+    kwargs = {}
+    if images:
+        kwargs["images"] = [images]
+    if audio:
+        sampling_rate = processor.feature_extractor.sampling_rate
+        kwargs["audio"] = load_audio(audio, sampling_rate)
+    outputs = processor(text=[text], return_tensors="pt", **kwargs)
+
+    sample = {}
+    for key, value in outputs.items():
+        if key in MEDIA_KEYS:
+            sample[key] = value
+        elif isinstance(value, torch.Tensor):
+            sample[key] = value[0]
+    return sample
+
+
+def process_multimodal_with_budget(processor, build, text: str, images, audio, max_length: int, reserved: int = 0):
+    """Process build(text, media_placeholders) so that it fits max_length - reserved tokens.
+
+    The processed sequence is never truncated (that would cut through expanded media placeholders);
+    when it is too long, the text is shortened by the overflow and processed again.
+    """
+    images = load_images(images)
+    if audio:
+        audio = load_audio(audio, processor.feature_extractor.sampling_rate)
+    media = format_media_prefix(processor, len(images), len(audio or []))
+    sample = process_multimodal(processor, build(text, media), images, audio)
+    overflow = len(sample["input_ids"]) + reserved - max_length
+    if overflow > 0:
+        tokenizer = get_tokenizer(processor)
+        text_ids = tokenizer(text, add_special_tokens=False)["input_ids"]
+        text = tokenizer.decode(text_ids[: max(0, len(text_ids) - overflow)])
+        sample = process_multimodal(processor, build(text, media), images, audio)
+    return sample
+
+
+def _pad_trailing(tensor: torch.Tensor, shape, value) -> torch.Tensor:
+    pad = []
+    for dim in reversed(range(1, tensor.dim())):
+        pad.extend((0, shape[dim] - tensor.shape[dim]))
+    if not any(pad):
+        return tensor
+    return torch.nn.functional.pad(tensor, pad, value=value)
+
+
+def collate_media(samples: list[dict]) -> dict:
+    """Concatenate media tensors of a batch along dim 0 (in batch order), padding trailing dims if needed."""
+    batch = {}
+    for key in MEDIA_KEYS:
+        tensors = [sample[key] for sample in samples if sample.get(key) is not None]
+        if not tensors:
+            continue
+        shape = [max(tensor.shape[dim] for tensor in tensors) for dim in range(tensors[0].dim())]
+        value = MEDIA_PAD_VALUES.get(key, False if tensors[0].dtype == torch.bool else 0)
+        batch[key] = torch.cat([_pad_trailing(tensor, shape, value) for tensor in tensors], dim=0)
+    return batch
+
+
+def freeze_media_encoders(model) -> None:
+    """Freeze the vision / audio towers and their projectors of a multi-modal backbone."""
+    for name in ("visual", "vision_tower", "audio_tower", "embed_vision", "embed_audio"):
+        module = getattr(model, name, None)
+        if module is not None:
+            module.requires_grad_(False)
diff --git a/gliclass/pipeline.py b/gliclass/pipeline.py
index 3298f2a..44a5233 100644
--- a/gliclass/pipeline.py
+++ b/gliclass/pipeline.py
@@ -8,6 +8,7 @@
 from .model import GLiClassModel, GLiClassBiEncoder
 from .utils import retrieval_augmented_text
 from .data_processing import format_decoder_kv_labels, format_decoder_kv_context, format_decoder_kv_sequence
+from .multimodal import collate_media, get_tokenizer, is_processor, process_multimodal_with_budget
 
 
 def flatten_hierarchical_labels(
@@ -186,7 +187,12 @@ def __init__(
     ):
         self.model = model
         if isinstance(tokenizer, str):
-            self.tokenizer = AutoTokenizer.from_pretrained(tokenizer)
+            if getattr(model.config, "multimodal", False):
+                from transformers import AutoProcessor
+
+                self.tokenizer = AutoProcessor.from_pretrained(tokenizer)
+            else:
+                self.tokenizer = AutoTokenizer.from_pretrained(tokenizer)
         else:
             self.tokenizer = tokenizer
         self.max_classes = max_classes
@@ -250,6 +256,15 @@ def _normalize_classification_types(
         normalized = self._normalize_classification_type(classification_type)
         return [normalized] * num_texts
 
+    @staticmethod
+    def _normalize_media(media, num_texts: int) -> List[List]:
+        """Per-text lists of media items; a single text may pass its items as a flat list."""
+        if num_texts == 1 and not any(isinstance(item, (list, tuple)) or item is None for item in media):
+            media = [media]
+        if len(media) != num_texts:
+            raise ValueError("images / audio must provide one list of items per text.")
+        return [[] if item is None else list(item) if isinstance(item, (list, tuple)) else [item] for item in media]
+
     def _normalize_adapter_ids(self, adapter_ids: str | List[str] | None, num_texts: int) -> List[str] | None:
         if adapter_ids is None:
             return None
@@ -449,6 +464,8 @@ def __call__(
         prompt: str | List[str] | None = None,
         return_hierarchical: bool = False,
         adapter_ids: str | List[str] | None = None,
+        images: List | None = None,
+        audio: List | None = None,
     ):
         """
         Perform zero-shot classification.
@@ -466,6 +483,8 @@ def __call__(
             prompt: Task description - string (same for all) or list (per-text)
             return_hierarchical: If True, return hierarchical structure with all scores
             adapter_ids: Optional LoRA adapter id for all texts or one adapter id per text.
+            images: Multi-modal decoder-kv only: per text, a list of images (paths, URLs or PIL images).
+            audio: Multi-modal decoder-kv only: per text, a list of audio clips (paths, URLs or arrays).
 
         Returns:
             List of classification results or hierarchical dict structure.
@@ -473,6 +492,11 @@ def __call__(
         original_labels = labels
 
         texts = self._normalize_texts(texts)
+        media = {
+            key: self._normalize_media(value, len(texts))
+            for key, value in (("images", images), ("audio", audio))
+            if value is not None
+        }
         thresholds = self._normalize_thresholds(threshold, len(texts))
         classification_types = self._normalize_classification_types(classification_type, len(texts))
         adapter_ids = self._normalize_adapter_ids(adapter_ids, len(texts))
@@ -508,8 +532,10 @@ def __call__(
             batch_prompt = self._get_batch_prompt(prompt, idx, len(batch_texts))
             batch_adapter_ids = adapter_ids[idx : idx + len(batch_texts)] if adapter_ids is not None else None
 
+            batch_media = {key: value[idx : idx + len(batch_texts)] for key, value in media.items()}
+
             tokenized_inputs = self.prepare_inputs(
-                batch_texts, batch_labels, same_labels, examples=batch_examples, prompt=batch_prompt
+                batch_texts, batch_labels, same_labels, examples=batch_examples, prompt=batch_prompt, **batch_media
             )
             max_num_classes = self._resolve_max_num_classes(batch_labels, same_labels)
             model_output = self.model(
@@ -812,15 +838,21 @@ def _build_context_and_labels(self, texts, labels, same_labels, examples, prompt
             )
         return contexts, label_sequences
 
-    def prepare_inputs(self, texts, labels, same_labels=False, examples=None, prompt=None):
+    def prepare_inputs(self, texts, labels, same_labels=False, examples=None, prompt=None, images=None, audio=None):
+        tokenizer = get_tokenizer(self.tokenizer)
         contexts, label_sequences = self._build_context_and_labels(texts, labels, same_labels, examples, prompt)
 
-        label_token_ids = self.tokenizer(label_sequences, add_special_tokens=False)["input_ids"]
+        label_token_ids = tokenizer(label_sequences, add_special_tokens=False)["input_ids"]
+
+        if images is not None or audio is not None:
+            return self._prepare_multimodal_inputs(
+                texts, labels, same_labels, examples, prompt, label_token_ids, images, audio
+            )
 
         input_ids_list = []
         for context, label_ids in zip(contexts, label_token_ids, strict=True):
             context_budget = max(1, self.max_length - len(label_ids))
-            context_ids = self.tokenizer(
+            context_ids = tokenizer(
                 context,
                 truncation=True,
                 max_length=context_budget,
@@ -828,27 +860,55 @@ def prepare_inputs(self, texts, labels, same_labels=False, examples=None, prompt
             )["input_ids"]
             input_ids_list.append(context_ids + label_ids)
 
-        max_len = max(len(ids) for ids in input_ids_list)
-        pad_token_id = self.tokenizer.pad_token_id
+        return self._pad_inputs({"input_ids": input_ids_list})
+
+    def _prepare_multimodal_inputs(self, texts, labels, same_labels, examples, prompt, label_token_ids, images, audio):
+        """Process [prompt][examples]<media>text with the processor and append the label section ids."""
+        if not is_processor(self.tokenizer):
+            raise ValueError("images / audio require building the pipeline with the model's processor (AutoProcessor).")
+
+        sequences = {"input_ids": [], "mm_token_type_ids": []}
+        samples = []
+        for i, (text, label_ids) in enumerate(zip(texts, label_token_ids, strict=True)):
+            examples_text = self._format_examples_for_input(self._get_text_examples(examples, i))
+            text_prompt = self._format_prompt(prompt, i) or ""
+            sample = process_multimodal_with_budget(
+                self.tokenizer,
+                lambda text, media: format_decoder_kv_context(text, text_prompt, examples_text, media),
+                text,
+                images[i] if images else [],
+                audio[i] if audio else [],
+                self.max_length,
+                reserved=len(label_ids),
+            )
+            context_ids = sample["input_ids"].tolist()
+            mm_token_type_ids = sample.get("mm_token_type_ids")
+            mm_token_type_ids = [0] * len(context_ids) if mm_token_type_ids is None else mm_token_type_ids.tolist()
+            sequences["input_ids"].append(context_ids + label_ids)
+            sequences["mm_token_type_ids"].append(mm_token_type_ids + [0] * len(label_ids))
+            samples.append(sample)
+
+        inputs = self._pad_inputs(sequences)
+        inputs.update({key: value.to(self.device) for key, value in collate_media(samples).items()})
+        return inputs
+
+    def _pad_inputs(self, sequences: dict) -> dict:
+        """Right-pad id sequences (the model locates the label section from the attention mask's length)."""
+        tokenizer = get_tokenizer(self.tokenizer)
+        pad_token_id = tokenizer.pad_token_id
         if pad_token_id is None:
-            pad_token_id = self.tokenizer.eos_token_id or 0
-        padding_side = getattr(self.tokenizer, "padding_side", "right")
-
-        padded_input_ids = []
-        attention_masks = []
-        for ids in input_ids_list:
-            pad_len = max_len - len(ids)
-            if padding_side == "left":
-                padded_input_ids.append([pad_token_id] * pad_len + ids)
-                attention_masks.append([0] * pad_len + [1] * len(ids))
-            else:
-                padded_input_ids.append(ids + [pad_token_id] * pad_len)
-                attention_masks.append([1] * len(ids) + [0] * pad_len)
+            pad_token_id = tokenizer.eos_token_id or 0
 
-        return {
-            "input_ids": torch.tensor(padded_input_ids, device=self.device),
-            "attention_mask": torch.tensor(attention_masks, device=self.device),
+        input_ids_list = sequences["input_ids"]
+        max_len = max(len(ids) for ids in input_ids_list)
+        inputs = {
+            "input_ids": [ids + [pad_token_id] * (max_len - len(ids)) for ids in input_ids_list],
+            "attention_mask": [[1] * len(ids) + [0] * (max_len - len(ids)) for ids in input_ids_list],
         }
+        for key, values in sequences.items():
+            if key != "input_ids":
+                inputs[key] = [ids + [0] * (max_len - len(ids)) for ids in values]
+        return {key: torch.tensor(value, device=self.device) for key, value in inputs.items()}
 
 
 class ZeroShotClassificationPipeline:
@@ -1007,6 +1067,8 @@ def __call__(
         prompt: str | List[str] | None = None,
         return_hierarchical: bool = False,
         adapter_ids: str | List[str] | None = None,
+        images: List | None = None,
+        audio: List | None = None,
     ):
         """
         Perform zero-shot classification.
@@ -1027,10 +1089,13 @@ def __call__(
             prompt: Task description - string or list of strings (per-text)
             return_hierarchical: If True, return structure matching input labels
             adapter_ids: Optional LoRA adapter id for all texts or one adapter id per text.
+            images: Multi-modal decoder-kv only: per text, a list of images (paths, URLs or PIL images).
+            audio: Multi-modal decoder-kv only: per text, a list of audio clips (paths, URLs or arrays).
 
         Returns:
             List of predictions (flat) or hierarchical dicts with all scores
         """
+        media = {key: value for key, value in (("images", images), ("audio", audio)) if value is not None}
         return self.pipe(
             texts,
             labels,
@@ -1042,6 +1107,7 @@ def __call__(
             prompt=prompt,
             return_hierarchical=return_hierarchical,
             adapter_ids=adapter_ids,
+            **media,
         )
 
 
diff --git a/gliclass/training.py b/gliclass/training.py
index 746181e..e00630e 100644
--- a/gliclass/training.py
+++ b/gliclass/training.py
@@ -273,6 +273,9 @@ def consolidate(self, dataset: Dataset, alpha: float = 0.5):
 class TrainingArguments(transformers.TrainingArguments):
     cache_dir: str | None = field(default=None)
     optim: str = field(default="adamw_torch")
+    # GLiClass models take **kwargs, so the forward signature cannot tell which columns are unused
+    # (plain-dict samples, e.g. multi-modal ones, would otherwise lose every key)
+    remove_unused_columns: bool = field(default=False)
     others_lr: float | None = None
     others_weight_decay: float | None = 0.0
     recurrent_lr: float | None = field(

From 67933e8f674d77cd2f3acef5a0d37f2320092917 Mon Sep 17 00:00:00 2001
From: Ingvar <igor.stepanov2000@gmail.com>
Date: Fri, 2 Oct 2026 16:57:59 +0300
Subject: [PATCH 05/12] make training support multi-modal data

---
 train.py | 90 ++++++++++++++++++++++++++++++++++++++++++++++++--------
 1 file changed, 78 insertions(+), 12 deletions(-)

diff --git a/train.py b/train.py
index 94e3b65..8e3336f 100644
--- a/train.py
+++ b/train.py
@@ -6,7 +6,7 @@
 
 from sklearn.metrics import precision_recall_fscore_support, accuracy_score
 import transformers
-from transformers import AutoTokenizer, AutoConfig
+from transformers import AutoTokenizer, AutoConfig, AutoProcessor
 from torch.utils.data import WeightedRandomSampler
 from packaging import version
 
@@ -16,6 +16,7 @@
 from gliclass import GLiClassModelConfig, GLiClassModel
 from gliclass.training import TrainingArguments, Trainer
 from gliclass.data_processing import DataCollatorWithPadding, GLiClassDataset, AugmentationConfig
+from gliclass.multimodal import get_tokenizer
 
 class CustomTrainer(Trainer):
     """Trainer with weighted random sampling support."""
@@ -89,6 +90,24 @@ def load_dataset(data_path: str) -> list:
     return data
 
 
+def recurrent_kwargs(args):
+    """Recurrent reasoning settings for the decoder-kv scorer; unset flags keep config/checkpoint values."""
+    kwargs = dict(
+        recurrent_steps=args.recurrent_steps,
+        recurrent_min_steps=args.recurrent_min_steps,
+        recurrent_inference_max_steps=args.recurrent_inference_max_steps,
+        recurrent_halt_threshold=args.recurrent_halt_threshold,
+        recurrent_improvement_coef=args.recurrent_improvement_coef,
+        recurrent_improvement_margin=args.recurrent_improvement_margin,
+        recurrent_bptt_steps=args.recurrent_bptt_steps,
+        recurrent_read_text=args.recurrent_read_text,
+        recurrent_confidence_coef=args.recurrent_confidence_coef,
+        recurrent_confidence_mode=args.recurrent_confidence_mode,
+        recurrent_confidence_margin=args.recurrent_confidence_margin,
+    )
+    return {key: value for key, value in kwargs.items() if value is not None}
+
+
 def main(args):
     device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')
 
@@ -98,11 +117,16 @@ def main(args):
             args.model_name, 
             focal_loss_alpha=args.focal_loss_alpha,
             focal_loss_gamma=args.focal_loss_gamma,
-            focal_loss_reduction=args.focal_loss_reduction
+            focal_loss_reduction=args.focal_loss_reduction,
+            **recurrent_kwargs(args),
         )
-        tokenizer = AutoTokenizer.from_pretrained(args.model_name)
+        # multi-modal decoder-kv: the processor tokenizes text and prepares images / audio
+        tokenizer_class = AutoProcessor if model.config.multimodal else AutoTokenizer
+        tokenizer = tokenizer_class.from_pretrained(args.model_name)
     else:
-        tokenizer = AutoTokenizer.from_pretrained(args.encoder_model_name)
+        tokenizer_class = AutoProcessor if args.multimodal else AutoTokenizer
+        tokenizer = tokenizer_class.from_pretrained(args.encoder_model_name)
+        text_tokenizer = get_tokenizer(tokenizer)
         encoder_config = AutoConfig.from_pretrained(args.encoder_model_name)
 
         label_model_config = None
@@ -114,9 +138,9 @@ def main(args):
             encoder_model=args.encoder_model_name,
             label_model_name=args.label_model_name,
             label_model_config=label_model_config,
-            class_token_index=len(tokenizer),
-            text_token_index=len(tokenizer)+1,
-            example_token_index=len(tokenizer)+2,
+            class_token_index=len(text_tokenizer),
+            text_token_index=len(text_tokenizer)+1,
+            example_token_index=len(text_tokenizer)+2,
             pooling_strategy=args.pooler_type,
             class_token_pooling=args.class_token_pooling,
             scorer_type=args.scorer_type,
@@ -135,14 +159,20 @@ def main(args):
             shuffle_labels=args.shuffle_labels,
             dropout=args.dropout,
             use_segment_embeddings=args.use_segment_embeddings,
+            multimodal=args.multimodal,
+            freeze_media_encoders=not args.train_media_encoders,
+            **recurrent_kwargs(args),
         )
+        if args.architecture_type == 'decoder-kv':
+            # <<SEP>> is added right after <<LABEL>> below
+            glicalss_config.sep_token_index = len(text_tokenizer) + 1
 
         model = GLiClassModel(glicalss_config, from_pretrained=True).to(dtype=torch.float32)
 
-        if args.architecture_type in {'uni-encoder', 'bi-encoder-fused', 'encoder-decoder'}:
+        if args.architecture_type in {'uni-encoder', 'bi-encoder-fused', 'encoder-decoder', 'decoder-kv'}:
             new_words = ["<<LABEL>>", "<<SEP>>", "<<EXAMPLE>>"]
-            tokenizer.add_tokens(new_words, special_tokens=True)
-            model.resize_token_embeddings(len(tokenizer))
+            text_tokenizer.add_tokens(new_words, special_tokens=True)
+            model.resize_token_embeddings(len(text_tokenizer))
 
     model.to(device)
 
@@ -224,8 +254,12 @@ def main(args):
         weight_decay=args.encoder_weight_decay,
         others_lr=args.others_lr,
         others_weight_decay=args.others_weight_decay,
+        recurrent_lr=args.recurrent_lr,
         lr_scheduler_type=args.lr_scheduler_type,
-        warmup_ratio=args.warmup_ratio,
+        optim=args.optim,
+        # transformers v5 dropped warmup_ratio; a float warmup_steps is a ratio there
+        **({"warmup_steps": args.warmup_ratio} if version.parse(transformers.__version__) >= version.parse("5.0.0")
+           else {"warmup_ratio": args.warmup_ratio}),
         per_device_train_batch_size=args.batch_size,
         per_device_eval_batch_size=args.batch_size,
         gradient_accumulation_steps=args.gradient_accumulation_steps,
@@ -233,10 +267,11 @@ def main(args):
         save_steps=args.save_steps,
         save_total_limit=args.save_total_limit,
         dataloader_num_workers=args.num_workers,
-        logging_steps=100,
+        logging_steps=args.logging_steps,
         use_cpu=False,
         report_to="none",
         fp16=args.fp16,
+        bf16=args.bf16,
         # EWC parameters
         use_ewc=args.use_ewc,
         ewc_lambda=args.ewc_lambda,
@@ -321,6 +356,33 @@ def compute_metrics_fn(p):
     parser.add_argument('--dropout', type=float, default=0.3)
     parser.add_argument('--shuffle_labels', type=bool, default=True)
     parser.add_argument('--use_segment_embeddings', type=bool, default=False)
+    parser.add_argument('--multimodal', action='store_true',
+                        help='decoder-kv: keep the vision / audio encoders of Qwen3.5 / Gemma 4; '
+                             'data items may then carry "images" and "audio" lists')
+    parser.add_argument('--train_media_encoders', action='store_true',
+                        help='Also train the vision / audio encoders (frozen by default)')
+
+    # Recurrent hidden reasoning (decoder-kv scorer)
+    parser.add_argument('--recurrent_steps', type=int, default=None,
+                        help='Max scorer reasoning steps during training (default 1 = disabled)')
+    parser.add_argument('--recurrent_min_steps', type=int, default=None,
+                        help='Min training steps; depth is sampled from [min, max] per batch (default 1)')
+    parser.add_argument('--recurrent_inference_max_steps', type=int, default=None,
+                        help='Max reasoning steps at inference (default: recurrent_steps)')
+    parser.add_argument('--recurrent_halt_threshold', type=float, default=None,
+                        help='Stop inference recurrence once label probabilities change less than this')
+    parser.add_argument('--recurrent_improvement_coef', type=float, default=None)
+    parser.add_argument('--recurrent_improvement_margin', type=float, default=None)
+    parser.add_argument('--recurrent_bptt_steps', type=int, default=None)
+    parser.add_argument('--recurrent_read_text', action='store_true', default=None,
+                        help='Let every recurrent step cross-attend to the text (default: labels only)')
+    parser.add_argument('--recurrent_confidence_coef', type=float, default=None,
+                        help='Weight of the confidence (low-entropy) loss on recurrent steps (default 0 = off)')
+    parser.add_argument('--recurrent_confidence_mode', type=str, default=None, choices=['relative', 'absolute'])
+    parser.add_argument('--recurrent_confidence_margin', type=float, default=None,
+                        help='relative mode: each step must cut entropy by this fraction vs the previous step')
+    parser.add_argument('--recurrent_lr', type=float, default=None,
+                        help='Learning rate for the recurrent reasoning cell (default: others_lr)')
 
     # Training arguments
     parser.add_argument('--num_epochs', type=int, default=3)
@@ -332,11 +394,15 @@ def compute_metrics_fn(p):
     parser.add_argument('--others_weight_decay', type=float, default=0.01)
     parser.add_argument('--warmup_ratio', type=float, default=0.05)
     parser.add_argument('--lr_scheduler_type', type=str, default='linear')
+    parser.add_argument('--optim', type=str, default='adamw_torch',
+                        help='transformers optimizer name, e.g. adafactor to fit larger backbones in memory')
     parser.add_argument('--max_length', type=int, default=1024)
     parser.add_argument('--save_steps', type=int, default=1000)
     parser.add_argument('--save_total_limit', type=int, default=3)
     parser.add_argument('--num_workers', type=int, default=12)
     parser.add_argument('--fp16', type=bool, default=False)
+    parser.add_argument('--bf16', action='store_true')
+    parser.add_argument('--logging_steps', type=int, default=100)
     
     # Augmentation parameters
     parser.add_argument('--enable_augmentation', type=bool, default=True)

From dc5e333e8a6744d6e674f80fcfc7d198ccb26c6a Mon Sep 17 00:00:00 2001
From: Ingvar <igor.stepanov2000@gmail.com>
Date: Fri, 2 Oct 2026 16:58:21 +0300
Subject: [PATCH 06/12] make calibration work with multi-modal models

---
 calibrate.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/calibrate.py b/calibrate.py
index c1135b3..7295f66 100644
--- a/calibrate.py
+++ b/calibrate.py
@@ -8,6 +8,8 @@
 Example:
     python calibrate.py --model_path models/checkpoint-1000 --data_path data/val.json
         --save_path models/checkpoint-1000-calibrated
+
+For multi-modal decoder-kv models, data items may carry "images" / "audio" lists as in training.
 """
 
 import os
@@ -18,7 +20,7 @@
 import argparse
 
 import torch
-from transformers import AutoTokenizer
+from transformers import AutoTokenizer, AutoProcessor
 from torch.utils.data import DataLoader
 
 from gliclass import GLiClassModel
@@ -57,7 +59,9 @@ def main(args):
     device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
 
     model = GLiClassModel.from_pretrained(args.model_path, torch_dtype=DTYPES[args.dtype]).to(device)
-    tokenizer = AutoTokenizer.from_pretrained(args.model_path)
+    # multi-modal decoder-kv: the processor also prepares the images / audio of each data item
+    tokenizer_class = AutoProcessor if getattr(model.config, "multimodal", False) else AutoTokenizer
+    tokenizer = tokenizer_class.from_pretrained(args.model_path)
     model.eval()
     model.requires_grad_(False)
 

From 3d2276239faedfd5094f6e2479d5f62b2d80c9d4 Mon Sep 17 00:00:00 2001
From: Ingvar <igor.stepanov2000@gmail.com>
Date: Fri, 2 Oct 2026 19:51:25 +0300
Subject: [PATCH 07/12] update tests

---
 tests/test_calibration.py | 277 +++++++++++++++++++++++++++++++++++
 tests/test_decoder_kv.py  | 294 ++++++++++++++++++++++++++++++++++++++
 tests/test_multimodal.py  |  94 ++++++++++++
 3 files changed, 665 insertions(+)
 create mode 100644 tests/test_calibration.py
 create mode 100644 tests/test_multimodal.py

diff --git a/tests/test_calibration.py b/tests/test_calibration.py
new file mode 100644
index 0000000..294e53c
--- /dev/null
+++ b/tests/test_calibration.py
@@ -0,0 +1,277 @@
+"""Tests for the post-hoc calibrator."""
+
+import math
+
+import torch
+import pytest
+from torch.utils.data import DataLoader
+
+from gliclass import GLiClassModel, GLiClassModelConfig
+from gliclass.pipeline import BaseZeroShotClassificationPipeline
+from gliclass.calibration import (
+    CalibrationExample,
+    GLiClassCalibrator,
+    fit_calibrator,
+    logit_features,
+    calibration_loss,
+    fit_global_scaling,
+    evaluate_calibrator,
+    collect_calibration_examples,
+)
+
+
+@pytest.mark.parametrize("use_bias", [False, True])
+def test_calibrator_starts_at_global_scaling_and_ignores_padding(use_bias):
+    calibrator = GLiClassCalibrator(8, hidden_size=16, beta_max=3.0, use_bias=use_bias)
+    calibrator.reset_to_global_scaling(0.7, -1.5)
+    mask = torch.tensor([[1, 1, 0], [1, 0, 0]])
+    beta, bias = calibrator(torch.randn(2, 8), torch.randn(2, 3, 8), torch.randn(2, 3), mask)
+
+    torch.testing.assert_close(beta[mask.bool()], torch.full((3,), 0.7))
+    torch.testing.assert_close(bias[mask.bool()], torch.full((3,), -1.5 if use_bias else 0.0))
+    torch.testing.assert_close(beta[~mask.bool()], torch.ones(3))
+    torch.testing.assert_close(bias[~mask.bool()], torch.zeros(3))
+
+
+def test_beta_is_bounded_and_never_flips_decisions_without_bias():
+    torch.manual_seed(0)
+    calibrator = GLiClassCalibrator(8, hidden_size=16, beta_max=2.0)
+    torch.nn.init.normal_(calibrator.mlp[-1].weight, std=1.0)
+    logits = torch.randn(4, 5) * 5
+    beta, bias = calibrator(torch.randn(4, 8), torch.randn(4, 5, 8), logits, torch.ones(4, 5))
+
+    assert (bias == 0).all()
+    assert (beta > 0).all() and (beta < 2.0).all()
+    assert torch.equal(torch.sign(beta * logits), torch.sign(logits))
+
+
+def test_logit_features_ignore_padded_labels():
+    logits = torch.tensor([[2.0, -1.0, 100.0]])
+    mask = torch.tensor([[1, 1, 0]])
+    features = logit_features(logits, mask)
+
+    assert features[0, 2].abs().sum() == 0
+    torch.testing.assert_close(features[0, :2, 2], torch.tensor([0.0, -3.0]))  # z - max over real labels
+    torch.testing.assert_close(features[0, :2, 5], torch.full((2,), math.log(2)))
+
+
+def test_nll_pushes_beta_down_for_wrong_pairs_and_up_for_correct_ones():
+    logits = torch.tensor([[3.0, 3.0]])
+    targets = torch.tensor([[0.0, 1.0]])  # first pair is confidently wrong
+    beta = torch.ones(1, 2, requires_grad=True)
+    loss, _ = calibration_loss(
+        beta, torch.zeros(1, 2), logits, targets, torch.ones(1, 2, dtype=torch.bool), 1.0, prior_weight=0.0
+    )
+    loss.backward()
+
+    assert beta.grad[0, 0] > 0  # gradient descent lowers beta -> probability moves to 0.5
+    assert beta.grad[0, 1] < 0
+
+
+def make_synthetic_examples(num_texts, seed):
+    """Logits that are right on 'easy' labels and sign-flipped on labels marked hard in their representation."""
+    generator = torch.Generator().manual_seed(seed)
+    examples = []
+    for _ in range(num_texts):
+        num_labels = int(torch.randint(2, 6, (1,), generator=generator))
+        targets = torch.randint(0, 2, (num_labels,), generator=generator).float()
+        hard = torch.rand(num_labels, generator=generator) < 0.3
+        logits = (targets * 2 - 1) * 4.0
+        logits = torch.where(hard, -logits, logits)
+        label_repr = torch.randn(num_labels, 8, generator=generator) * 0.1
+        label_repr[:, 0] = hard.float() * 3
+        examples.append(CalibrationExample(torch.randn(8, generator=generator), label_repr, logits, targets))
+    return examples
+
+
+def test_fit_learns_to_lower_confidence_where_backbone_is_wrong():
+    torch.manual_seed(0)
+    train, held_out = make_synthetic_examples(400, seed=1), make_synthetic_examples(100, seed=2)
+    base_beta, _ = fit_global_scaling(train)
+    calibrator = fit_calibrator(
+        GLiClassCalibrator(8, hidden_size=32),
+        train,
+        held_out,
+        "cpu",
+        base_beta=base_beta,
+        num_epochs=30,
+        lr=3e-3,
+        log_fn=lambda *_: None,
+    )
+    report = evaluate_calibrator(calibrator, held_out, "cpu", base_beta)
+
+    assert report["calibrated"]["nll"] < report["global_scaling"]["nll"] < report["uncalibrated"]["nll"]
+    assert report["calibrated"]["accuracy"] == pytest.approx(report["uncalibrated"]["accuracy"])
+
+    with torch.no_grad():
+        example = held_out[0]
+        hard = example.label_repr[:, 0] > 1
+        beta = calibrator(
+            example.text_repr[None], example.label_repr[None], example.logits[None], torch.ones(1, len(hard))
+        )[0][0]
+    if hard.any() and (~hard).any():
+        assert beta[hard].max() < beta[~hard].min()
+
+
+def make_overconfident_examples(num_texts, seed):
+    """A backbone that calls most negatives positive: 1 positive among 8 labels, negatives at logit ~ +2,
+    positives at ~ +5. Ranking is informative, but every logit is positive.
+    """
+    generator = torch.Generator().manual_seed(seed)
+    examples = []
+    for _ in range(num_texts):
+        targets = torch.zeros(8)
+        targets[int(torch.randint(0, 8, (1,), generator=generator))] = 1.0
+        logits = 2.0 + 3.0 * targets + torch.randn(8, generator=generator) * 0.7
+        label_repr = torch.randn(8, 8, generator=generator) * 0.1
+        examples.append(CalibrationExample(torch.randn(8, generator=generator), label_repr, logits, targets))
+    return examples
+
+
+def test_global_platt_scaling_recovers_negative_bias():
+    beta, bias = fit_global_scaling(make_overconfident_examples(300, seed=0), fit_bias=True)
+    assert beta > 0 and bias < 0
+    _, no_bias = fit_global_scaling(make_overconfident_examples(300, seed=0), fit_bias=False)
+    assert no_bias == 0.0
+
+
+def test_bias_fixes_base_rate_miscalibration_that_scaling_cannot():
+    torch.manual_seed(0)
+    train, held_out = make_overconfident_examples(400, seed=1), make_overconfident_examples(100, seed=2)
+    reports = {}
+    for use_bias in (False, True):
+        base_beta, base_bias = fit_global_scaling(train, fit_bias=use_bias)
+        calibrator = fit_calibrator(
+            GLiClassCalibrator(8, hidden_size=32, use_bias=use_bias),
+            train,
+            held_out,
+            "cpu",
+            base_beta=base_beta,
+            base_bias=base_bias,
+            num_epochs=10,
+            lr=3e-3,
+            log_fn=lambda *_: None,
+        )
+        reports[use_bias] = evaluate_calibrator(calibrator, held_out, "cpu", base_beta, base_bias)
+
+    # without the bias every logit stays positive, so accuracy is stuck at the positive rate
+    assert reports[False]["calibrated"]["accuracy"] == pytest.approx(reports[False]["uncalibrated"]["accuracy"])
+    assert reports[True]["calibrated"]["accuracy"] > 0.9
+    assert reports[True]["calibrated"]["nll"] < 0.5 * reports[False]["calibrated"]["nll"]
+
+
+@pytest.fixture
+def tiny_uni_encoder():
+    torch.manual_seed(0)
+    encoder_config = {
+        "model_type": "deberta-v2",
+        "vocab_size": 64,
+        "hidden_size": 32,
+        "num_hidden_layers": 1,
+        "num_attention_heads": 2,
+        "intermediate_size": 64,
+        "max_position_embeddings": 64,
+    }
+    config = GLiClassModelConfig(
+        encoder_config=encoder_config,
+        encoder_model="tiny-deberta",  # never loaded: the backbone is built from encoder_config
+        class_token_index=60,
+        text_token_index=61,
+        architecture_type="uni-encoder",
+        problem_type="multi_label_classification",
+        max_num_classes=4,
+    )
+    model = GLiClassModel(config).eval()
+    # two texts, the second with only two labels
+    input_ids = torch.tensor(
+        [
+            [60, 5, 60, 6, 60, 7, 61, 10, 11, 12],
+            [60, 5, 60, 6, 61, 13, 14, 0, 0, 0],
+        ]
+    )
+    attention_mask = (input_ids != 0).long()
+    return model, {"input_ids": input_ids, "attention_mask": attention_mask, "max_num_classes": 3}
+
+
+def test_model_returns_calibrated_and_initial_logits(tiny_uni_encoder, tmp_path):
+    model, inputs = tiny_uni_encoder
+    with torch.no_grad():
+        raw = model(**inputs).logits
+    assert model(**inputs).uncalibrated_logits is None
+
+    calibrator = model.add_calibrator(hidden_size=16, use_bias=True)
+    calibrator.reset_to_global_scaling(0.5, -1.0)
+    with torch.no_grad():
+        output = model(**inputs)
+
+    torch.testing.assert_close(output.uncalibrated_logits, raw)
+    torch.testing.assert_close(output.class_mask, torch.tensor([[1, 1, 1], [1, 1, 0]]))
+    torch.testing.assert_close(output.logits[0], 0.5 * raw[0] - 1.0)
+    torch.testing.assert_close(output.calibration_biases[1], torch.tensor([-1.0, -1.0, 0.0]))
+    assert output.text_embeddings is None and output.class_embeddings is None
+
+    model.save_pretrained(tmp_path)
+    reloaded = GLiClassModel.from_pretrained(tmp_path).eval()
+    assert reloaded.calibrator is not None and reloaded.calibrator.use_bias
+    with torch.no_grad():
+        torch.testing.assert_close(reloaded(**inputs).logits, output.logits)
+
+
+def test_collect_calibration_examples_uses_real_labels_only(tiny_uni_encoder):
+    model, inputs = tiny_uni_encoder
+    batch = dict(inputs, labels=torch.tensor([[1.0, 0.0, 1.0], [0.0, 1.0, 0.0]]))
+    examples = collect_calibration_examples(model, DataLoader([batch], batch_size=None), "cpu", show_progress=False)
+
+    assert [len(ex.logits) for ex in examples] == [3, 2]
+    torch.testing.assert_close(examples[1].targets, torch.tensor([0.0, 1.0]))
+    assert examples[0].label_repr.shape == (3, model.config.encoder_config.hidden_size)
+
+
+def test_postprocess_reports_initial_score():
+    predictions, all_scores = BaseZeroShotClassificationPipeline._postprocess_logits(
+        torch.tensor([0.5, -2.0]),
+        ["a", "b"],
+        "multi-label",
+        0.5,
+        initial_logits=torch.tensor([3.0, -4.0]),
+    )
+    assert predictions == [
+        {"label": "a", "score": pytest.approx(torch.sigmoid(torch.tensor(0.5)).item()),
+         "initial_score": pytest.approx(torch.sigmoid(torch.tensor(3.0)).item())}
+    ]
+    assert set(all_scores) == {"a", "b"}
+
+    predictions, _ = BaseZeroShotClassificationPipeline._postprocess_logits(
+        torch.tensor([0.5, -2.0]), ["a", "b"], "multi-label", 0.5
+    )
+    assert "initial_score" not in predictions[0]
+
+
+@pytest.mark.parametrize("halt_threshold", [None, 1e-3, 1.1])
+def test_decoder_kv_representations_match_final_logits(halt_threshold):
+    from types import SimpleNamespace
+
+    from gliclass.scorers import DecoderKVScorer
+
+    torch.manual_seed(0)
+    config = SimpleNamespace(
+        hidden_size=64,
+        scorer_encoder_num_layers=1,
+        scorer_mlp_hidden_size=32,
+        dropout=0.0,
+        normalize_features=False,
+        sep_token_index=98,
+        class_token_index=99,
+        problem_type="multi_label_classification",
+        recurrent_steps=3,
+    )
+    scorer = DecoderKVScorer(config).eval()
+    hidden = torch.randn(2, 5, 64)
+    input_ids = torch.tensor([[10, 99, 11, 99, 98], [12, 99, 98, 0, 0]])
+    steps, _, (text_repr, label_repr) = scorer.recurrent_forward(
+        hidden, input_ids, input_ids.ne(0).long(), num_steps=4, halt_threshold=halt_threshold,
+        return_representations=True,
+    )
+
+    combined = torch.cat([text_repr.unsqueeze(1).expand_as(label_repr), label_repr], dim=-1)
+    torch.testing.assert_close(scorer.mlp(combined).squeeze(-1), steps[-1])
diff --git a/tests/test_decoder_kv.py b/tests/test_decoder_kv.py
index d907a03..210e969 100644
--- a/tests/test_decoder_kv.py
+++ b/tests/test_decoder_kv.py
@@ -3,6 +3,7 @@
 from types import SimpleNamespace
 
 import torch
+import pytest
 
 from gliclass.model import GLiClassOutput, GLiClassDecoderKV
 
@@ -81,3 +82,296 @@ def test_extract_label_section_preserves_fallback_and_gradients():
     padded_hidden.sum().backward()
     torch.testing.assert_close(hidden_states.grad[:, :3], torch.ones_like(hidden_states[:, :3]))
     torch.testing.assert_close(hidden_states.grad[:, 3:], torch.zeros_like(hidden_states[:, 3:]))
+
+
+def make_recurrent_stub(**overrides):
+    config = SimpleNamespace(
+        problem_type="multi_label_classification",
+        focal_loss_alpha=-1,
+        focal_loss_gamma=-1,
+        ignore_index=-100,
+        recurrent_steps=4,
+        recurrent_min_steps=2,
+        recurrent_inference=True,
+        recurrent_inference_max_steps=6,
+        recurrent_halt_threshold=0.01,
+        recurrent_improvement_coef=0.0,
+        recurrent_improvement_margin=0.0,
+    )
+    vars(config).update(overrides)
+    stub = SimpleNamespace(config=config, training=True, scorer=SimpleNamespace(recurrent=True))
+    for name in ("_resolve_recurrence", "_per_sample_loss", "_per_sample_entropy", "get_recurrent_loss"):
+        setattr(stub, name, getattr(GLiClassDecoderKV, name).__get__(stub))
+    return stub
+
+
+def make_scorer_config(**overrides):
+    config = SimpleNamespace(
+        hidden_size=64,
+        scorer_encoder_num_layers=1,
+        scorer_mlp_hidden_size=32,
+        dropout=0.0,
+        normalize_features=False,
+        sep_token_index=98,
+        class_token_index=99,
+        problem_type="multi_label_classification",
+        recurrent_steps=1,
+    )
+    vars(config).update(overrides)
+    return config
+
+
+def make_scorer_inputs():
+    hidden = torch.randn(2, 5, 64)
+    input_ids = torch.tensor([[10, 99, 11, 99, 98], [12, 99, 98, 0, 0]])
+    return hidden, input_ids, input_ids.ne(0).long()
+
+
+def test_recurrent_scorer_first_step_matches_plain_scorer():
+    from gliclass.scorers import DecoderKVScorer
+
+    torch.manual_seed(0)
+    plain = DecoderKVScorer(make_scorer_config()).eval()
+    recurrent = DecoderKVScorer(make_scorer_config(recurrent_steps=3)).eval()
+    recurrent.load_state_dict(plain.state_dict(), strict=False)
+    hidden, input_ids, mask = make_scorer_inputs()
+
+    expected = plain(hidden, input_ids, mask)
+    steps, steps_taken = recurrent.recurrent_forward(hidden, input_ids, mask, num_steps=3, return_all_steps=True)
+
+    assert len(steps) == 3
+    torch.testing.assert_close(steps[0], expected)
+    torch.testing.assert_close(recurrent(hidden, input_ids, mask, num_steps=3), steps[-1])
+    assert steps_taken.tolist() == [3, 3]
+
+
+def test_recurrent_scorer_extrapolates_beyond_training_depth():
+    from gliclass.scorers import DecoderKVScorer
+
+    scorer = DecoderKVScorer(make_scorer_config(recurrent_steps=2)).eval()
+    hidden, input_ids, mask = make_scorer_inputs()
+    steps, steps_taken = scorer.recurrent_forward(hidden, input_ids, mask, num_steps=10, return_all_steps=True)
+    assert len(steps) == 10
+    assert steps_taken.tolist() == [10, 10]
+    assert all(torch.isfinite(logits).all() for logits in steps)
+
+
+def test_recurrent_scorer_halts_when_probabilities_stop_changing():
+    from gliclass.scorers import DecoderKVScorer
+
+    scorer = DecoderKVScorer(make_scorer_config(recurrent_steps=2)).eval()
+    hidden, input_ids, mask = make_scorer_inputs()
+    # probabilities can never change by more than 1, so every example halts after step 2
+    steps, steps_taken = scorer.recurrent_forward(
+        hidden, input_ids, mask, num_steps=10, halt_threshold=1.1, return_all_steps=True
+    )
+    assert len(steps) == 2
+    assert steps_taken.tolist() == [2, 2]
+
+
+def test_plain_scorer_rejects_multiple_steps():
+    from gliclass.scorers import DecoderKVScorer
+
+    scorer = DecoderKVScorer(make_scorer_config())
+    hidden, input_ids, mask = make_scorer_inputs()
+    with pytest.raises(ValueError):
+        scorer(hidden, input_ids, mask, num_steps=2)
+
+
+def test_resolve_recurrence_training_inference_and_overrides():
+    stub = make_recurrent_stub()
+    sampled = {stub._resolve_recurrence()[0] for _ in range(200)}
+    assert sampled == {2, 3, 4}
+    assert stub._resolve_recurrence()[1] is None
+    assert stub._resolve_recurrence(use_recurrence=False) == (1, None)
+
+    stub.training = False
+    assert stub._resolve_recurrence() == (6, 0.01)
+    assert stub._resolve_recurrence(max_recurrent_steps=12, halt_threshold=0.0) == (12, 0.0)
+    assert stub._resolve_recurrence(use_recurrence=False) == (1, None)
+    stub.config.recurrent_inference = False
+    assert stub._resolve_recurrence() == (1, None)
+    assert stub._resolve_recurrence(use_recurrence=True) == (6, 0.01)
+
+    stub.scorer.recurrent = False
+    assert stub._resolve_recurrence(use_recurrence=True) == (1, None)
+
+
+def test_improvement_term_starts_at_step_three_and_spares_earlier_steps():
+    labels = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
+    # later steps are worse than earlier ones, so the hinge is active
+    values = [
+        [[2.0, -2.0], [-2.0, 2.0]],
+        [[1.0, -1.0], [-1.0, 1.0]],
+        [[0.0, 0.0], [0.0, 0.0]],
+    ]
+
+    def step_grads(coef, num_steps):
+        stub = make_recurrent_stub(recurrent_improvement_coef=coef)
+        step_logits = [torch.tensor(value, requires_grad=True) for value in values[:num_steps]]
+        loss, step_losses, _ = stub.get_recurrent_loss(step_logits, labels)
+        loss.backward()
+        assert step_losses.shape == (num_steps,)
+        return [logits.grad for logits in step_logits]
+
+    # two steps: no improvement term at all
+    for with_hinge, without in zip(step_grads(1.0, 2), step_grads(0.0, 2)):
+        torch.testing.assert_close(with_hinge, without)
+
+    with_hinge, without = step_grads(1.0, 3), step_grads(0.0, 3)
+    torch.testing.assert_close(with_hinge[0], without[0])
+    torch.testing.assert_close(with_hinge[1], without[1])  # baseline of the step-3 hinge is stop-gradient
+    assert with_hinge[2].abs().sum() > without[2].abs().sum()
+
+
+def test_recurrent_loss_ignores_padded_labels():
+    stub = make_recurrent_stub()
+    labels = torch.tensor([[1.0, -100.0]])
+    step_logits = [torch.tensor([[2.0, 50.0]]), torch.tensor([[3.0, -50.0]])]
+    _, step_losses, _ = stub.get_recurrent_loss(step_logits, labels)
+    expected = torch.nn.functional.binary_cross_entropy_with_logits(
+        torch.tensor([2.0, 3.0]), torch.ones(2), reduction="none"
+    )
+    torch.testing.assert_close(step_losses, expected)
+
+
+def test_read_text_scorer_requires_and_uses_text():
+    from gliclass.scorers import DecoderKVScorer
+
+    torch.manual_seed(0)
+    plain = DecoderKVScorer(make_scorer_config()).eval()
+    scorer = DecoderKVScorer(make_scorer_config(recurrent_steps=3, recurrent_read_text=True)).eval()
+    scorer.load_state_dict(plain.state_dict(), strict=False)
+    hidden, input_ids, mask = make_scorer_inputs()
+    text = torch.randn(2, 7, 64)
+    text_mask = torch.tensor([[1] * 7, [1] * 4 + [0] * 3])
+
+    with pytest.raises(ValueError):
+        scorer(hidden, input_ids, mask, num_steps=2)
+
+    steps, _ = scorer.recurrent_forward(
+        hidden,
+        input_ids,
+        mask,
+        num_steps=3,
+        return_all_steps=True,
+        text_hidden_states=text,
+        text_attention_mask=text_mask,
+    )
+    torch.testing.assert_close(steps[0], plain(hidden, input_ids, mask))
+
+    # recurrent steps depend on the text; padded text positions do not
+    other_text = text.clone()
+    other_text[0] += torch.randn(7, 64)
+    other_text[1, 4:] += 100.0
+    changed, _ = scorer.recurrent_forward(
+        hidden,
+        input_ids,
+        mask,
+        num_steps=3,
+        return_all_steps=True,
+        text_hidden_states=other_text,
+        text_attention_mask=text_mask,
+    )
+    assert not torch.allclose(changed[-1][0], steps[-1][0])
+    torch.testing.assert_close(changed[-1][1], steps[-1][1])
+
+
+def test_extract_label_section_returns_text_section():
+    model = make_decoder_kv_stub()
+    hidden_states = torch.arange(2 * 10 * 3, dtype=torch.float32).reshape(2, 10, 3)
+    input_ids = torch.tensor(
+        [
+            [10, 98, 11, 98, 20, 99, 98, 0, 0, 99],
+            [98, 12, 98, 21, 22, 99, 23, 99, 98, 0],
+        ]
+    )
+    attention_mask = torch.tensor(
+        [
+            [1, 1, 1, 1, 1, 1, 1, 0, 0, 0],
+            [1, 1, 1, 1, 1, 1, 1, 1, 1, 0],
+        ]
+    )
+
+    *_, text_hidden, text_mask = GLiClassDecoderKV._extract_label_section(
+        model, hidden_states, input_ids, attention_mask, return_text_section=True
+    )
+
+    torch.testing.assert_close(text_mask, attention_mask.new_tensor([[1, 1, 1, 1], [1, 1, 1, 0]]))
+    torch.testing.assert_close(text_hidden[0], hidden_states[0, :4])
+    torch.testing.assert_close(text_hidden[1, :3], hidden_states[1, :3])
+
+
+def test_recurrent_lr_moves_reasoning_cell_into_own_groups():
+    from gliclass.training import Trainer
+
+    model = torch.nn.Module()
+    model.backbone = torch.nn.Linear(2, 2)
+    model.reasoning_cell = torch.nn.Linear(2, 2)
+    groups = [
+        {"params": [model.backbone.weight, model.reasoning_cell.weight], "weight_decay": 0.01, "lr": 1e-5},
+        {"params": [model.backbone.bias, model.reasoning_cell.bias], "weight_decay": 0.0, "lr": 1e-5},
+    ]
+    stub = SimpleNamespace(args=SimpleNamespace(recurrent_lr=1e-3))
+
+    split = Trainer._split_recurrent_groups(stub, model, groups)
+
+    assert [len(group["params"]) for group in split] == [1, 1, 1, 1]
+    assert [group["lr"] for group in split] == [1e-5, 1e-3, 1e-5, 1e-3]
+    assert [group["weight_decay"] for group in split] == [0.01, 0.01, 0.0, 0.0]
+    assert split[1]["params"][0] is model.reasoning_cell.weight
+
+
+def test_per_sample_entropy_is_binary_entropy_over_real_labels():
+    stub = make_recurrent_stub()
+    logits = torch.tensor([[0.0, 30.0, 5.0]])
+    valid = torch.tensor([[True, True, False]])
+    entropy = stub._per_sample_entropy(logits, valid)
+    # p=0.5 -> log 2, p~1 -> 0; the padded third label is ignored
+    torch.testing.assert_close(entropy, torch.tensor([torch.log(torch.tensor(2.0)).item() / 2]), atol=1e-6, rtol=0)
+
+
+def test_relative_confidence_loss_sharpens_later_step_only():
+    labels = torch.tensor([[1.0, 0.0]])
+    valid = torch.ones(1, 2, dtype=torch.bool)
+
+    def grads(coef, mode="relative"):
+        stub = make_recurrent_stub(recurrent_confidence_coef=coef, recurrent_confidence_mode=mode)
+        # step 2 is less certain than step 1, so the relative hinge is active
+        step_logits = [torch.tensor([[2.0, -2.0]], requires_grad=True), torch.tensor([[0.5, -0.5]], requires_grad=True)]
+        loss, _, entropies = stub.get_recurrent_loss(step_logits, labels, valid)
+        loss.backward()
+        return [logits.grad for logits in step_logits], entropies
+
+    without, entropies = grads(0.0)
+    with_conf, _ = grads(1.0)
+    assert entropies[1] > entropies[0]
+    # earlier step is a stop-gradient baseline
+    torch.testing.assert_close(with_conf[0], without[0])
+    # gradient descent moves step-2 logits further from 0 (probabilities away from 0.5)
+    extra = with_conf[1] - without[1]
+    assert extra[0, 0] < 0 and extra[0, 1] > 0
+
+    absolute, _ = grads(1.0, mode="absolute")
+    torch.testing.assert_close(absolute[0], without[0])
+    assert (absolute[1] - without[1])[0, 0] < 0
+
+
+def test_confidence_loss_inactive_when_later_step_already_more_certain():
+    stub = make_recurrent_stub(recurrent_confidence_coef=1.0)
+    labels = torch.tensor([[1.0, 0.0]])
+    step_logits = [torch.tensor([[0.5, -0.5]]), torch.tensor([[3.0, -3.0]])]
+    with_conf, _, _ = stub.get_recurrent_loss(step_logits, labels)
+    stub.config.recurrent_confidence_coef = 0.0
+    without, _, _ = stub.get_recurrent_loss(step_logits, labels)
+    torch.testing.assert_close(with_conf, without)
+
+
+def test_single_label_losses_accept_class_index_targets():
+    stub = make_recurrent_stub(problem_type="single_label_classification", recurrent_confidence_coef=0.0)
+    stub.get_loss = GLiClassDecoderKV.get_loss.__get__(stub)
+    labels = torch.tensor([2, 0])
+    logits = torch.tensor([[0.0, 1.0, 3.0], [2.0, 0.0, 0.0]])
+    torch.testing.assert_close(stub.get_loss(logits, labels), torch.nn.functional.cross_entropy(logits, labels))
+    _, step_losses, _ = stub.get_recurrent_loss([logits, logits], labels)
+    torch.testing.assert_close(step_losses, torch.nn.functional.cross_entropy(logits, labels).repeat(2))
diff --git a/tests/test_multimodal.py b/tests/test_multimodal.py
new file mode 100644
index 0000000..911a0c9
--- /dev/null
+++ b/tests/test_multimodal.py
@@ -0,0 +1,94 @@
+"""Tests for gliclass.multimodal helpers (no model downloads)."""
+
+import pytest
+import torch
+
+from gliclass.data_processing import DataCollatorWithPadding, format_decoder_kv_sequence
+from gliclass.multimodal import collate_media, format_media_prefix, media_placeholder
+from gliclass.pipeline import BaseZeroShotClassificationPipeline
+
+
+class QwenLikeProcessor:
+    image_token = "<|image_pad|>"
+    vision_start_token = "<|vision_start|>"
+    vision_end_token = "<|vision_end|>"
+
+
+class GemmaLikeProcessor:
+    image_token = "<|image|>"
+    audio_token = "<|audio|>"
+
+
+def test_media_placeholders_per_backbone():
+    assert media_placeholder(QwenLikeProcessor(), "image") == "<|vision_start|><|image_pad|><|vision_end|>"
+    assert media_placeholder(GemmaLikeProcessor(), "image") == "<|image|>"
+    assert format_media_prefix(GemmaLikeProcessor(), num_images=2, num_audio=1) == "<|image|><|image|><|audio|>"
+    # no audio placeholder is needed (or required) when there is no audio
+    assert format_media_prefix(QwenLikeProcessor(), num_images=1) == "<|vision_start|><|image_pad|><|vision_end|>"
+    with pytest.raises(ValueError):
+        format_media_prefix(QwenLikeProcessor(), num_audio=1)
+
+
+def test_media_goes_into_context_before_text():
+    sequence = format_decoder_kv_sequence("text", ["a", "b"], prompt="P:", media="<|image|>")
+    assert sequence == "P:<|image|>text<<SEP>>a<<LABEL>>b<<LABEL>><<SEP>>"
+
+
+def test_collate_media_concatenates_items_in_batch_order():
+    samples = [
+        {"pixel_values": torch.ones(4, 3), "image_grid_thw": torch.tensor([[1, 2, 2]])},
+        {},
+        {"pixel_values": torch.full((2, 3), 2.0), "image_grid_thw": torch.tensor([[1, 1, 2]])},
+    ]
+    batch = collate_media(samples)
+    assert batch["pixel_values"].shape == (6, 3)
+    assert batch["pixel_values"][4:].eq(2).all()
+    assert batch["image_grid_thw"].tolist() == [[1, 2, 2], [1, 1, 2]]
+
+
+def test_collate_media_pads_trailing_dims():
+    samples = [
+        {
+            "input_features": torch.ones(1, 5, 2),
+            "input_features_mask": torch.ones(1, 5, dtype=torch.bool),
+            "image_position_ids": torch.zeros(1, 2, 2, dtype=torch.long),
+        },
+        {
+            "input_features": torch.ones(2, 3, 2),
+            "input_features_mask": torch.ones(2, 3, dtype=torch.bool),
+            "image_position_ids": torch.zeros(1, 4, 2, dtype=torch.long),
+        },
+    ]
+    batch = collate_media(samples)
+    assert batch["input_features"].shape == (3, 5, 2)
+    assert batch["input_features_mask"].sum(dim=1).tolist() == [5, 3, 3]
+    assert batch["image_position_ids"].shape == (2, 4, 2)
+    assert batch["image_position_ids"][0, 2:].eq(-1).all()
+
+
+def test_collator_handles_samples_with_and_without_media():
+    collator = DataCollatorWithPadding(device="cpu")
+    batch = collator(
+        [
+            {"input_ids": torch.tensor([1, 2]), "mm_token_type_ids": torch.tensor([0, 0]), "labels": torch.tensor([1.0])},
+            {
+                "input_ids": torch.tensor([1, 9, 9]),
+                "mm_token_type_ids": torch.tensor([0, 1, 1]),
+                "labels": torch.tensor([0.0]),
+                "pixel_values": torch.ones(8, 4),
+                "image_grid_thw": torch.tensor([[1, 2, 4]]),
+            },
+        ]
+    )
+    assert batch["input_ids"].shape == (2, 3)
+    assert batch["mm_token_type_ids"].tolist() == [[0, 0, 0], [0, 1, 1]]
+    assert batch["pixel_values"].shape == (8, 4)
+    assert batch["image_grid_thw"].tolist() == [[1, 2, 4]]
+
+
+def test_normalize_media_accepts_per_text_and_flat_lists():
+    normalize = BaseZeroShotClassificationPipeline._normalize_media
+    assert normalize(["img1", "img2"], 1) == [["img1", "img2"]]
+    assert normalize([["img1"], None, "img2"], 3) == [["img1"], [], ["img2"]]
+    with pytest.raises(ValueError):
+        normalize([["img1"]], 2)

From a9fb1d198037e4ed44f6a0b61cac6aa6044dff84 Mon Sep 17 00:00:00 2001
From: Ingvar <igor.stepanov2000@gmail.com>
Date: Fri, 2 Oct 2026 19:56:44 +0300
Subject: [PATCH 08/12] fix ruff errors

---
 gliclass/config.py          | 2 +-
 gliclass/data_processing.py | 4 ++--
 gliclass/pipeline.py        | 6 ++++--
 3 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/gliclass/config.py b/gliclass/config.py
index ecb4b87..6ebdd6f 100644
--- a/gliclass/config.py
+++ b/gliclass/config.py
@@ -20,7 +20,7 @@ class GLiClassModelConfig(PretrainedConfig):
     model_type = "GLiClass"
     is_composition = True
 
-    def __init__(
+    def __init__(  # noqa: PLR0917 - HF config, arguments mirror the saved config fields
         self,
         encoder_config=None,
         encoder_model=None,
diff --git a/gliclass/data_processing.py b/gliclass/data_processing.py
index 1da30e1..7a20d85 100644
--- a/gliclass/data_processing.py
+++ b/gliclass/data_processing.py
@@ -6,7 +6,7 @@
 from torch.utils.data import Dataset
 from torch.nn.utils.rnn import pad_sequence
 
-from .multimodal import MEDIA_KEYS, collate_media, is_processor, process_multimodal_with_budget
+from .multimodal import MEDIA_KEYS, is_processor, collate_media, process_multimodal_with_budget
 
 
 def format_decoder_kv_context(text: str, prompt: str = "", examples: str = "", media: str = "") -> str:
@@ -495,7 +495,7 @@ def _resolve_max_num_classes(self, batch):
 
     def __call__(self, batch):
         # Media tensors are indexed by image / audio item and only present in samples that have media
-        keys = [key for key in batch[0].keys() if key not in MEDIA_KEYS]
+        keys = [key for key in batch[0] if key not in MEDIA_KEYS]
         padded_batch = {key: [] for key in keys}
 
         for key in keys:
diff --git a/gliclass/pipeline.py b/gliclass/pipeline.py
index 44a5233..6095d5b 100644
--- a/gliclass/pipeline.py
+++ b/gliclass/pipeline.py
@@ -7,8 +7,8 @@
 
 from .model import GLiClassModel, GLiClassBiEncoder
 from .utils import retrieval_augmented_text
+from .multimodal import is_processor, collate_media, get_tokenizer, process_multimodal_with_budget
 from .data_processing import format_decoder_kv_labels, format_decoder_kv_context, format_decoder_kv_sequence
-from .multimodal import collate_media, get_tokenizer, is_processor, process_multimodal_with_budget
 
 
 def flatten_hierarchical_labels(
@@ -874,7 +874,9 @@ def _prepare_multimodal_inputs(self, texts, labels, same_labels, examples, promp
             text_prompt = self._format_prompt(prompt, i) or ""
             sample = process_multimodal_with_budget(
                 self.tokenizer,
-                lambda text, media: format_decoder_kv_context(text, text_prompt, examples_text, media),
+                lambda text, media, text_prompt=text_prompt, examples_text=examples_text: format_decoder_kv_context(
+                    text, text_prompt, examples_text, media
+                ),
                 text,
                 images[i] if images else [],
                 audio[i] if audio else [],

From 2412a8507268f76572640cdc0bc582129403c48c Mon Sep 17 00:00:00 2001
From: Ingvar <igor.stepanov2000@gmail.com>
Date: Sat, 3 Oct 2026 09:52:02 +0300
Subject: [PATCH 09/12] fix ruff error in calibration file

---
 gliclass/calibration.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/gliclass/calibration.py b/gliclass/calibration.py
index 34b3054..b7bbbd4 100644
--- a/gliclass/calibration.py
+++ b/gliclass/calibration.py
@@ -355,9 +355,7 @@ def eval_nll():
             for start in range(0, len(eval_examples), 256):
                 batch = collate_calibration_examples(eval_examples[start : start + 256])
                 batch = {key: value.to(device) for key, value in batch.items()}
-                beta, bias = calibrator(
-                    batch["text_repr"], batch["label_repr"], batch["logits"], batch["label_mask"]
-                )
+                beta, bias = calibrator(batch["text_repr"], batch["label_repr"], batch["logits"], batch["label_mask"])
                 _, nll = calibration_loss(
                     beta, bias, batch["logits"], batch["targets"], batch["label_mask"], base_beta, base_bias
                 )

From 0b5fd9f01f8a870af3bfb138fed599b33124b64b Mon Sep 17 00:00:00 2001
From: Ingvarstep <igor.stepanov2000@gmail.com>
Date: Tue, 6 Oct 2026 18:47:57 +0000
Subject: [PATCH 10/12] fix multi-modal training issues

---
 gliclass/config.py          |  7 +++
 gliclass/data_processing.py | 18 ++++++-
 gliclass/model.py           | 96 ++++++++++++++++++++++++-------------
 gliclass/multimodal.py      | 27 +++++++++--
 gliclass/scorers.py         | 25 ++++++++++
 gliclass/training.py        | 11 ++++-
 6 files changed, 144 insertions(+), 40 deletions(-)

diff --git a/gliclass/config.py b/gliclass/config.py
index 6ebdd6f..a75047d 100644
--- a/gliclass/config.py
+++ b/gliclass/config.py
@@ -60,6 +60,7 @@ def __init__(  # noqa: PLR0917 - HF config, arguments mirror the saved config fi
         dropout=0.1,
         use_segment_embeddings=False,
         scorer_encoder_num_layers=2,
+        scorer_full_sequence=False,
         recurrent_steps=1,
         recurrent_min_steps=1,
         recurrent_inference=True,
@@ -116,6 +117,9 @@ def __init__(  # noqa: PLR0917 - HF config, arguments mirror the saved config fi
 
         if hidden_size is None:
             self.hidden_size = text_config.hidden_size
+            # EmbeddingGemma 2 projects its last hidden states from hidden_size to embedding_dim
+            if text_config.model_type == "embedding_gemma2_text":
+                self.hidden_size = text_config.embedding_dim
         else:
             self.hidden_size = hidden_size
 
@@ -170,6 +174,9 @@ def __init__(  # noqa: PLR0917 - HF config, arguments mirror the saved config fi
         self.use_segment_embeddings = use_segment_embeddings
 
         self.scorer_encoder_num_layers = scorer_encoder_num_layers
+        # decoder-kv: False → the scorer encoder sees only the label section (label1<<LABEL>>...<<SEP>>);
+        # True → it runs over the whole sequence ([prompt][examples]text<<SEP>>labels...<<SEP>>, media included).
+        self.scorer_full_sequence = scorer_full_sequence
 
         # Recurrent hidden reasoning for the decoder-kv scorer (recurrent_steps=1 disables it).
         # Training depth is sampled uniformly from [recurrent_min_steps, recurrent_steps].
diff --git a/gliclass/data_processing.py b/gliclass/data_processing.py
index 7a20d85..71830f5 100644
--- a/gliclass/data_processing.py
+++ b/gliclass/data_processing.py
@@ -224,6 +224,8 @@ def augment(self, item):
 
 
 class GLiClassDataset(Dataset):
+    MAX_MEDIA_RETRIES = 10
+
     def __init__(
         self,
         examples,
@@ -426,8 +428,20 @@ def __len__(self):
         return len(self._data)
 
     def __getitem__(self, idx):
-        example = self._data[idx]
-
+        # A sample whose image / audio file cannot be loaded is replaced by a random other sample
+        for _ in range(self.MAX_MEDIA_RETRIES):
+            example = self._data[idx]
+            if not (example.get("images") or example.get("audio")):
+                return self._prepare(example)
+            try:
+                return self._prepare(example)
+            except Exception as error:
+                media = (example.get("images") or []) + (example.get("audio") or [])
+                print(f"Skipping sample {idx}, could not load its media {media}: {error!r}")
+                idx = random.randrange(len(self._data))
+        raise RuntimeError(f"Could not load media for {self.MAX_MEDIA_RETRIES} samples in a row")
+
+    def _prepare(self, example):
         example = self.augmenter.augment(example)
 
         if self.architecture_type == "uni-encoder":
diff --git a/gliclass/model.py b/gliclass/model.py
index 688b4df..d81d7a4 100644
--- a/gliclass/model.py
+++ b/gliclass/model.py
@@ -995,6 +995,8 @@ def __init__(self, config: GLiClassModelConfig, from_pretrained=False):
             config.encoder_config = AutoConfig.from_pretrained(config.encoder_model_name)
 
         config_name = config.encoder_config.__class__.__name__
+        # EmbeddingGemma 2 is a bidirectional encoder without a KV cache: the whole sequence is encoded in one pass
+        self.bidirectional = config_name == "EmbeddingGemma2Config"
 
         if getattr(config, "multimodal", False):
             # Base multi-modal model (no LM head): encodes images / audio and merges them into the sequence
@@ -1002,8 +1004,18 @@ def __init__(self, config: GLiClassModelConfig, from_pretrained=False):
                 from transformers.models.qwen3_5 import Qwen3_5Model as ModelClass
             elif config_name == "Gemma4Config":
                 from transformers.models.gemma4 import Gemma4Model as ModelClass
+            elif config_name == "EmbeddingGemma2Config":
+                from transformers.models.embedding_gemma2 import EmbeddingGemma2Model as ModelClass
             else:
-                raise ValueError(f"multimodal decoder-kv requires Qwen3.5 or Gemma 4. Got: {config_name}")
+                raise ValueError(
+                    f"multimodal decoder-kv requires Qwen3.5, Gemma 4 or EmbeddingGemma 2. Got: {config_name}"
+                )
+        elif config_name == "EmbeddingGemma2Config":
+            from transformers.models.embedding_gemma2 import EmbeddingGemma2Model as ModelClass
+
+            # text only: skip the vision / audio towers (their weights are ignored on load)
+            config.encoder_config.vision_config = None
+            config.encoder_config.audio_config = None
         elif config_name == "Qwen3_5TextConfig":
             from transformers.models.qwen3_5 import Qwen3_5TextModel
 
@@ -1018,9 +1030,12 @@ def __init__(self, config: GLiClassModelConfig, from_pretrained=False):
 
             ModelClass = Qwen3Model
         else:
-            raise ValueError(f"decoder-kv architecture requires Qwen3 or Qwen3.5. Got: {config_name}")
+            raise ValueError(f"decoder-kv architecture requires Qwen3, Qwen3.5 or EmbeddingGemma 2. Got: {config_name}")
 
-        if from_pretrained:
+        if from_pretrained and self.bidirectional:
+            # pass the config so a text-only one (no vision / audio config) skips the media towers
+            self.decoder_model = ModelClass.from_pretrained(config.encoder_model_name, config=config.encoder_config)
+        elif from_pretrained:
             self.decoder_model = ModelClass.from_pretrained(config.encoder_model_name)
         else:
             self.decoder_model = ModelClass(config.encoder_config)
@@ -1032,10 +1047,12 @@ def __init__(self, config: GLiClassModelConfig, from_pretrained=False):
             if current_vocab != config.vocab_size:
                 self.decoder_model.resize_token_embeddings(config.vocab_size)
 
-        if getattr(config, "multimodal", False) and getattr(config, "freeze_media_encoders", True):
-            from .multimodal import freeze_media_encoders
+        if getattr(config, "multimodal", False):
+            from .multimodal import freeze_media_encoders, match_media_feature_dtype
 
-            freeze_media_encoders(self.decoder_model)
+            match_media_feature_dtype(self.decoder_model)
+            if getattr(config, "freeze_media_encoders", True):
+                freeze_media_encoders(self.decoder_model)
 
         from .scorers import DecoderKVScorer
 
@@ -1047,7 +1064,7 @@ def __init__(self, config: GLiClassModelConfig, from_pretrained=False):
 
         self.sep_token_id = config.sep_token_index
 
-    def get_loss(self, logits, labels):
+    def get_loss(self, logits, labels, valid_labels=None):
         loss = None
         if labels is not None:
             # Sequence truncation may cut label tokens → align labels to logits
@@ -1055,24 +1072,7 @@ def get_loss(self, logits, labels):
             # single-label targets are class indices (batch,); only per-label targets need aligning
             if labels.dim() > 1 and labels.shape[-1] != num_labels:
                 labels = labels[:, :num_labels]
-
-            if self.config.problem_type == "multi_label_classification":
-                from .loss_functions import focal_loss_with_logits
-
-                reduction = self.config.focal_loss_reduction or "none"
-                all_losses = focal_loss_with_logits(
-                    logits,
-                    labels,
-                    self.config.focal_loss_alpha,
-                    self.config.focal_loss_gamma,
-                    reduction,
-                )
-                loss = all_losses.mean()
-            elif self.config.problem_type == "single_label_classification":
-                loss_fct = nn.CrossEntropyLoss()
-                loss = loss_fct(logits.view(-1, num_labels), labels.view(-1))
-            else:
-                raise NotImplementedError(f"{self.config.problem_type} is not implemented.")
+            loss = self._per_sample_loss(logits, labels, valid_labels).mean()
         return loss
 
     def _resolve_recurrence(self, use_recurrence=None, max_recurrent_steps=None, halt_threshold=None):
@@ -1098,8 +1098,13 @@ def _resolve_recurrence(self, use_recurrence=None, max_recurrent_steps=None, hal
             halt_threshold = self.config.recurrent_halt_threshold
         return max_steps, halt_threshold
 
-    def _per_sample_loss(self, logits, labels):
-        """Unreduced loss per example, shape (batch,)."""
+    def _per_sample_loss(self, logits, labels, valid_labels=None):
+        """Unreduced loss per example, shape (batch,).
+
+        valid_labels (batch, num_labels) marks the example's real labels. Slots past them only exist because
+        the batch is padded to its longest label list, so they are left out; otherwise the loss would depend
+        on which examples share a batch.
+        """
         if self.config.problem_type == "multi_label_classification":
             from .loss_functions import focal_loss_with_logits
 
@@ -1111,8 +1116,12 @@ def _per_sample_loss(self, logits, labels):
                 "none",
             )
             valid = labels.ne(self.config.ignore_index)
-            return all_losses.sum(-1) / valid.sum(-1).clamp_min(1)
+            if valid_labels is not None:
+                valid = valid & valid_labels
+            return (all_losses * valid).sum(-1) / valid.sum(-1).clamp_min(1)
         elif self.config.problem_type == "single_label_classification":
+            if valid_labels is not None:
+                logits = logits.masked_fill(~valid_labels, torch.finfo(logits.dtype).min)
             return nn.functional.cross_entropy(
                 logits, labels.view(-1), ignore_index=self.config.ignore_index, reduction="none"
             )
@@ -1157,7 +1166,9 @@ def get_recurrent_loss(self, step_logits, labels, valid_labels=None):
         if valid_labels is None:
             valid_labels = torch.ones_like(step_logits[-1], dtype=torch.bool)
 
-        per_sample = torch.stack([self._per_sample_loss(logits, labels) for logits in step_logits])  # (T, B)
+        per_sample = torch.stack(
+            [self._per_sample_loss(logits, labels, valid_labels) for logits in step_logits]
+        )  # (T, B)
         step_losses = per_sample.mean(dim=1)
         loss = step_losses.mean()
 
@@ -1261,6 +1272,13 @@ def _extract_label_section(
         text_mask = (text_positions.unsqueeze(0) < text_lengths.unsqueeze(1)).to(attention_mask.dtype)
         return padded_hidden, padded_ids, label_mask, text_hidden, text_mask
 
+    def _require_kv_cache(self):
+        if self.bidirectional:
+            raise ValueError(
+                f"{type(self.decoder_model).__name__} is a bidirectional encoder without a KV cache; "
+                "use the classic (non-streaming) decoder-kv pipeline."
+            )
+
     def update_decoder_cache(
         self,
         input_ids: torch.Tensor,
@@ -1269,6 +1287,7 @@ def update_decoder_cache(
         past_key_values=None,
     ):
         """Extend a text-only decoder cache without running the scorer."""
+        self._require_kv_cache()
         return self.decoder_model(
             input_ids=input_ids,
             attention_mask=attention_mask,
@@ -1298,6 +1317,7 @@ def classify_from_decoder_cache(
         With config.recurrent_read_text, pass the decoder's last hidden states of the cached text
         (text_hidden_states / text_attention_mask); the KV cache alone does not contain them.
         """
+        self._require_kv_cache()
         decoder_outputs = self.decoder_model(
             input_ids=input_ids,
             attention_mask=attention_mask,
@@ -1370,12 +1390,20 @@ def forward(
         """
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
+        if self.bidirectional:
+            if past_key_values is not None or use_cache:
+                self._require_kv_cache()
+            # media positions come from the placeholder token ids; padding mm_token_type_ids are not accepted
+            kwargs.pop("mm_token_type_ids", None)
+            cache_kwargs = {}
+        else:
+            cache_kwargs = {"past_key_values": past_key_values, "use_cache": use_cache}
+
         decoder_outputs = self.decoder_model(
             input_ids=input_ids,
             attention_mask=attention_mask,
-            past_key_values=past_key_values,
-            use_cache=use_cache,
             return_dict=True,
+            **cache_kwargs,
             **kwargs,
         )
 
@@ -1408,11 +1436,11 @@ def forward(
             self.num_labels = logits.shape[-1]
 
         recurrent_losses = recurrent_entropies = None
+        valid_labels = self.scorer._valid_label_mask(label_ids, label_mask, logits.shape[-1])
         if len(step_logits) > 1 and labels is not None:
-            valid_labels = self.scorer._valid_label_mask(label_ids, label_mask, logits.shape[-1])
             loss, recurrent_losses, recurrent_entropies = self.get_recurrent_loss(step_logits, labels, valid_labels)
         else:
-            loss = self.get_loss(logits, labels)
+            loss = self.get_loss(logits, labels, valid_labels)
 
         if not return_dict:
             output = (logits,)
@@ -1432,7 +1460,7 @@ def forward(
             recurrent_num_steps=steps_taken if num_steps > 1 else None,
             text_embeddings=text_repr if output_text_embeddings else None,
             class_embeddings=label_repr if output_class_embeddings else None,
-            class_mask=self.scorer._valid_label_mask(label_ids, label_mask, logits.shape[-1]).long(),
+            class_mask=valid_labels.long(),
         )
 
 
diff --git a/gliclass/multimodal.py b/gliclass/multimodal.py
index 0e2aefd..52df99d 100644
--- a/gliclass/multimodal.py
+++ b/gliclass/multimodal.py
@@ -5,10 +5,12 @@
     [prompt][examples]<image><image><audio>text<<SEP>>label1<<LABEL>>...<<SEP>>
 
 The HF processor expands every placeholder into the item's soft tokens and returns the media tensors;
-the multi-modal backbone (Qwen3_5Model, Gemma4Model) encodes them and writes the features into those
-positions. Placeholders are matched to media in order, so the i-th image placeholder is the i-th image.
+the multi-modal backbone (Qwen3_5Model, Gemma4Model, EmbeddingGemma2Model) encodes them and writes the features
+into those positions. Placeholders are matched to media in order, so the i-th image placeholder is the i-th image.
 """
 
+import os
+
 import torch
 from transformers.processing_utils import ProcessorMixin
 
@@ -74,7 +76,12 @@ def load_audio(audio, sampling_rate: int):
     """Load audio given as paths, URLs or 1D arrays (arrays must already be at `sampling_rate`)."""
     from transformers.audio_utils import load_audio as _load_audio
 
-    return [_load_audio(item, sampling_rate=sampling_rate) for item in audio or []]
+    audio = audio or []
+    for item in audio:
+        # transformers returns a missing local path unchanged, which only fails later inside feature extraction
+        if isinstance(item, str) and not item.startswith(("http://", "https://")) and not os.path.isfile(item):
+            raise FileNotFoundError(f"Audio file not found: {item}")
+    return [_load_audio(item, sampling_rate=sampling_rate) for item in audio]
 
 
 def process_multimodal(processor, text: str, images=None, audio=None) -> dict:
@@ -148,3 +155,17 @@ def freeze_media_encoders(model) -> None:
         module = getattr(model, name, None)
         if module is not None:
             module.requires_grad_(False)
+
+
+def match_media_feature_dtype(model) -> None:
+    """Cast projected image / audio features to the text embedding dtype before they are merged.
+
+    Under bf16 autocast the projectors return bf16 while the fp32 text embeddings stay fp32. Gemma 4 casts
+    image features itself but not audio features, so masked_scatter fails on the mixed dtypes.
+    """
+    for name in ("embed_vision", "embed_audio"):
+        module = getattr(model, name, None)
+        if module is not None:
+            module.register_forward_hook(
+                lambda _module, _inputs, output: output.to(model.get_input_embeddings().weight.dtype)
+            )
diff --git a/gliclass/scorers.py b/gliclass/scorers.py
index cc7de5c..957e4d4 100644
--- a/gliclass/scorers.py
+++ b/gliclass/scorers.py
@@ -1,7 +1,25 @@
+import os
+
 import torch
 from torch import nn
 
 from .ops import attn_padded
+from .utils import is_module_available
+
+IS_FLASHDEBERTA = is_module_available("flashdeberta")
+
+if IS_FLASHDEBERTA:
+    from flashdeberta import FlashDisentangledSelfAttention
+    from transformers.models.deberta_v2.modeling_deberta_v2 import DisentangledSelfAttention
+
+    class ScorerFlashSelfAttention(FlashDisentangledSelfAttention):
+        """FlashDeBERTa attention for the scorer encoder; its Triton kernels need CUDA, so other devices
+        (e.g. scoring on CPU) fall back to the eager DeBERTa attention with the same weights."""
+
+        def forward(self, hidden_states, *args, **kwargs):
+            if hidden_states.is_cuda:
+                return super().forward(hidden_states, *args, **kwargs)
+            return DisentangledSelfAttention.forward(self, hidden_states, *args, **kwargs)
 
 
 class ScorerWeightedDot(nn.Module):
@@ -260,6 +278,10 @@ def __init__(self, config, **kwargs):
             max_relative_positions=512,
         )
         self.scorer_encoder = DebertaV2Encoder(encoder_config)
+        if os.environ.get("USE_FLASHDEBERTA", "") and IS_FLASHDEBERTA:
+            print("Using FlashDeberta backend for the scorer encoder.")
+            for layer in self.scorer_encoder.layer:
+                layer.attention.self = ScorerFlashSelfAttention(encoder_config)
 
         self.text_projector = nn.Linear(config.hidden_size, config.hidden_size)
         self.label_projector = nn.Linear(config.hidden_size, config.hidden_size)
@@ -278,6 +300,9 @@ def __init__(self, config, **kwargs):
 
         self.epsilon = 1e-8
 
+        # False: the encoder sees only the label section; True: the whole sequence (text included)
+        self.full_sequence = getattr(config, "scorer_full_sequence", False)
+
         # Recurrent reasoning over scorer_encoder (see RecurrentReasoningCell). Step 1 is the plain scorer.
         self.recurrent = getattr(config, "recurrent_steps", 1) > 1
         self.read_text = self.recurrent and getattr(config, "recurrent_read_text", False)
diff --git a/gliclass/training.py b/gliclass/training.py
index e00630e..d81eb1a 100644
--- a/gliclass/training.py
+++ b/gliclass/training.py
@@ -557,8 +557,17 @@ def create_optimizer(self):
         if self.optimizer is None:
             decay_parameters = get_parameter_names(opt_model, ALL_LAYERNORM_LAYERS)
             decay_parameters = [name for name in decay_parameters if "bias" not in name]
+            # e.g. a partly frozen embedding matrix: decay would also shrink the frozen rows
+            no_decay = getattr(opt_model, "no_weight_decay_param_names", set())
+            decay_parameters = [name for name in decay_parameters if name not in no_decay]
             if self.args.others_lr is not None:
-                encoder_parameters = [name for name, _ in opt_model.named_parameters() if "encoder" in name]
+                # the pretrained backbone gets learning_rate, new heads get others_lr; decoder-kv names its
+                # backbone decoder_model, and its scorer's scorer_encoder is a new head, not the backbone
+                encoder_parameters = [
+                    name
+                    for name, _ in opt_model.named_parameters()
+                    if ("encoder" in name or "decoder_model" in name) and "scorer" not in name
+                ]
                 optimizer_grouped_parameters = [
                     {
                         "params": [

From 5d8c51a189d8674c74fd8aba675ea221fdb64676 Mon Sep 17 00:00:00 2001
From: Ingvarstep <igor.stepanov2000@gmail.com>
Date: Fri, 9 Oct 2026 05:48:07 +0000
Subject: [PATCH 11/12] improve multi-modal training

---
 gliclass/config.py   | 13 +++++++++++--
 gliclass/model.py    |  5 +++++
 gliclass/training.py | 12 ++++++++++++
 3 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/gliclass/config.py b/gliclass/config.py
index a75047d..aff07e7 100644
--- a/gliclass/config.py
+++ b/gliclass/config.py
@@ -80,6 +80,7 @@ def __init__(  # noqa: PLR0917 - HF config, arguments mirror the saved config fi
         calibrator_use_bias=False,
         multimodal=False,
         freeze_media_encoders=True,
+        use_embedding_projection=None,
         **kwargs,
     ):
         if isinstance(encoder_config, dict):
@@ -115,10 +116,18 @@ def __init__(  # noqa: PLR0917 - HF config, arguments mirror the saved config fi
         # Multi-modal backbones have a composite config; text attributes live in its text config
         text_config = self.encoder_config.get_text_config()
 
+        # EmbeddingGemma 2 ends with embedding_projection (hidden_size -> embedding_dim), a head meant for
+        # mean pooling + L2 normalization that scales token states ~30x; fed to the scorer, it blows up the
+        # gradients. By default the scorer reads the final-norm states instead. Older checkpoints were trained
+        # with the projection: without the flag in their config, it is kept when hidden_size == embedding_dim.
+        is_embedding_gemma = text_config.model_type == "embedding_gemma2_text"
+        if use_embedding_projection is None:
+            use_embedding_projection = is_embedding_gemma and hidden_size == getattr(text_config, "embedding_dim", None)
+        self.use_embedding_projection = use_embedding_projection
+
         if hidden_size is None:
             self.hidden_size = text_config.hidden_size
-            # EmbeddingGemma 2 projects its last hidden states from hidden_size to embedding_dim
-            if text_config.model_type == "embedding_gemma2_text":
+            if is_embedding_gemma and use_embedding_projection:
                 self.hidden_size = text_config.embedding_dim
         else:
             self.hidden_size = hidden_size
diff --git a/gliclass/model.py b/gliclass/model.py
index d81d7a4..75db568 100644
--- a/gliclass/model.py
+++ b/gliclass/model.py
@@ -1042,6 +1042,11 @@ def __init__(self, config: GLiClassModelConfig, from_pretrained=False):
             # multi-modal backbones build their sub-models in the config's dtype (e.g. bf16); match the scorer
             self.decoder_model.to(torch.get_default_dtype())
 
+        if self.bidirectional and not getattr(config, "use_embedding_projection", True):
+            # scorer reads the final-norm hidden states (see GLiClassModelConfig.use_embedding_projection)
+            text_model = getattr(self.decoder_model, "language_model", self.decoder_model)
+            text_model.embedding_projection = nn.Identity()
+
         if config.vocab_size is not None and hasattr(self.decoder_model, "resize_token_embeddings"):
             current_vocab = self.decoder_model.config.get_text_config().vocab_size
             if current_vocab != config.vocab_size:
diff --git a/gliclass/training.py b/gliclass/training.py
index d81eb1a..10c332e 100644
--- a/gliclass/training.py
+++ b/gliclass/training.py
@@ -625,7 +625,19 @@ def create_optimizer(self):
             if getattr(self.args, "recurrent_lr", None) is not None:
                 optimizer_grouped_parameters = self._split_recurrent_groups(opt_model, optimizer_grouped_parameters)
 
+            optimizer_grouped_parameters = [group for group in optimizer_grouped_parameters if group["params"]]
+
             optimizer_cls, optimizer_kwargs = Trainer.get_optimizer_cls_and_kwargs(self.args)
+            if self.is_deepspeed_enabled and self.accelerator.state.deepspeed_plugin.hf_ds_config.is_offload():
+                # ZeRO offload keeps optimizer states on CPU, which needs DeepSpeed's CPU implementation of Adam
+                from deepspeed.ops.adam import DeepSpeedCPUAdam
+                optimizer_cls = DeepSpeedCPUAdam
+                optimizer_kwargs = {
+                    "lr": self.args.learning_rate,
+                    "betas": (self.args.adam_beta1, self.args.adam_beta2),
+                    "eps": self.args.adam_epsilon,
+                    "adamw_mode": True,
+                }
 
             self.optimizer = optimizer_cls(optimizer_grouped_parameters, **optimizer_kwargs)
 

From 7e9f31382b77d6942f56c99f045300c9add0c8f1 Mon Sep 17 00:00:00 2001
From: Ingvarstep <igor.stepanov2000@gmail.com>
Date: Fri, 9 Oct 2026 14:23:18 +0000
Subject: [PATCH 12/12] optimize training

---
 gliclass/data_processing.py |  55 +++++++-
 gliclass/training.py        |  33 ++++-
 train.py                    | 251 ++++++++++++++++++++++++++++++++----
 3 files changed, 315 insertions(+), 24 deletions(-)

diff --git a/gliclass/data_processing.py b/gliclass/data_processing.py
index 71830f5..b880311 100644
--- a/gliclass/data_processing.py
+++ b/gliclass/data_processing.py
@@ -6,7 +6,7 @@
 from torch.utils.data import Dataset
 from torch.nn.utils.rnn import pad_sequence
 
-from .multimodal import MEDIA_KEYS, is_processor, collate_media, process_multimodal_with_budget
+from .multimodal import MEDIA_KEYS, is_processor, collate_media, get_tokenizer, process_multimodal_with_budget
 
 
 def format_decoder_kv_context(text: str, prompt: str = "", examples: str = "", media: str = "") -> str:
@@ -266,6 +266,59 @@ def __init__(
     def get_diversity(self):
         return [item.get("_diversity", {}).get("overall_diversity", 0.5) for item in self._data]
 
+    def estimate_lengths(self, batch_size=1000):
+        """Approximate token length of every sample (text + labels + media tokens), capped at max_length.
+
+        Used for length-grouped batching; much cheaper than preprocessing every sample (no image decoding).
+        """
+        if getattr(self, "_estimated_lengths", None) is not None:
+            return self._estimated_lengths
+        tokenizer = get_tokenizer(self.tokenizer)
+        lengths = []
+        for start in range(0, len(self._data), batch_size):
+            chunk = self._data[start : start + batch_size]
+            # longer texts are truncated to max_length anyway; bound the tokenizer work
+            texts = [
+                " ".join(
+                    [
+                        str(example.get("prompt", "")),
+                        str(example["text"])[: self.max_length * 16],
+                        *map(str, example.get("all_labels", [])),
+                    ]
+                )
+                for example in chunk
+            ]
+            ids = tokenizer(texts, add_special_tokens=False)["input_ids"]
+            for example, text_ids in zip(chunk, ids):
+                length = len(text_ids) + 2 * len(example.get("all_labels", [])) + 4
+                length += sum(self._estimate_image_tokens(image) for image in example.get("images") or [])
+                length += self.ESTIMATED_AUDIO_TOKENS * len(example.get("audio") or [])
+                lengths.append(min(length, self.max_length))
+        self._estimated_lengths = lengths
+        return lengths
+
+    ESTIMATED_AUDIO_TOKENS = 256
+
+    def _estimate_image_tokens(self, image, default=256):
+        """Image tokens from the image size; reads only the file header. Exact for Qwen-VL image processors."""
+        image_processor = getattr(self.tokenizer, "image_processor", None)
+        if image_processor is None or not hasattr(image_processor, "get_number_of_image_patches"):
+            return default
+        try:
+            if isinstance(image, str):
+                from PIL import Image
+
+                with Image.open(image) as opened:
+                    size = opened.size
+            elif hasattr(image, "size"):
+                size = image.size
+            else:
+                return default
+            patches = image_processor.get_number_of_image_patches(size[1], size[0], {})
+        except Exception:
+            return default
+        return patches // getattr(image_processor, "merge_size", 1) ** 2
+
     def collect_dataset_labels(self):
         dataset_labels = set()
         for example in self._data:
diff --git a/gliclass/training.py b/gliclass/training.py
index 10c332e..581da3c 100644
--- a/gliclass/training.py
+++ b/gliclass/training.py
@@ -415,6 +415,29 @@ def train(self, *args, **kwargs):
         self._maybe_initialize_ewc()
         return super().train(*args, **kwargs)
 
+    def _get_train_sampler(self, train_dataset=None):
+        """batch_rebalance: use the dataset's cheap length estimates instead of preprocessing every sample."""
+        train_dataset = train_dataset if train_dataset is not None else self.train_dataset
+        if (
+            getattr(self.args, "train_sampling_strategy", None) != "batch_rebalance"
+            or not hasattr(train_dataset, "estimate_lengths")
+        ):
+            return super()._get_train_sampler(train_dataset)
+        from transformers.trainer_pt_utils import BatchRebalanceSampler
+
+        lengths = train_dataset.estimate_lengths()
+        world_size = max(1, self.args.world_size)
+        grad_accum = self.args.gradient_accumulation_steps
+        return BatchRebalanceSampler(
+            lengths=lengths,
+            effective_batch_size=self.args.train_batch_size * grad_accum * world_size,
+            dp_size=world_size,
+            grad_accum=grad_accum,
+            rank=self.args.process_index if world_size > 1 else 0,
+            seed=self.args.data_seed if self.args.data_seed is not None else self.args.seed,
+            drop_last=self.args.dataloader_drop_last,
+        )
+
     def training_step(self, model, inputs, *args, **kwargs) -> torch.Tensor:
         """
         Perform a training step on a batch of inputs.
@@ -445,7 +468,6 @@ def training_step(self, model, inputs, *args, **kwargs) -> torch.Tensor:
                 loss = self.compute_loss(model, inputs)
 
             del inputs
-            torch.cuda.empty_cache()
 
             kwargs = {}
 
@@ -641,6 +663,15 @@ def create_optimizer(self):
 
             self.optimizer = optimizer_cls(optimizer_grouped_parameters, **optimizer_kwargs)
 
+            if "bitsandbytes" in str(optimizer_cls) and optimizer_kwargs.get("optim_bits") == 8:
+                # mirror transformers.Trainer.create_optimizer: keep embedding optimizer states in 32-bit for stability
+                import bitsandbytes
+
+                manager = bitsandbytes.optim.GlobalOptimManager.get_instance()
+                for module in opt_model.modules():
+                    if isinstance(module, nn.Embedding):
+                        manager.register_module_override(module, "weight", {"optim_bits": 32})
+
         return self.optimizer
 
     def _split_recurrent_groups(self, model, groups):
diff --git a/train.py b/train.py
index 8e3336f..3392d5c 100644
--- a/train.py
+++ b/train.py
@@ -3,6 +3,7 @@
 import numpy as np
 import argparse
 import json
+import math
 
 from sklearn.metrics import precision_recall_fscore_support, accuracy_score
 import transformers
@@ -12,6 +13,7 @@
 
 import random
 import torch
+from torch import nn
 
 from gliclass import GLiClassModelConfig, GLiClassModel
 from gliclass.training import TrainingArguments, Trainer
@@ -27,7 +29,7 @@ def __init__(self, *args, use_weighted_sampling=False, **kwargs):
     
     def _get_train_sampler(self, train_dataset) -> torch.utils.data.Sampler:
         if not self.use_weighted_sampling:
-            return super()._get_train_sampler()
+            return super()._get_train_sampler(train_dataset)
         
         weights = train_dataset.get_diversity()
         return WeightedRandomSampler(
@@ -36,6 +38,121 @@ def _get_train_sampler(self, train_dataset) -> torch.utils.data.Sampler:
             replacement=True
         )
     
+def log_encoder_backend(model):
+    """Print which encoder class was loaded, e.g. to confirm the FlashDeBERTa backend is in use."""
+    encoder = getattr(model.model, 'encoder_model', None)
+    if encoder is None:
+        return
+    if hasattr(encoder, 'get_base_model'):  # PEFT wrapper
+        encoder = encoder.get_base_model()
+    name = type(encoder).__name__
+    print(f'Encoder backend: {name}')
+    if name == 'FlashDebertaV2Model':
+        print('FlashDeBERTa is used: attention runs on FlashDeBERTa kernels.')
+    elif name == 'DebertaV2Model' and os.environ.get('USE_FLASHDEBERTA'):
+        print('WARNING: USE_FLASHDEBERTA is set but the standard DebertaV2Model was loaded '
+              '(is the flashdeberta package installed in this Python environment?).')
+
+SPECIAL_TOKEN_INIT_WORDS = {"<<LABEL>>": "label", "<<SEP>>": "text", "<<EXAMPLE>>": "example"}
+
+def init_special_token_embeddings(embeddings, text_tokenizer, new_ids, num_pretrained_rows):
+    """Initialize the rows of newly added special tokens from the embedding of a descriptive word.
+
+    Rows of new tokens are either mean-initialized by resize_token_embeddings or, when the backbone
+    has spare vocabulary rows (e.g. deberta-v2-xxlarge), left as never-trained pretrained rows that
+    look like [PAD]. Either way <<LABEL>>, <<SEP>> and <<EXAMPLE>> start nearly identical to each other
+    and with a much smaller norm than real tokens. Seeding each from a distinct word ("label", "text",
+    "example") and rescaling to the median norm of the pretrained vocabulary keeps them distinguishable.
+    """
+    weight = embeddings.weight
+    with torch.no_grad():
+        target_norm = weight[:num_pretrained_rows].float().norm(dim=-1).median()
+        for token, token_id in new_ids.items():
+            word = SPECIAL_TOKEN_INIT_WORDS.get(token, token.strip("<>").lower())
+            word_ids = [i for i in text_tokenizer(word, add_special_tokens=False)['input_ids']
+                        if i < num_pretrained_rows]
+            if not word_ids:
+                continue
+            vector = weight[word_ids].float().mean(dim=0)
+            weight[token_id] = (vector * target_norm / vector.norm()).to(weight.dtype)
+            print(f'Initialized {token} (id {token_id}) from "{word}" embedding, norm {target_norm:.3f}')
+
+def freeze_decoder(model, train_last_layers=None):
+    """decoder-kv: freeze the decoder backbone, train only the scorer.
+
+    Without train_last_layers the rows of the special tokens (<<LABEL>>, <<SEP>>, <<EXAMPLE>>) are new, so they
+    stay trainable: their gradients pass through a row mask and the embedding matrix is kept out of weight decay,
+    so every pretrained row stays exactly as loaded.
+
+    With train_last_layers (a fraction) the last layers of the decoder and its final norm train too. The input of
+    the first trainable layer is detached, so the backward pass stops there instead of running through all frozen
+    layers; the special token rows then stay at their initial values.
+    """
+    decoder = getattr(model.model, 'decoder_model', None)
+    if decoder is None:
+        raise ValueError('--freeze_decoder / --train_last_layers require --architecture_type decoder-kv')
+    decoder.requires_grad_(False)
+
+    if train_last_layers is not None:
+        if not 0 < train_last_layers <= 1:
+            raise ValueError(f'--train_last_layers must be in (0, 1], got {train_last_layers}')
+        text_model = getattr(decoder, 'language_model', decoder)
+        layers = text_model.layers
+        num_trainable = math.ceil(train_last_layers * len(layers))
+        first_trainable = len(layers) - num_trainable
+        for layer in layers[first_trainable:]:
+            layer.requires_grad_(True)
+        if getattr(text_model, 'norm', None) is not None:
+            text_model.norm.requires_grad_(True)
+
+        # gradient checkpointing makes the embedding output require grad (enable_input_require_grads),
+        # which would otherwise backpropagate through all frozen layers for nothing
+        def detach_hidden_states(_module, args, kwargs):
+            if args:
+                return (args[0].detach(), *args[1:]), kwargs
+            return args, {**kwargs, 'hidden_states': kwargs['hidden_states'].detach()}
+
+        layers[first_trainable].register_forward_pre_hook(detach_hidden_states, with_kwargs=True)
+        trainable = sum(param.numel() for param in model.parameters() if param.requires_grad)
+        total = sum(param.numel() for param in model.parameters())
+        print(f'Decoder frozen except its last {num_trainable} of {len(layers)} layers '
+              f'(layers {first_trainable}-{len(layers) - 1}) and final norm: '
+              f'{trainable / 1e6:.1f}M of {total / 1e6:.1f}M parameters receive gradients')
+        return
+
+    config = model.config
+    special_ids = sorted({index for index in (config.class_token_index, config.text_token_index,
+                                              config.example_token_index, getattr(config, 'sep_token_index', None))
+                          if index is not None and index >= 0})
+    embeddings = decoder.get_input_embeddings().weight
+    special_ids = [index for index in special_ids if index < embeddings.shape[0]]
+    if special_ids:
+        mask = torch.zeros(embeddings.shape[0], 1, device=embeddings.device, dtype=embeddings.dtype)
+        mask[special_ids] = 1.0
+        embeddings.requires_grad_(True)
+        embeddings.register_hook(lambda grad: grad * mask)
+        model.no_weight_decay_param_names = {name for name, param in model.named_parameters() if param is embeddings}
+
+    trainable = sum(param.numel() for param in model.parameters() if param.requires_grad and param is not embeddings)
+    trainable += len(special_ids) * embeddings.shape[1]
+    total = sum(param.numel() for param in model.parameters())
+    print(f'Decoder frozen (special token rows {special_ids} stay trainable): '
+          f'{trainable / 1e6:.1f}M of {total / 1e6:.1f}M parameters receive gradients')
+
+def cast_frozen_to_bf16(model):
+    """Store frozen decoder linear / embedding weights in bf16.
+
+    Under bf16 autocast every fp32 weight is cast to bf16 on each forward (twice with checkpoint recompute);
+    frozen weights never need fp32. Norms, convolutions, small fp32 parameters and trainable layers stay fp32.
+    """
+    decoder = model.model.decoder_model
+    saved = 0
+    for module in decoder.modules():
+        if isinstance(module, (nn.Linear, nn.Embedding)) and not any(p.requires_grad for p in module.parameters()):
+            saved += sum(p.numel() for p in module.parameters() if p.dtype == torch.float32) * 2
+            module.to(torch.bfloat16)
+    print(f'Frozen decoder weights stored in bf16: {saved / 1e9:.1f} GB saved')
+
 def compute_metrics(p, problem_type='multi_label_classification'):
     """Compute evaluation metrics.
     
@@ -76,23 +193,46 @@ def compute_metrics(p, problem_type='multi_label_classification'):
         raise NotImplementedError(f"{problem_type} is not implemented.")
 
 
-def load_dataset(data_path: str) -> list:
-    """Load dataset from JSON file.
+def load_dataset(data_path) -> list:
+    """Load dataset from one or several JSON / JSONL files.
     
     Args:
-        data_path: Path to JSON data file
+        data_path: Path to a JSON (list of samples) or JSONL (one sample per line) file,
+            or a list of such paths whose samples are concatenated
         
     Returns:
         List of data samples
     """
-    with open(data_path, 'r') as f:
-        data = json.load(f)
+    paths = [data_path] if isinstance(data_path, str) else data_path
+    data = []
+    for path in paths:
+        with open(path, 'r') as f:
+            if path.endswith('.jsonl'):
+                part = [json.loads(line) for line in f if line.strip()]
+            else:
+                part = json.load(f)
+        if len(paths) > 1:
+            print(f'Loaded {len(part)} samples from {path}')
+        data.extend(part)
     return data
 
 
+def has_media(example) -> bool:
+    return bool(example.get('images') or example.get('audio'))
+
+
+def filter_text_only(data, name) -> list:
+    """Drop samples with images / audio, which a text-only model cannot consume."""
+    filtered = [example for example in data if not has_media(example)]
+    if len(filtered) < len(data):
+        print(f'Removed {len(data) - len(filtered)} multi-modal samples from {name} (model is not multi-modal)')
+    return filtered
+
+
 def recurrent_kwargs(args):
-    """Recurrent reasoning settings for the decoder-kv scorer; unset flags keep config/checkpoint values."""
+    """Decoder-kv scorer settings (recurrent reasoning, scorer input); unset flags keep config/checkpoint values."""
     kwargs = dict(
+        scorer_full_sequence=args.scorer_full_sequence,
         recurrent_steps=args.recurrent_steps,
         recurrent_min_steps=args.recurrent_min_steps,
         recurrent_inference_max_steps=args.recurrent_inference_max_steps,
@@ -161,6 +301,7 @@ def main(args):
             use_segment_embeddings=args.use_segment_embeddings,
             multimodal=args.multimodal,
             freeze_media_encoders=not args.train_media_encoders,
+            use_embedding_projection=args.use_embedding_projection,
             **recurrent_kwargs(args),
         )
         if args.architecture_type == 'decoder-kv':
@@ -171,10 +312,20 @@ def main(args):
 
         if args.architecture_type in {'uni-encoder', 'bi-encoder-fused', 'encoder-decoder', 'decoder-kv'}:
             new_words = ["<<LABEL>>", "<<SEP>>", "<<EXAMPLE>>"]
+            num_pretrained_rows = len(text_tokenizer)
+            added = [word for word in new_words if word not in text_tokenizer.get_vocab()]
             text_tokenizer.add_tokens(new_words, special_tokens=True)
-            model.resize_token_embeddings(len(text_tokenizer))
+            embeddings = model.resize_token_embeddings(len(text_tokenizer))
+            init_special_token_embeddings(embeddings, text_tokenizer,
+                                          {word: text_tokenizer.convert_tokens_to_ids(word) for word in added},
+                                          num_pretrained_rows)
 
     model.to(device)
+    log_encoder_backend(model)
+    if args.freeze_decoder or args.train_last_layers is not None:
+        freeze_decoder(model, args.train_last_layers)
+        if args.train_last_layers is not None and args.bf16:
+            cast_frozen_to_bf16(model)
 
     # Get labels tokenizer if needed
     if model.config.label_model_name is not None:
@@ -186,14 +337,27 @@ def main(args):
 
     # Load current training data
     data = load_dataset(args.data_path)
+    multimodal = model.config.multimodal
+    if not multimodal:
+        data = filter_text_only(data, 'training data')
     print(f'Dataset size: {len(data)}')
     
     random.shuffle(data)    
     print('Dataset is shuffled...')
 
-    train_data = data[:int(len(data) * 0.9)]
-    test_data = data[int(len(data) * 0.9):]
-    print('Dataset is splitted...')
+    if args.val_data_path is not None:
+        train_data = data
+        test_data = load_dataset(args.val_data_path)
+        if not multimodal:
+            test_data = filter_text_only(test_data, 'validation data')
+        random.shuffle(test_data)
+        print(f'Validation dataset size: {len(test_data)}')
+    else:
+        train_data = data[:int(len(data) * 0.9)]
+        test_data = data[int(len(data) * 0.9):]
+        print('Dataset is splitted...')
+    if args.max_eval_samples is not None:
+        test_data = test_data[:args.max_eval_samples]
 
     # Create augmentation config with all parameters
     augment_config = AugmentationConfig(
@@ -231,6 +395,8 @@ def main(args):
     if args.use_ewc and args.prev_data_path is not None:
         print(f'Loading previous dataset for EWC from: {args.prev_data_path}')
         prev_data = load_dataset(args.prev_data_path)
+        if not multimodal:
+            prev_data = filter_text_only(prev_data, 'previous data')
         print(f'Previous dataset size: {len(prev_data)}')
         
         # Use a subset if specified
@@ -257,6 +423,10 @@ def main(args):
         recurrent_lr=args.recurrent_lr,
         lr_scheduler_type=args.lr_scheduler_type,
         optim=args.optim,
+        optim_args=args.optim_args,
+        # transformers v4 has no train_sampling_strategy; only pass it when it is not the default
+        **({"train_sampling_strategy": args.train_sampling_strategy}
+           if args.train_sampling_strategy != 'random' else {}),
         # transformers v5 dropped warmup_ratio; a float warmup_steps is a ratio there
         **({"warmup_steps": args.warmup_ratio} if version.parse(transformers.__version__) >= version.parse("5.0.0")
            else {"warmup_ratio": args.warmup_ratio}),
@@ -266,12 +436,17 @@ def main(args):
         num_train_epochs=args.num_epochs,
         save_steps=args.save_steps,
         save_total_limit=args.save_total_limit,
+        save_only_model=args.save_only_model,
+        eval_strategy="steps" if args.val_data_path is not None else "no",
+        eval_steps=args.eval_steps or args.save_steps,
         dataloader_num_workers=args.num_workers,
         logging_steps=args.logging_steps,
         use_cpu=False,
         report_to="none",
         fp16=args.fp16,
         bf16=args.bf16,
+        deepspeed=args.deepspeed,
+        gradient_checkpointing=args.gradient_checkpointing,
         # EWC parameters
         use_ewc=args.use_ewc,
         ewc_lambda=args.ewc_lambda,
@@ -333,10 +508,12 @@ def compute_metrics_fn(p):
     # Path arguments
     parser.add_argument('--save_path', type=str, default='models/',
                         help='Path to save trained model')
-    parser.add_argument('--data_path', type=str, default='data/zero-cats.json',
-                        help='Path to training data JSON file')
-    parser.add_argument('--prev_data_path', type=str, default=None,
-                        help='Path to previous task data for EWC (required if use_ewc=True)')
+    parser.add_argument('--data_path', type=str, nargs='+', default=['data/zero-cats.json'],
+                        help='Path(s) to training data JSON file(s); multiple files are concatenated')
+    parser.add_argument('--val_data_path', type=str, nargs='+', default=None,
+                        help='Validation data JSON file(s); replaces the 90/10 split of data_path and enables evaluation')
+    parser.add_argument('--prev_data_path', type=str, nargs='+', default=None,
+                        help='Path(s) to previous task data for EWC (required if use_ewc=True)')
     parser.add_argument('--labels_desc_path', type=str, default = None)
 
     # Model architecture arguments
@@ -357,11 +534,24 @@ def compute_metrics_fn(p):
     parser.add_argument('--shuffle_labels', type=bool, default=True)
     parser.add_argument('--use_segment_embeddings', type=bool, default=False)
     parser.add_argument('--multimodal', action='store_true',
-                        help='decoder-kv: keep the vision / audio encoders of Qwen3.5 / Gemma 4; '
+                        help='decoder-kv: keep the vision / audio encoders of Qwen3.5 / Gemma 4 / EmbeddingGemma 2; '
                              'data items may then carry "images" and "audio" lists')
     parser.add_argument('--train_media_encoders', action='store_true',
                         help='Also train the vision / audio encoders (frozen by default)')
-
+    parser.add_argument('--use_embedding_projection', action='store_true',
+                        help='EmbeddingGemma 2: feed the scorer the output of embedding_projection instead of the '
+                             'final-norm hidden states (the projection scales them ~30x and destabilizes training)')
+    parser.add_argument('--freeze_decoder', action='store_true',
+                        help='decoder-kv: freeze the decoder backbone and train only the scorer '
+                             '(and the embeddings of the added special tokens)')
+    parser.add_argument('--train_last_layers', type=float, nargs='?', const=0.3, default=None,
+                        help='decoder-kv: freeze the decoder except the given fraction of its last layers '
+                             '(and the final norm); 0.3 when passed without a value. The backward pass stops at '
+                             'the first trainable layer and, with --bf16, frozen weights are stored in bf16')
+
+    parser.add_argument('--scorer_full_sequence', action='store_true', default=None,
+                        help='decoder-kv: run the scorer encoder over the whole sequence (text + labels) '
+                             'instead of only the label section')
     # Recurrent hidden reasoning (decoder-kv scorer)
     parser.add_argument('--recurrent_steps', type=int, default=None,
                         help='Max scorer reasoning steps during training (default 1 = disabled)')
@@ -388,24 +578,41 @@ def compute_metrics_fn(p):
     parser.add_argument('--num_epochs', type=int, default=3)
     parser.add_argument('--batch_size', type=int, default=8)
     parser.add_argument('--gradient_accumulation_steps', type=int, default=1)
-    parser.add_argument('--encoder_lr', type=float, default=1e-5)
-    parser.add_argument('--others_lr', type=float, default=3e-5)
+    parser.add_argument('--encoder_lr', type=float, default=2e-5)
+    parser.add_argument('--others_lr', type=float, default=5e-5)
     parser.add_argument('--encoder_weight_decay', type=float, default=0.01)
     parser.add_argument('--others_weight_decay', type=float, default=0.01)
     parser.add_argument('--warmup_ratio', type=float, default=0.05)
     parser.add_argument('--lr_scheduler_type', type=str, default='linear')
     parser.add_argument('--optim', type=str, default='adamw_torch',
-                        help='transformers optimizer name, e.g. adafactor to fit larger backbones in memory')
-    parser.add_argument('--max_length', type=int, default=1024)
+                        help='transformers optimizer name, e.g. adafactor, adamw_bnb_8bit or adamw_torch_4bit '
+                             'to fit larger backbones in memory')
+    parser.add_argument('--optim_args', type=str, default=None,
+                        help='extra optimizer args as "k1=v1,k2=v2", e.g. block_size=256 for adamw_torch_4bit')
+    parser.add_argument('--train_sampling_strategy', type=str, default='random',
+                        choices=['random', 'sequential', 'group_by_length', 'batch_rebalance'],
+                        help='batch_rebalance: sort each optimizer step\'s samples by (estimated) length and split '
+                             'them into variable-size micro-batches of balanced cost, cutting padding')
+    parser.add_argument('--max_length', type=int, default=2048)
     parser.add_argument('--save_steps', type=int, default=1000)
     parser.add_argument('--save_total_limit', type=int, default=3)
+    parser.add_argument('--save_only_model', action='store_true',
+                        help='Save only model weights in checkpoints (no optimizer / scheduler / RNG state; training cannot be resumed from them)')
+    parser.add_argument('--eval_steps', type=int, default=None,
+                        help='Evaluate every N steps when val_data_path is set (default: save_steps)')
+    parser.add_argument('--max_eval_samples', type=int, default=5000,
+                        help='Evaluate on at most this many validation examples')
     parser.add_argument('--num_workers', type=int, default=12)
     parser.add_argument('--fp16', type=bool, default=False)
     parser.add_argument('--bf16', action='store_true')
+    parser.add_argument('--deepspeed', type=str, default=None,
+                        help='Path to a DeepSpeed config JSON, e.g. configs/deepspeed/zero2_offload.json')
+    parser.add_argument('--gradient_checkpointing', action='store_true',
+                        help='Recompute activations in the backward pass to save memory')
     parser.add_argument('--logging_steps', type=int, default=100)
     
     # Augmentation parameters
-    parser.add_argument('--enable_augmentation', type=bool, default=True)
+    parser.add_argument('--enable_augmentation', type=bool, default=False)
     parser.add_argument('--random_label_removal_prob', type=float, default=0.05)
     parser.add_argument('--random_label_addition_prob', type=float, default=0.05)
     parser.add_argument('--random_text_addition_prob', type=float, default=0.05)