From 72d22c8f39269a4029307ad22794b7835cc82b79 Mon Sep 17 00:00:00 2001
From: BioMikeUkr <mykhayloshtopko@gmail.com>
Date: Fri, 10 Jul 2026 12:52:25 +0000
Subject: [PATCH 01/25] caching utils added

---
 gliclass/streaming/cache.py | 351 ++++++++++++++++++++++++++++++++++++
 1 file changed, 351 insertions(+)
 create mode 100644 gliclass/streaming/cache.py

diff --git a/gliclass/streaming/cache.py b/gliclass/streaming/cache.py
new file mode 100644
index 0000000..a7941ba
--- /dev/null
+++ b/gliclass/streaming/cache.py
@@ -0,0 +1,351 @@
+"""
+KV cache management for GLiClass streaming classification.
+
+Includes:
+  - CacheState / DynamicKVCacheManager  (low-level per-session ops)
+  - BatchedKVHelper                      (stack / unstack heterogeneous KV caches for batched inference)
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+import torch
+
+
+# ---------------------------------------------------------------------------
+# CacheState
+# ---------------------------------------------------------------------------
+
+@dataclass
+class CacheState:
+    past_key_values: Any
+    input_ids: torch.Tensor
+    attention_mask: torch.Tensor
+    position_ids: torch.Tensor | None = None
+    cached_length: int = 0
+    session_id: str | None = None
+    metadata: dict = field(default_factory=dict)
+
+    def to(self, device: torch.device) -> "CacheState":
+        return CacheState(
+            past_key_values=_move_past_kv(self.past_key_values, device),
+            input_ids=self.input_ids.to(device),
+            attention_mask=self.attention_mask.to(device),
+            position_ids=self.position_ids.to(device) if self.position_ids is not None else None,
+            cached_length=self.cached_length,
+            session_id=self.session_id,
+            metadata=self.metadata.copy(),
+        )
+
+
+def _move_past_kv(past_kv, device):
+    if past_kv is None:
+        return None
+    if hasattr(past_kv, "to"):
+        return past_kv.to(device)
+    if isinstance(past_kv, (tuple, list)):
+        return type(past_kv)(_move_past_kv(item, device) for item in past_kv)
+    return past_kv
+
+
+# ---------------------------------------------------------------------------
+def create_empty_cache(session_id: str | None = None, device=None) -> CacheState:
+    kw = {"device": device} if device is not None else {}
+    return CacheState(
+        past_key_values=None,
+        input_ids=torch.empty(0, dtype=torch.long, **kw),
+        attention_mask=torch.empty(0, dtype=torch.long, **kw),
+        position_ids=None,
+        cached_length=0,
+        session_id=session_id,
+        metadata={},
+    )
+
+
+def truncate_cache(cache_state: CacheState, max_length: int) -> CacheState:
+    if cache_state.cached_length <= max_length or cache_state.past_key_values is None:
+        return cache_state
+    truncated = _crop_past_kv(cache_state.past_key_values, max_length)
+    if truncated is None:
+        return cache_state
+    return CacheState(
+        past_key_values=truncated,
+        input_ids=cache_state.input_ids[-max_length:],
+        attention_mask=cache_state.attention_mask[-max_length:],
+        position_ids=cache_state.position_ids[-max_length:] if cache_state.position_ids is not None else None,
+        cached_length=max_length,
+        session_id=cache_state.session_id,
+        metadata=cache_state.metadata.copy(),
+    )
+
+
+def _deep_copy_past_kv(past_kv):
+    if past_kv is None:
+        return None
+    from transformers.cache_utils import DynamicCache
+    if isinstance(past_kv, DynamicCache):
+        new = DynamicCache()
+        for i, layer in enumerate(past_kv.layers):
+            if layer.is_initialized:
+                new.update(layer.keys.clone(), layer.values.clone(), i)
+        return new
+    if hasattr(past_kv, "clone"):
+        return past_kv.clone()
+    if isinstance(past_kv, (tuple, list)):
+        return type(past_kv)(_deep_copy_past_kv(x) for x in past_kv)
+    return past_kv
+
+
+def _crop_past_kv(past_kv, max_length: int):
+    """Crop past_key_values to last max_length tokens. Returns None if not supported."""
+    if past_kv is None:
+        return None
+    from transformers.cache_utils import DynamicCache
+    if isinstance(past_kv, DynamicCache):
+        new = DynamicCache()
+        for i, layer in enumerate(past_kv.layers):
+            if layer.is_initialized:
+                new.update(layer.keys[:, :, -max_length:, :].clone(),
+                           layer.values[:, :, -max_length:, :].clone(), i)
+        return new
+    return None
+
+
+# ---------------------------------------------------------------------------
+# BatchedKVHelper  – stack / unstack heterogeneous KV caches
+# ---------------------------------------------------------------------------
+
+class BatchedKVHelper:
+    """
+    Utilities for batching decoder forwards across sessions with different
+    KV cache lengths.
+
+    Stacking strategy:
+      - KV caches are prepend-padded with zeros to max_cached_len so that
+        position embeddings (RoPE) baked into each K/V tensor remain correct.
+      - The attention_mask covers [prepend_zeros | real_cached | new_tokens | new_pad].
+      - position_ids for new tokens are set per-session: [cached_len_i .. cached_len_i + new_len_i).
+    """
+
+    @staticmethod
+    def stack_for_update(
+        caches: list[CacheState],
+        new_input_ids: list[torch.Tensor],      # one per session, already on device
+        new_attention_masks: list[torch.Tensor],
+        device: torch.device,
+    ) -> dict:
+        """
+        Prepare a batched decoder forward for the cache-update stage.
+
+        Returns a dict with keys:
+          input_ids, attention_mask, position_ids, past_key_values,
+          cached_lengths, new_lengths, max_cached_len
+        """
+        batch_size = len(caches)
+        cached_lengths = [c.cached_length for c in caches]
+        new_lengths = [ids.shape[-1] for ids in new_input_ids]
+        max_cached = max(cached_lengths)
+        max_new = max(new_lengths)
+
+        # --- pad new tokens (right-pad) ---
+        padded_ids = new_input_ids[0].new_zeros(batch_size, max_new)
+        padded_new_mask = new_input_ids[0].new_zeros(batch_size, max_new)
+        position_ids = new_input_ids[0].new_zeros(batch_size, max_new)
+
+        for i, (ids, mask, clen, nlen) in enumerate(
+            zip(new_input_ids, new_attention_masks, cached_lengths, new_lengths)
+        ):
+            padded_ids[i, :nlen] = ids[0] if ids.dim() == 2 else ids
+            padded_new_mask[i, :nlen] = mask[0] if mask.dim() == 2 else mask
+            position_ids[i, :nlen] = torch.arange(clen, clen + nlen, device=device)
+
+        # --- build full attention mask: [prepend_zeros | cached | new | new_pad] ---
+        full_mask = torch.zeros(batch_size, max_cached + max_new, dtype=torch.long, device=device)
+        for i, (clen, nlen) in enumerate(zip(cached_lengths, new_lengths)):
+            pad = max_cached - clen
+            full_mask[i, pad : pad + clen] = 1          # real cached tokens
+            full_mask[i, max_cached : max_cached + nlen] = 1  # real new tokens
+
+        # --- stack past_key_values (prepend-pad to max_cached) ---
+        stacked_past_kv = None
+        if max_cached > 0:
+            stacked_past_kv = BatchedKVHelper._stack_past_kv(caches, max_cached, device)
+
+        return {
+            "input_ids": padded_ids,
+            "attention_mask": full_mask,
+            "position_ids": position_ids,
+            "past_key_values": stacked_past_kv,
+            "cached_lengths": cached_lengths,
+            "new_lengths": new_lengths,
+            "max_cached_len": max_cached,
+        }
+
+    @staticmethod
+    def unstack_after_update(
+        new_past_kv,
+        stacked_info: dict,
+        old_caches: list[CacheState],
+        new_input_ids: list[torch.Tensor],
+        new_attention_masks: list[torch.Tensor],
+    ) -> list[CacheState]:
+        """
+        Unstack per-session CacheState from a batched decoder output.
+
+        new_past_kv has shape [..., max_cached + max_new, ...] along seq dim.
+        We slice [max_cached - clen_i : max_cached + new_len_i] to recover
+        only the real tokens for each session.
+        """
+        cached_lengths = stacked_info["cached_lengths"]
+        new_lengths = stacked_info["new_lengths"]
+        max_cached = stacked_info["max_cached_len"]
+
+        results = []
+        for i, (old, clen, nlen) in enumerate(zip(old_caches, cached_lengths, new_lengths)):
+            real_start = max_cached - clen
+            real_end = max_cached + nlen
+            sliced_kv = BatchedKVHelper._slice_past_kv(new_past_kv, i, real_start, real_end)
+            new_clen = clen + nlen
+
+            # rebuild input_ids / attention_mask for CacheState
+            new_ids = new_input_ids[i]
+            if new_ids.dim() == 2:
+                new_ids = new_ids[0]
+            new_mask = new_attention_masks[i]
+            if new_mask.dim() == 2:
+                new_mask = new_mask[0]
+
+            if clen > 0:
+                full_ids = torch.cat([old.input_ids, new_ids[:nlen]], dim=0)
+                full_mask = torch.cat([old.attention_mask, new_mask[:nlen]], dim=0)
+            else:
+                full_ids = new_ids[:nlen]
+                full_mask = new_mask[:nlen]
+
+            results.append(CacheState(
+                past_key_values=sliced_kv,
+                input_ids=full_ids,
+                attention_mask=full_mask,
+                position_ids=None,
+                cached_length=new_clen,
+                session_id=old.session_id,
+                metadata=old.metadata.copy(),
+            ))
+
+        return results
+
+    @staticmethod
+    def stack_for_classify(
+        caches: list[CacheState],
+        label_ids: list[torch.Tensor],
+        label_masks: list[torch.Tensor],
+        device: torch.device,
+    ) -> dict:
+        """
+        Prepare a batched decoder forward for the classification stage.
+
+        Label tokens are right-padded. KV caches are prepend-padded.
+        Returns input_ids, attention_mask, position_ids, past_key_values,
+        plus metadata needed for slicing scorer inputs afterwards.
+        """
+        batch_size = len(caches)
+        cached_lengths = [c.cached_length for c in caches]
+        label_lengths = [ids.shape[-1] for ids in label_ids]
+        max_cached = max(cached_lengths)
+        max_label = max(label_lengths)
+
+        padded_label_ids = label_ids[0].new_zeros(batch_size, max_label)
+        padded_label_mask = label_ids[0].new_zeros(batch_size, max_label)
+        position_ids = label_ids[0].new_zeros(batch_size, max_label)
+
+        for i, (ids, mask, clen, llen) in enumerate(
+            zip(label_ids, label_masks, cached_lengths, label_lengths)
+        ):
+            flat_ids = ids[0] if ids.dim() == 2 else ids
+            flat_mask = mask[0] if mask.dim() == 2 else mask
+            padded_label_ids[i, :llen] = flat_ids
+            padded_label_mask[i, :llen] = flat_mask
+            position_ids[i, :llen] = torch.arange(clen, clen + llen, device=device)
+
+        full_mask = torch.zeros(batch_size, max_cached + max_label, dtype=torch.long, device=device)
+        for i, (clen, llen) in enumerate(zip(cached_lengths, label_lengths)):
+            pad = max_cached - clen
+            full_mask[i, pad : pad + clen] = 1
+            full_mask[i, max_cached : max_cached + llen] = 1
+
+        stacked_past_kv = None
+        if max_cached > 0:
+            stacked_past_kv = BatchedKVHelper._stack_past_kv(caches, max_cached, device)
+
+        return {
+            "input_ids": padded_label_ids,
+            "attention_mask": full_mask,          # full: cache + labels (for decoder)
+            "label_mask": padded_label_mask,      # labels only (for scorer)
+            "position_ids": position_ids,
+            "past_key_values": stacked_past_kv,
+            "label_lengths": label_lengths,
+            "max_label_len": max_label,
+        }
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def _stack_past_kv(caches: list[CacheState], max_cached: int, device: torch.device):
+        """Prepend-pad each session's KV to max_cached and stack along batch dim."""
+        from transformers.cache_utils import DynamicCache
+
+        non_empty = [c for c in caches if c.past_key_values is not None]
+        num_layers = len(non_empty[0].past_key_values.layers)
+
+        # Get reference shape from first non-empty layer
+        ref_layer = non_empty[0].past_key_values.layers[0]
+        ref_k = ref_layer.keys  # [1, heads, seq, head_dim]
+
+        stacked = DynamicCache()
+        for layer_idx in range(num_layers):
+            keys, values = [], []
+            for c in caches:
+                kv = c.past_key_values
+                if kv is None or not kv.layers[layer_idx].is_initialized:
+                    k = torch.zeros(1, ref_k.shape[1], max_cached, ref_k.shape[3],
+                                    device=device, dtype=ref_k.dtype)
+                    v = torch.zeros_like(k)
+                else:
+                    layer = kv.layers[layer_idx]
+                    k = layer.keys.to(device)   # [1, heads, clen, head_dim]
+                    v = layer.values.to(device)
+                    clen = k.shape[2]
+                    pad = max_cached - clen
+                    if pad > 0:
+                        zeros_k = torch.zeros(1, k.shape[1], pad, k.shape[3], device=device, dtype=k.dtype)
+                        zeros_v = torch.zeros_like(zeros_k)
+                        k = torch.cat([zeros_k, k], dim=2)
+                        v = torch.cat([zeros_v, v], dim=2)
+                keys.append(k)
+                values.append(v)
+
+            stacked.update(torch.cat(keys, dim=0), torch.cat(values, dim=0), layer_idx)
+
+        return stacked
+
+    @staticmethod
+    def _slice_past_kv(past_kv, batch_idx: int, start: int, end: int):
+        """Extract single-session KV from a batched past_key_values."""
+        if past_kv is None:
+            return None
+
+        from transformers.cache_utils import DynamicCache
+        sliced = DynamicCache()
+        for i, layer in enumerate(past_kv.layers):
+            if not layer.is_initialized:
+                continue
+            sliced.update(
+                layer.keys[batch_idx : batch_idx + 1, :, start:end, :],
+                layer.values[batch_idx : batch_idx + 1, :, start:end, :],
+                i,
+            )
+        return sliced

From 2937d33d94a3af5ad48ccba43b87ec97d499b109 Mon Sep 17 00:00:00 2001
From: BioMikeUkr <mykhayloshtopko@gmail.com>
Date: Fri, 10 Jul 2026 12:52:47 +0000
Subject: [PATCH 02/25] streaming pipeline added

---
 gliclass/streaming/pipeline.py | 443 +++++++++++++++++++++++++++++++++
 1 file changed, 443 insertions(+)
 create mode 100644 gliclass/streaming/pipeline.py

diff --git a/gliclass/streaming/pipeline.py b/gliclass/streaming/pipeline.py
new file mode 100644
index 0000000..22eaa83
--- /dev/null
+++ b/gliclass/streaming/pipeline.py
@@ -0,0 +1,443 @@
+"""
+StreamingPipeline – batched multi-session streaming classification.
+
+Flow per __call__:
+  1. Cleanup sessions absent from current batch.
+  2. Create CacheState for new session_ids.
+  3. Stage 1 – batch update KV caches with new text (skip empty text).
+  4. Stage 2 – resolve strategies; batch classify sessions that triggered.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from .cache import BatchedKVHelper, CacheState, create_empty_cache, truncate_cache
+from .types import SessionInput, SessionOutput
+
+
+class StreamingPipeline:
+    def __init__(
+        self,
+        model,
+        tokenizer,
+        device: str | torch.device = "cpu",
+        max_cache_len: int | None = None,
+        batch_size: int | None = None,
+        offload_to_cpu: bool = False,
+        use_pinned_memory: bool = False,
+        score_on_cpu: bool = False,
+    ):
+        self.model = model
+        self.tokenizer = tokenizer
+        self.device = torch.device(device) if isinstance(device, str) else device
+        self.max_cache_len = max_cache_len
+        self.batch_size = batch_size
+        self.offload_to_cpu = offload_to_cpu
+        self.use_pinned_memory = use_pinned_memory and offload_to_cpu and self.device.type == "cuda"
+        self.score_on_cpu = score_on_cpu
+        self._caches: dict[str, CacheState] = {}
+
+        scorer_target = "cpu" if score_on_cpu else self.device
+        self.model.model.scorer.to(scorer_target)
+
+    # ------------------------------------------------------------------
+    # Public API
+    # ------------------------------------------------------------------
+
+    def __call__(self, inputs: list[SessionInput]) -> list[SessionOutput]:
+        self._cleanup_dead_sessions(inputs)
+        self._ensure_sessions(inputs)
+
+        if self.batch_size is None or len(inputs) <= self.batch_size:
+            return self._process_batch(inputs)
+
+        outputs = [None] * len(inputs)
+        for start in range(0, len(inputs), self.batch_size):
+            chunk = inputs[start : start + self.batch_size]
+            chunk_ids = {inp.session_id for inp in chunk}
+
+            if self.offload_to_cpu:
+                self._load_to_device(chunk_ids, self.device)
+
+            for i, out in enumerate(self._process_batch(chunk)):
+                outputs[start + i] = out
+
+            if self.offload_to_cpu:
+                self._load_to_device(chunk_ids, torch.device("cpu"))
+
+        return outputs
+
+    def _process_batch(self, inputs: list[SessionInput]) -> list[SessionOutput]:
+        tokens_added = self._stage_update(inputs)
+        if self.max_cache_len is not None:
+            self._enforce_max_cache_len()
+        return self._stage_classify(inputs, tokens_added)
+
+    def _load_to_device(self, session_ids: set[str], device: torch.device) -> None:
+        to_cpu = device.type == "cpu"
+        for sid in session_ids:
+            cache = self._caches.get(sid)
+            if cache is None or cache.past_key_values is None:
+                continue
+            if to_cpu and self.use_pinned_memory:
+                self._caches[sid] = _move_cache_pinned(cache)
+            else:
+                self._caches[sid] = _move_cache_nonblocking(cache, device)
+
+        # synchronize before using on GPU so non_blocking transfers complete
+        if not to_cpu and self.use_pinned_memory and self.device.type == "cuda":
+            torch.cuda.synchronize(self.device)
+
+    # ------------------------------------------------------------------
+    # Session management
+    # ------------------------------------------------------------------
+
+    def _cleanup_dead_sessions(self, inputs: list[SessionInput]) -> None:
+        active = {inp.session_id for inp in inputs}
+        for sid in list(self._caches):
+            if sid not in active:
+                del self._caches[sid]
+
+    def _enforce_max_cache_len(self) -> None:
+        for sid, cache in self._caches.items():
+            if cache.cached_length > self.max_cache_len:
+                self._caches[sid] = truncate_cache(cache, self.max_cache_len)
+
+    def _ensure_sessions(self, inputs: list[SessionInput]) -> None:
+        for inp in inputs:
+            if inp.session_id not in self._caches:
+                self._caches[inp.session_id] = create_empty_cache(
+                    session_id=inp.session_id, device=self.device
+                )
+
+    # ------------------------------------------------------------------
+    # Stage 1: update KV caches
+    # ------------------------------------------------------------------
+
+    def _stage_update(self, inputs: list[SessionInput]) -> list[int]:
+        """
+        Update KV cache for each session. Returns tokens_added per input.
+        Sessions with empty text are skipped (tokens_added = 0).
+        """
+        tokens_added = [0] * len(inputs)
+        to_update = [(i, inp) for i, inp in enumerate(inputs) if inp.text]
+
+        if not to_update:
+            return tokens_added
+
+        indices = [i for i, _ in to_update]
+        active_inputs = [inp for _, inp in to_update]
+
+        # Tokenize new text per session
+        tokenized = [
+            self.tokenizer(inp.text, return_tensors="pt", add_special_tokens=False)
+            for inp in active_inputs
+        ]
+        new_ids = [t["input_ids"].to(self.device) for t in tokenized]
+        new_masks = [t["attention_mask"].to(self.device) for t in tokenized]
+
+        caches = [self._caches[inp.session_id] for inp in active_inputs]
+
+        # Check if any session has a non-empty cache (need KV stacking)
+        has_cache = any(c.cached_length > 0 for c in caches)
+
+        if has_cache:
+            stacked = BatchedKVHelper.stack_for_update(caches, new_ids, new_masks, self.device)
+            decoder = self.model.model.decoder_model
+
+            with torch.no_grad():
+                out = decoder(
+                    input_ids=stacked["input_ids"],
+                    attention_mask=stacked["attention_mask"],
+                    position_ids=stacked["position_ids"],
+                    past_key_values=stacked["past_key_values"],
+                    use_cache=True,
+                    return_dict=True,
+                )
+
+            updated = BatchedKVHelper.unstack_after_update(
+                out.past_key_values, stacked, caches, new_ids, new_masks
+            )
+        else:
+            # All caches empty – simple batched forward without stacking
+            updated = self._update_from_scratch(caches, new_ids, new_masks)
+
+        for session_idx, cache, new_len in zip(
+            indices, updated, stacked["new_lengths"] if has_cache else [ids.shape[-1] for ids in new_ids]
+        ):
+            inp = inputs[session_idx]
+            self._caches[inp.session_id] = cache
+            tokens_added[session_idx] = new_len
+
+        return tokens_added
+
+    def _update_from_scratch(
+        self,
+        caches: list[CacheState],
+        new_ids: list[torch.Tensor],
+        new_masks: list[torch.Tensor],
+    ) -> list[CacheState]:
+        """Batched update when all caches are empty (no KV stacking needed)."""
+        max_len = max(ids.shape[-1] for ids in new_ids)
+        batch_size = len(caches)
+
+        padded_ids = new_ids[0].new_zeros(batch_size, max_len)
+        padded_masks = new_ids[0].new_zeros(batch_size, max_len)
+        position_ids = new_ids[0].new_zeros(batch_size, max_len)
+
+        new_lengths = []
+        for i, (ids, mask) in enumerate(zip(new_ids, new_masks)):
+            flat = ids[0] if ids.dim() == 2 else ids
+            flat_m = mask[0] if mask.dim() == 2 else mask
+            nlen = flat.shape[0]
+            padded_ids[i, :nlen] = flat
+            padded_masks[i, :nlen] = flat_m
+            position_ids[i, :nlen] = torch.arange(nlen, device=self.device)
+            new_lengths.append(nlen)
+
+        decoder = self.model.model.decoder_model
+        with torch.no_grad():
+            out = decoder(
+                input_ids=padded_ids,
+                attention_mask=padded_masks,
+                position_ids=position_ids,
+                use_cache=True,
+                return_dict=True,
+            )
+
+        results = []
+        for i, (cache, nlen) in enumerate(zip(caches, new_lengths)):
+            sliced_kv = BatchedKVHelper._slice_past_kv(out.past_key_values, i, 0, nlen)
+            flat_ids = new_ids[i][0] if new_ids[i].dim() == 2 else new_ids[i]
+            flat_mask = new_masks[i][0] if new_masks[i].dim() == 2 else new_masks[i]
+            results.append(CacheState(
+                past_key_values=sliced_kv,
+                input_ids=flat_ids[:nlen],
+                attention_mask=flat_mask[:nlen],
+                cached_length=nlen,
+                session_id=cache.session_id,
+                metadata=cache.metadata.copy(),
+            ))
+
+        return results
+
+    # ------------------------------------------------------------------
+    # Stage 2: classify triggered sessions
+    # ------------------------------------------------------------------
+
+    def _stage_classify(
+        self, inputs: list[SessionInput], tokens_added: list[int]
+    ) -> list[SessionOutput]:
+        # Resolve which sessions trigger classification
+        triggered_indices = []
+        for i, (inp, n_added) in enumerate(zip(inputs, tokens_added)):
+            if n_added == 0:
+                continue  # no new tokens → cache unchanged, skip classification
+            cache = self._caches[inp.session_id]
+            if inp.strategy.should_classify(n_added, cache.cached_length, inp.text):
+                triggered_indices.append(i)
+
+        outputs: list[SessionOutput | None] = [None] * len(inputs)
+
+        # Fill non-triggered outputs
+        for i, inp in enumerate(inputs):
+            if i not in triggered_indices:
+                cache = self._caches[inp.session_id]
+                outputs[i] = SessionOutput(
+                    session_id=inp.session_id,
+                    triggered=False,
+                    predictions=None,
+                    cached_length=cache.cached_length,
+                    tokens_added=tokens_added[i],
+                )
+
+        if not triggered_indices:
+            return outputs
+
+        # Prepare label sequences for triggered sessions
+        triggered_inputs = [inputs[i] for i in triggered_indices]
+        triggered_caches = [
+            inputs[i].strategy.get_window(self._caches[inputs[i].session_id])
+            for i in triggered_indices
+        ]
+
+        label_seqs = [
+            self._prepare_label_seq(inp.labels) for inp in triggered_inputs
+        ]
+        tokenized_labels = [
+            self.tokenizer(seq, return_tensors="pt", add_special_tokens=False)
+            for seq in label_seqs
+        ]
+        label_ids = [t["input_ids"].to(self.device) for t in tokenized_labels]
+        label_masks = [t["attention_mask"].to(self.device) for t in tokenized_labels]
+
+        stacked = BatchedKVHelper.stack_for_classify(
+            triggered_caches, label_ids, label_masks, self.device
+        )
+
+        decoder = self.model.model.decoder_model
+        with torch.no_grad():
+            out = decoder(
+                input_ids=stacked["input_ids"],
+                attention_mask=stacked["attention_mask"],
+                position_ids=stacked["position_ids"],
+                past_key_values=stacked["past_key_values"],
+                use_cache=False,
+                return_dict=True,
+            )
+
+        # Slice label hidden states and pass to scorer
+        hidden = out.last_hidden_state          # [batch, max_cached + max_label, hidden]
+        max_cached = stacked["past_key_values"] and max(c.cached_length for c in triggered_caches) or 0
+        label_lengths = stacked["label_lengths"]
+        max_label = stacked["max_label_len"]
+
+        # hidden for labels starts at index max_cached (prepend-padded KV)
+        # but we need the actual label slice per session
+        scorer = self.model.model.scorer
+        batch_logits = self._run_scorer_batched(
+            scorer, hidden, stacked["input_ids"], stacked["label_mask"],
+            label_lengths, max_label,
+        )
+
+        for local_idx, global_idx in enumerate(triggered_indices):
+            inp = inputs[global_idx]
+            cache = self._caches[inp.session_id]
+            logits = batch_logits[local_idx]               # [num_labels]
+            preds = self._decode_predictions(logits, inp.labels, inp.classification_type)
+            outputs[global_idx] = SessionOutput(
+                session_id=inp.session_id,
+                triggered=True,
+                predictions=preds,
+                cached_length=cache.cached_length,
+                tokens_added=tokens_added[global_idx],
+            )
+
+        return outputs
+
+    # ------------------------------------------------------------------
+    # Helpers
+    # ------------------------------------------------------------------
+
+    def _prepare_label_seq(self, labels: list[str]) -> str:
+        label_token = "<<LABEL>>"
+        sep_token = "<<SEP>>"
+        # Opening <<SEP>> matches training format: text<<SEP>>label1<<LABEL>>...<<SEP>>
+        # Cache contains only text; <<SEP>> boundary must be part of label sequence
+        parts = [sep_token]
+        for label in labels:
+            parts.append(label)
+            parts.append(label_token)
+        parts.append(sep_token)
+        return "".join(parts)
+
+    def _run_scorer_batched(
+        self,
+        scorer,
+        hidden: torch.Tensor,
+        label_ids: torch.Tensor,
+        label_mask: torch.Tensor,
+        label_lengths: list[int],
+        max_label: int,
+    ) -> list[torch.Tensor]:
+        """
+        hidden: [batch, max_cached + max_label, hidden_size]  (full decoder output)
+        We slice out only the label part [max_cached :] per session before scorer.
+        label_ids / label_mask are already label-only ([batch, max_label]).
+        """
+        max_cached_in_hidden = hidden.shape[1] - max_label
+        label_hidden = hidden[:, max_cached_in_hidden:, :]   # [batch, max_label, hidden]
+
+        if self.score_on_cpu:
+            label_hidden = label_hidden.cpu()
+            label_ids = label_ids.cpu()
+            label_mask = label_mask.cpu()
+
+        logits = scorer(
+            hidden_states=label_hidden,
+            input_ids=label_ids,
+            attention_mask=label_mask,
+        )  # [batch, num_labels]
+
+        return [logits[i] for i in range(logits.shape[0])]
+
+    def _decode_predictions(
+        self, logits: torch.Tensor, labels: list[str], classification_type: str
+    ) -> list[dict]:
+        logits = logits[:len(labels)]
+        if classification_type == "single-label":
+            scores = torch.softmax(logits, dim=-1)
+            best = int(scores.argmax().item())
+            return [{"label": labels[best], "score": float(scores[best].detach())}]
+        else:
+            probs = torch.sigmoid(logits)
+            return [
+                {"label": label, "score": float(prob.detach())}
+                for label, prob in zip(labels, probs)
+            ]
+
+
+def _move_cache_pinned(cache: "CacheState") -> "CacheState":
+    """Move KV cache to CPU pinned memory for fast async GPU uploads."""
+    from transformers.cache_utils import DynamicCache, DynamicLayer
+    from .cache import CacheState
+
+    past_kv = cache.past_key_values
+    if not isinstance(past_kv, DynamicCache):
+        return cache.to(torch.device("cpu"))
+
+    new_kv = DynamicCache()
+    for i, layer in enumerate(past_kv.layers):
+        if not layer.is_initialized:
+            continue
+        pinned_k = layer.keys.cpu().pin_memory()
+        pinned_v = layer.values.cpu().pin_memory()
+        # bypass update() to avoid torch.cat losing pinned property
+        while len(new_kv.layers) <= i:
+            new_kv.layers.append(DynamicLayer())
+        new_layer = new_kv.layers[i]
+        new_layer.dtype = pinned_k.dtype
+        new_layer.device = pinned_k.device
+        new_layer.keys = pinned_k
+        new_layer.values = pinned_v
+        new_layer.is_initialized = True
+
+    return CacheState(
+        past_key_values=new_kv,
+        input_ids=cache.input_ids.cpu(),
+        attention_mask=cache.attention_mask.cpu(),
+        position_ids=cache.position_ids.cpu() if cache.position_ids is not None else None,
+        cached_length=cache.cached_length,
+        session_id=cache.session_id,
+        metadata=cache.metadata.copy(),
+    )
+
+
+def _move_cache_nonblocking(cache: "CacheState", device: torch.device) -> "CacheState":
+    """Move KV cache to device with non_blocking=True (async when src is pinned)."""
+    from transformers.cache_utils import DynamicCache
+    from .cache import CacheState
+
+    past_kv = cache.past_key_values
+    if not isinstance(past_kv, DynamicCache):
+        return cache.to(device)
+
+    new_kv = DynamicCache()
+    for i, layer in enumerate(past_kv.layers):
+        if layer.is_initialized:
+            new_kv.update(
+                layer.keys.to(device, non_blocking=True),
+                layer.values.to(device, non_blocking=True),
+                i,
+            )
+
+    return CacheState(
+        past_key_values=new_kv,
+        input_ids=cache.input_ids.to(device, non_blocking=True),
+        attention_mask=cache.attention_mask.to(device, non_blocking=True),
+        position_ids=cache.position_ids.to(device, non_blocking=True) if cache.position_ids is not None else None,
+        cached_length=cache.cached_length,
+        session_id=cache.session_id,
+        metadata=cache.metadata.copy(),
+    )

From 9e1aa2bdfe3bcf67b5f25f3cfd4f6ec82e1e0616 Mon Sep 17 00:00:00 2001
From: BioMikeUkr <mykhayloshtopko@gmail.com>
Date: Fri, 10 Jul 2026 12:53:13 +0000
Subject: [PATCH 03/25] classification triggers added

---
 gliclass/streaming/strategies.py | 100 +++++++++++++++++++++++++++++++
 1 file changed, 100 insertions(+)
 create mode 100644 gliclass/streaming/strategies.py

diff --git a/gliclass/streaming/strategies.py b/gliclass/streaming/strategies.py
new file mode 100644
index 0000000..c4fde4c
--- /dev/null
+++ b/gliclass/streaming/strategies.py
@@ -0,0 +1,100 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from .cache import CacheState
+
+
+class ClassificationStrategy(ABC):
+    """
+    Base class for streaming classification strategies.
+
+    Controls two things:
+      - should_classify: whether to trigger classification on this call
+      - get_window:      which slice of the cache to classify over
+    """
+
+    @abstractmethod
+    def should_classify(self, tokens_added: int, cached_length: int, text: str) -> bool: ...
+
+    def get_window(self, cache_state: "CacheState") -> "CacheState":
+        """Return the cache slice to classify over. Default: full context."""
+        return cache_state
+
+
+class EveryChunkStrategy(ClassificationStrategy):
+    """Classify on every non-empty chunk."""
+
+    def should_classify(self, tokens_added: int, cached_length: int, text: str) -> bool:
+        return tokens_added > 0
+
+
+class EveryNTokensStrategy(ClassificationStrategy):
+    """Classify after accumulating at least N new tokens."""
+
+    def __init__(self, n: int):
+        self.n = n
+        self._accumulated = 0
+
+    def should_classify(self, tokens_added: int, cached_length: int, text: str) -> bool:
+        self._accumulated += tokens_added
+        if self._accumulated >= self.n:
+            self._accumulated = 0
+            return True
+        return False
+
+
+class OnDelimiterStrategy(ClassificationStrategy):
+    """Classify when the incoming text contains a delimiter."""
+
+    def __init__(self, delimiter: str):
+        self.delimiter = delimiter
+
+    def should_classify(self, tokens_added: int, cached_length: int, text: str) -> bool:
+        return self.delimiter in text
+
+
+class NeverStrategy(ClassificationStrategy):
+    """Only update cache, never classify. Useful for pre-filling context."""
+
+    def should_classify(self, tokens_added: int, cached_length: int, text: str) -> bool:
+        return False
+
+
+class SlidingWindowStrategy(ClassificationStrategy):
+    """
+    Classify on every chunk, but only over the last window_size cached tokens.
+
+    Combine with EveryNTokensStrategy behaviour by subclassing or wrapping.
+    """
+
+    def __init__(self, window_size: int):
+        self.window_size = window_size
+
+    def should_classify(self, tokens_added: int, cached_length: int, text: str) -> bool:
+        return tokens_added > 0
+
+    def get_window(self, cache_state: "CacheState") -> "CacheState":
+        from .cache import truncate_cache
+        return truncate_cache(cache_state, self.window_size)
+
+
+class ComposedStrategy(ClassificationStrategy):
+    """
+    Combine a trigger strategy with a window strategy.
+
+    Example:
+        ComposedStrategy(trigger=EveryNTokensStrategy(200), window=SlidingWindowStrategy(512))
+    """
+
+    def __init__(self, trigger: ClassificationStrategy, window: ClassificationStrategy):
+        self.trigger = trigger
+        self.window = window
+
+    def should_classify(self, tokens_added: int, cached_length: int, text: str) -> bool:
+        return self.trigger.should_classify(tokens_added, cached_length, text)
+
+    def get_window(self, cache_state: "CacheState") -> "CacheState":
+        return self.window.get_window(cache_state)

From 0bd779c2c4fb1e200daf4d090a04f5811955cf7e Mon Sep 17 00:00:00 2001
From: BioMikeUkr <mykhayloshtopko@gmail.com>
Date: Fri, 10 Jul 2026 12:53:41 +0000
Subject: [PATCH 04/25] io models added

---
 gliclass/streaming/types.py | 26 ++++++++++++++++++++++++++
 1 file changed, 26 insertions(+)
 create mode 100644 gliclass/streaming/types.py

diff --git a/gliclass/streaming/types.py b/gliclass/streaming/types.py
new file mode 100644
index 0000000..68de682
--- /dev/null
+++ b/gliclass/streaming/types.py
@@ -0,0 +1,26 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from .strategies import ClassificationStrategy
+
+
+@dataclass
+class SessionInput:
+    session_id: str
+    text: str
+    labels: list[str]
+    strategy: "ClassificationStrategy"
+    classification_type: str = "multi-label"  # "multi-label" or "single-label"
+
+
+@dataclass
+class SessionOutput:
+    session_id: str
+    triggered: bool
+    predictions: list[dict] | None
+    cached_length: int
+    tokens_added: int
+    metadata: dict = field(default_factory=dict)

From ecb5accda828eaf8d9bd172f8ac3cd5d23cfd3aa Mon Sep 17 00:00:00 2001
From: BioMikeUkr <mykhayloshtopko@gmail.com>
Date: Fri, 10 Jul 2026 12:54:15 +0000
Subject: [PATCH 05/25] streaming implemented

---
 gliclass/streaming/__init__.py | 12 ++++++++++++
 1 file changed, 12 insertions(+)
 create mode 100644 gliclass/streaming/__init__.py

diff --git a/gliclass/streaming/__init__.py b/gliclass/streaming/__init__.py
new file mode 100644
index 0000000..84e3dde
--- /dev/null
+++ b/gliclass/streaming/__init__.py
@@ -0,0 +1,12 @@
+from .cache import CacheState, BatchedKVHelper, create_empty_cache, truncate_cache
+from .types import SessionInput, SessionOutput
+from .strategies import (
+    ClassificationStrategy,
+    EveryChunkStrategy,
+    EveryNTokensStrategy,
+    OnDelimiterStrategy,
+    NeverStrategy,
+    SlidingWindowStrategy,
+    ComposedStrategy,
+)
+from .pipeline import StreamingPipeline

From fd7a4f204cf4c8dc63ba56b55b961df9203813da Mon Sep 17 00:00:00 2001
From: BioMikeUkr <mykhayloshtopko@gmail.com>
Date: Fri, 10 Jul 2026 12:58:48 +0000
Subject: [PATCH 06/25] ruff checks

---
 gliclass/streaming/__init__.py   | 12 ++--
 gliclass/streaming/cache.py      | 52 +++++++++---------
 gliclass/streaming/pipeline.py   | 94 ++++++++++++++------------------
 gliclass/streaming/strategies.py |  7 ++-
 gliclass/streaming/types.py      |  6 +-
 5 files changed, 80 insertions(+), 91 deletions(-)

diff --git a/gliclass/streaming/__init__.py b/gliclass/streaming/__init__.py
index 84e3dde..e27b343 100644
--- a/gliclass/streaming/__init__.py
+++ b/gliclass/streaming/__init__.py
@@ -1,12 +1,12 @@
-from .cache import CacheState, BatchedKVHelper, create_empty_cache, truncate_cache
+from .cache import CacheState, BatchedKVHelper, truncate_cache, create_empty_cache
 from .types import SessionInput, SessionOutput
+from .pipeline import StreamingPipeline
 from .strategies import (
-    ClassificationStrategy,
+    NeverStrategy,
+    ComposedStrategy,
     EveryChunkStrategy,
-    EveryNTokensStrategy,
     OnDelimiterStrategy,
-    NeverStrategy,
+    EveryNTokensStrategy,
     SlidingWindowStrategy,
-    ComposedStrategy,
+    ClassificationStrategy,
 )
-from .pipeline import StreamingPipeline
diff --git a/gliclass/streaming/cache.py b/gliclass/streaming/cache.py
index a7941ba..e81f3ed 100644
--- a/gliclass/streaming/cache.py
+++ b/gliclass/streaming/cache.py
@@ -8,16 +8,16 @@
 
 from __future__ import annotations
 
-from dataclasses import dataclass, field
 from typing import Any
+from dataclasses import field, dataclass
 
 import torch
 
-
 # ---------------------------------------------------------------------------
 # CacheState
 # ---------------------------------------------------------------------------
 
+
 @dataclass
 class CacheState:
     past_key_values: Any
@@ -28,7 +28,7 @@ class CacheState:
     session_id: str | None = None
     metadata: dict = field(default_factory=dict)
 
-    def to(self, device: torch.device) -> "CacheState":
+    def to(self, device: torch.device) -> CacheState:
         return CacheState(
             past_key_values=_move_past_kv(self.past_key_values, device),
             input_ids=self.input_ids.to(device),
@@ -85,6 +85,7 @@ def _deep_copy_past_kv(past_kv):
     if past_kv is None:
         return None
     from transformers.cache_utils import DynamicCache
+
     if isinstance(past_kv, DynamicCache):
         new = DynamicCache()
         for i, layer in enumerate(past_kv.layers):
@@ -103,20 +104,21 @@ def _crop_past_kv(past_kv, max_length: int):
     if past_kv is None:
         return None
     from transformers.cache_utils import DynamicCache
+
     if isinstance(past_kv, DynamicCache):
         new = DynamicCache()
         for i, layer in enumerate(past_kv.layers):
             if layer.is_initialized:
-                new.update(layer.keys[:, :, -max_length:, :].clone(),
-                           layer.values[:, :, -max_length:, :].clone(), i)
+                new.update(layer.keys[:, :, -max_length:, :].clone(), layer.values[:, :, -max_length:, :].clone(), i)
         return new
     return None
 
 
 # ---------------------------------------------------------------------------
-# BatchedKVHelper  – stack / unstack heterogeneous KV caches
+# BatchedKVHelper  - stack / unstack heterogeneous KV caches
 # ---------------------------------------------------------------------------
 
+
 class BatchedKVHelper:
     """
     Utilities for batching decoder forwards across sessions with different
@@ -132,7 +134,7 @@ class BatchedKVHelper:
     @staticmethod
     def stack_for_update(
         caches: list[CacheState],
-        new_input_ids: list[torch.Tensor],      # one per session, already on device
+        new_input_ids: list[torch.Tensor],  # one per session, already on device
         new_attention_masks: list[torch.Tensor],
         device: torch.device,
     ) -> dict:
@@ -165,7 +167,7 @@ def stack_for_update(
         full_mask = torch.zeros(batch_size, max_cached + max_new, dtype=torch.long, device=device)
         for i, (clen, nlen) in enumerate(zip(cached_lengths, new_lengths)):
             pad = max_cached - clen
-            full_mask[i, pad : pad + clen] = 1          # real cached tokens
+            full_mask[i, pad : pad + clen] = 1  # real cached tokens
             full_mask[i, max_cached : max_cached + nlen] = 1  # real new tokens
 
         # --- stack past_key_values (prepend-pad to max_cached) ---
@@ -224,15 +226,17 @@ def unstack_after_update(
                 full_ids = new_ids[:nlen]
                 full_mask = new_mask[:nlen]
 
-            results.append(CacheState(
-                past_key_values=sliced_kv,
-                input_ids=full_ids,
-                attention_mask=full_mask,
-                position_ids=None,
-                cached_length=new_clen,
-                session_id=old.session_id,
-                metadata=old.metadata.copy(),
-            ))
+            results.append(
+                CacheState(
+                    past_key_values=sliced_kv,
+                    input_ids=full_ids,
+                    attention_mask=full_mask,
+                    position_ids=None,
+                    cached_length=new_clen,
+                    session_id=old.session_id,
+                    metadata=old.metadata.copy(),
+                )
+            )
 
         return results
 
@@ -260,9 +264,7 @@ def stack_for_classify(
         padded_label_mask = label_ids[0].new_zeros(batch_size, max_label)
         position_ids = label_ids[0].new_zeros(batch_size, max_label)
 
-        for i, (ids, mask, clen, llen) in enumerate(
-            zip(label_ids, label_masks, cached_lengths, label_lengths)
-        ):
+        for i, (ids, mask, clen, llen) in enumerate(zip(label_ids, label_masks, cached_lengths, label_lengths)):
             flat_ids = ids[0] if ids.dim() == 2 else ids
             flat_mask = mask[0] if mask.dim() == 2 else mask
             padded_label_ids[i, :llen] = flat_ids
@@ -281,8 +283,8 @@ def stack_for_classify(
 
         return {
             "input_ids": padded_label_ids,
-            "attention_mask": full_mask,          # full: cache + labels (for decoder)
-            "label_mask": padded_label_mask,      # labels only (for scorer)
+            "attention_mask": full_mask,  # full: cache + labels (for decoder)
+            "label_mask": padded_label_mask,  # labels only (for scorer)
             "position_ids": position_ids,
             "past_key_values": stacked_past_kv,
             "label_lengths": label_lengths,
@@ -311,12 +313,11 @@ def _stack_past_kv(caches: list[CacheState], max_cached: int, device: torch.devi
             for c in caches:
                 kv = c.past_key_values
                 if kv is None or not kv.layers[layer_idx].is_initialized:
-                    k = torch.zeros(1, ref_k.shape[1], max_cached, ref_k.shape[3],
-                                    device=device, dtype=ref_k.dtype)
+                    k = torch.zeros(1, ref_k.shape[1], max_cached, ref_k.shape[3], device=device, dtype=ref_k.dtype)
                     v = torch.zeros_like(k)
                 else:
                     layer = kv.layers[layer_idx]
-                    k = layer.keys.to(device)   # [1, heads, clen, head_dim]
+                    k = layer.keys.to(device)  # [1, heads, clen, head_dim]
                     v = layer.values.to(device)
                     clen = k.shape[2]
                     pad = max_cached - clen
@@ -339,6 +340,7 @@ def _slice_past_kv(past_kv, batch_idx: int, start: int, end: int):
             return None
 
         from transformers.cache_utils import DynamicCache
+
         sliced = DynamicCache()
         for i, layer in enumerate(past_kv.layers):
             if not layer.is_initialized:
diff --git a/gliclass/streaming/pipeline.py b/gliclass/streaming/pipeline.py
index 22eaa83..88e054e 100644
--- a/gliclass/streaming/pipeline.py
+++ b/gliclass/streaming/pipeline.py
@@ -1,18 +1,18 @@
 """
-StreamingPipeline – batched multi-session streaming classification.
+StreamingPipeline - batched multi-session streaming classification.
 
 Flow per __call__:
   1. Cleanup sessions absent from current batch.
   2. Create CacheState for new session_ids.
-  3. Stage 1 – batch update KV caches with new text (skip empty text).
-  4. Stage 2 – resolve strategies; batch classify sessions that triggered.
+  3. Stage 1 - batch update KV caches with new text (skip empty text).
+  4. Stage 2 - resolve strategies; batch classify sessions that triggered.
 """
 
 from __future__ import annotations
 
 import torch
 
-from .cache import BatchedKVHelper, CacheState, create_empty_cache, truncate_cache
+from .cache import CacheState, BatchedKVHelper, truncate_cache, create_empty_cache
 from .types import SessionInput, SessionOutput
 
 
@@ -107,9 +107,7 @@ def _enforce_max_cache_len(self) -> None:
     def _ensure_sessions(self, inputs: list[SessionInput]) -> None:
         for inp in inputs:
             if inp.session_id not in self._caches:
-                self._caches[inp.session_id] = create_empty_cache(
-                    session_id=inp.session_id, device=self.device
-                )
+                self._caches[inp.session_id] = create_empty_cache(session_id=inp.session_id, device=self.device)
 
     # ------------------------------------------------------------------
     # Stage 1: update KV caches
@@ -130,10 +128,7 @@ def _stage_update(self, inputs: list[SessionInput]) -> list[int]:
         active_inputs = [inp for _, inp in to_update]
 
         # Tokenize new text per session
-        tokenized = [
-            self.tokenizer(inp.text, return_tensors="pt", add_special_tokens=False)
-            for inp in active_inputs
-        ]
+        tokenized = [self.tokenizer(inp.text, return_tensors="pt", add_special_tokens=False) for inp in active_inputs]
         new_ids = [t["input_ids"].to(self.device) for t in tokenized]
         new_masks = [t["attention_mask"].to(self.device) for t in tokenized]
 
@@ -156,11 +151,9 @@ def _stage_update(self, inputs: list[SessionInput]) -> list[int]:
                     return_dict=True,
                 )
 
-            updated = BatchedKVHelper.unstack_after_update(
-                out.past_key_values, stacked, caches, new_ids, new_masks
-            )
+            updated = BatchedKVHelper.unstack_after_update(out.past_key_values, stacked, caches, new_ids, new_masks)
         else:
-            # All caches empty – simple batched forward without stacking
+            # All caches empty - simple batched forward without stacking
             updated = self._update_from_scratch(caches, new_ids, new_masks)
 
         for session_idx, cache, new_len in zip(
@@ -211,14 +204,16 @@ def _update_from_scratch(
             sliced_kv = BatchedKVHelper._slice_past_kv(out.past_key_values, i, 0, nlen)
             flat_ids = new_ids[i][0] if new_ids[i].dim() == 2 else new_ids[i]
             flat_mask = new_masks[i][0] if new_masks[i].dim() == 2 else new_masks[i]
-            results.append(CacheState(
-                past_key_values=sliced_kv,
-                input_ids=flat_ids[:nlen],
-                attention_mask=flat_mask[:nlen],
-                cached_length=nlen,
-                session_id=cache.session_id,
-                metadata=cache.metadata.copy(),
-            ))
+            results.append(
+                CacheState(
+                    past_key_values=sliced_kv,
+                    input_ids=flat_ids[:nlen],
+                    attention_mask=flat_mask[:nlen],
+                    cached_length=nlen,
+                    session_id=cache.session_id,
+                    metadata=cache.metadata.copy(),
+                )
+            )
 
         return results
 
@@ -226,9 +221,7 @@ def _update_from_scratch(
     # Stage 2: classify triggered sessions
     # ------------------------------------------------------------------
 
-    def _stage_classify(
-        self, inputs: list[SessionInput], tokens_added: list[int]
-    ) -> list[SessionOutput]:
+    def _stage_classify(self, inputs: list[SessionInput], tokens_added: list[int]) -> list[SessionOutput]:
         # Resolve which sessions trigger classification
         triggered_indices = []
         for i, (inp, n_added) in enumerate(zip(inputs, tokens_added)):
@@ -258,23 +251,15 @@ def _stage_classify(
         # Prepare label sequences for triggered sessions
         triggered_inputs = [inputs[i] for i in triggered_indices]
         triggered_caches = [
-            inputs[i].strategy.get_window(self._caches[inputs[i].session_id])
-            for i in triggered_indices
+            inputs[i].strategy.get_window(self._caches[inputs[i].session_id]) for i in triggered_indices
         ]
 
-        label_seqs = [
-            self._prepare_label_seq(inp.labels) for inp in triggered_inputs
-        ]
-        tokenized_labels = [
-            self.tokenizer(seq, return_tensors="pt", add_special_tokens=False)
-            for seq in label_seqs
-        ]
+        label_seqs = [self._prepare_label_seq(inp.labels) for inp in triggered_inputs]
+        tokenized_labels = [self.tokenizer(seq, return_tensors="pt", add_special_tokens=False) for seq in label_seqs]
         label_ids = [t["input_ids"].to(self.device) for t in tokenized_labels]
         label_masks = [t["attention_mask"].to(self.device) for t in tokenized_labels]
 
-        stacked = BatchedKVHelper.stack_for_classify(
-            triggered_caches, label_ids, label_masks, self.device
-        )
+        stacked = BatchedKVHelper.stack_for_classify(triggered_caches, label_ids, label_masks, self.device)
 
         decoder = self.model.model.decoder_model
         with torch.no_grad():
@@ -288,8 +273,8 @@ def _stage_classify(
             )
 
         # Slice label hidden states and pass to scorer
-        hidden = out.last_hidden_state          # [batch, max_cached + max_label, hidden]
-        max_cached = stacked["past_key_values"] and max(c.cached_length for c in triggered_caches) or 0
+        hidden = out.last_hidden_state  # [batch, max_cached + max_label, hidden]
+        max_cached = max(c.cached_length for c in triggered_caches) if stacked["past_key_values"] else 0  # noqa: F841
         label_lengths = stacked["label_lengths"]
         max_label = stacked["max_label_len"]
 
@@ -297,14 +282,18 @@ def _stage_classify(
         # but we need the actual label slice per session
         scorer = self.model.model.scorer
         batch_logits = self._run_scorer_batched(
-            scorer, hidden, stacked["input_ids"], stacked["label_mask"],
-            label_lengths, max_label,
+            scorer,
+            hidden,
+            stacked["input_ids"],
+            stacked["label_mask"],
+            label_lengths,
+            max_label,
         )
 
         for local_idx, global_idx in enumerate(triggered_indices):
             inp = inputs[global_idx]
             cache = self._caches[inp.session_id]
-            logits = batch_logits[local_idx]               # [num_labels]
+            logits = batch_logits[local_idx]  # [num_labels]
             preds = self._decode_predictions(logits, inp.labels, inp.classification_type)
             outputs[global_idx] = SessionOutput(
                 session_id=inp.session_id,
@@ -347,7 +336,7 @@ def _run_scorer_batched(
         label_ids / label_mask are already label-only ([batch, max_label]).
         """
         max_cached_in_hidden = hidden.shape[1] - max_label
-        label_hidden = hidden[:, max_cached_in_hidden:, :]   # [batch, max_label, hidden]
+        label_hidden = hidden[:, max_cached_in_hidden:, :]  # [batch, max_label, hidden]
 
         if self.score_on_cpu:
             label_hidden = label_hidden.cpu()
@@ -362,25 +351,21 @@ def _run_scorer_batched(
 
         return [logits[i] for i in range(logits.shape[0])]
 
-    def _decode_predictions(
-        self, logits: torch.Tensor, labels: list[str], classification_type: str
-    ) -> list[dict]:
-        logits = logits[:len(labels)]
+    def _decode_predictions(self, logits: torch.Tensor, labels: list[str], classification_type: str) -> list[dict]:
+        logits = logits[: len(labels)]
         if classification_type == "single-label":
             scores = torch.softmax(logits, dim=-1)
             best = int(scores.argmax().item())
             return [{"label": labels[best], "score": float(scores[best].detach())}]
         else:
             probs = torch.sigmoid(logits)
-            return [
-                {"label": label, "score": float(prob.detach())}
-                for label, prob in zip(labels, probs)
-            ]
+            return [{"label": label, "score": float(prob.detach())} for label, prob in zip(labels, probs)]
 
 
-def _move_cache_pinned(cache: "CacheState") -> "CacheState":
+def _move_cache_pinned(cache: CacheState) -> CacheState:
     """Move KV cache to CPU pinned memory for fast async GPU uploads."""
     from transformers.cache_utils import DynamicCache, DynamicLayer
+
     from .cache import CacheState
 
     past_kv = cache.past_key_values
@@ -414,9 +399,10 @@ def _move_cache_pinned(cache: "CacheState") -> "CacheState":
     )
 
 
-def _move_cache_nonblocking(cache: "CacheState", device: torch.device) -> "CacheState":
+def _move_cache_nonblocking(cache: CacheState, device: torch.device) -> CacheState:
     """Move KV cache to device with non_blocking=True (async when src is pinned)."""
     from transformers.cache_utils import DynamicCache
+
     from .cache import CacheState
 
     past_kv = cache.past_key_values
diff --git a/gliclass/streaming/strategies.py b/gliclass/streaming/strategies.py
index c4fde4c..2cf716d 100644
--- a/gliclass/streaming/strategies.py
+++ b/gliclass/streaming/strategies.py
@@ -19,7 +19,7 @@ class ClassificationStrategy(ABC):
     @abstractmethod
     def should_classify(self, tokens_added: int, cached_length: int, text: str) -> bool: ...
 
-    def get_window(self, cache_state: "CacheState") -> "CacheState":
+    def get_window(self, cache_state: CacheState) -> CacheState:
         """Return the cache slice to classify over. Default: full context."""
         return cache_state
 
@@ -76,8 +76,9 @@ def __init__(self, window_size: int):
     def should_classify(self, tokens_added: int, cached_length: int, text: str) -> bool:
         return tokens_added > 0
 
-    def get_window(self, cache_state: "CacheState") -> "CacheState":
+    def get_window(self, cache_state: CacheState) -> CacheState:
         from .cache import truncate_cache
+
         return truncate_cache(cache_state, self.window_size)
 
 
@@ -96,5 +97,5 @@ def __init__(self, trigger: ClassificationStrategy, window: ClassificationStrate
     def should_classify(self, tokens_added: int, cached_length: int, text: str) -> bool:
         return self.trigger.should_classify(tokens_added, cached_length, text)
 
-    def get_window(self, cache_state: "CacheState") -> "CacheState":
+    def get_window(self, cache_state: CacheState) -> CacheState:
         return self.window.get_window(cache_state)
diff --git a/gliclass/streaming/types.py b/gliclass/streaming/types.py
index 68de682..c73db1a 100644
--- a/gliclass/streaming/types.py
+++ b/gliclass/streaming/types.py
@@ -1,7 +1,7 @@
 from __future__ import annotations
 
-from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
+from dataclasses import field, dataclass
 
 if TYPE_CHECKING:
     from .strategies import ClassificationStrategy
@@ -12,7 +12,7 @@ class SessionInput:
     session_id: str
     text: str
     labels: list[str]
-    strategy: "ClassificationStrategy"
+    strategy: ClassificationStrategy
     classification_type: str = "multi-label"  # "multi-label" or "single-label"
 
 

From 6737b9f21c83a02e8c755a7741a47a097ab33b18 Mon Sep 17 00:00:00 2001
From: BioMikeUkr <mykhayloshtopko@gmail.com>
Date: Fri, 10 Jul 2026 12:59:18 +0000
Subject: [PATCH 07/25] new scorer param added

---
 gliclass/config.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/gliclass/config.py b/gliclass/config.py
index 8698977..dd11728 100644
--- a/gliclass/config.py
+++ b/gliclass/config.py
@@ -59,6 +59,7 @@ def __init__(
         embed_class_token=True,
         dropout=0.1,
         use_segment_embeddings=False,
+        scorer_encoder_num_layers=2,
         **kwargs,
     ):
         if isinstance(encoder_config, dict):
@@ -145,4 +146,7 @@ def __init__(
         self.pad_token_id = self.encoder_config.pad_token_id
         self.dropout = dropout
         self.use_segment_embeddings = use_segment_embeddings
-        super().__init__(**kwargs)
+
+        self.scorer_encoder_num_layers = scorer_encoder_num_layers
+
+        super().__init__(problem_type=problem_type, **kwargs)

From 3b90edfe2ebf1508fc408b9d58cd3863304f3b16 Mon Sep 17 00:00:00 2001
From: BioMikeUkr <mykhayloshtopko@gmail.com>
Date: Fri, 10 Jul 2026 13:01:00 +0000
Subject: [PATCH 08/25] new schema for decoder-kv model

---
 gliclass/data_processing.py | 36 ++++++++++++++++++++++++++++++++++++
 1 file changed, 36 insertions(+)

diff --git a/gliclass/data_processing.py b/gliclass/data_processing.py
index 8cc361e..700c5f5 100644
--- a/gliclass/data_processing.py
+++ b/gliclass/data_processing.py
@@ -292,6 +292,40 @@ def tokenize_and_prepare_labels_for_uniencoder(self, example):
         tokenized_inputs["input_texts"] = example["text"]
         return tokenized_inputs
 
+    def tokenize_and_prepare_labels_for_decoder_kv(self, example):
+        """
+        Prepare input for decoder-kv architecture.
+
+        Format: [prompt][examples]text<<SEP>>label1<<LABEL>>label2<<LABEL>>...<<SEP>>
+        """
+        if self.shuffle_labels:
+            random.shuffle(example["all_labels"])
+
+        prompt = example.get("prompt", "")
+        examples_text = self.format_examples(example)
+        text = str(example["text"])
+
+        input_parts = []
+        if prompt:
+            input_parts.append(prompt)
+        input_parts.append(examples_text)
+        input_parts.append(text)
+        input_parts.append(self.sep_token)
+
+        for label in example["all_labels"]:
+            input_parts.append(label)
+            input_parts.append(self.label_token)
+        input_parts.append(self.sep_token)
+
+        input_text = "".join(input_parts)
+        label2idx = {label: idx for idx, label in enumerate(example["all_labels"])}
+
+        tokenized_inputs = self.tokenize(input_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"]
+        return tokenized_inputs
+
     def tokenize_and_prepare_labels_for_encoder_decoder(self, example):
         if self.shuffle_labels:
             random.shuffle(example["all_labels"])
@@ -357,6 +391,8 @@ def __getitem__(self, idx):
             model_inputs = self.tokenize_and_prepare_labels_for_encoder_decoder(example)
         elif self.architecture_type in {"bi-encoder", "bi-encoder-fused"}:
             model_inputs = self.tokenize_and_prepare_labels_for_biencoder(example)
+        elif self.architecture_type == "decoder-kv":
+            model_inputs = self.tokenize_and_prepare_labels_for_decoder_kv(example)
         else:
             raise NotImplementedError("This architecture type is not implemented.")
         return model_inputs

From 6de99ab903d79ce8e8a1dda150bf7fd5430c1b6d Mon Sep 17 00:00:00 2001
From: BioMikeUkr <mykhayloshtopko@gmail.com>
Date: Fri, 10 Jul 2026 13:01:31 +0000
Subject: [PATCH 09/25] decoder kv model added

---
 gliclass/model.py | 224 ++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 224 insertions(+)

diff --git a/gliclass/model.py b/gliclass/model.py
index bc0dd97..696a4ce 100644
--- a/gliclass/model.py
+++ b/gliclass/model.py
@@ -950,6 +950,224 @@ def forward(
         )
 
 
+class GLiClassDecoderKV(nn.Module):
+    """
+    Decoder-KV architecture with dynamic KV cache for streaming classification.
+
+    Sequence format: [prompt][examples]text<<SEP>>label1<<LABEL>>label2<<LABEL>>...<<SEP>>
+
+    Cached part: [prompt][examples]
+    New part each time: text<<SEP>>labels...<<SEP>>
+
+    Flow:
+    1. Decoder backbone (Qwen3) processes full sequence with past_key_values
+    2. Update KV cache ONLY with [prompt][examples]text part (before labels <<SEP>>)
+    3. Hidden states → DecoderKVScorer (bidirectional encoder + extraction + MLP)
+    """
+
+    def __init__(self, config: GLiClassModelConfig, from_pretrained=False):
+        super().__init__()
+        self.config = config
+
+        if config.encoder_config is None:
+            if config.encoder_model_name is None:
+                raise ValueError("You need to specify encoder_model_name for decoder backbone (Qwen3).")
+            config.encoder_config = AutoConfig.from_pretrained(config.encoder_model_name)
+
+        config_name = config.encoder_config.__class__.__name__
+
+        if config_name == "Qwen3_5TextConfig":
+            from transformers.models.qwen3_5 import Qwen3_5TextModel
+
+            ModelClass = Qwen3_5TextModel
+        elif config_name == "Qwen3_5Config":
+            from transformers.models.qwen3_5 import Qwen3_5TextModel
+
+            ModelClass = Qwen3_5TextModel
+            config.encoder_config = config.encoder_config.text_config
+        elif config_name == "Qwen3Config":
+            from transformers import Qwen3Model
+
+            ModelClass = Qwen3Model
+        else:
+            raise ValueError(f"decoder-kv architecture requires Qwen3 or Qwen3.5. Got: {config_name}")
+
+        if from_pretrained:
+            self.decoder_model = ModelClass.from_pretrained(config.encoder_model_name)
+        else:
+            self.decoder_model = ModelClass(config.encoder_config)
+
+        if config.vocab_size is not None and hasattr(self.decoder_model, "resize_token_embeddings"):
+            current_vocab = self.decoder_model.config.vocab_size
+            if current_vocab != config.vocab_size:
+                self.decoder_model.resize_token_embeddings(config.vocab_size)
+
+        from .scorers import DecoderKVScorer
+
+        self.scorer = DecoderKVScorer(config)
+
+        self.vocab_size = config.vocab_size
+        self.pad_token_id = config.pad_token_id if config.pad_token_id is not None else -1
+        self.num_labels = -1
+
+        self.sep_token_id = config.sep_token_index
+
+    def get_loss(self, logits, labels):
+        loss = None
+        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:
+                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.")
+        return loss
+
+    def _extract_label_section(
+        self,
+        hidden_states: torch.Tensor,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Extract label section hidden states from full-sequence decoder output.
+
+        Training sequence: [prompt][examples<<SEP>>]text<<SEP>>label1<<LABEL>>...<<SEP>>
+        Scorer receives:   label1<<LABEL>>...<<SEP>>   (matches classify() inference format)
+
+        Start pos = first token after last <<SEP>> before first <<LABEL>>.
+        End pos = last real token (inclusive, determined by attention_mask).
+
+        Returns:
+            padded_hidden: (batch, max_label_len, hidden_size)
+            padded_ids:    (batch, max_label_len)
+            label_mask:    (batch, max_label_len)
+        """
+        batch_size, _, hidden_size = hidden_states.shape
+        sep_id = self.sep_token_id
+        label_id = self.config.class_token_index
+
+        slices_h, slices_ids = [], []
+
+        for i in range(batch_size):
+            real_len = int(attention_mask[i].sum().item())
+            ids_i = input_ids[i, :real_len]
+
+            label_pos = (ids_i == label_id).nonzero(as_tuple=False)
+            sep_pos = (ids_i == sep_id).nonzero(as_tuple=False)
+
+            if label_pos.numel() == 0 or sep_pos.numel() == 0:
+                # Fallback: use everything up to real_len
+                slices_h.append(hidden_states[i, :real_len])
+                slices_ids.append(ids_i)
+                continue
+
+            first_label = label_pos[0].item()
+            # last <<SEP>> strictly before first <<LABEL>>
+            seps_before = sep_pos[sep_pos < first_label]
+
+            if seps_before.numel() == 0:
+                slices_h.append(hidden_states[i, :real_len])
+                slices_ids.append(ids_i)
+                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])
+
+        max_len = max(s.shape[0] for s in slices_h)
+        padded_hidden = hidden_states.new_zeros(batch_size, max_len, hidden_size)
+        padded_ids = input_ids.new_zeros(batch_size, max_len)
+        label_mask = attention_mask.new_zeros(batch_size, max_len)
+
+        for i, (h, ids) in enumerate(zip(slices_h, slices_ids)):
+            length = h.shape[0]
+            padded_hidden[i, :length] = h
+            padded_ids[i, :length] = ids
+            label_mask[i, :length] = 1
+
+        return padded_hidden, padded_ids, label_mask
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        past_key_values=None,
+        labels: torch.Tensor | None = None,
+        return_dict: bool | None = None,
+        use_cache: bool = False,
+        **kwargs,
+    ):
+        """Forward pass.
+
+        Args:
+            input_ids: Full sequence [prompt][examples]text<<SEP>>labels...<<SEP>>
+            attention_mask: Attention mask
+            past_key_values: Cached KV for [prompt][examples] (optional)
+            labels: Classification labels
+            return_dict: Return dict output
+            use_cache: Whether to return updated cache
+
+        Returns:
+            GLiClassOutput with logits, loss, and optionally past_key_values
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        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,
+            **kwargs,
+        )
+
+        hidden_states = decoder_outputs.last_hidden_state
+
+        label_hidden, label_ids, label_mask = self._extract_label_section(hidden_states, input_ids, attention_mask)
+
+        logits = self.scorer(
+            hidden_states=label_hidden,
+            input_ids=label_ids,
+            attention_mask=label_mask,
+        )
+
+        if labels is not None:
+            self.num_labels = logits.shape[-1]
+
+        loss = self.get_loss(logits, labels)
+
+        if not return_dict:
+            output = (logits,)
+            if use_cache:
+                output += (decoder_outputs.past_key_values,)
+            return (loss, *output) if loss is not None else output
+
+        return GLiClassOutput(
+            loss=loss,
+            logits=logits,
+            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,
+        )
+
+
 class GLiClassModel(GLiClassPreTrainedModel):
     def __init__(self, config, from_pretrained=False):
         super().__init__(config)
@@ -963,6 +1181,8 @@ def __init__(self, config, from_pretrained=False):
             self.model = GLiClassEncoderDecoder(config, from_pretrained)
         elif config.architecture_type == "encoder-decoder-cls":
             self.model = GLiClassEncoderDecoderCLS(config, from_pretrained)
+        elif config.architecture_type == "decoder-kv":
+            self.model = GLiClassDecoderKV(config, from_pretrained)
         self.post_init()
 
     def get_input_embeddings(self):
@@ -1006,6 +1226,8 @@ def tie_weights(self, recompute_mapping=True, missing_keys=None):
             encoder_model = self.model.encoder_decoder_model
         elif self.config.architecture_type in {"bi-encoder", "bi-encoder-fused"}:
             encoder_model = self.model.encoder_model
+        elif self.config.architecture_type == "decoder-kv":
+            encoder_model = self.model.decoder_model
         else:
             raise NotImplementedError("Tie weights is not implemented for this architecture type")
 
@@ -1045,6 +1267,8 @@ def resize_token_embeddings(self, new_num_tokens: int | None = None, pad_to_mult
             model_embeds = self.model.encoder_decoder_model.resize_token_embeddings(new_num_tokens, pad_to_multiple_of)
         elif self.config.architecture_type in {"bi-encoder-fused"}:
             model_embeds = self.model.encoder_model.resize_token_embeddings(new_num_tokens, pad_to_multiple_of)
+        elif self.config.architecture_type == "decoder-kv":
+            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

From a093f442607d2b9197b826b071d7bc2d98516cd3 Mon Sep 17 00:00:00 2001
From: BioMikeUkr <mykhayloshtopko@gmail.com>
Date: Fri, 10 Jul 2026 13:02:03 +0000
Subject: [PATCH 10/25] decoder kv model support added

---
 gliclass/pipeline.py | 53 ++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 53 insertions(+)

diff --git a/gliclass/pipeline.py b/gliclass/pipeline.py
index 6d05cf0..0fc9d41 100644
--- a/gliclass/pipeline.py
+++ b/gliclass/pipeline.py
@@ -760,6 +760,55 @@ def prepare_inputs(self, texts, labels, same_labels=False, examples=None, prompt
         return tokenized_inputs
 
 
+class DecoderKVZeroShotClassificationPipeline(BaseZeroShotClassificationPipeline):
+    """Classic (non-streaming) pipeline for decoder-kv architecture.
+
+    Formats input as: [prompt][examples<<SEP>>]text<<SEP>>label1<<LABEL>>...<<SEP>>
+    Runs a single batched forward pass per batch — no KV cache reuse.
+    """
+
+    def prepare_input(self, text: str, labels: list[str], examples=None, prompt: str | None = None) -> str:
+        parts = []
+        if prompt:
+            parts.append(prompt)
+        if examples:
+            parts.append(self._format_examples_for_input(examples))
+        parts.append(text)
+        parts.append(self.sep_token)
+        for label in labels:
+            parts.append(label)
+            parts.append(self.label_token)
+        parts.append(self.sep_token)
+        return "".join(parts)
+
+    def prepare_inputs(self, texts, labels, same_labels=False, examples=None, prompt=None):
+        inputs = []
+        if same_labels:
+            for i, text in enumerate(texts):
+                inputs.append(
+                    self.prepare_input(
+                        text,
+                        labels,
+                        self._get_text_examples(examples, i),
+                        self._format_prompt(prompt, i),
+                    )
+                )
+        else:
+            for i, (text, labels_) in enumerate(zip(texts, labels)):
+                inputs.append(
+                    self.prepare_input(
+                        text,
+                        labels_,
+                        self._get_text_examples(examples, i),
+                        self._format_prompt(prompt, i),
+                    )
+                )
+
+        return self.tokenizer(
+            inputs, truncation=True, max_length=self.max_length, padding="longest", return_tensors="pt"
+        ).to(self.device)
+
+
 class ZeroShotClassificationPipeline:
     """
     Main pipeline class for zero-shot classification with GLiClass models.
@@ -880,6 +929,10 @@ def __init__(
             self.pipe = BiEncoderZeroShotClassificationPipeline(
                 model, tokenizer, max_classes, max_length, classification_type, device, progress_bar, label_separator
             )
+        elif model.config.architecture_type == "decoder-kv":
+            self.pipe = DecoderKVZeroShotClassificationPipeline(
+                model, tokenizer, max_classes, max_length, classification_type, device, progress_bar, label_separator
+            )
         else:
             raise NotImplementedError("This architecture is not implemented")
 

From fedb4a8bf79744ffdc0ea731847fd5fc6f8cad45 Mon Sep 17 00:00:00 2001
From: BioMikeUkr <mykhayloshtopko@gmail.com>
Date: Fri, 10 Jul 2026 13:02:36 +0000
Subject: [PATCH 11/25] Encoder scorer added

---
 gliclass/scorers.py | 137 ++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 137 insertions(+)

diff --git a/gliclass/scorers.py b/gliclass/scorers.py
index b33749c..0a93cbf 100644
--- a/gliclass/scorers.py
+++ b/gliclass/scorers.py
@@ -169,10 +169,147 @@ 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 DecoderKVScorer(nn.Module):
+    """
+    Scorer for decoder-kv architecture with built-in bidirectional encoder and representation extraction.
+
+    Sequence format: [prompt][examples]text<<SEP>>label1<<LABEL>>label2<<LABEL>>...<<SEP>>
+
+    Flow:
+    1. Takes hidden states from decoder backbone
+    2. Applies bidirectional scorer_encoder (DebertaV2Encoder without embeddings)
+    3. Extracts text repr from last <<SEP>> before labels
+    4. Extracts label repr from each <<LABEL>> token
+    5. Computes scores via MLP (concat text + label)
+    """
+
+    def __init__(self, config, **kwargs):
+        super().__init__()
+        self.config = config
+
+        from transformers import DebertaV2Config
+        from transformers.models.deberta_v2.modeling_deberta_v2 import DebertaV2Encoder
+
+        num_layers = getattr(config, "scorer_encoder_num_layers", 2)
+        num_heads = max(1, config.hidden_size // 64)
+
+        encoder_config = DebertaV2Config(
+            hidden_size=config.hidden_size,
+            num_hidden_layers=num_layers,
+            num_attention_heads=num_heads,
+            intermediate_size=config.hidden_size * 4,
+            relative_attention=True,
+            pos_att_type=["p2c", "c2p"],
+            max_relative_positions=512,
+        )
+        self.scorer_encoder = DebertaV2Encoder(encoder_config)
+
+        self.text_projector = nn.Linear(config.hidden_size, config.hidden_size)
+        self.label_projector = nn.Linear(config.hidden_size, config.hidden_size)
+        self.dropout = nn.Dropout(config.dropout)
+
+        mlp_hidden_size = getattr(config, "scorer_mlp_hidden_size", 1024)
+        total_input_size = config.hidden_size * 2
+
+        self.mlp = nn.Sequential(
+            nn.Linear(total_input_size, mlp_hidden_size),
+            nn.ReLU(),
+            nn.Linear(mlp_hidden_size, mlp_hidden_size // 2),
+            nn.ReLU(),
+            nn.Linear(mlp_hidden_size // 2, 1),
+        )
+
+        self.epsilon = 1e-8
+
+    def _extract_representations(self, hidden_states, input_ids, attention_mask):
+        """
+        Extract text and label representations from hidden states.
+
+        Sequence: text<<SEP>>label1<<LABEL>>label2<<LABEL>><<SEP>>
+        Text repr: LAST <<SEP>> token (after all labels) - sees both text and labels
+        Label repr: each <<LABEL>> token
+        """
+        batch_size, _, hidden_size = hidden_states.shape
+
+        sep_token_id = self.config.sep_token_index
+        label_token_id = self.config.class_token_index
+
+        sep_mask = input_ids == sep_token_id
+        label_mask = input_ids == label_token_id
+
+        sep_positions = torch.where(sep_mask)
+        label_positions = torch.where(label_mask)
+
+        text_repr = torch.zeros(batch_size, hidden_size, device=hidden_states.device, dtype=hidden_states.dtype)
+
+        for batch_idx in range(batch_size):
+            batch_sep_positions = sep_positions[1][sep_positions[0] == batch_idx]
+            if len(batch_sep_positions) > 0:
+                # Take LAST SEP (after all labels) so text repr sees both text and labels
+                last_sep_pos = batch_sep_positions[-1]
+                text_repr[batch_idx] = hidden_states[batch_idx, last_sep_pos]
+
+        num_labels_per_batch = torch.zeros(batch_size, dtype=torch.long, device=hidden_states.device)
+        for batch_idx in range(batch_size):
+            num_labels_per_batch[batch_idx] = (label_positions[0] == batch_idx).sum()
+
+        max_labels = num_labels_per_batch.max().item()
+        label_repr = torch.zeros(
+            batch_size,
+            max_labels,
+            hidden_size,
+            device=hidden_states.device,
+            dtype=hidden_states.dtype,
+        )
+
+        for batch_idx in range(batch_size):
+            batch_label_positions = label_positions[1][label_positions[0] == batch_idx]
+            for label_idx, pos in enumerate(batch_label_positions):
+                if label_idx < max_labels:
+                    label_repr[batch_idx, label_idx] = hidden_states[batch_idx, pos]
+
+        return text_repr, label_repr
+
+    def forward(self, hidden_states, input_ids, attention_mask, **kwargs):
+        """Forward pass.
+
+        Args:
+            hidden_states: (batch_size, seq_length, hidden_size) from decoder backbone
+            input_ids: (batch_size, seq_length)
+            attention_mask: (batch_size, seq_length)
+
+        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
+
+        text_repr, label_repr = self._extract_representations(contextualized_hidden_states, input_ids, attention_mask)
+
+        text_repr = self.text_projector(text_repr)
+        text_repr = self.dropout(text_repr)
+
+        label_repr = self.label_projector(label_repr)
+
+        if self.config.normalize_features:
+            text_repr = text_repr / (text_repr.norm(p=2, dim=-1, keepdim=True) + self.epsilon)
+            label_repr = label_repr / (label_repr.norm(p=2, dim=-1, keepdim=True) + self.epsilon)
+
+        batch_size, num_labels, dim = label_repr.shape
+        text_repr_expanded = text_repr.unsqueeze(1).expand(batch_size, num_labels, dim)
+        combined_rep = torch.cat([text_repr_expanded, label_repr], dim=-1)
+
+        logits = self.mlp(combined_rep).squeeze(-1)
+
+        return logits
+
+
 SCORER2OBJECT = {
     "weighted-dot": ScorerWeightedDot,
     "simple": ScorerDot,
     "mlp": MLPScorer,
     "hopfield": HopfieldScorer,
     "cross-attn": CrossAttnScorer,
+    "decoder-kv": DecoderKVScorer,
 }

From c944bf5894ca35f7390e22b87f41a5b4e4bed508 Mon Sep 17 00:00:00 2001
From: BioMikeUkr <mykhayloshtopko@gmail.com>
Date: Fri, 10 Jul 2026 13:04:27 +0000
Subject: [PATCH 12/25] save_optimizer_state argument added; useful for
 trainung large models

---
 gliclass/training.py | 23 +++++++++++++++++++++++
 1 file changed, 23 insertions(+)

diff --git a/gliclass/training.py b/gliclass/training.py
index 0be9a93..3b73b36 100644
--- a/gliclass/training.py
+++ b/gliclass/training.py
@@ -290,6 +290,13 @@ class TrainingArguments(transformers.TrainingArguments):
     )
     ewc_gamma: float = field(default=0.95, metadata={"help": "Decay factor for Online EWC."})
 
+    save_optimizer_state: bool = field(
+        default=True,
+        metadata={
+            "help": "Save optimizer/scheduler state in checkpoints. Set False for large models to save disk space."
+        },  # noqa: E501
+    )
+
 
 class Trainer(transformers.Trainer):
     """Extended Trainer with EWC support for continual learning."""
@@ -495,6 +502,22 @@ def prediction_step(
             print(f"An error occurred during prediction step: {e!s}")
             return (None, None, None)
 
+    def _save_checkpoint(self, model, trial=None, **kwargs):
+        super()._save_checkpoint(model, trial=trial, **kwargs)
+        if not self.args.save_optimizer_state:
+            import glob
+
+            checkpoint_folder = f"checkpoint-{self.state.global_step}"
+            output_dir = os.path.join(self.args.output_dir, checkpoint_folder)
+            # remove optimizer / scheduler / scaler / accelerator state files
+            for pattern in ["optimizer.pt", "scheduler.pt", "scaler.pt", "optimizer_*", "random_states_*"]:
+                for f in glob.glob(os.path.join(output_dir, pattern)):
+                    os.remove(f)
+            # accelerator saves a subfolder for optimizer state
+            accel_state = os.path.join(output_dir, "optimizer.bin")
+            if os.path.exists(accel_state):
+                os.remove(accel_state)
+
     def create_optimizer(self):
         """
         Setup the optimizer.

From cf44e45c318a14f5e7e91ad3e3397856db1b87d9 Mon Sep 17 00:00:00 2001
From: BioMikeUkr <mykhayloshtopko@gmail.com>
Date: Fri, 10 Jul 2026 13:04:48 +0000
Subject: [PATCH 13/25] streaming added

---
 gliclass/__init__.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/gliclass/__init__.py b/gliclass/__init__.py
index ce42569..8dfaa6b 100644
--- a/gliclass/__init__.py
+++ b/gliclass/__init__.py
@@ -1,4 +1,4 @@
-from .model import GLiClassModel, GLiClassBiEncoder, GLiClassUniEncoder, GLiClassEncoderDecoderCLS
+from .model import GLiClassModel, GLiClassBiEncoder, GLiClassDecoderKV, GLiClassUniEncoder, GLiClassEncoderDecoderCLS
 from .config import GLiClassModelConfig
 from .pipeline import (
     ZeroShotClassificationPipeline,
@@ -6,7 +6,7 @@
     ZeroShotClassificationWithChunkingPipeline,
 )
 
-__version__ = "0.1.19"
+__version__ = "0.1.20"
 
 # Serve module (optional import)
 try:

From b12806cf5f8ce5bb2657aafca8cbbd04358f82f5 Mon Sep 17 00:00:00 2001
From: BioMikeUkr <mykhayloshtopko@gmail.com>
Date: Fri, 10 Jul 2026 13:06:28 +0000
Subject: [PATCH 14/25] streaming docs created

---
 docs/streaming.md | 365 ++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 365 insertions(+)
 create mode 100644 docs/streaming.md

diff --git a/docs/streaming.md b/docs/streaming.md
new file mode 100644
index 0000000..4519cce
--- /dev/null
+++ b/docs/streaming.md
@@ -0,0 +1,365 @@
+# GLiClass Streaming Classification
+
+`gliclass.streaming` provides incremental, multi-session text classification over a decoder-KV model. Instead of classifying a complete document at once, it processes text in arbitrary chunks, maintains a persistent KV cache per session, and triggers classification only when a pluggable strategy decides it is appropriate.
+
+---
+
+## Contents
+
+- [Architecture overview](#architecture-overview)
+- [KV cache management](#kv-cache-management)
+- [Session lifecycle](#session-lifecycle)
+- [Classification strategies](#classification-strategies)
+- [StreamingPipeline reference](#streamingpipeline-reference)
+- [Data types](#data-types)
+- [Usage examples](#usage-examples)
+- [Memory and performance considerations](#memory-and-performance-considerations)
+
+---
+
+## Architecture overview
+
+GLiClass Decoder-KV separates the model into two components:
+
+- **`decoder_model`** — a causal language model (e.g. Qwen3.5) that encodes text into a sequence of hidden states and accumulates a KV cache.
+- **`scorer`** — a lightweight cross-attention encoder that reads the label hidden states produced by the decoder and outputs per-label logits.
+
+In standard (non-streaming) inference, both components run on the full text each time. In streaming mode, the decoder runs once per chunk to *extend* the KV cache; the scorer runs only when a strategy decides classification should fire. The result is:
+
+```
+chunk_1 → decoder → KV₁
+chunk_2 → decoder (KV₁ as prefix) → KV₂
+chunk_3 → decoder (KV₂ as prefix) → KV₃  ← strategy triggers → scorer → predictions
+chunk_4 → ...
+```
+
+Each forward pass of the decoder is O(n_chunk × n_cache) rather than O(n_total²), because the attention over cached tokens is handled by the KV cache mechanism and never recomputed.
+
+### Two-stage per call
+
+Every `StreamingPipeline.__call__` runs two sequential stages over the input batch:
+
+**Stage 1 — cache update.**  
+New text is tokenized per session, then all sessions in the batch are forwarded through `decoder_model` together. The resulting `past_key_values` are sliced per session and stored in `_caches`.
+
+**Stage 2 — classification.**  
+Sessions whose strategy returns `True` from `should_classify` are forwarded again through `decoder_model` with their KV cache as prefix and the label token sequence as input. The decoder's `last_hidden_state` over label positions is then passed to `scorer`, which produces logits. Sigmoid (multi-label) or softmax (single-label) is applied to produce probabilities.
+
+Sessions that did not trigger classification return `triggered=False` with `predictions=None`.
+
+---
+
+## KV cache management
+
+### CacheState
+
+`CacheState` is the per-session container that travels through the pipeline:
+
+```python
+@dataclass
+class CacheState:
+    past_key_values: Any          # transformers DynamicCache or None
+    input_ids: torch.Tensor       # flat 1-D tensor of all cached token ids
+    attention_mask: torch.Tensor  # flat 1-D tensor of all cached attention mask values
+    position_ids: torch.Tensor | None
+    cached_length: int            # total number of cached tokens
+    session_id: str | None
+    metadata: dict                # arbitrary user data
+```
+
+A fresh empty cache is created via `create_empty_cache(session_id, device)`.
+
+### Batched stacking with heterogeneous cache lengths
+
+When multiple sessions exist in a batch and have different numbers of cached tokens, their KV caches cannot be naively stacked: different sequence lengths mean different position embeddings baked into the K/V tensors.
+
+`BatchedKVHelper` resolves this with **prepend-padding**:
+
+- KV caches shorter than `max_cached_len` are zero-padded on the **left** (prepend), not the right.
+- The attention mask marks these prepended zeros as not-attended, so the model ignores them.
+- `position_ids` for the new tokens in each session are set to `[cached_len_i, ..., cached_len_i + new_len_i)`, so RoPE offsets remain correct regardless of padding.
+
+This means a batch of sessions with cache lengths `[150, 400, 300]` is padded to `max_cached=400`, with sessions 0 and 2 receiving 250 and 100 prepend zeros respectively.
+
+After the forward pass, the updated KV is sliced back per session using the real token range `[max_cached - clen_i : max_cached + new_len_i]`.
+
+### Cache truncation
+
+`truncate_cache(cache_state, max_length)` crops the KV to the **last** `max_length` tokens. Only `DynamicCache` is supported for cropping; other cache formats are returned unchanged. The pipeline applies this automatically after each stage-1 update when `max_cache_len` is set.
+
+Truncation discards the oldest context first. This is a simple sliding-window policy; for more selective eviction (e.g. keeping the first few tokens as a persistent prefix), extend or replace the default behavior by subclassing `StreamingPipeline._enforce_max_cache_len`.
+
+### Device offloading
+
+When running more sessions than fit in GPU memory simultaneously, `offload_to_cpu=True` keeps inactive session caches in CPU RAM and moves them to GPU only for the active sub-batch.
+
+With `use_pinned_memory=True`, KV caches are allocated in CPU pinned (page-locked) memory, enabling asynchronous DMA transfers to the GPU via `non_blocking=True`. A `torch.cuda.synchronize` call before the forward pass ensures transfers are complete before the kernel launches.
+
+---
+
+## Session lifecycle
+
+Sessions are identified by a string `session_id`. The pipeline maintains a dict `_caches: dict[str, CacheState]` internally.
+
+**Creation.** A session is created implicitly when a `session_id` appears for the first time in an input batch. `create_empty_cache` is called and stored.
+
+**Updating.** On every call, the session's cache is extended with new tokens if `text` is non-empty. Empty text causes the update stage to be skipped for that session (`tokens_added = 0`), which also skips classification regardless of strategy.
+
+**Cleanup.** Sessions that are absent from the current input batch are deleted. The pipeline does not keep stale sessions alive across calls: any `session_id` not present in the current `inputs` list will have its cache freed. If a session should remain alive without receiving new text, pass it with `text=""`.
+
+**Manual reset.** Delete the cache entry directly:
+```python
+del pipeline._caches["my_session"]
+```
+
+---
+
+## Classification strategies
+
+Strategies implement `ClassificationStrategy` and control:
+
+1. **`should_classify(tokens_added, cached_length, text) -> bool`** — whether to trigger classification this call.
+2. **`get_window(cache_state) -> CacheState`** — which portion of the cache to classify over (default: full cache).
+
+### Built-in strategies
+
+#### `EveryChunkStrategy`
+
+Classifies on every call that adds at least one token. Use when latency is acceptable and you want the most up-to-date prediction at every step.
+
+```python
+EveryChunkStrategy()
+```
+
+#### `EveryNTokensStrategy`
+
+Accumulates tokens across calls and classifies once every `n` tokens. The counter resets after each trigger. Useful for throttling classification frequency over a fast-arriving stream.
+
+```python
+EveryNTokensStrategy(n=100)
+```
+
+Note: the internal `_accumulated` counter is per-instance. If the same strategy object is shared across sessions, the counter is shared too. Create one instance per session or use `ComposedStrategy` with per-session instantiation.
+
+#### `OnDelimiterStrategy`
+
+Classifies when the incoming text chunk contains a specific string, such as a sentence boundary or a paragraph marker.
+
+```python
+OnDelimiterStrategy(delimiter="\n\n")   # paragraph boundary
+OnDelimiterStrategy(delimiter=". ")    # sentence boundary (approximate)
+```
+
+The delimiter check is a plain substring search on the raw text string, not on tokens.
+
+#### `NeverStrategy`
+
+Never triggers classification. Use this for a pre-filling phase where you want to build up a long context cache before any classification runs.
+
+```python
+NeverStrategy()
+```
+
+Switch to a different strategy object after pre-filling is complete.
+
+#### `SlidingWindowStrategy`
+
+Classifies on every chunk (same trigger condition as `EveryChunkStrategy`), but passes only the last `window_size` cached tokens to the scorer via `get_window`. This limits the attention span of the classification pass without discarding the full KV cache.
+
+```python
+SlidingWindowStrategy(window_size=512)
+```
+
+Note that this does not truncate the stored KV cache — the full cache is retained for future updates. Only the classification forward sees the windowed slice.
+
+#### `ComposedStrategy`
+
+Combines a **trigger strategy** (controls `should_classify`) with a **window strategy** (controls `get_window`). This is the standard way to pair a frequency policy with a context window policy.
+
+```python
+ComposedStrategy(
+    trigger=EveryNTokensStrategy(200),
+    window=SlidingWindowStrategy(512),
+)
+```
+
+Classifies every 200 new tokens, but uses only the last 512 cached tokens for scoring.
+
+### Custom strategies
+
+Subclass `ClassificationStrategy` and implement `should_classify`. Optionally override `get_window` for windowed classification.
+
+```python
+class SentenceCountStrategy(ClassificationStrategy):
+    def __init__(self, every_n_sentences: int):
+        self.n = every_n_sentences
+        self._count = 0
+
+    def should_classify(self, tokens_added, cached_length, text):
+        self._count += text.count(". ")
+        if self._count >= self.n:
+            self._count = 0
+            return True
+        return False
+```
+
+---
+
+## StreamingPipeline reference
+
+```python
+StreamingPipeline(
+    model,
+    tokenizer,
+    device="cpu",
+    max_cache_len=None,
+    batch_size=None,
+    offload_to_cpu=False,
+    use_pinned_memory=False,
+    score_on_cpu=False,
+)
+```
+
+| Parameter | Type | Description |
+|---|---|---|
+| `model` | `GLiClassModel` | Loaded GLiClass decoder-KV model in eval mode. |
+| `tokenizer` | `PreTrainedTokenizer` | Matching tokenizer with `<<LABEL>>`, `<<SEP>>`, `<<EXAMPLE>>` special tokens. |
+| `device` | `str \| torch.device` | Device for decoder forward passes. |
+| `max_cache_len` | `int \| None` | Hard cap on cached tokens per session. Oldest tokens are dropped when exceeded. `None` means unbounded. |
+| `batch_size` | `int \| None` | Maximum sessions per decoder forward. When `len(inputs) > batch_size`, sessions are chunked. With `offload_to_cpu=True`, only the active chunk is on GPU at once. |
+| `offload_to_cpu` | `bool` | Move inactive session caches to CPU RAM between sub-batches. |
+| `use_pinned_memory` | `bool` | Allocate CPU caches in pinned memory for async GPU uploads. Only effective when `offload_to_cpu=True` and `device` is CUDA. |
+| `score_on_cpu` | `bool` | Run the scorer on CPU. Useful when GPU VRAM is the bottleneck and the scorer is small relative to the decoder. |
+
+### `__call__(inputs: list[SessionInput]) -> list[SessionOutput]`
+
+Processes a batch of sessions. Returns one `SessionOutput` per input in the same order. Sessions not present in the current batch have their caches deleted.
+
+---
+
+## Data types
+
+### `SessionInput`
+
+```python
+@dataclass
+class SessionInput:
+    session_id: str                    # unique identifier for this session
+    text: str                          # new text chunk to append (empty = no update)
+    labels: list[str]                  # candidate label strings
+    strategy: ClassificationStrategy   # per-session strategy instance
+    classification_type: str           # "multi-label" (sigmoid) or "single-label" (softmax)
+```
+
+### `SessionOutput`
+
+```python
+@dataclass
+class SessionOutput:
+    session_id: str
+    triggered: bool                    # True if classification ran this call
+    predictions: list[dict] | None     # [{"label": str, "score": float}, ...] or None
+    cached_length: int                 # total cached tokens after this call
+    tokens_added: int                  # tokens added in stage 1 (0 if text was empty)
+    metadata: dict                     # pass-through from CacheState.metadata
+```
+
+For `multi-label`, `predictions` contains one entry per label with a sigmoid probability. For `single-label`, it contains only the top-1 label with its softmax probability.
+
+---
+
+## Usage examples
+
+### Minimal single-session loop
+
+```python
+import torch
+from transformers import AutoTokenizer
+from gliclass import GLiClassModel
+from gliclass.streaming import StreamingPipeline, SessionInput, EveryNTokensStrategy
+
+tokenizer = AutoTokenizer.from_pretrained("path/to/model")
+model = GLiClassModel.from_pretrained("path/to/model")
+model = model.to("cuda", dtype=torch.bfloat16).eval()
+
+pipeline = StreamingPipeline(model, tokenizer, device="cuda", max_cache_len=1024)
+
+labels = ["science", "politics", "finance", "sports"]
+strategy = EveryNTokensStrategy(n=50)
+
+chunks = ["The Fed raised interest rates", " by 25 basis points", " amid inflation concerns."]
+
+for chunk in chunks:
+    outputs = pipeline([
+        SessionInput(
+            session_id="doc_001",
+            text=chunk,
+            labels=labels,
+            strategy=strategy,
+            classification_type="multi-label",
+        )
+    ])
+    out = outputs[0]
+    print(f"cached={out.cached_length} +{out.tokens_added} triggered={out.triggered}")
+    if out.triggered:
+        for p in sorted(out.predictions, key=lambda x: -x["score"]):
+            print(f"  {p['label']}: {p['score']:.3f}")
+```
+
+### Multiple concurrent sessions
+
+```python
+session_strategies = {
+    "session_A": EveryNTokensStrategy(n=100),
+    "session_B": OnDelimiterStrategy(delimiter="\n"),
+    "session_C": ComposedStrategy(
+        trigger=EveryNTokensStrategy(50),
+        window=SlidingWindowStrategy(256),
+    ),
+}
+
+batch = [
+    SessionInput("session_A", chunk_a, labels, session_strategies["session_A"]),
+    SessionInput("session_B", chunk_b, labels, session_strategies["session_B"]),
+    SessionInput("session_C", chunk_c, labels, session_strategies["session_C"]),
+]
+
+outputs = pipeline(batch)
+```
+
+All three sessions are updated and classified (if triggered) in a single batched forward pass per stage.
+
+### Pre-filling a long context before classification
+
+```python
+from gliclass.streaming import NeverStrategy, EveryChunkStrategy
+
+# Phase 1: ingest background context without any classification overhead
+for paragraph in background_paragraphs:
+    pipeline([SessionInput("doc", paragraph, labels, NeverStrategy())])
+
+# Phase 2: stream the main document and classify on every chunk
+for chunk in main_document_chunks:
+    outputs = pipeline([SessionInput("doc", chunk, labels, EveryChunkStrategy())])
+    if outputs[0].triggered:
+        handle(outputs[0].predictions)
+```
+
+### Memory-constrained deployment with CPU offload
+
+```python
+pipeline = StreamingPipeline(
+    model=model,
+    tokenizer=tokenizer,
+    device="cuda",
+    max_cache_len=512,
+    batch_size=8,           # process 8 sessions per GPU forward
+    offload_to_cpu=True,    # keep inactive caches in RAM
+    use_pinned_memory=True, # async DMA transfers
+    score_on_cpu=True,      # scorer stays in CPU RAM
+)
+```
+
+With 32 concurrent sessions and `batch_size=8`, the GPU holds at most 8 KV caches at a time. Inactive caches are evicted to CPU pinned memory and uploaded asynchronously before their sub-batch runs.
+
+---
\ No newline at end of file

From 314ec62fb6495e43afcbacf1709819e4036e9558 Mon Sep 17 00:00:00 2001
From: BioMikeUkr <mykhayloshtopko@gmail.com>
Date: Fri, 10 Jul 2026 13:06:59 +0000
Subject: [PATCH 15/25] streaming tests added

---
 tests/test_streaming.py | 502 ++++++++++++++++++++++++++++++++++++++++
 1 file changed, 502 insertions(+)
 create mode 100644 tests/test_streaming.py

diff --git a/tests/test_streaming.py b/tests/test_streaming.py
new file mode 100644
index 0000000..4030bc2
--- /dev/null
+++ b/tests/test_streaming.py
@@ -0,0 +1,502 @@
+"""Tests for gliclass.streaming — strategies, cache management, BatchedKVHelper."""
+
+import pytest
+import torch
+from transformers.cache_utils import DynamicCache
+
+from gliclass.streaming.cache import (
+    CacheState,
+    BatchedKVHelper,
+    create_empty_cache,
+    truncate_cache,
+)
+from gliclass.streaming.strategies import (
+    ClassificationStrategy,
+    EveryChunkStrategy,
+    EveryNTokensStrategy,
+    OnDelimiterStrategy,
+    NeverStrategy,
+    SlidingWindowStrategy,
+    ComposedStrategy,
+)
+from gliclass.streaming.types import SessionInput, SessionOutput
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+def _make_dynamic_cache(num_layers: int, num_heads: int, seq_len: int, head_dim: int, batch: int = 1) -> DynamicCache:
+    cache = DynamicCache()
+    for i in range(num_layers):
+        k = torch.randn(batch, num_heads, seq_len, head_dim)
+        v = torch.randn(batch, num_heads, seq_len, head_dim)
+        cache.update(k, v, i)
+    return cache
+
+
+def _make_cache_state(seq_len: int, num_layers: int = 2, num_heads: int = 2, head_dim: int = 8) -> CacheState:
+    past_kv = _make_dynamic_cache(num_layers, num_heads, seq_len, head_dim)
+    return CacheState(
+        past_key_values=past_kv,
+        input_ids=torch.ones(seq_len, dtype=torch.long),
+        attention_mask=torch.ones(seq_len, dtype=torch.long),
+        cached_length=seq_len,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Strategy tests
+# ---------------------------------------------------------------------------
+
+class TestEveryChunkStrategy:
+    def test_triggers_on_any_positive_tokens(self):
+        s = EveryChunkStrategy()
+        assert s.should_classify(1, 10, "text") is True
+        assert s.should_classify(50, 200, "text") is True
+
+    def test_does_not_trigger_on_zero_tokens(self):
+        s = EveryChunkStrategy()
+        assert s.should_classify(0, 10, "text") is False
+
+    def test_get_window_returns_full_cache(self):
+        s = EveryChunkStrategy()
+        cache = _make_cache_state(20)
+        result = s.get_window(cache)
+        assert result is cache
+
+
+class TestEveryNTokensStrategy:
+    def test_triggers_after_n_tokens(self):
+        s = EveryNTokensStrategy(n=10)
+        assert s.should_classify(5, 5, "") is False
+        assert s.should_classify(5, 10, "") is True
+
+    def test_counter_resets_after_trigger(self):
+        s = EveryNTokensStrategy(n=10)
+        s.should_classify(10, 10, "")   # trigger, reset
+        assert s._accumulated == 0
+        assert s.should_classify(5, 15, "") is False
+
+    def test_accumulates_across_calls(self):
+        s = EveryNTokensStrategy(n=30)
+        assert s.should_classify(10, 10, "") is False
+        assert s.should_classify(10, 20, "") is False
+        assert s.should_classify(10, 30, "") is True
+
+    def test_exact_n(self):
+        s = EveryNTokensStrategy(n=7)
+        assert s.should_classify(7, 7, "") is True
+
+    def test_overshoot_still_triggers(self):
+        s = EveryNTokensStrategy(n=5)
+        assert s.should_classify(20, 20, "") is True
+
+    def test_independent_instances(self):
+        s1 = EveryNTokensStrategy(n=10)
+        s2 = EveryNTokensStrategy(n=10)
+        s1.should_classify(8, 8, "")
+        assert s2._accumulated == 0
+
+
+class TestOnDelimiterStrategy:
+    def test_triggers_when_delimiter_present(self):
+        s = OnDelimiterStrategy(delimiter=".")
+        assert s.should_classify(5, 5, "End of sentence.") is True
+
+    def test_does_not_trigger_without_delimiter(self):
+        s = OnDelimiterStrategy(delimiter=".")
+        assert s.should_classify(5, 5, "No period here") is False
+
+    def test_newline_delimiter(self):
+        s = OnDelimiterStrategy(delimiter="\n\n")
+        assert s.should_classify(3, 3, "para1\n\npara2") is True
+        assert s.should_classify(3, 3, "para1\npara2") is False
+
+    def test_empty_text(self):
+        s = OnDelimiterStrategy(delimiter=".")
+        assert s.should_classify(0, 0, "") is False
+
+    def test_delimiter_in_middle(self):
+        s = OnDelimiterStrategy(delimiter=";")
+        assert s.should_classify(5, 5, "part1; part2") is True
+
+
+class TestNeverStrategy:
+    def test_never_triggers(self):
+        s = NeverStrategy()
+        assert s.should_classify(100, 1000, "text") is False
+        assert s.should_classify(0, 0, "") is False
+
+    def test_get_window_returns_full_cache(self):
+        s = NeverStrategy()
+        cache = _make_cache_state(10)
+        assert s.get_window(cache) is cache
+
+
+class TestSlidingWindowStrategy:
+    def test_triggers_on_positive_tokens(self):
+        s = SlidingWindowStrategy(window_size=50)
+        assert s.should_classify(1, 10, "text") is True
+
+    def test_does_not_trigger_on_zero_tokens(self):
+        s = SlidingWindowStrategy(window_size=50)
+        assert s.should_classify(0, 10, "text") is False
+
+    def test_window_truncates_cache(self):
+        s = SlidingWindowStrategy(window_size=10)
+        cache = _make_cache_state(30)
+        windowed = s.get_window(cache)
+        assert windowed.cached_length == 10
+
+    def test_window_larger_than_cache_returns_full(self):
+        s = SlidingWindowStrategy(window_size=100)
+        cache = _make_cache_state(20)
+        windowed = s.get_window(cache)
+        assert windowed.cached_length == 20
+
+    def test_original_cache_unmodified(self):
+        s = SlidingWindowStrategy(window_size=5)
+        cache = _make_cache_state(20)
+        s.get_window(cache)
+        assert cache.cached_length == 20
+
+
+class TestComposedStrategy:
+    def test_trigger_delegated_to_trigger_strategy(self):
+        trigger = EveryNTokensStrategy(n=10)
+        window = SlidingWindowStrategy(window_size=50)
+        s = ComposedStrategy(trigger=trigger, window=window)
+        assert s.should_classify(5, 5, "") is False
+        assert s.should_classify(5, 10, "") is True
+
+    def test_window_delegated_to_window_strategy(self):
+        trigger = EveryChunkStrategy()
+        window = SlidingWindowStrategy(window_size=10)
+        s = ComposedStrategy(trigger=trigger, window=window)
+        cache = _make_cache_state(30)
+        windowed = s.get_window(cache)
+        assert windowed.cached_length == 10
+
+    def test_never_trigger_with_every_chunk_window(self):
+        s = ComposedStrategy(trigger=NeverStrategy(), window=SlidingWindowStrategy(20))
+        assert s.should_classify(100, 100, "text") is False
+
+
+class TestCustomStrategy:
+    def test_custom_strategy_interface(self):
+        class CountStrategy(ClassificationStrategy):
+            def __init__(self, n):
+                self.n = n
+                self.calls = 0
+
+            def should_classify(self, tokens_added, cached_length, text):
+                self.calls += 1
+                return self.calls >= self.n
+
+        s = CountStrategy(n=3)
+        assert s.should_classify(1, 1, "") is False
+        assert s.should_classify(1, 2, "") is False
+        assert s.should_classify(1, 3, "") is True
+
+
+# ---------------------------------------------------------------------------
+# CacheState tests
+# ---------------------------------------------------------------------------
+
+class TestCacheState:
+    def test_create_empty_cache(self):
+        cache = create_empty_cache(session_id="test")
+        assert cache.cached_length == 0
+        assert cache.past_key_values is None
+        assert cache.session_id == "test"
+        assert cache.input_ids.shape == (0,)
+
+    def test_create_empty_cache_with_device(self):
+        device = torch.device("cpu")
+        cache = create_empty_cache(session_id="s", device=device)
+        assert cache.input_ids.device.type == "cpu"
+
+    def test_to_device(self):
+        cache = _make_cache_state(10)
+        moved = cache.to(torch.device("cpu"))
+        assert moved.input_ids.device.type == "cpu"
+        assert moved.cached_length == 10
+        assert moved.session_id == cache.session_id
+
+    def test_metadata_preserved_on_move(self):
+        cache = _make_cache_state(5)
+        cache.metadata = {"key": "value"}
+        moved = cache.to(torch.device("cpu"))
+        assert moved.metadata == {"key": "value"}
+
+
+class TestTruncateCache:
+    def test_truncate_reduces_length(self):
+        cache = _make_cache_state(50)
+        truncated = truncate_cache(cache, max_length=20)
+        assert truncated.cached_length == 20
+
+    def test_truncate_keeps_last_tokens(self):
+        cache = _make_cache_state(10)
+        cache.input_ids = torch.arange(10, dtype=torch.long)
+        truncated = truncate_cache(cache, max_length=5)
+        assert truncated.cached_length == 5
+        assert truncated.input_ids.tolist() == [5, 6, 7, 8, 9]
+
+    def test_truncate_attention_mask(self):
+        cache = _make_cache_state(10)
+        truncated = truncate_cache(cache, max_length=4)
+        assert truncated.attention_mask.shape == (4,)
+
+    def test_no_truncate_when_within_limit(self):
+        cache = _make_cache_state(10)
+        result = truncate_cache(cache, max_length=20)
+        assert result is cache
+
+    def test_no_truncate_on_empty_cache(self):
+        cache = create_empty_cache()
+        result = truncate_cache(cache, max_length=10)
+        assert result is cache
+
+    def test_truncate_kv_shape(self):
+        num_layers, num_heads, seq_len, head_dim = 2, 2, 30, 8
+        cache = _make_cache_state(seq_len, num_layers, num_heads, head_dim)
+        truncated = truncate_cache(cache, max_length=10)
+        for layer in truncated.past_key_values.layers:
+            if layer.is_initialized:
+                assert layer.keys.shape[2] == 10
+
+    def test_session_id_preserved(self):
+        cache = _make_cache_state(10)
+        cache.session_id = "my_session"
+        truncated = truncate_cache(cache, max_length=5)
+        assert truncated.session_id == "my_session"
+
+    def test_metadata_preserved(self):
+        cache = _make_cache_state(10)
+        cache.metadata = {"info": 42}
+        truncated = truncate_cache(cache, max_length=5)
+        assert truncated.metadata == {"info": 42}
+
+
+# ---------------------------------------------------------------------------
+# BatchedKVHelper tests
+# ---------------------------------------------------------------------------
+
+class TestBatchedKVHelperStack:
+    @pytest.fixture
+    def device(self):
+        return torch.device("cpu")
+
+    def _make_inputs(self, lengths, device):
+        ids = [torch.ones(l, dtype=torch.long, device=device) for l in lengths]
+        masks = [torch.ones(l, dtype=torch.long, device=device) for l in lengths]
+        return ids, masks
+
+    def test_stack_homogeneous_caches(self, device):
+        caches = [_make_cache_state(10) for _ in range(3)]
+        ids, masks = self._make_inputs([5, 5, 5], device)
+        result = BatchedKVHelper.stack_for_update(caches, ids, masks, device)
+
+        assert result["input_ids"].shape == (3, 5)
+        assert result["max_cached_len"] == 10
+        assert result["cached_lengths"] == [10, 10, 10]
+        assert result["new_lengths"] == [5, 5, 5]
+
+    def test_stack_heterogeneous_caches(self, device):
+        caches = [_make_cache_state(5), _make_cache_state(20), _make_cache_state(10)]
+        ids, masks = self._make_inputs([3, 3, 3], device)
+        result = BatchedKVHelper.stack_for_update(caches, ids, masks, device)
+
+        assert result["max_cached_len"] == 20
+        assert result["input_ids"].shape == (3, 3)
+
+    def test_stack_with_empty_cache(self, device):
+        empty = create_empty_cache(device=device)
+        filled = _make_cache_state(10)
+        caches = [empty, filled]
+        ids, masks = self._make_inputs([4, 4], device)
+        result = BatchedKVHelper.stack_for_update(caches, ids, masks, device)
+
+        assert result["cached_lengths"] == [0, 10]
+        assert result["max_cached_len"] == 10
+
+    def test_attention_mask_shape(self, device):
+        caches = [_make_cache_state(8), _make_cache_state(12)]
+        ids, masks = self._make_inputs([5, 3], device)
+        result = BatchedKVHelper.stack_for_update(caches, ids, masks, device)
+
+        # full mask = max_cached + max_new
+        assert result["attention_mask"].shape == (2, 12 + 5)
+
+    def test_position_ids_per_session(self, device):
+        caches = [_make_cache_state(5), _make_cache_state(10)]
+        ids, masks = self._make_inputs([3, 3], device)
+        result = BatchedKVHelper.stack_for_update(caches, ids, masks, device)
+
+        pos = result["position_ids"]
+        # session 0: positions [5, 6, 7]
+        assert pos[0, :3].tolist() == [5, 6, 7]
+        # session 1: positions [10, 11, 12]
+        assert pos[1, :3].tolist() == [10, 11, 12]
+
+    def test_all_empty_caches_no_stacked_kv(self, device):
+        caches = [create_empty_cache(device=device) for _ in range(2)]
+        ids, masks = self._make_inputs([4, 4], device)
+        result = BatchedKVHelper.stack_for_update(caches, ids, masks, device)
+
+        assert result["past_key_values"] is None
+        assert result["max_cached_len"] == 0
+
+    def test_padded_new_tokens_are_zero(self, device):
+        caches = [_make_cache_state(5), _make_cache_state(5)]
+        ids = [torch.ones(6, dtype=torch.long, device=device),
+               torch.ones(3, dtype=torch.long, device=device)]
+        masks = [torch.ones(6, dtype=torch.long, device=device),
+                 torch.ones(3, dtype=torch.long, device=device)]
+        result = BatchedKVHelper.stack_for_update(caches, ids, masks, device)
+
+        # session 1 has 3 tokens, padded to 6 — positions 3,4,5 should be 0
+        assert result["input_ids"][1, 3:].sum().item() == 0
+
+
+class TestBatchedKVHelperUnstack:
+    @pytest.fixture
+    def device(self):
+        return torch.device("cpu")
+
+    def test_unstack_restores_cache_lengths(self, device):
+        caches = [_make_cache_state(5), _make_cache_state(10)]
+        ids = [torch.ones(3, dtype=torch.long, device=device),
+               torch.ones(3, dtype=torch.long, device=device)]
+        masks = [torch.ones(3, dtype=torch.long, device=device),
+                 torch.ones(3, dtype=torch.long, device=device)]
+
+        stacked = BatchedKVHelper.stack_for_update(caches, ids, masks, device)
+
+        # simulate decoder output with updated KV
+        num_layers, num_heads, head_dim = 2, 2, 8
+        max_total = stacked["max_cached_len"] + max(stacked["new_lengths"])
+        new_kv = _make_dynamic_cache(num_layers, num_heads, max_total, head_dim, batch=2)
+
+        updated = BatchedKVHelper.unstack_after_update(new_kv, stacked, caches, ids, masks)
+
+        assert updated[0].cached_length == 5 + 3
+        assert updated[1].cached_length == 10 + 3
+
+    def test_unstack_preserves_session_ids(self, device):
+        c1 = _make_cache_state(5)
+        c1.session_id = "alice"
+        c2 = _make_cache_state(8)
+        c2.session_id = "bob"
+        caches = [c1, c2]
+
+        ids = [torch.ones(2, dtype=torch.long, device=device),
+               torch.ones(2, dtype=torch.long, device=device)]
+        masks = ids[:]
+
+        stacked = BatchedKVHelper.stack_for_update(caches, ids, masks, device)
+        max_total = stacked["max_cached_len"] + 2
+        new_kv = _make_dynamic_cache(2, 2, max_total, 8, batch=2)
+        updated = BatchedKVHelper.unstack_after_update(new_kv, stacked, caches, ids, masks)
+
+        assert updated[0].session_id == "alice"
+        assert updated[1].session_id == "bob"
+
+    def test_unstack_input_ids_concatenated(self, device):
+        cache = _make_cache_state(3)
+        cache.input_ids = torch.tensor([1, 2, 3], dtype=torch.long)
+        cache.attention_mask = torch.ones(3, dtype=torch.long)
+        caches = [cache]
+
+        new_ids = [torch.tensor([4, 5], dtype=torch.long, device=device)]
+        new_masks = [torch.ones(2, dtype=torch.long, device=device)]
+
+        stacked = BatchedKVHelper.stack_for_update(caches, new_ids, new_masks, device)
+        new_kv = _make_dynamic_cache(2, 2, stacked["max_cached_len"] + 2, 8, batch=1)
+        updated = BatchedKVHelper.unstack_after_update(new_kv, stacked, caches, new_ids, new_masks)
+
+        assert updated[0].input_ids.tolist() == [1, 2, 3, 4, 5]
+
+
+class TestBatchedKVHelperClassify:
+    @pytest.fixture
+    def device(self):
+        return torch.device("cpu")
+
+    def test_stack_for_classify_shapes(self, device):
+        caches = [_make_cache_state(10), _make_cache_state(15)]
+        label_ids = [torch.ones(8, dtype=torch.long, device=device),
+                     torch.ones(6, dtype=torch.long, device=device)]
+        label_masks = [torch.ones(8, dtype=torch.long, device=device),
+                       torch.ones(6, dtype=torch.long, device=device)]
+
+        result = BatchedKVHelper.stack_for_classify(caches, label_ids, label_masks, device)
+
+        assert result["input_ids"].shape == (2, 8)   # max_label = 8
+        assert result["attention_mask"].shape == (2, 15 + 8)   # max_cached + max_label
+        assert result["label_lengths"] == [8, 6]
+        assert result["max_label_len"] == 8
+
+    def test_stack_for_classify_position_ids(self, device):
+        caches = [_make_cache_state(5), _make_cache_state(12)]
+        label_ids = [torch.ones(4, dtype=torch.long, device=device),
+                     torch.ones(4, dtype=torch.long, device=device)]
+        label_masks = label_ids[:]
+
+        result = BatchedKVHelper.stack_for_classify(caches, label_ids, label_masks, device)
+
+        pos = result["position_ids"]
+        assert pos[0, :4].tolist() == [5, 6, 7, 8]
+        assert pos[1, :4].tolist() == [12, 13, 14, 15]
+
+    def test_label_mask_is_label_only(self, device):
+        caches = [_make_cache_state(10)]
+        label_ids = [torch.ones(5, dtype=torch.long, device=device)]
+        label_masks = [torch.ones(5, dtype=torch.long, device=device)]
+
+        result = BatchedKVHelper.stack_for_classify(caches, label_ids, label_masks, device)
+
+        # label_mask covers only label tokens (not cache)
+        assert result["label_mask"].shape == (1, 5)
+
+
+# ---------------------------------------------------------------------------
+# Data types tests
+# ---------------------------------------------------------------------------
+
+class TestDataTypes:
+    def test_session_input_defaults(self):
+        from gliclass.streaming.strategies import EveryChunkStrategy
+        inp = SessionInput(
+            session_id="s1",
+            text="hello",
+            labels=["a", "b"],
+            strategy=EveryChunkStrategy(),
+        )
+        assert inp.classification_type == "multi-label"
+
+    def test_session_output_not_triggered(self):
+        out = SessionOutput(
+            session_id="s1",
+            triggered=False,
+            predictions=None,
+            cached_length=10,
+            tokens_added=5,
+        )
+        assert out.predictions is None
+        assert not out.triggered
+        assert out.metadata == {}
+
+    def test_session_output_triggered(self):
+        preds = [{"label": "positive", "score": 0.9}]
+        out = SessionOutput(
+            session_id="s1",
+            triggered=True,
+            predictions=preds,
+            cached_length=50,
+            tokens_added=10,
+        )
+        assert out.triggered
+        assert out.predictions[0]["label"] == "positive"

From 824d93893fa63201cfb6fcbcc7fc6b58154609e0 Mon Sep 17 00:00:00 2001
From: BioMikeUkr <mykhayloshtopko@gmail.com>
Date: Fri, 10 Jul 2026 13:07:43 +0000
Subject: [PATCH 16/25] readme updated; streaming

---
 README.md | 32 ++++++++++++++++++++++++++++++++
 1 file changed, 32 insertions(+)

diff --git a/README.md b/README.md
index 1a829cb..c4847b1 100644
--- a/README.md
+++ b/README.md
@@ -341,6 +341,38 @@ config = GLiClassModelConfig(
 
 ---
 
+### Streaming Classification
+
+GLiClass supports incremental, multi-session text classification over a decoder-KV model. Instead of re-encoding the full document on every update, it maintains a persistent KV cache per session and runs the scorer only when a pluggable strategy decides classification should fire.
+
+```bash
+pip install gliclass[streaming]
+```
+
+```python
+from gliclass.streaming import StreamingPipeline, SessionInput, EveryNTokensStrategy
+
+pipeline = StreamingPipeline(model, tokenizer, device="cuda", max_cache_len=1024)
+strategy = EveryNTokensStrategy(n=50)
+
+for chunk in text_chunks:
+    outputs = pipeline([SessionInput(
+        session_id="doc_001",
+        text=chunk,
+        labels=["science", "politics", "finance"],
+        strategy=strategy,
+        classification_type="multi-label",
+    )])
+    if outputs[0].triggered:
+        print(outputs[0].predictions)
+```
+
+Built-in strategies: `EveryChunkStrategy`, `EveryNTokensStrategy`, `OnDelimiterStrategy`, `SlidingWindowStrategy`, `ComposedStrategy`, `NeverStrategy`.
+
+For full documentation on session management, KV cache internals, batching, CPU offloading, and custom strategies, see [docs/streaming.md](docs/streaming.md).
+
+---
+
 ### Flash Attention Backends
 
 GLiClass supports optional flash attention backends for faster inference.

From 683ef0386145ea90068edabce7cd80ea4dfa054d Mon Sep 17 00:00:00 2001
From: BioMikeUkr <mykhayloshtopko@gmail.com>
Date: Fri, 10 Jul 2026 13:08:13 +0000
Subject: [PATCH 17/25] streaming added

---
 pyproject.toml | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/pyproject.toml b/pyproject.toml
index a24f8e8..8545e52 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -36,6 +36,10 @@ dependencies = [
 dynamic = ["version"]
 
 [project.optional-dependencies]
+streaming = [
+    "torch>=2.0.0",
+    "transformers>=5.0.0",
+]
 serve = [
     "ray[serve]>=2.0.0",
     "requests>=2.31.0",

From 8626b67a51c958baea092fed92201640f5bb4bf6 Mon Sep 17 00:00:00 2001
From: BioMikeUkr <mykhayloshtopko@gmail.com>
Date: Fri, 10 Jul 2026 13:14:23 +0000
Subject: [PATCH 18/25] sessions man. added

---
 gliclass/streaming/pipeline.py | 22 ++++++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/gliclass/streaming/pipeline.py b/gliclass/streaming/pipeline.py
index 88e054e..087f268 100644
--- a/gliclass/streaming/pipeline.py
+++ b/gliclass/streaming/pipeline.py
@@ -93,6 +93,28 @@ def _load_to_device(self, session_ids: set[str], device: torch.device) -> None:
     # Session management
     # ------------------------------------------------------------------
 
+    def delete_session(self, *session_ids: str) -> list[str]:
+        """Delete one or more sessions and free their KV caches.
+
+        Returns the list of session IDs that were actually deleted
+        (unknown IDs are silently ignored).
+        """
+        deleted = []
+        for sid in session_ids:
+            if sid in self._caches:
+                del self._caches[sid]
+                deleted.append(sid)
+        return deleted
+
+    def clear_sessions(self) -> None:
+        """Delete all active sessions and free all KV caches."""
+        self._caches.clear()
+
+    @property
+    def active_sessions(self) -> list[str]:
+        """Return the list of currently active session IDs."""
+        return list(self._caches.keys())
+
     def _cleanup_dead_sessions(self, inputs: list[SessionInput]) -> None:
         active = {inp.session_id for inp in inputs}
         for sid in list(self._caches):

From 3f62c4dffb198a150cdb52363d496b912aa485de Mon Sep 17 00:00:00 2001
From: BioMikeUkr <mykhayloshtopko@gmail.com>
Date: Fri, 10 Jul 2026 13:15:35 +0000
Subject: [PATCH 19/25] sessions management examplrs added

---
 docs/streaming.md | 78 +++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 75 insertions(+), 3 deletions(-)

diff --git a/docs/streaming.md b/docs/streaming.md
index 4519cce..855c015 100644
--- a/docs/streaming.md
+++ b/docs/streaming.md
@@ -105,15 +105,62 @@ Sessions are identified by a string `session_id`. The pipeline maintains a dict
 
 **Updating.** On every call, the session's cache is extended with new tokens if `text` is non-empty. Empty text causes the update stage to be skipped for that session (`tokens_added = 0`), which also skips classification regardless of strategy.
 
-**Cleanup.** Sessions that are absent from the current input batch are deleted. The pipeline does not keep stale sessions alive across calls: any `session_id` not present in the current `inputs` list will have its cache freed. If a session should remain alive without receiving new text, pass it with `text=""`.
+**Cleanup.** Sessions that are absent from the current input batch are deleted. The pipeline does not keep stale sessions alive across calls: any `session_id` not present in the current `inputs` list will have its cache freed.
 
-**Manual reset.** Delete the cache entry directly:
+**Manual reset.** Use the explicit API:
 ```python
-del pipeline._caches["my_session"]
+pipeline.delete_session("my_session")          # delete one
+pipeline.delete_session("session_a", "session_b")  # delete several at once
+pipeline.clear_sessions()                      # delete all
+print(pipeline.active_sessions)                # ["session_c", ...]
 ```
 
 ---
 
+### Session survival rule
+
+> **A session lives only as long as its `session_id` appears in every `__call__`.**
+
+The cleanup runs at the very beginning of each call, before any processing. Any session absent from the current batch is immediately deleted — its entire KV cache is freed with no warning and no recovery.
+
+This means:
+
+**If you have a long-running session but no new text to add on a given call, you must still include it in the batch with `text=""`:**
+
+```python
+# Session "doc_A" has no new text this tick, but must stay alive
+pipeline([
+    SessionInput("doc_A", "", labels, strategy),   # keeps cache, no update
+    SessionInput("doc_B", new_chunk, labels, strategy),
+])
+```
+
+Passing `text=""` preserves the cache exactly as it is — no tokens are added, no classification is triggered regardless of strategy.
+
+**If you omit a session entirely, its cache is gone:**
+
+```python
+# Round 1 — both sessions created
+pipeline([SessionInput("doc_A", chunk1, ...), SessionInput("doc_B", chunk1, ...)])
+
+# Round 2 — doc_A is absent → its cache is deleted immediately
+pipeline([SessionInput("doc_B", chunk2, ...)])
+
+# Round 3 — doc_A reappears, but starts from scratch (empty cache)
+pipeline([SessionInput("doc_A", chunk2, ...), SessionInput("doc_B", chunk3, ...)])
+```
+
+**Typical patterns:**
+
+| Situation | What to pass |
+|---|---|
+| New text available | `text=chunk` |
+| Session idle (no new text), keep alive | `text=""` |
+| Session finished, free its memory | Omit it from the next call |
+| Explicit reset mid-stream | `pipeline.delete_session("session_id")`, then pass as new |
+
+---
+
 ## Classification strategies
 
 Strategies implement `ClassificationStrategy` and control:
@@ -235,6 +282,31 @@ StreamingPipeline(
 
 Processes a batch of sessions. Returns one `SessionOutput` per input in the same order. Sessions not present in the current batch have their caches deleted.
 
+### `delete_session(*session_ids: str) -> list[str]`
+
+Explicitly delete one or more sessions and free their KV caches. Returns the list of IDs that were actually deleted; unknown IDs are silently ignored.
+
+```python
+pipeline.delete_session("doc_001")
+pipeline.delete_session("doc_001", "doc_002", "doc_003")
+```
+
+### `clear_sessions() -> None`
+
+Delete all active sessions and free all KV caches at once.
+
+```python
+pipeline.clear_sessions()
+```
+
+### `active_sessions -> list[str]`
+
+Read-only property. Returns the list of currently active session IDs.
+
+```python
+print(pipeline.active_sessions)   # ["doc_001", "doc_002"]
+```
+
 ---
 
 ## Data types

From 8820bd50f30449461d820eb11a25ddcd45fb5b60 Mon Sep 17 00:00:00 2001
From: BioMikeUkr <mykhayloshtopko@gmail.com>
Date: Fri, 10 Jul 2026 13:25:59 +0000
Subject: [PATCH 20/25] ruff check

---
 gliclass/training.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/gliclass/training.py b/gliclass/training.py
index 3b73b36..4a88e02 100644
--- a/gliclass/training.py
+++ b/gliclass/training.py
@@ -294,7 +294,7 @@ class TrainingArguments(transformers.TrainingArguments):
         default=True,
         metadata={
             "help": "Save optimizer/scheduler state in checkpoints. Set False for large models to save disk space."
-        },  # noqa: E501
+        },
     )
 
 

From 0cb33fe833f2897fb49c51f455ef226489cfcfd5 Mon Sep 17 00:00:00 2001
From: Ingvar <igor.stepanov2000@gmail.com>
Date: Sat, 11 Jul 2026 12:29:40 +0300
Subject: [PATCH 21/25] vectorize _extract_representations

---
 gliclass/scorers.py | 42 ++++++++++++++++--------------------------
 1 file changed, 16 insertions(+), 26 deletions(-)

diff --git a/gliclass/scorers.py b/gliclass/scorers.py
index 0a93cbf..584e7c0 100644
--- a/gliclass/scorers.py
+++ b/gliclass/scorers.py
@@ -229,44 +229,34 @@ def _extract_representations(self, hidden_states, input_ids, attention_mask):
         Text repr: LAST <<SEP>> token (after all labels) - sees both text and labels
         Label repr: each <<LABEL>> token
         """
-        batch_size, _, hidden_size = hidden_states.shape
+        batch_size, seq_length, hidden_size = hidden_states.shape
 
         sep_token_id = self.config.sep_token_index
         label_token_id = self.config.class_token_index
 
-        sep_mask = input_ids == sep_token_id
-        label_mask = input_ids == label_token_id
+        valid_mask = attention_mask.bool()
+        sep_mask = input_ids.eq(sep_token_id) & valid_mask
+        label_mask = input_ids.eq(label_token_id) & valid_mask
 
-        sep_positions = torch.where(sep_mask)
-        label_positions = torch.where(label_mask)
-
-        text_repr = torch.zeros(batch_size, hidden_size, device=hidden_states.device, dtype=hidden_states.dtype)
-
-        for batch_idx in range(batch_size):
-            batch_sep_positions = sep_positions[1][sep_positions[0] == batch_idx]
-            if len(batch_sep_positions) > 0:
-                # Take LAST SEP (after all labels) so text repr sees both text and labels
-                last_sep_pos = batch_sep_positions[-1]
-                text_repr[batch_idx] = hidden_states[batch_idx, last_sep_pos]
-
-        num_labels_per_batch = torch.zeros(batch_size, dtype=torch.long, device=hidden_states.device)
-        for batch_idx in range(batch_size):
-            num_labels_per_batch[batch_idx] = (label_positions[0] == batch_idx).sum()
+        positions = torch.arange(seq_length, device=input_ids.device).expand(batch_size, -1)
+        last_sep_positions = positions.masked_fill(~sep_mask, -1).amax(dim=1)
+        has_sep = last_sep_positions.ge(0)
+        batch_indices = torch.arange(batch_size, device=input_ids.device)
+        text_repr = hidden_states[batch_indices, last_sep_positions.clamp_min(0)]
+        text_repr = text_repr * has_sep.unsqueeze(-1)
 
+        num_labels_per_batch = label_mask.sum(dim=1)
         max_labels = num_labels_per_batch.max().item()
-        label_repr = torch.zeros(
+        label_repr = hidden_states.new_zeros(
             batch_size,
             max_labels,
             hidden_size,
-            device=hidden_states.device,
-            dtype=hidden_states.dtype,
         )
 
-        for batch_idx in range(batch_size):
-            batch_label_positions = label_positions[1][label_positions[0] == batch_idx]
-            for label_idx, pos in enumerate(batch_label_positions):
-                if label_idx < max_labels:
-                    label_repr[batch_idx, label_idx] = hidden_states[batch_idx, pos]
+        label_batch_indices, label_token_indices = label_mask.nonzero(as_tuple=True)
+        label_ranks = label_mask.long().cumsum(dim=1) - 1
+        label_indices = label_ranks[label_batch_indices, label_token_indices]
+        label_repr[label_batch_indices, label_indices] = hidden_states[label_batch_indices, label_token_indices]
 
         return text_repr, label_repr
 

From 92de6a794de14d9687e2fc68dfec3a59d9f22fd8 Mon Sep 17 00:00:00 2001
From: Ingvar <igor.stepanov2000@gmail.com>
Date: Sat, 11 Jul 2026 18:22:12 +0300
Subject: [PATCH 22/25] vectorize _extract_label_section gliclass decoder

---
 gliclass/model.py | 20 ++++++++++----------
 1 file changed, 10 insertions(+), 10 deletions(-)

diff --git a/gliclass/model.py b/gliclass/model.py
index 696a4ce..f46b208 100644
--- a/gliclass/model.py
+++ b/gliclass/model.py
@@ -7,6 +7,7 @@
 import torch
 import transformers
 from torch import nn
+from torch.nn.utils.rnn import pad_sequence
 from packaging import version
 from transformers import AutoModel, AutoConfig, PreTrainedModel
 from transformers.utils import logging
@@ -1059,7 +1060,7 @@ def _extract_label_section(
             padded_ids:    (batch, max_label_len)
             label_mask:    (batch, max_label_len)
         """
-        batch_size, _, hidden_size = hidden_states.shape
+        batch_size = hidden_states.shape[0]
         sep_id = self.sep_token_id
         label_id = self.config.class_token_index
 
@@ -1092,16 +1093,15 @@ def _extract_label_section(
             slices_h.append(hidden_states[i, start:real_len])
             slices_ids.append(ids_i[start:real_len])
 
-        max_len = max(s.shape[0] for s in slices_h)
-        padded_hidden = hidden_states.new_zeros(batch_size, max_len, hidden_size)
-        padded_ids = input_ids.new_zeros(batch_size, max_len)
-        label_mask = attention_mask.new_zeros(batch_size, max_len)
+        padded_hidden = pad_sequence(slices_h, batch_first=True)
+        padded_ids = pad_sequence(slices_ids, batch_first=True)
 
-        for i, (h, ids) in enumerate(zip(slices_h, slices_ids)):
-            length = h.shape[0]
-            padded_hidden[i, :length] = h
-            padded_ids[i, :length] = ids
-            label_mask[i, :length] = 1
+        section_lengths = torch.tensor(
+            [section.shape[0] for section in slices_h],
+            device=attention_mask.device,
+        )
+        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
 

From db7eb3f20fdf404abcf71f230851e35c504d8407 Mon Sep 17 00:00:00 2001
From: Ingvar <igor.stepanov2000@gmail.com>
Date: Sun, 12 Jul 2026 09:55:20 +0300
Subject: [PATCH 23/25] align design of streaming pipeline with the repo

---
 docs/streaming.md              | 533 +++++++++---------------
 gliclass/__init__.py           |   1 +
 gliclass/data_processing.py    |  57 ++-
 gliclass/model.py              |  62 ++-
 gliclass/pipeline.py           |  75 ++--
 gliclass/streaming/__init__.py |  18 +-
 gliclass/streaming/cache.py    |  24 +-
 gliclass/streaming/pipeline.py | 716 ++++++++++++++++++---------------
 gliclass/streaming/types.py    |  26 --
 tests/test_decoder_kv.py       |  83 ++++
 tests/test_scorers.py          |  56 +++
 tests/test_streaming.py        | 255 +++++++++---
 12 files changed, 1079 insertions(+), 827 deletions(-)
 delete mode 100644 gliclass/streaming/types.py
 create mode 100644 tests/test_decoder_kv.py

diff --git a/docs/streaming.md b/docs/streaming.md
index 855c015..529c08f 100644
--- a/docs/streaming.md
+++ b/docs/streaming.md
@@ -1,437 +1,274 @@
-# GLiClass Streaming Classification
+# Streaming classification
 
-`gliclass.streaming` provides incremental, multi-session text classification over a decoder-KV model. Instead of classifying a complete document at once, it processes text in arbitrary chunks, maintains a persistent KV cache per session, and triggers classification only when a pluggable strategy decides it is appropriate.
+`StreamingZeroShotClassificationPipeline` incrementally classifies text with a decoder-KV model. It keeps one text-only KV cache per active session and evaluates label sequences without adding those labels to the persistent cache.
 
----
+## Installation
 
-## Contents
-
-- [Architecture overview](#architecture-overview)
-- [KV cache management](#kv-cache-management)
-- [Session lifecycle](#session-lifecycle)
-- [Classification strategies](#classification-strategies)
-- [StreamingPipeline reference](#streamingpipeline-reference)
-- [Data types](#data-types)
-- [Usage examples](#usage-examples)
-- [Memory and performance considerations](#memory-and-performance-considerations)
-
----
-
-## Architecture overview
-
-GLiClass Decoder-KV separates the model into two components:
-
-- **`decoder_model`** — a causal language model (e.g. Qwen3.5) that encodes text into a sequence of hidden states and accumulates a KV cache.
-- **`scorer`** — a lightweight cross-attention encoder that reads the label hidden states produced by the decoder and outputs per-label logits.
-
-In standard (non-streaming) inference, both components run on the full text each time. In streaming mode, the decoder runs once per chunk to *extend* the KV cache; the scorer runs only when a strategy decides classification should fire. The result is:
-
-```
-chunk_1 → decoder → KV₁
-chunk_2 → decoder (KV₁ as prefix) → KV₂
-chunk_3 → decoder (KV₂ as prefix) → KV₃  ← strategy triggers → scorer → predictions
-chunk_4 → ...
+```bash
+pip install gliclass[streaming]
 ```
 
-Each forward pass of the decoder is O(n_chunk × n_cache) rather than O(n_total²), because the attention over cached tokens is handled by the KV cache mechanism and never recomputed.
+## Basic usage
 
-### Two-stage per call
-
-Every `StreamingPipeline.__call__` runs two sequential stages over the input batch:
-
-**Stage 1 — cache update.**  
-New text is tokenized per session, then all sessions in the batch are forwarded through `decoder_model` together. The resulting `past_key_values` are sliced per session and stored in `_caches`.
-
-**Stage 2 — classification.**  
-Sessions whose strategy returns `True` from `should_classify` are forwarded again through `decoder_model` with their KV cache as prefix and the label token sequence as input. The decoder's `last_hidden_state` over label positions is then passed to `scorer`, which produces logits. Sigmoid (multi-label) or softmax (single-label) is applied to produce probabilities.
-
-Sessions that did not trigger classification return `triggered=False` with `predictions=None`.
-
----
+```python
+from gliclass.streaming import EveryNTokensStrategy, StreamingZeroShotClassificationPipeline
 
-## KV cache management
+pipeline = StreamingZeroShotClassificationPipeline(
+    model,
+    tokenizer,
+    device="cuda",
+    max_cache_len=1024,
+    default_strategy=EveryNTokensStrategy(50),
+)
 
-### CacheState
+for chunk in text_chunks:
+    output = pipeline(
+        chunk,
+        ["science", "politics", "finance"],
+        session_ids="document-1",
+        threshold=0.5,
+    )[0]
+    if output["triggered"]:
+        print(output["predictions"])
+```
 
-`CacheState` is the per-session container that travels through the pipeline:
+The call returns one dictionary per input text:
 
 ```python
-@dataclass
-class CacheState:
-    past_key_values: Any          # transformers DynamicCache or None
-    input_ids: torch.Tensor       # flat 1-D tensor of all cached token ids
-    attention_mask: torch.Tensor  # flat 1-D tensor of all cached attention mask values
-    position_ids: torch.Tensor | None
-    cached_length: int            # total number of cached tokens
-    session_id: str | None
-    metadata: dict                # arbitrary user data
+{
+    "session_id": "document-1",
+    "triggered": True,
+    "predictions": [{"label": "science", "score": 0.91}],
+    "cached_length": 128,
+    "tokens_added": 17,
+}
 ```
 
-A fresh empty cache is created via `create_empty_cache(session_id, device)`.
-
-### Batched stacking with heterogeneous cache lengths
-
-When multiple sessions exist in a batch and have different numbers of cached tokens, their KV caches cannot be naively stacked: different sequence lengths mean different position embeddings baked into the K/V tensors.
-
-`BatchedKVHelper` resolves this with **prepend-padding**:
-
-- KV caches shorter than `max_cached_len` are zero-padded on the **left** (prepend), not the right.
-- The attention mask marks these prepended zeros as not-attended, so the model ignores them.
-- `position_ids` for the new tokens in each session are set to `[cached_len_i, ..., cached_len_i + new_len_i)`, so RoPE offsets remain correct regardless of padding.
+## Batched sessions
 
-This means a batch of sessions with cache lengths `[150, 400, 300]` is padded to `max_cached=400`, with sessions 0 and 2 receiving 250 and 100 prepend zeros respectively.
-
-After the forward pass, the updated KV is sliced back per session using the real token range `[max_cached - clen_i : max_cached + new_len_i]`.
-
-### Cache truncation
-
-`truncate_cache(cache_state, max_length)` crops the KV to the **last** `max_length` tokens. Only `DynamicCache` is supported for cropping; other cache formats are returned unchanged. The pipeline applies this automatically after each stage-1 update when `max_cache_len` is set.
-
-Truncation discards the oldest context first. This is a simple sliding-window policy; for more selective eviction (e.g. keeping the first few tokens as a persistent prefix), extend or replace the default behavior by subclassing `StreamingPipeline._enforce_max_cache_len`.
-
-### Device offloading
-
-When running more sessions than fit in GPU memory simultaneously, `offload_to_cpu=True` keeps inactive session caches in CPU RAM and moves them to GPU only for the active sub-batch.
+```python
+outputs = pipeline(
+    texts=[chunk_a, chunk_b],
+    labels=["positive", "negative"],
+    session_ids=["session-a", "session-b"],
+    strategies=[strategy_a, strategy_b],
+    batch_size=8,
+)
+```
 
-With `use_pinned_memory=True`, KV caches are allocated in CPU pinned (page-locked) memory, enabling asynchronous DMA transfers to the GPU via `non_blocking=True`. A `torch.cuda.synchronize` call before the forward pass ensures transfers are complete before the kernel launches.
+`texts`, `session_ids`, and a per-session `strategies` list must have matching lengths. Labels may be shared, specified per input, or provided as a hierarchical dictionary like the regular zero-shot pipeline. A single strategy is copied independently for every new session; a strategy list assigns one strategy to each new session.
 
----
+Session IDs must be unique within a call.
 
 ## Session lifecycle
 
-Sessions are identified by a string `session_id`. The pipeline maintains a dict `_caches: dict[str, CacheState]` internally.
-
-**Creation.** A session is created implicitly when a `session_id` appears for the first time in an input batch. `create_empty_cache` is called and stored.
-
-**Updating.** On every call, the session's cache is extended with new tokens if `text` is non-empty. Empty text causes the update stage to be skipped for that session (`tokens_added = 0`), which also skips classification regardless of strategy.
-
-**Cleanup.** Sessions that are absent from the current input batch are deleted. The pipeline does not keep stale sessions alive across calls: any `session_id` not present in the current `inputs` list will have its cache freed.
+The session IDs in each top-level call are the complete active set. A cached session is deleted automatically when it is absent from the next call.
 
-**Manual reset.** Use the explicit API:
 ```python
-pipeline.delete_session("my_session")          # delete one
-pipeline.delete_session("session_a", "session_b")  # delete several at once
-pipeline.clear_sessions()                      # delete all
-print(pipeline.active_sessions)                # ["session_c", ...]
-```
-
----
-
-### Session survival rule
-
-> **A session lives only as long as its `session_id` appears in every `__call__`.**
+pipeline([chunk_a, chunk_b], labels, session_ids=["a", "b"])
+pipeline([next_b, chunk_c], labels, session_ids=["b", "c"])
 
-The cleanup runs at the very beginning of each call, before any processing. Any session absent from the current batch is immediately deleted — its entire KV cache is freed with no warning and no recovery.
+# Session "a" has been deleted; "b" was extended; "c" was created.
+assert pipeline.active_sessions == ["b", "c"]
+```
 
-This means:
+Cleanup is performed once for the complete call, before internal `batch_size` splitting. Therefore, sessions are not accidentally removed merely because they belong to different internal sub-batches.
 
-**If you have a long-running session but no new text to add on a given call, you must still include it in the batch with `text=""`:**
+To retain a session without adding text, include it with an empty string:
 
 ```python
-# Session "doc_A" has no new text this tick, but must stay alive
-pipeline([
-    SessionInput("doc_A", "", labels, strategy),   # keeps cache, no update
-    SessionInput("doc_B", new_chunk, labels, strategy),
-])
+pipeline(
+    texts=["", next_chunk_b],
+    labels=labels,
+    session_ids=["a", "b"],
+)
 ```
 
-Passing `text=""` preserves the cache exactly as it is — no tokens are added, no classification is triggered regardless of strategy.
+Session `a` remains active, reports `tokens_added=0`, and is not classified. Omitting `a` would delete it.
 
-**If you omit a session entirely, its cache is gone:**
+Manual lifecycle methods are also available:
 
 ```python
-# Round 1 — both sessions created
-pipeline([SessionInput("doc_A", chunk1, ...), SessionInput("doc_B", chunk1, ...)])
-
-# Round 2 — doc_A is absent → its cache is deleted immediately
-pipeline([SessionInput("doc_B", chunk2, ...)])
-
-# Round 3 — doc_A reappears, but starts from scratch (empty cache)
-pipeline([SessionInput("doc_A", chunk2, ...), SessionInput("doc_B", chunk3, ...)])
+pipeline.delete_session("b")
+pipeline.clear_sessions()
+print(pipeline.active_sessions)
 ```
 
-**Typical patterns:**
-
-| Situation | What to pass |
-|---|---|
-| New text available | `text=chunk` |
-| Session idle (no new text), keep alive | `text=""` |
-| Session finished, free its memory | Omit it from the next call |
-| Explicit reset mid-stream | `pipeline.delete_session("session_id")`, then pass as new |
-
----
-
-## Classification strategies
-
-Strategies implement `ClassificationStrategy` and control:
-
-1. **`should_classify(tokens_added, cached_length, text) -> bool`** — whether to trigger classification this call.
-2. **`get_window(cache_state) -> CacheState`** — which portion of the cache to classify over (default: full cache).
-
-### Built-in strategies
-
-#### `EveryChunkStrategy`
-
-Classifies on every call that adds at least one token. Use when latency is acceptable and you want the most up-to-date prediction at every step.
+Calling the pipeline with empty `texts` and `session_ids` clears all sessions:
 
 ```python
-EveryChunkStrategy()
+pipeline([], [], session_ids=[])
 ```
 
-#### `EveryNTokensStrategy`
+## Prompt and examples
 
-Accumulates tokens across calls and classifies once every `n` tokens. The counter resets after each trigger. Useful for throttling classification frequency over a fast-arriving stream.
+Prompt and few-shot examples become a one-time prefix when a session is created:
 
 ```python
-EveryNTokensStrategy(n=100)
+pipeline(
+    first_chunk,
+    labels,
+    session_ids="document-1",
+    prompt="Classify this document: ",
+    examples=examples,
+)
 ```
 
-Note: the internal `_accumulated` counter is per-instance. If the same strategy object is shared across sessions, the counter is shared too. Create one instance per session or use `ComposedStrategy` with per-session instantiation.
+The prefix is not appended again on later calls. Repeating the same prompt/examples is allowed. Once a session contains cached tokens, changing its prompt or examples raises an error; delete the session first when a different prefix is required.
 
-#### `OnDelimiterStrategy`
+If a new session has an empty text but a non-empty prompt or examples, the prefix is still cached and counted in `tokens_added`. Prefix-only updates never trigger classification.
 
-Classifies when the incoming text chunk contains a specific string, such as a sentence boundary or a paragraph marker.
+## Classification strategies
 
-```python
-OnDelimiterStrategy(delimiter="\n\n")   # paragraph boundary
-OnDelimiterStrategy(delimiter=". ")    # sentence boundary (approximate)
-```
+Available strategies:
 
-The delimiter check is a plain substring search on the raw text string, not on tokens.
+- `EveryChunkStrategy`: classify every non-empty update.
+- `EveryNTokensStrategy(n)`: classify after accumulating `n` tokens.
+- `OnDelimiterStrategy(delimiter)`: classify when the incoming text contains the delimiter.
+- `NeverStrategy`: update the cache without classifying.
+- `SlidingWindowStrategy(window_size)`: classify every chunk using only the newest cached tokens.
+- `ComposedStrategy(trigger, window)`: combine a trigger with a classification window.
 
-#### `NeverStrategy`
+Each session owns an independent copy of its strategy. Stateful counters are never shared across sessions, even when one strategy is supplied as the pipeline default. `default_strategy=None` selects `EveryChunkStrategy`.
 
-Never triggers classification. Use this for a pre-filling phase where you want to build up a long context cache before any classification runs.
+The `strategies` call argument initializes new sessions. Passing another strategy for an already-active session does not reset its state. Replace an active strategy explicitly with:
 
 ```python
-NeverStrategy()
+pipeline.set_session_strategy("document-1", EveryChunkStrategy())
 ```
 
-Switch to a different strategy object after pre-filling is complete.
+## Two-stage execution
 
-#### `SlidingWindowStrategy`
+Each call has two stages:
 
-Classifies on every chunk (same trigger condition as `EveryChunkStrategy`), but passes only the last `window_size` cached tokens to the scorer via `get_window`. This limits the attention span of the classification pass without discarding the full KV cache.
+1. New text is passed to `model.update_decoder_cache(...)` with `use_cache=True`. The resulting text-only cache replaces the cache for that session.
+2. Triggered sessions pass `<<SEP>>label<<LABEL>>...<<SEP>>` to `model.classify_from_decoder_cache(...)`. This stage uses the text cache as context but does not save label KV tensors.
 
-```python
-SlidingWindowStrategy(window_size=512)
-```
+Existing sessions with empty text receive `tokens_added=0` and are not classified. A new prefix-only session may add prompt/example tokens, but it is also not classified until a non-empty text chunk arrives.
+
+## Heterogeneous batching
 
-Note that this does not truncate the stored KV cache — the full cache is retained for future updates. Only the classification forward sees the windowed slice.
+Sessions may have different cache lengths. Before a decoder call, shorter caches are padded on the left to the longest cache in the batch. The attention mask hides padding, and explicit position IDs preserve each session's logical positions. After updating, the batched KV tensors are copied back into independent per-session caches.
 
-#### `ComposedStrategy`
+## Cache truncation
 
-Combines a **trigger strategy** (controls `should_classify`) with a **window strategy** (controls `get_window`). This is the standard way to pair a frequency policy with a context window policy.
+Set `max_cache_len` to retain only the newest KV entries:
 
 ```python
-ComposedStrategy(
-    trigger=EveryNTokensStrategy(200),
-    window=SlidingWindowStrategy(512),
+pipeline = StreamingZeroShotClassificationPipeline(
+    model,
+    tokenizer,
+    max_cache_len=512,
 )
 ```
 
-Classifies every 200 new tokens, but uses only the last 512 cached tokens for scoring.
+The physical cache length is capped, while the next absolute position ID continues increasing. This avoids reusing RoPE positions after truncation. Cropping currently requires the Transformers `DynamicCache` format used by supported decoder-KV models.
 
-### Custom strategies
-
-Subclass `ClassificationStrategy` and implement `should_classify`. Optionally override `get_window` for windowed classification.
+## CPU offload
 
 ```python
-class SentenceCountStrategy(ClassificationStrategy):
-    def __init__(self, every_n_sentences: int):
-        self.n = every_n_sentences
-        self._count = 0
-
-    def should_classify(self, tokens_added, cached_length, text):
-        self._count += text.count(". ")
-        if self._count >= self.n:
-            self._count = 0
-            return True
-        return False
+pipeline = StreamingZeroShotClassificationPipeline(
+    model,
+    tokenizer,
+    device="cuda",
+    offload_to_cpu=True,
+    use_pinned_memory=True,
+    score_on_cpu=True,
+)
 ```
 
----
+With CPU offload enabled, caches are loaded for the active internal batch and moved back to CPU afterward. Pinned memory enables non-blocking GPU transfers when CUDA is used.
 
-## StreamingPipeline reference
+## Constructor
 
 ```python
-StreamingPipeline(
+StreamingZeroShotClassificationPipeline(
     model,
     tokenizer,
-    device="cpu",
+    max_classes=25,
+    max_length=1024,
+    classification_type="multi-label",
+    device="cuda:0",
+    progress_bar=False,
+    label_separator=".",
     max_cache_len=None,
-    batch_size=None,
+    default_batch_size=8,
+    default_strategy=None,
     offload_to_cpu=False,
     use_pinned_memory=False,
     score_on_cpu=False,
 )
 ```
 
-| Parameter | Type | Description |
-|---|---|---|
-| `model` | `GLiClassModel` | Loaded GLiClass decoder-KV model in eval mode. |
-| `tokenizer` | `PreTrainedTokenizer` | Matching tokenizer with `<<LABEL>>`, `<<SEP>>`, `<<EXAMPLE>>` special tokens. |
-| `device` | `str \| torch.device` | Device for decoder forward passes. |
-| `max_cache_len` | `int \| None` | Hard cap on cached tokens per session. Oldest tokens are dropped when exceeded. `None` means unbounded. |
-| `batch_size` | `int \| None` | Maximum sessions per decoder forward. When `len(inputs) > batch_size`, sessions are chunked. With `offload_to_cpu=True`, only the active chunk is on GPU at once. |
-| `offload_to_cpu` | `bool` | Move inactive session caches to CPU RAM between sub-batches. |
-| `use_pinned_memory` | `bool` | Allocate CPU caches in pinned memory for async GPU uploads. Only effective when `offload_to_cpu=True` and `device` is CUDA. |
-| `score_on_cpu` | `bool` | Run the scorer on CPU. Useful when GPU VRAM is the bottleneck and the scorer is small relative to the decoder. |
-
-### `__call__(inputs: list[SessionInput]) -> list[SessionOutput]`
-
-Processes a batch of sessions. Returns one `SessionOutput` per input in the same order. Sessions not present in the current batch have their caches deleted.
-
-### `delete_session(*session_ids: str) -> list[str]`
-
-Explicitly delete one or more sessions and free their KV caches. Returns the list of IDs that were actually deleted; unknown IDs are silently ignored.
-
-```python
-pipeline.delete_session("doc_001")
-pipeline.delete_session("doc_001", "doc_002", "doc_003")
-```
-
-### `clear_sessions() -> None`
-
-Delete all active sessions and free all KV caches at once.
-
-```python
-pipeline.clear_sessions()
-```
-
-### `active_sessions -> list[str]`
-
-Read-only property. Returns the list of currently active session IDs.
-
-```python
-print(pipeline.active_sessions)   # ["doc_001", "doc_002"]
-```
-
----
-
-## Data types
-
-### `SessionInput`
-
-```python
-@dataclass
-class SessionInput:
-    session_id: str                    # unique identifier for this session
-    text: str                          # new text chunk to append (empty = no update)
-    labels: list[str]                  # candidate label strings
-    strategy: ClassificationStrategy   # per-session strategy instance
-    classification_type: str           # "multi-label" (sigmoid) or "single-label" (softmax)
-```
-
-### `SessionOutput`
-
-```python
-@dataclass
-class SessionOutput:
-    session_id: str
-    triggered: bool                    # True if classification ran this call
-    predictions: list[dict] | None     # [{"label": str, "score": float}, ...] or None
-    cached_length: int                 # total cached tokens after this call
-    tokens_added: int                  # tokens added in stage 1 (0 if text was empty)
-    metadata: dict                     # pass-through from CacheState.metadata
-```
-
-For `multi-label`, `predictions` contains one entry per label with a sigmoid probability. For `single-label`, it contains only the top-1 label with its softmax probability.
-
----
-
-## Usage examples
+The model must use `architecture_type="decoder-kv"`.
 
-### Minimal single-session loop
-
-```python
-import torch
-from transformers import AutoTokenizer
-from gliclass import GLiClassModel
-from gliclass.streaming import StreamingPipeline, SessionInput, EveryNTokensStrategy
-
-tokenizer = AutoTokenizer.from_pretrained("path/to/model")
-model = GLiClassModel.from_pretrained("path/to/model")
-model = model.to("cuda", dtype=torch.bfloat16).eval()
-
-pipeline = StreamingPipeline(model, tokenizer, device="cuda", max_cache_len=1024)
-
-labels = ["science", "politics", "finance", "sports"]
-strategy = EveryNTokensStrategy(n=50)
-
-chunks = ["The Fed raised interest rates", " by 25 basis points", " amid inflation concerns."]
-
-for chunk in chunks:
-    outputs = pipeline([
-        SessionInput(
-            session_id="doc_001",
-            text=chunk,
-            labels=labels,
-            strategy=strategy,
-            classification_type="multi-label",
-        )
-    ])
-    out = outputs[0]
-    print(f"cached={out.cached_length} +{out.tokens_added} triggered={out.triggered}")
-    if out.triggered:
-        for p in sorted(out.predictions, key=lambda x: -x["score"]):
-            print(f"  {p['label']}: {p['score']:.3f}")
-```
-
-### Multiple concurrent sessions
+| Parameter | Behavior |
+|---|---|
+| `model` | A loaded GLiClass decoder-KV model. The pipeline moves it to `device` and enables evaluation mode. |
+| `tokenizer` | Matching tokenizer instance or pretrained tokenizer name. |
+| `max_classes` | Inherited compatibility setting; decoder-KV streaming does not currently split labels using it. |
+| `max_length` | Inherited pipeline configuration value. Streaming cache length is controlled separately by `max_cache_len`. |
+| `classification_type` | Default `"multi-label"` or `"single-label"` behavior. |
+| `device` | Decoder execution device. Falls back to CPU when CUDA is requested but unavailable. |
+| `progress_bar` | Inherited pipeline option; streaming calls do not display a progress bar. |
+| `label_separator` | Separator used to flatten and rebuild hierarchical labels. |
+| `max_cache_len` | Maximum physically retained tokens per session; `None` is unbounded. Must be positive when provided. |
+| `default_batch_size` | Internal session batch size when `batch_size` is omitted from a call. Must be positive. |
+| `default_strategy` | Strategy copied into each new session; `None` means `EveryChunkStrategy`. |
+| `offload_to_cpu` | Move each processed internal batch's caches back to CPU. |
+| `use_pinned_memory` | Use pinned CPU cache memory for CUDA transfers; effective only with CUDA offload. |
+| `score_on_cpu` | Move the decoder-KV scorer and its inputs to CPU. |
+
+## Call interface
 
 ```python
-session_strategies = {
-    "session_A": EveryNTokensStrategy(n=100),
-    "session_B": OnDelimiterStrategy(delimiter="\n"),
-    "session_C": ComposedStrategy(
-        trigger=EveryNTokensStrategy(50),
-        window=SlidingWindowStrategy(256),
-    ),
-}
-
-batch = [
-    SessionInput("session_A", chunk_a, labels, session_strategies["session_A"]),
-    SessionInput("session_B", chunk_b, labels, session_strategies["session_B"]),
-    SessionInput("session_C", chunk_c, labels, session_strategies["session_C"]),
-]
-
-outputs = pipeline(batch)
+pipeline(
+    texts,
+    labels,
+    *,
+    session_ids,
+    threshold=0.5,
+    batch_size=None,
+    classification_type=None,
+    strategies=None,
+    examples=None,
+    prompt=None,
+    return_hierarchical=False,
+)
 ```
 
-All three sessions are updated and classified (if triggered) in a single batched forward pass per stage.
-
-### Pre-filling a long context before classification
-
-```python
-from gliclass.streaming import NeverStrategy, EveryChunkStrategy
-
-# Phase 1: ingest background context without any classification overhead
-for paragraph in background_paragraphs:
-    pipeline([SessionInput("doc", paragraph, labels, NeverStrategy())])
-
-# Phase 2: stream the main document and classify on every chunk
-for chunk in main_document_chunks:
-    outputs = pipeline([SessionInput("doc", chunk, labels, EveryChunkStrategy())])
-    if outputs[0].triggered:
-        handle(outputs[0].predictions)
-```
+Multi-label predictions are filtered by `threshold`. Single-label classification applies softmax and returns the highest-scoring label. `return_hierarchical=True` reconstructs the same hierarchical output format as the regular zero-shot pipeline.
 
-### Memory-constrained deployment with CPU offload
+| Argument | Accepted values |
+|---|---|
+| `texts` | One chunk string or a list of chunk strings. |
+| `labels` | Shared label list, per-text label lists, one hierarchical dictionary, or per-text hierarchical dictionaries. At least one label is required whenever the call contains sessions, including sessions with empty chunks. |
+| `session_ids` | One non-empty ID or a list matching `texts`. IDs must be unique within the call. |
+| `threshold` | One multi-label threshold or a list matching `texts`. |
+| `batch_size` | Positive internal batch size override; `None` uses `default_batch_size`. |
+| `classification_type` | Default, one supported mode, or one mode per text. Aliases such as `"single"` and `"multi_label"` are normalized. |
+| `strategies` | Default, one strategy copied for new sessions, or one strategy per input. Use `set_session_strategy()` for active sessions. |
+| `examples` | Shared examples or one examples list per input. Applied once when a session cache is empty. |
+| `prompt` | Shared prompt or one prompt per input. Applied once when a session cache is empty. |
+| `return_hierarchical` | Rebuild prediction scores into the supplied hierarchical label structure. |
+
+## Return values
+
+The pipeline always returns a list in input order. Each item contains:
+
+| Field | Meaning |
+|---|---|
+| `session_id` | Session associated with this result. |
+| `triggered` | Whether classification ran in this call. |
+| `predictions` | Flat prediction list, hierarchical score dictionary, or `None` when not triggered. |
+| `cached_length` | Physical text-cache length after update and optional truncation. |
+| `tokens_added` | Number of context tokens added in this call, including a new session's one-time prefix. |
 
-```python
-pipeline = StreamingPipeline(
-    model=model,
-    tokenizer=tokenizer,
-    device="cuda",
-    max_cache_len=512,
-    batch_size=8,           # process 8 sessions per GPU forward
-    offload_to_cpu=True,    # keep inactive caches in RAM
-    use_pinned_memory=True, # async DMA transfers
-    score_on_cpu=True,      # scorer stays in CPU RAM
-)
-```
+For multi-label classification, a triggered call may legitimately return an empty prediction list when every score is below `threshold`.
 
-With 32 concurrent sessions and `batch_size=8`, the GPU holds at most 8 KV caches at a time. Inactive caches are evicted to CPU pinned memory and uploaded asynchronously before their sub-batch runs.
+## Session methods
 
----
\ No newline at end of file
+- `active_sessions` returns active IDs in insertion order.
+- `delete_session(*session_ids)` removes cache and strategy state and returns the IDs that actually existed.
+- `clear_sessions()` removes all cache and strategy state.
+- `set_session_strategy(session_id, strategy)` deep-copies a new strategy into an active session and raises `KeyError` for an unknown ID.
diff --git a/gliclass/__init__.py b/gliclass/__init__.py
index 8dfaa6b..1981be3 100644
--- a/gliclass/__init__.py
+++ b/gliclass/__init__.py
@@ -5,6 +5,7 @@
     BiEncoderZeroShotClassificationPipeline,
     ZeroShotClassificationWithChunkingPipeline,
 )
+from .streaming import StreamingZeroShotClassificationPipeline
 
 __version__ = "0.1.20"
 
diff --git a/gliclass/data_processing.py b/gliclass/data_processing.py
index 700c5f5..f9bfc12 100644
--- a/gliclass/data_processing.py
+++ b/gliclass/data_processing.py
@@ -7,6 +7,42 @@
 from torch.nn.utils.rnn import pad_sequence
 
 
+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_labels(
+    labels: list[str],
+    *,
+    label_token: str = "<<LABEL>>",
+    sep_token: str = "<<SEP>>",
+) -> str:
+    """Format the transient label portion of a decoder-KV sequence."""
+    parts = [sep_token]
+    for label in labels:
+        parts.extend((label, label_token))
+    parts.append(sep_token)
+    return "".join(parts)
+
+
+def format_decoder_kv_sequence(
+    text: str,
+    labels: list[str],
+    *,
+    prompt: str = "",
+    examples: str = "",
+    label_token: str = "<<LABEL>>",
+    sep_token: str = "<<SEP>>",
+) -> str:
+    """Format a complete decoder-KV sequence for training or classic inference."""
+    return format_decoder_kv_context(text, prompt, examples) + format_decoder_kv_labels(
+        labels,
+        label_token=label_token,
+        sep_token=sep_token,
+    )
+
+
 @dataclass
 class AugmentationConfig:
     """Configuration for data augmentation."""
@@ -305,19 +341,14 @@ def tokenize_and_prepare_labels_for_decoder_kv(self, example):
         examples_text = self.format_examples(example)
         text = str(example["text"])
 
-        input_parts = []
-        if prompt:
-            input_parts.append(prompt)
-        input_parts.append(examples_text)
-        input_parts.append(text)
-        input_parts.append(self.sep_token)
-
-        for label in example["all_labels"]:
-            input_parts.append(label)
-            input_parts.append(self.label_token)
-        input_parts.append(self.sep_token)
-
-        input_text = "".join(input_parts)
+        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,
+        )
         label2idx = {label: idx for idx, label in enumerate(example["all_labels"])}
 
         tokenized_inputs = self.tokenize(input_text)
diff --git a/gliclass/model.py b/gliclass/model.py
index f46b208..8d8702b 100644
--- a/gliclass/model.py
+++ b/gliclass/model.py
@@ -1,15 +1,15 @@
 import os
 import warnings
-from typing import Tuple
+from typing import Any, Tuple
 from pathlib import Path
 from dataclasses import dataclass
 
 import torch
 import transformers
 from torch import nn
-from torch.nn.utils.rnn import pad_sequence
 from packaging import version
 from transformers import AutoModel, AutoConfig, PreTrainedModel
+from torch.nn.utils.rnn import pad_sequence
 from transformers.utils import logging
 from transformers.modeling_outputs import SequenceClassifierOutput
 
@@ -61,6 +61,7 @@
 class GLiClassOutput(SequenceClassifierOutput):
     text_embeddings: torch.Tensor | None = None
     class_embeddings: torch.Tensor | None = None
+    past_key_values: Any | None = None
 
 
 class GLiClassPreTrainedModel(PreTrainedModel):
@@ -1105,6 +1106,50 @@ def _extract_label_section(
 
         return padded_hidden, padded_ids, label_mask
 
+    def update_decoder_cache(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        position_ids: torch.Tensor | None = None,
+        past_key_values=None,
+    ):
+        """Extend a text-only decoder cache without running the scorer."""
+        return self.decoder_model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            use_cache=True,
+            return_dict=True,
+        )
+
+    def classify_from_decoder_cache(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        label_mask: torch.Tensor,
+        position_ids: torch.Tensor,
+        past_key_values,
+    ) -> GLiClassOutput:
+        """Classify label tokens against a text cache without persisting labels."""
+        decoder_outputs = self.decoder_model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            use_cache=False,
+            return_dict=True,
+        )
+
+        label_hidden = decoder_outputs.last_hidden_state[:, -input_ids.shape[1] :, :]
+        scorer_device = next(self.scorer.parameters()).device
+        logits = self.scorer(
+            hidden_states=label_hidden.to(scorer_device),
+            input_ids=input_ids.to(scorer_device),
+            attention_mask=label_mask.to(scorer_device),
+        )
+        return GLiClassOutput(logits=logits)
+
     def forward(
         self,
         input_ids: torch.Tensor | None = None,
@@ -1165,6 +1210,7 @@ def forward(
             logits=logits,
             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,
         )
 
 
@@ -1281,3 +1327,15 @@ def forward(self, *args, **kwargs):
             kwargs.pop("adapter_ids", None)
         outputs = self.model(*args, **kwargs)
         return outputs
+
+    def update_decoder_cache(self, **kwargs):
+        """Extend a decoder-KV text cache."""
+        if self.config.architecture_type != "decoder-kv":
+            raise ValueError("Decoder cache updates require architecture_type='decoder-kv'.")
+        return self.model.update_decoder_cache(**kwargs)
+
+    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)
diff --git a/gliclass/pipeline.py b/gliclass/pipeline.py
index 0fc9d41..938754a 100644
--- a/gliclass/pipeline.py
+++ b/gliclass/pipeline.py
@@ -7,6 +7,7 @@
 
 from .model import GLiClassModel, GLiClassBiEncoder
 from .utils import retrieval_augmented_text
+from .data_processing import format_decoder_kv_sequence
 
 
 def flatten_hierarchical_labels(
@@ -258,6 +259,26 @@ def _normalize_adapter_ids(self, adapter_ids: str | List[str] | None, num_texts:
             raise ValueError("Length of adapter_ids list must match number of texts.")
         return adapter_ids
 
+    @staticmethod
+    def _postprocess_logits(
+        logits: torch.Tensor,
+        labels: List[str],
+        classification_type: str,
+        threshold: float,
+    ) -> tuple[List[Dict[str, Any]], Dict[str, float]]:
+        """Convert one row of logits into predictions and a complete score map."""
+        logits = logits[: len(labels)]
+        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]
+        return predictions, all_scores
+
     def _process_labels(
         self, labels: List[str] | Dict[str, Any] | List[List[str]] | List[Dict[str, Any]]
     ) -> List[str] | List[List[str]]:
@@ -484,7 +505,6 @@ def __call__(
                 adapter_ids=batch_adapter_ids,
             )
             logits = model_output.logits
-            probs = torch.sigmoid(logits)
 
             for i in range(len(batch_texts)):
                 global_idx = idx + i
@@ -496,29 +516,15 @@ def __call__(
                 else:
                     curr_labels = batch_labels[i]
 
-                if item_classification_type == "single-label":
-                    score = torch.softmax(logits[i][: len(curr_labels)], dim=-1)
-
-                    if return_hierarchical:
-                        all_scores = {curr_labels[j]: score[j].item() for j in range(len(curr_labels))}
-                        all_scores_list.append(all_scores)
-
-                    pred_label = curr_labels[torch.argmax(score).item()]
-                    results.append([{"label": pred_label, "score": score.max().item()}])
-                elif item_classification_type == "multi-label":
-                    text_results = []
-
-                    if return_hierarchical:
-                        all_scores = {curr_labels[j]: probs[i][j].item() for j in range(len(curr_labels))}
-                        all_scores_list.append(all_scores)
-
-                    for j, prob in enumerate(probs[i][: len(curr_labels)]):
-                        score = prob.item()
-                        if score >= item_threshold:
-                            text_results.append({"label": curr_labels[j], "score": score})
-                    results.append(text_results)
-                else:
-                    raise ValueError("Unsupported classification type: choose 'single-label' or 'multi-label'")
+                predictions, all_scores = self._postprocess_logits(
+                    logits[i],
+                    curr_labels,
+                    item_classification_type,
+                    item_threshold,
+                )
+                results.append(predictions)
+                if return_hierarchical:
+                    all_scores_list.append(all_scores)
 
         if return_hierarchical:
             hierarchical_results = []
@@ -768,18 +774,15 @@ class DecoderKVZeroShotClassificationPipeline(BaseZeroShotClassificationPipeline
     """
 
     def prepare_input(self, text: str, labels: list[str], examples=None, prompt: str | None = None) -> str:
-        parts = []
-        if prompt:
-            parts.append(prompt)
-        if examples:
-            parts.append(self._format_examples_for_input(examples))
-        parts.append(text)
-        parts.append(self.sep_token)
-        for label in labels:
-            parts.append(label)
-            parts.append(self.label_token)
-        parts.append(self.sep_token)
-        return "".join(parts)
+        examples_text = self._format_examples_for_input(examples) if examples else ""
+        return format_decoder_kv_sequence(
+            text,
+            labels,
+            prompt=prompt or "",
+            examples=examples_text,
+            label_token=self.label_token,
+            sep_token=self.sep_token,
+        )
 
     def prepare_inputs(self, texts, labels, same_labels=False, examples=None, prompt=None):
         inputs = []
diff --git a/gliclass/streaming/__init__.py b/gliclass/streaming/__init__.py
index e27b343..822b4ed 100644
--- a/gliclass/streaming/__init__.py
+++ b/gliclass/streaming/__init__.py
@@ -1,6 +1,5 @@
 from .cache import CacheState, BatchedKVHelper, truncate_cache, create_empty_cache
-from .types import SessionInput, SessionOutput
-from .pipeline import StreamingPipeline
+from .pipeline import StreamingZeroShotClassificationPipeline
 from .strategies import (
     NeverStrategy,
     ComposedStrategy,
@@ -10,3 +9,18 @@
     SlidingWindowStrategy,
     ClassificationStrategy,
 )
+
+__all__ = [
+    "BatchedKVHelper",
+    "CacheState",
+    "ClassificationStrategy",
+    "ComposedStrategy",
+    "EveryChunkStrategy",
+    "EveryNTokensStrategy",
+    "NeverStrategy",
+    "OnDelimiterStrategy",
+    "SlidingWindowStrategy",
+    "StreamingZeroShotClassificationPipeline",
+    "create_empty_cache",
+    "truncate_cache",
+]
diff --git a/gliclass/streaming/cache.py b/gliclass/streaming/cache.py
index e81f3ed..99db8f6 100644
--- a/gliclass/streaming/cache.py
+++ b/gliclass/streaming/cache.py
@@ -25,9 +25,15 @@ class CacheState:
     attention_mask: torch.Tensor
     position_ids: torch.Tensor | None = None
     cached_length: int = 0
+    next_position_id: int | None = None
     session_id: str | None = None
     metadata: dict = field(default_factory=dict)
 
+    @property
+    def next_position(self) -> int:
+        """Return the absolute position for the next appended token."""
+        return self.cached_length if self.next_position_id is None else self.next_position_id
+
     def to(self, device: torch.device) -> CacheState:
         return CacheState(
             past_key_values=_move_past_kv(self.past_key_values, device),
@@ -35,6 +41,7 @@ def to(self, device: torch.device) -> CacheState:
             attention_mask=self.attention_mask.to(device),
             position_ids=self.position_ids.to(device) if self.position_ids is not None else None,
             cached_length=self.cached_length,
+            next_position_id=self.next_position_id,
             session_id=self.session_id,
             metadata=self.metadata.copy(),
         )
@@ -59,6 +66,7 @@ def create_empty_cache(session_id: str | None = None, device=None) -> CacheState
         attention_mask=torch.empty(0, dtype=torch.long, **kw),
         position_ids=None,
         cached_length=0,
+        next_position_id=0,
         session_id=session_id,
         metadata={},
     )
@@ -76,6 +84,7 @@ def truncate_cache(cache_state: CacheState, max_length: int) -> CacheState:
         attention_mask=cache_state.attention_mask[-max_length:],
         position_ids=cache_state.position_ids[-max_length:] if cache_state.position_ids is not None else None,
         cached_length=max_length,
+        next_position_id=cache_state.next_position,
         session_id=cache_state.session_id,
         metadata=cache_state.metadata.copy(),
     )
@@ -156,12 +165,13 @@ def stack_for_update(
         padded_new_mask = new_input_ids[0].new_zeros(batch_size, max_new)
         position_ids = new_input_ids[0].new_zeros(batch_size, max_new)
 
-        for i, (ids, mask, clen, nlen) in enumerate(
+        for i, (ids, mask, _clen, nlen) in enumerate(
             zip(new_input_ids, new_attention_masks, cached_lengths, new_lengths)
         ):
             padded_ids[i, :nlen] = ids[0] if ids.dim() == 2 else ids
             padded_new_mask[i, :nlen] = mask[0] if mask.dim() == 2 else mask
-            position_ids[i, :nlen] = torch.arange(clen, clen + nlen, device=device)
+            next_position = caches[i].next_position
+            position_ids[i, :nlen] = torch.arange(next_position, next_position + nlen, device=device)
 
         # --- build full attention mask: [prepend_zeros | cached | new | new_pad] ---
         full_mask = torch.zeros(batch_size, max_cached + max_new, dtype=torch.long, device=device)
@@ -233,6 +243,7 @@ def unstack_after_update(
                     attention_mask=full_mask,
                     position_ids=None,
                     cached_length=new_clen,
+                    next_position_id=old.next_position + nlen,
                     session_id=old.session_id,
                     metadata=old.metadata.copy(),
                 )
@@ -264,12 +275,13 @@ def stack_for_classify(
         padded_label_mask = label_ids[0].new_zeros(batch_size, max_label)
         position_ids = label_ids[0].new_zeros(batch_size, max_label)
 
-        for i, (ids, mask, clen, llen) in enumerate(zip(label_ids, label_masks, cached_lengths, label_lengths)):
+        for i, (ids, mask, _clen, llen) in enumerate(zip(label_ids, label_masks, cached_lengths, label_lengths)):
             flat_ids = ids[0] if ids.dim() == 2 else ids
             flat_mask = mask[0] if mask.dim() == 2 else mask
             padded_label_ids[i, :llen] = flat_ids
             padded_label_mask[i, :llen] = flat_mask
-            position_ids[i, :llen] = torch.arange(clen, clen + llen, device=device)
+            next_position = caches[i].next_position
+            position_ids[i, :llen] = torch.arange(next_position, next_position + llen, device=device)
 
         full_mask = torch.zeros(batch_size, max_cached + max_label, dtype=torch.long, device=device)
         for i, (clen, llen) in enumerate(zip(cached_lengths, label_lengths)):
@@ -346,8 +358,8 @@ def _slice_past_kv(past_kv, batch_idx: int, start: int, end: int):
             if not layer.is_initialized:
                 continue
             sliced.update(
-                layer.keys[batch_idx : batch_idx + 1, :, start:end, :],
-                layer.values[batch_idx : batch_idx + 1, :, start:end, :],
+                layer.keys[batch_idx : batch_idx + 1, :, start:end, :].clone(),
+                layer.values[batch_idx : batch_idx + 1, :, start:end, :].clone(),
                 i,
             )
         return sliced
diff --git a/gliclass/streaming/pipeline.py b/gliclass/streaming/pipeline.py
index 087f268..c18310a 100644
--- a/gliclass/streaming/pipeline.py
+++ b/gliclass/streaming/pipeline.py
@@ -1,190 +1,307 @@
-"""
-StreamingPipeline - batched multi-session streaming classification.
-
-Flow per __call__:
-  1. Cleanup sessions absent from current batch.
-  2. Create CacheState for new session_ids.
-  3. Stage 1 - batch update KV caches with new text (skip empty text).
-  4. Stage 2 - resolve strategies; batch classify sessions that triggered.
-"""
+"""Stateful, batched streaming classification for decoder-KV models."""
 
 from __future__ import annotations
 
+import copy
+from typing import Any
+
 import torch
 
 from .cache import CacheState, BatchedKVHelper, truncate_cache, create_empty_cache
-from .types import SessionInput, SessionOutput
+from ..pipeline import DecoderKVZeroShotClassificationPipeline, build_hierarchical_output
+from .strategies import EveryChunkStrategy, ClassificationStrategy
+from ..data_processing import format_decoder_kv_labels, format_decoder_kv_context
+
 
+class StreamingZeroShotClassificationPipeline(DecoderKVZeroShotClassificationPipeline):
+    """Incrementally classify multiple decoder-KV sessions with persistent KV caches."""
 
-class StreamingPipeline:
     def __init__(
         self,
         model,
         tokenizer,
-        device: str | torch.device = "cpu",
+        max_classes: int = 25,
+        max_length: int = 1024,
+        classification_type: str = "multi-label",
+        device: str | torch.device = "cuda:0",
+        progress_bar: bool = False,
+        label_separator: str = ".",
         max_cache_len: int | None = None,
-        batch_size: int | None = None,
+        default_batch_size: int = 8,
+        default_strategy: ClassificationStrategy | None = None,
         offload_to_cpu: bool = False,
         use_pinned_memory: bool = False,
         score_on_cpu: bool = False,
     ):
-        self.model = model
-        self.tokenizer = tokenizer
-        self.device = torch.device(device) if isinstance(device, str) else device
+        if model.config.architecture_type != "decoder-kv":
+            raise ValueError("Streaming classification requires architecture_type='decoder-kv'.")
+
+        super().__init__(
+            model,
+            tokenizer,
+            max_classes=max_classes,
+            max_length=max_length,
+            classification_type=classification_type,
+            device=device,
+            progress_bar=progress_bar,
+            label_separator=label_separator,
+        )
+        if default_batch_size <= 0:
+            raise ValueError("default_batch_size must be greater than zero.")
+        if max_cache_len is not None and max_cache_len <= 0:
+            raise ValueError("max_cache_len must be greater than zero when provided.")
+        if default_strategy is not None and not isinstance(default_strategy, ClassificationStrategy):
+            raise TypeError("default_strategy must implement ClassificationStrategy.")
+
         self.max_cache_len = max_cache_len
-        self.batch_size = batch_size
+        self.default_batch_size = default_batch_size
+        self.default_strategy = default_strategy or EveryChunkStrategy()
         self.offload_to_cpu = offload_to_cpu
         self.use_pinned_memory = use_pinned_memory and offload_to_cpu and self.device.type == "cuda"
         self.score_on_cpu = score_on_cpu
         self._caches: dict[str, CacheState] = {}
+        self._session_strategies: dict[str, ClassificationStrategy] = {}
 
-        scorer_target = "cpu" if score_on_cpu else self.device
-        self.model.model.scorer.to(scorer_target)
+        if score_on_cpu:
+            self.model.model.scorer.to("cpu")
 
-    # ------------------------------------------------------------------
-    # Public API
-    # ------------------------------------------------------------------
+    @torch.no_grad()
+    def __call__(
+        self,
+        texts: str | list[str],
+        labels: list[str] | dict[str, Any] | list[list[str]] | list[dict[str, Any]],
+        *,
+        session_ids: str | list[str],
+        threshold: float | list[float] = 0.5,
+        batch_size: int | None = None,
+        classification_type: str | list[str] | None = None,
+        strategies: ClassificationStrategy | list[ClassificationStrategy] | None = None,
+        examples: list[dict[str, Any]] | list[list[dict[str, Any]]] | None = None,
+        prompt: str | list[str] | None = None,
+        return_hierarchical: bool = False,
+    ) -> list[dict[str, Any]]:
+        """Append text chunks and classify sessions selected by their strategies."""
+        normalized = self._normalize_streaming_inputs(
+            texts=texts,
+            labels=labels,
+            session_ids=session_ids,
+            threshold=threshold,
+            classification_type=classification_type,
+            strategies=strategies,
+            examples=examples,
+            prompt=prompt,
+            return_hierarchical=return_hierarchical,
+        )
 
-    def __call__(self, inputs: list[SessionInput]) -> list[SessionOutput]:
-        self._cleanup_dead_sessions(inputs)
-        self._ensure_sessions(inputs)
+        batch_size = self.default_batch_size if batch_size is None else batch_size
+        if batch_size <= 0:
+            raise ValueError("batch_size must be greater than zero.")
 
-        if self.batch_size is None or len(inputs) <= self.batch_size:
-            return self._process_batch(inputs)
+        active_ids = set(normalized["session_ids"])
+        self._cleanup_absent_sessions(active_ids)
+        if not active_ids:
+            return []
 
-        outputs = [None] * len(inputs)
-        for start in range(0, len(inputs), self.batch_size):
-            chunk = inputs[start : start + self.batch_size]
-            chunk_ids = {inp.session_id for inp in chunk}
+        outputs: list[dict[str, Any] | None] = [None] * len(normalized["texts"])
+        for start in range(0, len(normalized["texts"]), batch_size):
+            end = start + batch_size
+            batch = {key: value[start:end] for key, value in normalized.items()}
+            batch_ids = set(batch["session_ids"])
+            self._ensure_sessions(batch)
 
             if self.offload_to_cpu:
-                self._load_to_device(chunk_ids, self.device)
+                self._load_to_device(batch_ids, self.device)
 
-            for i, out in enumerate(self._process_batch(chunk)):
-                outputs[start + i] = out
+            batch_outputs = self._process_batch(batch)
+            outputs[start:end] = batch_outputs
 
             if self.offload_to_cpu:
-                self._load_to_device(chunk_ids, torch.device("cpu"))
+                self._load_to_device(batch_ids, torch.device("cpu"))
 
         return outputs
 
-    def _process_batch(self, inputs: list[SessionInput]) -> list[SessionOutput]:
-        tokens_added = self._stage_update(inputs)
-        if self.max_cache_len is not None:
-            self._enforce_max_cache_len()
-        return self._stage_classify(inputs, tokens_added)
-
-    def _load_to_device(self, session_ids: set[str], device: torch.device) -> None:
-        to_cpu = device.type == "cpu"
-        for sid in session_ids:
-            cache = self._caches.get(sid)
-            if cache is None or cache.past_key_values is None:
-                continue
-            if to_cpu and self.use_pinned_memory:
-                self._caches[sid] = _move_cache_pinned(cache)
-            else:
-                self._caches[sid] = _move_cache_nonblocking(cache, device)
-
-        # synchronize before using on GPU so non_blocking transfers complete
-        if not to_cpu and self.use_pinned_memory and self.device.type == "cuda":
-            torch.cuda.synchronize(self.device)
-
-    # ------------------------------------------------------------------
-    # Session management
-    # ------------------------------------------------------------------
+    @property
+    def active_sessions(self) -> list[str]:
+        """Return active session IDs in insertion order."""
+        return list(self._caches)
 
     def delete_session(self, *session_ids: str) -> list[str]:
-        """Delete one or more sessions and free their KV caches.
-
-        Returns the list of session IDs that were actually deleted
-        (unknown IDs are silently ignored).
-        """
+        """Delete sessions and all cache and strategy state associated with them."""
         deleted = []
-        for sid in session_ids:
-            if sid in self._caches:
-                del self._caches[sid]
-                deleted.append(sid)
+        for session_id in session_ids:
+            if session_id in self._caches:
+                del self._caches[session_id]
+                self._session_strategies.pop(session_id, None)
+                deleted.append(session_id)
         return deleted
 
     def clear_sessions(self) -> None:
-        """Delete all active sessions and free all KV caches."""
+        """Delete every active session."""
         self._caches.clear()
+        self._session_strategies.clear()
 
-    @property
-    def active_sessions(self) -> list[str]:
-        """Return the list of currently active session IDs."""
-        return list(self._caches.keys())
-
-    def _cleanup_dead_sessions(self, inputs: list[SessionInput]) -> None:
-        active = {inp.session_id for inp in inputs}
-        for sid in list(self._caches):
-            if sid not in active:
-                del self._caches[sid]
-
-    def _enforce_max_cache_len(self) -> None:
-        for sid, cache in self._caches.items():
-            if cache.cached_length > self.max_cache_len:
-                self._caches[sid] = truncate_cache(cache, self.max_cache_len)
-
-    def _ensure_sessions(self, inputs: list[SessionInput]) -> None:
-        for inp in inputs:
-            if inp.session_id not in self._caches:
-                self._caches[inp.session_id] = create_empty_cache(session_id=inp.session_id, device=self.device)
-
-    # ------------------------------------------------------------------
-    # Stage 1: update KV caches
-    # ------------------------------------------------------------------
-
-    def _stage_update(self, inputs: list[SessionInput]) -> list[int]:
-        """
-        Update KV cache for each session. Returns tokens_added per input.
-        Sessions with empty text are skipped (tokens_added = 0).
-        """
-        tokens_added = [0] * len(inputs)
-        to_update = [(i, inp) for i, inp in enumerate(inputs) if inp.text]
-
-        if not to_update:
-            return tokens_added
+    def set_session_strategy(self, session_id: str, strategy: ClassificationStrategy) -> None:
+        """Replace the classification strategy for an active session."""
+        if session_id not in self._caches:
+            raise KeyError(f"Unknown session_id: {session_id}")
+        self._session_strategies[session_id] = copy.deepcopy(strategy)
 
-        indices = [i for i, _ in to_update]
-        active_inputs = [inp for _, inp in to_update]
-
-        # Tokenize new text per session
-        tokenized = [self.tokenizer(inp.text, return_tensors="pt", add_special_tokens=False) for inp in active_inputs]
-        new_ids = [t["input_ids"].to(self.device) for t in tokenized]
-        new_masks = [t["attention_mask"].to(self.device) for t in tokenized]
-
-        caches = [self._caches[inp.session_id] for inp in active_inputs]
+    def _normalize_streaming_inputs(
+        self,
+        *,
+        texts,
+        labels,
+        session_ids,
+        threshold,
+        classification_type,
+        strategies,
+        examples,
+        prompt,
+        return_hierarchical,
+    ) -> dict[str, list]:
+        texts = self._normalize_texts(texts)
+        if isinstance(session_ids, str):
+            session_ids = [session_ids]
+        else:
+            session_ids = list(session_ids)
+
+        if len(session_ids) != len(texts):
+            raise ValueError("session_ids and texts must have the same length.")
+        if len(set(session_ids)) != len(session_ids):
+            raise ValueError("session_ids must be unique within a streaming call.")
+        if any(not isinstance(session_id, str) or not session_id for session_id in session_ids):
+            raise ValueError("Every session_id must be a non-empty string.")
+
+        thresholds = self._normalize_thresholds(threshold, len(texts))
+        classification_types = self._normalize_classification_types(classification_type, len(texts))
+        processed_labels = self._process_labels(labels)
+        if texts and not processed_labels:
+            raise ValueError("At least one label is required.")
+
+        same_labels = not texts or isinstance(processed_labels[0], str)
+        if same_labels:
+            labels_per_text = [processed_labels] * len(texts)
+            original_labels = [labels] * len(texts)
+        else:
+            if len(processed_labels) != len(texts):
+                raise ValueError("Per-text labels and texts must have the same length.")
+            labels_per_text = processed_labels
+            original_labels = list(labels)
+
+        if strategies is None:
+            strategies_per_text = [None] * len(texts)
+        elif isinstance(strategies, ClassificationStrategy):
+            strategies_per_text = [strategies] * len(texts)
+        else:
+            strategies_per_text = list(strategies)
+            if len(strategies_per_text) != len(texts):
+                raise ValueError("strategies and texts must have the same length.")
+            if not all(isinstance(strategy, ClassificationStrategy) for strategy in strategies_per_text):
+                raise TypeError("Every strategy must implement ClassificationStrategy.")
+
+        prefixes = []
+        for index, session_id in enumerate(session_ids):
+            text_prompt = self._format_prompt(prompt, index)
+            text_examples = self._get_text_examples(examples, index)
+            examples_text = self._format_examples_for_input(text_examples) if text_examples else ""
+            prefix = format_decoder_kv_context("", text_prompt, examples_text)
+            existing = self._caches.get(session_id)
+            if (
+                existing is not None
+                and existing.cached_length > 0
+                and prefix
+                and prefix != existing.metadata.get("prefix", "")
+            ):
+                raise ValueError(
+                    f"Prompt or examples changed for active session {session_id!r}; "
+                    "delete the session before changing them."
+                )
+            prefixes.append(prefix)
+
+        return {
+            "texts": texts,
+            "labels": labels_per_text,
+            "original_labels": original_labels,
+            "session_ids": session_ids,
+            "thresholds": thresholds,
+            "classification_types": classification_types,
+            "strategies": strategies_per_text,
+            "prefixes": prefixes,
+            "return_hierarchical": [return_hierarchical] * len(texts),
+        }
+
+    def _cleanup_absent_sessions(self, active_ids: set[str]) -> None:
+        for session_id in list(self._caches):
+            if session_id not in active_ids:
+                del self._caches[session_id]
+                self._session_strategies.pop(session_id, None)
+
+    def _ensure_sessions(self, batch: dict[str, list]) -> None:
+        for session_id, strategy, prefix in zip(
+            batch["session_ids"], batch["strategies"], batch["prefixes"], strict=True
+        ):
+            if session_id in self._caches:
+                if self._caches[session_id].cached_length == 0 and prefix:
+                    self._caches[session_id].metadata["prefix"] = prefix
+                continue
+            cache = create_empty_cache(session_id=session_id, device=self.device)
+            cache.metadata["prefix"] = prefix
+            self._caches[session_id] = cache
+            self._session_strategies[session_id] = copy.deepcopy(strategy or self.default_strategy)
+
+    def _process_batch(self, batch: dict[str, list]) -> list[dict[str, Any]]:
+        texts = []
+        for session_id, text, prefix in zip(batch["session_ids"], batch["texts"], batch["prefixes"], strict=True):
+            cache = self._caches[session_id]
+            texts.append(f"{prefix}{text}" if cache.cached_length == 0 else text)
+
+        tokens_added = self._update_caches(batch["session_ids"], texts)
+        if self.max_cache_len is not None:
+            for session_id in batch["session_ids"]:
+                cache = self._caches[session_id]
+                if cache.cached_length > self.max_cache_len:
+                    self._caches[session_id] = truncate_cache(cache, self.max_cache_len)
+        return self._classify(batch, tokens_added)
+
+    def _update_caches(self, session_ids: list[str], texts: list[str]) -> list[int]:
+        tokens_added = [0] * len(texts)
+        active = [
+            (index, session_id, text)
+            for index, (session_id, text) in enumerate(zip(session_ids, texts, strict=True))
+            if text
+        ]
+        if not active:
+            return tokens_added
 
-        # Check if any session has a non-empty cache (need KV stacking)
-        has_cache = any(c.cached_length > 0 for c in caches)
+        indices = [item[0] for item in active]
+        active_ids = [item[1] for item in active]
+        tokenized = [self.tokenizer(item[2], return_tensors="pt", add_special_tokens=False) for item in active]
+        new_ids = [item["input_ids"].to(self.device) for item in tokenized]
+        new_masks = [item["attention_mask"].to(self.device) for item in tokenized]
+        caches = [self._caches[session_id] for session_id in active_ids]
 
-        if has_cache:
+        if any(cache.cached_length > 0 for cache in caches):
             stacked = BatchedKVHelper.stack_for_update(caches, new_ids, new_masks, self.device)
-            decoder = self.model.model.decoder_model
-
-            with torch.no_grad():
-                out = decoder(
-                    input_ids=stacked["input_ids"],
-                    attention_mask=stacked["attention_mask"],
-                    position_ids=stacked["position_ids"],
-                    past_key_values=stacked["past_key_values"],
-                    use_cache=True,
-                    return_dict=True,
-                )
-
-            updated = BatchedKVHelper.unstack_after_update(out.past_key_values, stacked, caches, new_ids, new_masks)
+            model_output = self.model.update_decoder_cache(
+                input_ids=stacked["input_ids"],
+                attention_mask=stacked["attention_mask"],
+                position_ids=stacked["position_ids"],
+                past_key_values=stacked["past_key_values"],
+            )
+            updated = BatchedKVHelper.unstack_after_update(
+                model_output.past_key_values,
+                stacked,
+                caches,
+                new_ids,
+                new_masks,
+            )
+            new_lengths = stacked["new_lengths"]
         else:
-            # All caches empty - simple batched forward without stacking
-            updated = self._update_from_scratch(caches, new_ids, new_masks)
-
-        for session_idx, cache, new_len in zip(
-            indices, updated, stacked["new_lengths"] if has_cache else [ids.shape[-1] for ids in new_ids]
-        ):
-            inp = inputs[session_idx]
-            self._caches[inp.session_id] = cache
-            tokens_added[session_idx] = new_len
+            updated, new_lengths = self._update_from_scratch(caches, new_ids, new_masks)
 
+        for index, session_id, cache, new_length in zip(indices, active_ids, updated, new_lengths, strict=True):
+            self._caches[session_id] = cache
+            tokens_added[index] = new_length
         return tokens_added
 
     def _update_from_scratch(
@@ -192,222 +309,168 @@ def _update_from_scratch(
         caches: list[CacheState],
         new_ids: list[torch.Tensor],
         new_masks: list[torch.Tensor],
-    ) -> list[CacheState]:
-        """Batched update when all caches are empty (no KV stacking needed)."""
-        max_len = max(ids.shape[-1] for ids in new_ids)
+    ) -> tuple[list[CacheState], list[int]]:
+        new_lengths = [ids.shape[-1] for ids in new_ids]
+        max_length = max(new_lengths)
         batch_size = len(caches)
-
-        padded_ids = new_ids[0].new_zeros(batch_size, max_len)
-        padded_masks = new_ids[0].new_zeros(batch_size, max_len)
-        position_ids = new_ids[0].new_zeros(batch_size, max_len)
-
-        new_lengths = []
-        for i, (ids, mask) in enumerate(zip(new_ids, new_masks)):
-            flat = ids[0] if ids.dim() == 2 else ids
-            flat_m = mask[0] if mask.dim() == 2 else mask
-            nlen = flat.shape[0]
-            padded_ids[i, :nlen] = flat
-            padded_masks[i, :nlen] = flat_m
-            position_ids[i, :nlen] = torch.arange(nlen, device=self.device)
-            new_lengths.append(nlen)
-
-        decoder = self.model.model.decoder_model
-        with torch.no_grad():
-            out = decoder(
-                input_ids=padded_ids,
-                attention_mask=padded_masks,
-                position_ids=position_ids,
-                use_cache=True,
-                return_dict=True,
+        padded_ids = new_ids[0].new_zeros(batch_size, max_length)
+        padded_masks = new_ids[0].new_zeros(batch_size, max_length)
+        position_ids = new_ids[0].new_zeros(batch_size, max_length)
+
+        for index, (ids, mask, cache, length) in enumerate(zip(new_ids, new_masks, caches, new_lengths, strict=True)):
+            flat_ids = ids[0] if ids.dim() == 2 else ids
+            flat_mask = mask[0] if mask.dim() == 2 else mask
+            padded_ids[index, :length] = flat_ids
+            padded_masks[index, :length] = flat_mask
+            position_ids[index, :length] = torch.arange(
+                cache.next_position,
+                cache.next_position + length,
+                device=self.device,
             )
 
+        model_output = self.model.update_decoder_cache(
+            input_ids=padded_ids,
+            attention_mask=padded_masks,
+            position_ids=position_ids,
+            past_key_values=None,
+        )
+
         results = []
-        for i, (cache, nlen) in enumerate(zip(caches, new_lengths)):
-            sliced_kv = BatchedKVHelper._slice_past_kv(out.past_key_values, i, 0, nlen)
-            flat_ids = new_ids[i][0] if new_ids[i].dim() == 2 else new_ids[i]
-            flat_mask = new_masks[i][0] if new_masks[i].dim() == 2 else new_masks[i]
+        for index, (cache, length) in enumerate(zip(caches, new_lengths, strict=True)):
+            flat_ids = new_ids[index][0] if new_ids[index].dim() == 2 else new_ids[index]
+            flat_mask = new_masks[index][0] if new_masks[index].dim() == 2 else new_masks[index]
             results.append(
                 CacheState(
-                    past_key_values=sliced_kv,
-                    input_ids=flat_ids[:nlen],
-                    attention_mask=flat_mask[:nlen],
-                    cached_length=nlen,
+                    past_key_values=BatchedKVHelper._slice_past_kv(
+                        model_output.past_key_values,
+                        index,
+                        0,
+                        length,
+                    ),
+                    input_ids=flat_ids[:length],
+                    attention_mask=flat_mask[:length],
+                    cached_length=length,
+                    next_position_id=cache.next_position + length,
                     session_id=cache.session_id,
                     metadata=cache.metadata.copy(),
                 )
             )
+        return results, new_lengths
 
-        return results
-
-    # ------------------------------------------------------------------
-    # Stage 2: classify triggered sessions
-    # ------------------------------------------------------------------
-
-    def _stage_classify(self, inputs: list[SessionInput], tokens_added: list[int]) -> list[SessionOutput]:
-        # Resolve which sessions trigger classification
+    def _classify(
+        self,
+        batch: dict[str, list],
+        tokens_added: list[int],
+    ) -> list[dict[str, Any]]:
         triggered_indices = []
-        for i, (inp, n_added) in enumerate(zip(inputs, tokens_added)):
-            if n_added == 0:
-                continue  # no new tokens → cache unchanged, skip classification
-            cache = self._caches[inp.session_id]
-            if inp.strategy.should_classify(n_added, cache.cached_length, inp.text):
-                triggered_indices.append(i)
-
-        outputs: list[SessionOutput | None] = [None] * len(inputs)
-
-        # Fill non-triggered outputs
-        for i, inp in enumerate(inputs):
-            if i not in triggered_indices:
-                cache = self._caches[inp.session_id]
-                outputs[i] = SessionOutput(
-                    session_id=inp.session_id,
-                    triggered=False,
-                    predictions=None,
-                    cached_length=cache.cached_length,
-                    tokens_added=tokens_added[i],
-                )
-
+        for index, (session_id, text, added) in enumerate(
+            zip(batch["session_ids"], batch["texts"], tokens_added, strict=True)
+        ):
+            cache = self._caches[session_id]
+            strategy = self._session_strategies[session_id]
+            if text and added > 0 and strategy.should_classify(added, cache.cached_length, text):
+                triggered_indices.append(index)
+
+        outputs = [
+            {
+                "session_id": session_id,
+                "triggered": False,
+                "predictions": None,
+                "cached_length": self._caches[session_id].cached_length,
+                "tokens_added": tokens_added[index],
+            }
+            for index, session_id in enumerate(batch["session_ids"])
+        ]
         if not triggered_indices:
             return outputs
 
-        # Prepare label sequences for triggered sessions
-        triggered_inputs = [inputs[i] for i in triggered_indices]
-        triggered_caches = [
-            inputs[i].strategy.get_window(self._caches[inputs[i].session_id]) for i in triggered_indices
-        ]
-
-        label_seqs = [self._prepare_label_seq(inp.labels) for inp in triggered_inputs]
-        tokenized_labels = [self.tokenizer(seq, return_tensors="pt", add_special_tokens=False) for seq in label_seqs]
-        label_ids = [t["input_ids"].to(self.device) for t in tokenized_labels]
-        label_masks = [t["attention_mask"].to(self.device) for t in tokenized_labels]
-
-        stacked = BatchedKVHelper.stack_for_classify(triggered_caches, label_ids, label_masks, self.device)
-
-        decoder = self.model.model.decoder_model
-        with torch.no_grad():
-            out = decoder(
-                input_ids=stacked["input_ids"],
-                attention_mask=stacked["attention_mask"],
-                position_ids=stacked["position_ids"],
-                past_key_values=stacked["past_key_values"],
-                use_cache=False,
-                return_dict=True,
+        triggered_caches = []
+        label_ids = []
+        label_masks = []
+        for index in triggered_indices:
+            session_id = batch["session_ids"][index]
+            cache = self._session_strategies[session_id].get_window(self._caches[session_id])
+            triggered_caches.append(cache)
+            sequence = format_decoder_kv_labels(
+                batch["labels"][index],
+                label_token=self.label_token,
+                sep_token=self.sep_token,
             )
+            tokenized = self.tokenizer(sequence, return_tensors="pt", add_special_tokens=False)
+            label_ids.append(tokenized["input_ids"].to(self.device))
+            label_masks.append(tokenized["attention_mask"].to(self.device))
 
-        # Slice label hidden states and pass to scorer
-        hidden = out.last_hidden_state  # [batch, max_cached + max_label, hidden]
-        max_cached = max(c.cached_length for c in triggered_caches) if stacked["past_key_values"] else 0  # noqa: F841
-        label_lengths = stacked["label_lengths"]
-        max_label = stacked["max_label_len"]
-
-        # hidden for labels starts at index max_cached (prepend-padded KV)
-        # but we need the actual label slice per session
-        scorer = self.model.model.scorer
-        batch_logits = self._run_scorer_batched(
-            scorer,
-            hidden,
-            stacked["input_ids"],
-            stacked["label_mask"],
-            label_lengths,
-            max_label,
+        stacked = BatchedKVHelper.stack_for_classify(triggered_caches, label_ids, label_masks, self.device)
+        model_output = self.model.classify_from_decoder_cache(
+            input_ids=stacked["input_ids"],
+            attention_mask=stacked["attention_mask"],
+            label_mask=stacked["label_mask"],
+            position_ids=stacked["position_ids"],
+            past_key_values=stacked["past_key_values"],
         )
 
-        for local_idx, global_idx in enumerate(triggered_indices):
-            inp = inputs[global_idx]
-            cache = self._caches[inp.session_id]
-            logits = batch_logits[local_idx]  # [num_labels]
-            preds = self._decode_predictions(logits, inp.labels, inp.classification_type)
-            outputs[global_idx] = SessionOutput(
-                session_id=inp.session_id,
-                triggered=True,
-                predictions=preds,
-                cached_length=cache.cached_length,
-                tokens_added=tokens_added[global_idx],
+        for local_index, batch_index in enumerate(triggered_indices):
+            labels = batch["labels"][batch_index]
+            logits = model_output.logits[local_index, : len(labels)]
+            predictions, all_scores = self._postprocess_logits(
+                logits,
+                labels,
+                batch["classification_types"][batch_index],
+                batch["thresholds"][batch_index],
             )
+            if batch["return_hierarchical"][batch_index]:
+                predictions = build_hierarchical_output(
+                    predictions,
+                    batch["original_labels"][batch_index],
+                    self.label_separator,
+                    all_scores,
+                )
 
+            session_id = batch["session_ids"][batch_index]
+            outputs[batch_index] = {
+                "session_id": session_id,
+                "triggered": True,
+                "predictions": predictions,
+                "cached_length": self._caches[session_id].cached_length,
+                "tokens_added": tokens_added[batch_index],
+            }
         return outputs
 
-    # ------------------------------------------------------------------
-    # Helpers
-    # ------------------------------------------------------------------
-
-    def _prepare_label_seq(self, labels: list[str]) -> str:
-        label_token = "<<LABEL>>"
-        sep_token = "<<SEP>>"
-        # Opening <<SEP>> matches training format: text<<SEP>>label1<<LABEL>>...<<SEP>>
-        # Cache contains only text; <<SEP>> boundary must be part of label sequence
-        parts = [sep_token]
-        for label in labels:
-            parts.append(label)
-            parts.append(label_token)
-        parts.append(sep_token)
-        return "".join(parts)
-
-    def _run_scorer_batched(
-        self,
-        scorer,
-        hidden: torch.Tensor,
-        label_ids: torch.Tensor,
-        label_mask: torch.Tensor,
-        label_lengths: list[int],
-        max_label: int,
-    ) -> list[torch.Tensor]:
-        """
-        hidden: [batch, max_cached + max_label, hidden_size]  (full decoder output)
-        We slice out only the label part [max_cached :] per session before scorer.
-        label_ids / label_mask are already label-only ([batch, max_label]).
-        """
-        max_cached_in_hidden = hidden.shape[1] - max_label
-        label_hidden = hidden[:, max_cached_in_hidden:, :]  # [batch, max_label, hidden]
-
-        if self.score_on_cpu:
-            label_hidden = label_hidden.cpu()
-            label_ids = label_ids.cpu()
-            label_mask = label_mask.cpu()
-
-        logits = scorer(
-            hidden_states=label_hidden,
-            input_ids=label_ids,
-            attention_mask=label_mask,
-        )  # [batch, num_labels]
-
-        return [logits[i] for i in range(logits.shape[0])]
-
-    def _decode_predictions(self, logits: torch.Tensor, labels: list[str], classification_type: str) -> list[dict]:
-        logits = logits[: len(labels)]
-        if classification_type == "single-label":
-            scores = torch.softmax(logits, dim=-1)
-            best = int(scores.argmax().item())
-            return [{"label": labels[best], "score": float(scores[best].detach())}]
-        else:
-            probs = torch.sigmoid(logits)
-            return [{"label": label, "score": float(prob.detach())} for label, prob in zip(labels, probs)]
+    def _load_to_device(self, session_ids: set[str], device: torch.device) -> None:
+        to_cpu = device.type == "cpu"
+        for session_id in session_ids:
+            cache = self._caches.get(session_id)
+            if cache is None or cache.past_key_values is None:
+                continue
+            if to_cpu and self.use_pinned_memory:
+                self._caches[session_id] = _move_cache_pinned(cache)
+            else:
+                self._caches[session_id] = _move_cache_nonblocking(cache, device)
+
+        if not to_cpu and self.use_pinned_memory and self.device.type == "cuda":
+            torch.cuda.synchronize(self.device)
 
 
 def _move_cache_pinned(cache: CacheState) -> CacheState:
-    """Move KV cache to CPU pinned memory for fast async GPU uploads."""
+    """Move a DynamicCache to CPU pinned memory."""
     from transformers.cache_utils import DynamicCache, DynamicLayer
 
-    from .cache import CacheState
-
-    past_kv = cache.past_key_values
-    if not isinstance(past_kv, DynamicCache):
+    if not isinstance(cache.past_key_values, DynamicCache):
         return cache.to(torch.device("cpu"))
 
     new_kv = DynamicCache()
-    for i, layer in enumerate(past_kv.layers):
+    for index, layer in enumerate(cache.past_key_values.layers):
         if not layer.is_initialized:
             continue
-        pinned_k = layer.keys.cpu().pin_memory()
-        pinned_v = layer.values.cpu().pin_memory()
-        # bypass update() to avoid torch.cat losing pinned property
-        while len(new_kv.layers) <= i:
+        pinned_keys = layer.keys.cpu().pin_memory()
+        pinned_values = layer.values.cpu().pin_memory()
+        while len(new_kv.layers) <= index:
             new_kv.layers.append(DynamicLayer())
-        new_layer = new_kv.layers[i]
-        new_layer.dtype = pinned_k.dtype
-        new_layer.device = pinned_k.device
-        new_layer.keys = pinned_k
-        new_layer.values = pinned_v
+        new_layer = new_kv.layers[index]
+        new_layer.dtype = pinned_keys.dtype
+        new_layer.device = pinned_keys.device
+        new_layer.keys = pinned_keys
+        new_layer.values = pinned_values
         new_layer.is_initialized = True
 
     return CacheState(
@@ -416,28 +479,26 @@ def _move_cache_pinned(cache: CacheState) -> CacheState:
         attention_mask=cache.attention_mask.cpu(),
         position_ids=cache.position_ids.cpu() if cache.position_ids is not None else None,
         cached_length=cache.cached_length,
+        next_position_id=cache.next_position_id,
         session_id=cache.session_id,
         metadata=cache.metadata.copy(),
     )
 
 
 def _move_cache_nonblocking(cache: CacheState, device: torch.device) -> CacheState:
-    """Move KV cache to device with non_blocking=True (async when src is pinned)."""
+    """Move a DynamicCache to a device, asynchronously when possible."""
     from transformers.cache_utils import DynamicCache
 
-    from .cache import CacheState
-
-    past_kv = cache.past_key_values
-    if not isinstance(past_kv, DynamicCache):
+    if not isinstance(cache.past_key_values, DynamicCache):
         return cache.to(device)
 
     new_kv = DynamicCache()
-    for i, layer in enumerate(past_kv.layers):
+    for index, layer in enumerate(cache.past_key_values.layers):
         if layer.is_initialized:
             new_kv.update(
                 layer.keys.to(device, non_blocking=True),
                 layer.values.to(device, non_blocking=True),
-                i,
+                index,
             )
 
     return CacheState(
@@ -446,6 +507,7 @@ def _move_cache_nonblocking(cache: CacheState, device: torch.device) -> CacheSta
         attention_mask=cache.attention_mask.to(device, non_blocking=True),
         position_ids=cache.position_ids.to(device, non_blocking=True) if cache.position_ids is not None else None,
         cached_length=cache.cached_length,
+        next_position_id=cache.next_position_id,
         session_id=cache.session_id,
         metadata=cache.metadata.copy(),
     )
diff --git a/gliclass/streaming/types.py b/gliclass/streaming/types.py
deleted file mode 100644
index c73db1a..0000000
--- a/gliclass/streaming/types.py
+++ /dev/null
@@ -1,26 +0,0 @@
-from __future__ import annotations
-
-from typing import TYPE_CHECKING
-from dataclasses import field, dataclass
-
-if TYPE_CHECKING:
-    from .strategies import ClassificationStrategy
-
-
-@dataclass
-class SessionInput:
-    session_id: str
-    text: str
-    labels: list[str]
-    strategy: ClassificationStrategy
-    classification_type: str = "multi-label"  # "multi-label" or "single-label"
-
-
-@dataclass
-class SessionOutput:
-    session_id: str
-    triggered: bool
-    predictions: list[dict] | None
-    cached_length: int
-    tokens_added: int
-    metadata: dict = field(default_factory=dict)
diff --git a/tests/test_decoder_kv.py b/tests/test_decoder_kv.py
new file mode 100644
index 0000000..d907a03
--- /dev/null
+++ b/tests/test_decoder_kv.py
@@ -0,0 +1,83 @@
+"""Tests for the decoder-KV architecture."""
+
+from types import SimpleNamespace
+
+import torch
+
+from gliclass.model import GLiClassOutput, GLiClassDecoderKV
+
+
+def make_decoder_kv_stub():
+    return SimpleNamespace(
+        sep_token_id=98,
+        config=SimpleNamespace(class_token_index=99),
+    )
+
+
+def test_gliclass_output_exposes_past_key_values():
+    cache = object()
+    output = GLiClassOutput(logits=torch.zeros(1, 1), past_key_values=cache)
+    assert output.past_key_values is cache
+
+
+def test_extract_label_section_pads_uneven_sections():
+    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],
+        ]
+    )
+
+    padded_hidden, padded_ids, label_mask = GLiClassDecoderKV._extract_label_section(
+        model,
+        hidden_states,
+        input_ids,
+        attention_mask,
+    )
+
+    expected_hidden = hidden_states.new_zeros(2, 6, 3)
+    expected_hidden[0, :3] = hidden_states[0, 4:7]
+    expected_hidden[1] = hidden_states[1, 3:9]
+    expected_ids = input_ids.new_zeros(2, 6)
+    expected_ids[0, :3] = input_ids[0, 4:7]
+    expected_ids[1] = input_ids[1, 3:9]
+    expected_mask = attention_mask.new_tensor(
+        [
+            [1, 1, 1, 0, 0, 0],
+            [1, 1, 1, 1, 1, 1],
+        ]
+    )
+
+    torch.testing.assert_close(padded_hidden, expected_hidden)
+    torch.testing.assert_close(padded_ids, expected_ids)
+    torch.testing.assert_close(label_mask, expected_mask)
+
+
+def test_extract_label_section_preserves_fallback_and_gradients():
+    model = make_decoder_kv_stub()
+    hidden_states = torch.randn(1, 5, 3, requires_grad=True)
+    input_ids = torch.tensor([[1, 99, 2, 0, 0]])
+    attention_mask = torch.tensor([[1, 1, 1, 0, 0]])
+
+    padded_hidden, padded_ids, label_mask = GLiClassDecoderKV._extract_label_section(
+        model,
+        hidden_states,
+        input_ids,
+        attention_mask,
+    )
+
+    torch.testing.assert_close(padded_hidden, hidden_states[:, :3])
+    torch.testing.assert_close(padded_ids, input_ids[:, :3])
+    torch.testing.assert_close(label_mask, attention_mask[:, :3])
+
+    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:]))
diff --git a/tests/test_scorers.py b/tests/test_scorers.py
index a999722..ca8189d 100644
--- a/tests/test_scorers.py
+++ b/tests/test_scorers.py
@@ -1,5 +1,7 @@
 """Tests for gliclass.scorers module."""
 
+from types import SimpleNamespace
+
 import pytest
 import torch
 
@@ -9,6 +11,7 @@
     MLPScorer,
     HopfieldScorer,
     CrossAttnScorer,
+    DecoderKVScorer,
 )
 
 
@@ -189,3 +192,56 @@ def test_eval_mode(self, scorer):
 
         assert scores.shape == (4, 10)
         assert not torch.isnan(scores).any()
+
+
+class TestDecoderKVScorerRepresentationExtraction:
+    @pytest.fixture
+    def scorer(self):
+        scorer = DecoderKVScorer.__new__(DecoderKVScorer)
+        torch.nn.Module.__init__(scorer)
+        scorer.config = SimpleNamespace(sep_token_index=98, class_token_index=99)
+        return scorer
+
+    def test_extracts_last_sep_and_labels_with_padding(self, scorer):
+        hidden_states = torch.arange(2 * 8 * 3, dtype=torch.float32).reshape(2, 8, 3)
+        input_ids = torch.tensor(
+            [
+                [98, 4, 99, 5, 99, 98, 0, 0],
+                [98, 6, 99, 98, 0, 99, 98, 0],
+            ]
+        )
+        attention_mask = torch.tensor(
+            [
+                [1, 1, 1, 1, 1, 1, 0, 0],
+                [1, 1, 1, 1, 0, 0, 0, 0],
+            ]
+        )
+
+        text_repr, label_repr = scorer._extract_representations(hidden_states, input_ids, attention_mask)
+
+        torch.testing.assert_close(text_repr, torch.stack([hidden_states[0, 5], hidden_states[1, 3]]))
+        torch.testing.assert_close(label_repr[0], torch.stack([hidden_states[0, 2], hidden_states[0, 4]]))
+        torch.testing.assert_close(label_repr[1, 0], hidden_states[1, 2])
+        torch.testing.assert_close(label_repr[1, 1], torch.zeros(3))
+
+    def test_missing_sep_returns_zero_representation(self, scorer):
+        hidden_states = torch.randn(1, 4, 3)
+        input_ids = torch.tensor([[1, 99, 2, 0]])
+        attention_mask = torch.tensor([[1, 1, 1, 0]])
+
+        text_repr, label_repr = scorer._extract_representations(hidden_states, input_ids, attention_mask)
+
+        torch.testing.assert_close(text_repr, torch.zeros_like(text_repr))
+        torch.testing.assert_close(label_repr[0, 0], hidden_states[0, 1])
+
+    def test_extracted_representations_preserve_gradients(self, scorer):
+        hidden_states = torch.randn(1, 5, 3, requires_grad=True)
+        input_ids = torch.tensor([[98, 1, 99, 2, 98]])
+        attention_mask = torch.ones_like(input_ids)
+
+        text_repr, label_repr = scorer._extract_representations(hidden_states, input_ids, attention_mask)
+        (text_repr.sum() + label_repr.sum()).backward()
+
+        assert hidden_states.grad is not None
+        torch.testing.assert_close(hidden_states.grad[0, 2], torch.ones(3))
+        torch.testing.assert_close(hidden_states.grad[0, 4], torch.ones(3))
diff --git a/tests/test_streaming.py b/tests/test_streaming.py
index 4030bc2..8408775 100644
--- a/tests/test_streaming.py
+++ b/tests/test_streaming.py
@@ -1,31 +1,34 @@
-"""Tests for gliclass.streaming — strategies, cache management, BatchedKVHelper."""
+"""Tests for streaming strategies, cache management, and classification."""
+
+from types import SimpleNamespace
 
-import pytest
 import torch
+import pytest
+from torch import nn
 from transformers.cache_utils import DynamicCache
 
 from gliclass.streaming.cache import (
     CacheState,
     BatchedKVHelper,
-    create_empty_cache,
     truncate_cache,
+    create_empty_cache,
 )
+from gliclass.streaming.pipeline import StreamingZeroShotClassificationPipeline
 from gliclass.streaming.strategies import (
-    ClassificationStrategy,
+    NeverStrategy,
+    ComposedStrategy,
     EveryChunkStrategy,
-    EveryNTokensStrategy,
     OnDelimiterStrategy,
-    NeverStrategy,
+    EveryNTokensStrategy,
     SlidingWindowStrategy,
-    ComposedStrategy,
+    ClassificationStrategy,
 )
-from gliclass.streaming.types import SessionInput, SessionOutput
-
 
 # ---------------------------------------------------------------------------
 # Helpers
 # ---------------------------------------------------------------------------
 
+
 def _make_dynamic_cache(num_layers: int, num_heads: int, seq_len: int, head_dim: int, batch: int = 1) -> DynamicCache:
     cache = DynamicCache()
     for i in range(num_layers):
@@ -42,6 +45,7 @@ def _make_cache_state(seq_len: int, num_layers: int = 2, num_heads: int = 2, hea
         input_ids=torch.ones(seq_len, dtype=torch.long),
         attention_mask=torch.ones(seq_len, dtype=torch.long),
         cached_length=seq_len,
+        next_position_id=seq_len,
     )
 
 
@@ -49,6 +53,7 @@ def _make_cache_state(seq_len: int, num_layers: int = 2, num_heads: int = 2, hea
 # Strategy tests
 # ---------------------------------------------------------------------------
 
+
 class TestEveryChunkStrategy:
     def test_triggers_on_any_positive_tokens(self):
         s = EveryChunkStrategy()
@@ -74,7 +79,7 @@ def test_triggers_after_n_tokens(self):
 
     def test_counter_resets_after_trigger(self):
         s = EveryNTokensStrategy(n=10)
-        s.should_classify(10, 10, "")   # trigger, reset
+        s.should_classify(10, 10, "")  # trigger, reset
         assert s._accumulated == 0
         assert s.should_classify(5, 15, "") is False
 
@@ -204,6 +209,7 @@ def should_classify(self, tokens_added, cached_length, text):
 # CacheState tests
 # ---------------------------------------------------------------------------
 
+
 class TestCacheState:
     def test_create_empty_cache(self):
         cache = create_empty_cache(session_id="test")
@@ -230,6 +236,11 @@ def test_metadata_preserved_on_move(self):
         moved = cache.to(torch.device("cpu"))
         assert moved.metadata == {"key": "value"}
 
+    def test_next_position_preserved_on_move(self):
+        cache = _make_cache_state(5)
+        cache.next_position_id = 12
+        assert cache.to(torch.device("cpu")).next_position == 12
+
 
 class TestTruncateCache:
     def test_truncate_reduces_length(self):
@@ -279,19 +290,27 @@ def test_metadata_preserved(self):
         truncated = truncate_cache(cache, max_length=5)
         assert truncated.metadata == {"info": 42}
 
+    def test_truncate_preserves_absolute_next_position(self):
+        cache = _make_cache_state(10)
+        cache.next_position_id = 20
+        truncated = truncate_cache(cache, max_length=4)
+        assert truncated.cached_length == 4
+        assert truncated.next_position == 20
+
 
 # ---------------------------------------------------------------------------
 # BatchedKVHelper tests
 # ---------------------------------------------------------------------------
 
+
 class TestBatchedKVHelperStack:
     @pytest.fixture
     def device(self):
         return torch.device("cpu")
 
     def _make_inputs(self, lengths, device):
-        ids = [torch.ones(l, dtype=torch.long, device=device) for l in lengths]
-        masks = [torch.ones(l, dtype=torch.long, device=device) for l in lengths]
+        ids = [torch.ones(length, dtype=torch.long, device=device) for length in lengths]
+        masks = [torch.ones(length, dtype=torch.long, device=device) for length in lengths]
         return ids, masks
 
     def test_stack_homogeneous_caches(self, device):
@@ -351,10 +370,8 @@ def test_all_empty_caches_no_stacked_kv(self, device):
 
     def test_padded_new_tokens_are_zero(self, device):
         caches = [_make_cache_state(5), _make_cache_state(5)]
-        ids = [torch.ones(6, dtype=torch.long, device=device),
-               torch.ones(3, dtype=torch.long, device=device)]
-        masks = [torch.ones(6, dtype=torch.long, device=device),
-                 torch.ones(3, dtype=torch.long, device=device)]
+        ids = [torch.ones(6, dtype=torch.long, device=device), torch.ones(3, dtype=torch.long, device=device)]
+        masks = [torch.ones(6, dtype=torch.long, device=device), torch.ones(3, dtype=torch.long, device=device)]
         result = BatchedKVHelper.stack_for_update(caches, ids, masks, device)
 
         # session 1 has 3 tokens, padded to 6 — positions 3,4,5 should be 0
@@ -368,10 +385,8 @@ def device(self):
 
     def test_unstack_restores_cache_lengths(self, device):
         caches = [_make_cache_state(5), _make_cache_state(10)]
-        ids = [torch.ones(3, dtype=torch.long, device=device),
-               torch.ones(3, dtype=torch.long, device=device)]
-        masks = [torch.ones(3, dtype=torch.long, device=device),
-                 torch.ones(3, dtype=torch.long, device=device)]
+        ids = [torch.ones(3, dtype=torch.long, device=device), torch.ones(3, dtype=torch.long, device=device)]
+        masks = [torch.ones(3, dtype=torch.long, device=device), torch.ones(3, dtype=torch.long, device=device)]
 
         stacked = BatchedKVHelper.stack_for_update(caches, ids, masks, device)
 
@@ -392,8 +407,7 @@ def test_unstack_preserves_session_ids(self, device):
         c2.session_id = "bob"
         caches = [c1, c2]
 
-        ids = [torch.ones(2, dtype=torch.long, device=device),
-               torch.ones(2, dtype=torch.long, device=device)]
+        ids = [torch.ones(2, dtype=torch.long, device=device), torch.ones(2, dtype=torch.long, device=device)]
         masks = ids[:]
 
         stacked = BatchedKVHelper.stack_for_update(caches, ids, masks, device)
@@ -419,6 +433,22 @@ def test_unstack_input_ids_concatenated(self, device):
 
         assert updated[0].input_ids.tolist() == [1, 2, 3, 4, 5]
 
+    def test_unstack_advances_absolute_position(self, device):
+        cache = _make_cache_state(3)
+        cache.next_position_id = 10
+        ids = [torch.ones(2, dtype=torch.long, device=device)]
+        masks = [torch.ones(2, dtype=torch.long, device=device)]
+        stacked = BatchedKVHelper.stack_for_update([cache], ids, masks, device)
+        new_kv = _make_dynamic_cache(2, 2, 5, 8, batch=1)
+        updated = BatchedKVHelper.unstack_after_update(new_kv, stacked, [cache], ids, masks)
+        assert updated[0].next_position == 12
+
+    def test_sliced_cache_owns_its_storage(self):
+        source = _make_dynamic_cache(1, 2, 8, 4, batch=2)
+        sliced = BatchedKVHelper._slice_past_kv(source, 0, 2, 6)
+        source.layers[0].keys.fill_(123)
+        assert not torch.all(sliced.layers[0].keys == 123)
+
 
 class TestBatchedKVHelperClassify:
     @pytest.fixture
@@ -427,22 +457,19 @@ def device(self):
 
     def test_stack_for_classify_shapes(self, device):
         caches = [_make_cache_state(10), _make_cache_state(15)]
-        label_ids = [torch.ones(8, dtype=torch.long, device=device),
-                     torch.ones(6, dtype=torch.long, device=device)]
-        label_masks = [torch.ones(8, dtype=torch.long, device=device),
-                       torch.ones(6, dtype=torch.long, device=device)]
+        label_ids = [torch.ones(8, dtype=torch.long, device=device), torch.ones(6, dtype=torch.long, device=device)]
+        label_masks = [torch.ones(8, dtype=torch.long, device=device), torch.ones(6, dtype=torch.long, device=device)]
 
         result = BatchedKVHelper.stack_for_classify(caches, label_ids, label_masks, device)
 
-        assert result["input_ids"].shape == (2, 8)   # max_label = 8
-        assert result["attention_mask"].shape == (2, 15 + 8)   # max_cached + max_label
+        assert result["input_ids"].shape == (2, 8)  # max_label = 8
+        assert result["attention_mask"].shape == (2, 15 + 8)  # max_cached + max_label
         assert result["label_lengths"] == [8, 6]
         assert result["max_label_len"] == 8
 
     def test_stack_for_classify_position_ids(self, device):
         caches = [_make_cache_state(5), _make_cache_state(12)]
-        label_ids = [torch.ones(4, dtype=torch.long, device=device),
-                     torch.ones(4, dtype=torch.long, device=device)]
+        label_ids = [torch.ones(4, dtype=torch.long, device=device), torch.ones(4, dtype=torch.long, device=device)]
         label_masks = label_ids[:]
 
         result = BatchedKVHelper.stack_for_classify(caches, label_ids, label_masks, device)
@@ -462,41 +489,135 @@ def test_label_mask_is_label_only(self, device):
         assert result["label_mask"].shape == (1, 5)
 
 
-# ---------------------------------------------------------------------------
-# Data types tests
-# ---------------------------------------------------------------------------
+class _Tokenizer:
+    def __call__(self, text, return_tensors="pt", add_special_tokens=False):
+        del return_tensors, add_special_tokens
+        ids = []
+        cursor = 0
+        special = {"<<SEP>>": 98, "<<LABEL>>": 99, "<<EXAMPLE>>": 97}
+        while cursor < len(text):
+            match = next((token for token in special if text.startswith(token, cursor)), None)
+            if match is not None:
+                ids.append(special[match])
+                cursor += len(match)
+            else:
+                ids.append((ord(text[cursor]) % 90) + 1)
+                cursor += 1
+        tensor = torch.tensor([ids], dtype=torch.long)
+        return {"input_ids": tensor, "attention_mask": torch.ones_like(tensor)}
+
+
+class _StreamingModel:
+    def __init__(self):
+        self.config = SimpleNamespace(architecture_type="decoder-kv", max_labels_alloc="dynamic")
+        self.device = torch.device("cpu")
+        self.model = SimpleNamespace(scorer=nn.Linear(1, 1))
+        self.position_history = []
+
+    def to(self, device):
+        self.device = torch.device(device)
+        return self
+
+    def eval(self):
+        return self
+
+    def update_decoder_cache(self, input_ids, attention_mask, position_ids, past_key_values):
+        del attention_mask
+        self.position_history.append(position_ids.clone())
+        batch, new_length = input_ids.shape
+        old_length = 0 if past_key_values is None else past_key_values.layers[0].keys.shape[2]
+        cache = _make_dynamic_cache(1, 1, old_length + new_length, 2, batch=batch)
+        return SimpleNamespace(past_key_values=cache)
+
+    def classify_from_decoder_cache(
+        self,
+        input_ids,
+        attention_mask,
+        label_mask,
+        position_ids,
+        past_key_values,
+    ):
+        del attention_mask, label_mask, position_ids, past_key_values
+        counts = input_ids.eq(99).sum(dim=1)
+        max_labels = int(counts.max().item())
+        logits = torch.zeros(input_ids.shape[0], max_labels)
+        if max_labels:
+            logits[:, 0] = 2.0
+        return SimpleNamespace(logits=logits)
+
+
+def _make_pipeline(**kwargs):
+    return StreamingZeroShotClassificationPipeline(
+        _StreamingModel(),
+        _Tokenizer(),
+        device="cpu",
+        **kwargs,
+    )
+
 
-class TestDataTypes:
-    def test_session_input_defaults(self):
-        from gliclass.streaming.strategies import EveryChunkStrategy
-        inp = SessionInput(
-            session_id="s1",
-            text="hello",
-            labels=["a", "b"],
-            strategy=EveryChunkStrategy(),
-        )
-        assert inp.classification_type == "multi-label"
-
-    def test_session_output_not_triggered(self):
-        out = SessionOutput(
-            session_id="s1",
-            triggered=False,
-            predictions=None,
-            cached_length=10,
-            tokens_added=5,
-        )
-        assert out.predictions is None
-        assert not out.triggered
-        assert out.metadata == {}
-
-    def test_session_output_triggered(self):
-        preds = [{"label": "positive", "score": 0.9}]
-        out = SessionOutput(
-            session_id="s1",
-            triggered=True,
-            predictions=preds,
-            cached_length=50,
-            tokens_added=10,
-        )
-        assert out.triggered
-        assert out.predictions[0]["label"] == "positive"
+class TestStreamingPipeline:
+    def test_new_api_and_automatic_session_cleanup(self):
+        pipeline = _make_pipeline()
+        first = pipeline(["ab", "cd"], ["x", "y"], session_ids=["a", "b"])
+        assert [item["session_id"] for item in first] == ["a", "b"]
+        assert pipeline.active_sessions == ["a", "b"]
+
+        pipeline("ef", ["x", "y"], session_ids="b")
+        assert pipeline.active_sessions == ["b"]
+        assert "a" not in pipeline._session_strategies
+
+    def test_duplicate_session_ids_are_rejected_before_cleanup(self):
+        pipeline = _make_pipeline()
+        pipeline("a", ["x"], session_ids="kept")
+        with pytest.raises(ValueError, match="unique"):
+            pipeline(["b", "c"], ["x"], session_ids=["dup", "dup"])
+        assert pipeline.active_sessions == ["kept"]
+
+    def test_invalid_batch_size_is_rejected_before_cleanup(self):
+        pipeline = _make_pipeline()
+        pipeline("a", ["x"], session_ids="kept")
+        with pytest.raises(ValueError, match="batch_size"):
+            pipeline("b", ["x"], session_ids="new", batch_size=0)
+        assert pipeline.active_sessions == ["kept"]
+
+    def test_multi_label_threshold_matches_regular_pipeline_behavior(self):
+        pipeline = _make_pipeline()
+        output = pipeline("a", ["first", "second"], session_ids="s", threshold=0.6)[0]
+        assert output["triggered"] is True
+        assert [item["label"] for item in output["predictions"]] == ["first"]
+
+    def test_strategy_state_is_independent_per_session(self):
+        pipeline = _make_pipeline(default_strategy=EveryNTokensStrategy(3))
+        first = pipeline(["a", "a"], ["x"], session_ids=["a", "b"])
+        assert [item["triggered"] for item in first] == [False, False]
+        second = pipeline(["aa", "a"], ["x"], session_ids=["a", "b"])
+        assert [item["triggered"] for item in second] == [True, False]
+
+    def test_prompt_is_cached_only_once(self):
+        pipeline = _make_pipeline()
+        first = pipeline("a", ["x"], session_ids="s", prompt="p")[0]
+        second = pipeline("b", ["x"], session_ids="s", prompt="p")[0]
+        assert first["tokens_added"] == 2
+        assert second["tokens_added"] == 1
+        assert second["cached_length"] == 3
+
+    def test_prefix_only_update_does_not_classify(self):
+        pipeline = _make_pipeline()
+        output = pipeline("", ["x"], session_ids="s", prompt="prefix")[0]
+        assert output["tokens_added"] == len("prefix")
+        assert output["triggered"] is False
+
+    def test_truncated_cache_keeps_monotonic_positions(self):
+        pipeline = _make_pipeline(max_cache_len=3)
+        pipeline("abcde", ["x"], session_ids="s")
+        assert pipeline._caches["s"].cached_length == 3
+        assert pipeline._caches["s"].next_position == 5
+        pipeline("fg", ["x"], session_ids="s")
+        assert pipeline.model.position_history[-1][0, :2].tolist() == [5, 6]
+        assert pipeline._caches["s"].next_position == 7
+
+    def test_empty_call_clears_sessions(self):
+        pipeline = _make_pipeline()
+        pipeline("a", ["x"], session_ids="s")
+        assert pipeline([], [], session_ids=[]) == []
+        assert pipeline.active_sessions == []

From 5f47e0b755fd3fdbf73ce02785954fe59987a81b Mon Sep 17 00:00:00 2001
From: Ingvarstep <igor.stepanov2000@gmail.com>
Date: Wed, 15 Jul 2026 12:47:51 +0000
Subject: [PATCH 24/25] fix issue with longer texts for decoder gliclass

---
 gliclass/pipeline.py | 74 +++++++++++++++++++++++++++++---------------
 1 file changed, 49 insertions(+), 25 deletions(-)

diff --git a/gliclass/pipeline.py b/gliclass/pipeline.py
index 938754a..9da574a 100644
--- a/gliclass/pipeline.py
+++ b/gliclass/pipeline.py
@@ -7,7 +7,7 @@
 
 from .model import GLiClassModel, GLiClassBiEncoder
 from .utils import retrieval_augmented_text
-from .data_processing import format_decoder_kv_sequence
+from .data_processing import format_decoder_kv_context, format_decoder_kv_labels, format_decoder_kv_sequence
 
 
 def flatten_hierarchical_labels(
@@ -784,32 +784,56 @@ def prepare_input(self, text: str, labels: list[str], examples=None, prompt: str
             sep_token=self.sep_token,
         )
 
+    def _build_context_and_labels(self, texts, labels, same_labels, examples, prompt):
+        contexts = []
+        label_sequences = []
+        for i, text in enumerate(texts):
+            item_labels = labels if same_labels else labels[i]
+            examples_text = self._format_examples_for_input(self._get_text_examples(examples, i))
+            text_prompt = self._format_prompt(prompt, i) or ""
+            contexts.append(format_decoder_kv_context(text, text_prompt, examples_text))
+            label_sequences.append(
+                format_decoder_kv_labels(item_labels, label_token=self.label_token, sep_token=self.sep_token)
+            )
+        return contexts, label_sequences
+
     def prepare_inputs(self, texts, labels, same_labels=False, examples=None, prompt=None):
-        inputs = []
-        if same_labels:
-            for i, text in enumerate(texts):
-                inputs.append(
-                    self.prepare_input(
-                        text,
-                        labels,
-                        self._get_text_examples(examples, i),
-                        self._format_prompt(prompt, i),
-                    )
-                )
-        else:
-            for i, (text, labels_) in enumerate(zip(texts, labels)):
-                inputs.append(
-                    self.prepare_input(
-                        text,
-                        labels_,
-                        self._get_text_examples(examples, i),
-                        self._format_prompt(prompt, i),
-                    )
-                )
+        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"]
+
+        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,
+                truncation=True,
+                max_length=context_budget,
+                add_special_tokens=False,
+            )["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
+        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)
 
-        return self.tokenizer(
-            inputs, truncation=True, max_length=self.max_length, padding="longest", return_tensors="pt"
-        ).to(self.device)
+        return {
+            "input_ids": torch.tensor(padded_input_ids, device=self.device),
+            "attention_mask": torch.tensor(attention_masks, device=self.device),
+        }
 
 
 class ZeroShotClassificationPipeline:

From f1956e30d622aadfdea0c5a50c918f3c0efe9ecf Mon Sep 17 00:00:00 2001
From: Ingvarstep <igor.stepanov2000@gmail.com>
Date: Fri, 17 Jul 2026 10:27:16 +0000
Subject: [PATCH 25/25] fix ruff error

---
 gliclass/pipeline.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/gliclass/pipeline.py b/gliclass/pipeline.py
index 9da574a..19150f6 100644
--- a/gliclass/pipeline.py
+++ b/gliclass/pipeline.py
@@ -7,7 +7,7 @@
 
 from .model import GLiClassModel, GLiClassBiEncoder
 from .utils import retrieval_augmented_text
-from .data_processing import format_decoder_kv_context, format_decoder_kv_labels, format_decoder_kv_sequence
+from .data_processing import format_decoder_kv_labels, format_decoder_kv_context, format_decoder_kv_sequence
 
 
 def flatten_hierarchical_labels(