Eval nan bug

Author: adaamko

squeez/encoder/chunking.py

@@ -0,0 +1,64 @@
+"""Helpers for splitting pathological long lines into token chunks.
+
+The encoder task stays line-level by default. Only lines whose tokenized
+length exceeds the configured per-line budget are split into chunked
+"pseudo-lines". Training assigns the original line label to every chunk.
+Inference aggregates chunk scores back to the original line index.
+"""
+
+from __future__ import annotations
+
+from transformers import PreTrainedTokenizer
+
+
+def encode_text(
+    tokenizer: PreTrainedTokenizer,
+    text: str,
+    truncation: bool = False,
+    max_length: int | None = None,
+) -> list[int]:
+    """Tokenize a single text span without special tokens or warning spam."""
+    encoded = tokenizer(
+        text,
+        add_special_tokens=False,
+        truncation=truncation,
+        max_length=max_length,
+        return_attention_mask=False,
+        return_token_type_ids=False,
+        verbose=False,
+    )
+    input_ids = encoded["input_ids"]
+    return input_ids if isinstance(input_ids, list) else list(input_ids)
+
+
+def chunk_output_lines(
+    tokenizer: PreTrainedTokenizer,
+    output_lines: list[str],
+    max_tokens_per_chunk: int,
+) -> tuple[list[list[int]], list[int]]:
+    """Tokenize output lines, splitting only oversized lines into chunks.
+
+    Returns:
+        chunk_token_ids: token ids for each pseudo-line/chunk
+        chunk_to_line: mapping from chunk index back to original line index
+    """
+    chunk_token_ids: list[list[int]] = []
+    chunk_to_line: list[int] = []
+
+    for line_idx, line in enumerate(output_lines):
+        token_ids = encode_text(tokenizer, line)
+        if not token_ids:
+            continue
+
+        if len(token_ids) <= max_tokens_per_chunk:
+            chunk_token_ids.append(token_ids)
+            chunk_to_line.append(line_idx)
+            continue
+
+        for start in range(0, len(token_ids), max_tokens_per_chunk):
+            chunk = token_ids[start : start + max_tokens_per_chunk]
+            if chunk:
+                chunk_token_ids.append(chunk)
+                chunk_to_line.append(line_idx)
+
+    return chunk_token_ids, chunk_to_line

squeez/encoder/dataset.py

@@ -23,6 +23,7 @@
 from torch.utils.data import Dataset
 from transformers import PreTrainedTokenizer
 
+from squeez.encoder.chunking import chunk_output_lines, encode_text
 from squeez.encoder.model import LINE_SEP_TOKEN
 
 logger = logging.getLogger(__name__)
@@ -106,6 +107,7 @@ def __init__(
         # covering one window of lines that fits within max_length.
         self._windows: list[tuple[list[int], list[list[int]], list[bool]]] = []
         n_expanded = 0
+        n_skipped_empty = 0
 
         for sample in raw_samples:
             task = sample["task"]
@@ -116,23 +118,20 @@ def __init__(
             line_labels = _match_lines(output_lines, relevant_lines)
 
             # Tokenize task, cap at half of max_length
-            task_ids = tokenizer.encode(
+            task_ids = encode_text(
+                tokenizer,
                 task,
-                add_special_tokens=False,
                 truncation=True,
                 max_length=self._max_task_tokens,
             )
 
-            # Tokenize each line
-            line_token_ids = [
-                tokenizer.encode(
-                    ln,
-                    add_special_tokens=False,
-                    truncation=True,
-                    max_length=self._max_line_tokens,
-                )
-                for ln in output_lines
-            ]
+            # Tokenize each line, chunking only pathological long lines.
+            line_token_ids, chunk_to_line = chunk_output_lines(
+                tokenizer,
+                output_lines,
+                max_tokens_per_chunk=self._max_line_tokens,
+            )
+            chunk_labels = [line_labels[line_idx] for line_idx in chunk_to_line]
 
             # overhead = [CLS] + task + [SEP] + ... + [SEP]
             prefix_len = 1 + len(task_ids) + 1
@@ -142,11 +141,15 @@ def __init__(
             windows = self._build_windows(line_token_ids, budget)
 
             for start, end in windows:
+                window_line_token_ids = line_token_ids[start:end]
+                if not any(window_line_token_ids):
+                    n_skipped_empty += 1
+                    continue
                 self._windows.append(
                     (
                         task_ids,
-                        line_token_ids[start:end],
-                        line_labels[start:end],
+                        window_line_token_ids,
+                        chunk_labels[start:end],
                     )
                 )
 
@@ -156,7 +159,8 @@ def __init__(
         logger.info(
             f"Loaded {len(raw_samples)} samples from {data_path} → "
             f"{len(self._windows)} windows "
-            f"({n_expanded} extra from sliding, max_length={max_length})"
+            f"({n_expanded} extra from sliding, {n_skipped_empty} empty windows skipped, "
+            f"max_length={max_length})"
         )
 
     # ------------------------------------------------------------------

squeez/encoder/model.py

@@ -18,6 +18,8 @@
 from transformers import AutoConfig, AutoModel, AutoTokenizer, PretrainedConfig, PreTrainedModel
 from transformers.modeling_outputs import TokenClassifierOutput
 
+from squeez.encoder.chunking import chunk_output_lines, encode_text
+
 logger = logging.getLogger(__name__)
 
 LINE_SEP_TOKEN = "[LINE_SEP]"
@@ -163,9 +165,9 @@ def extract(
         sep_id = tokenizer.sep_token_id
 
         # Tokenize task prefix (will be reused for every window)
-        task_ids = tokenizer.encode(
+        task_ids = encode_text(
+            tokenizer,
             task,
-            add_special_tokens=False,
             truncation=True,
             max_length=max(max_len - 3 - _MIN_LINE_BUDGET, 0),
         )
@@ -175,16 +177,14 @@ def extract(
         suffix_len = 1  # final SEP
         budget = max_len - prefix_len - suffix_len
 
-        # Tokenize each line
-        line_token_ids: list[list[int]] = []
-        for line in lines:
-            ids = tokenizer.encode(
-                line,
-                add_special_tokens=False,
-                truncation=True,
-                max_length=max(max_len - 4, 1),
-            )
-            line_token_ids.append(ids)
+        # Tokenize lines, chunking only pathological long lines.
+        line_token_ids, chunk_to_line = chunk_output_lines(
+            tokenizer,
+            lines,
+            max_tokens_per_chunk=max(max_len - 4, 1),
+        )
+        if not line_token_ids:
+            return []
 
         # Build windows
         windows = self._build_windows(line_token_ids, budget, window_overlap)
@@ -199,8 +199,9 @@ def extract(
 
             scores = self._predict_window(input_ids, attention_mask, line_sep_positions, sep_id)
             for i, score in enumerate(scores):
-                global_idx = start_idx + i
-                line_scores[global_idx] = max(line_scores[global_idx], score)
+                chunk_idx = start_idx + i
+                line_idx = chunk_to_line[chunk_idx]
+                line_scores[line_idx] = max(line_scores[line_idx], score)
 
         return [line for line, score in zip(lines, line_scores) if score >= threshold]
 

squeez/encoder/train.py

@@ -98,7 +98,6 @@ def train(
     data_collator = DataCollatorForTokenClassification(
         tokenizer=tokenizer,
         padding=True,
-        max_length=max_length,
     )
 
     # Training arguments

tests/test_encoder_chunking.py

@@ -0,0 +1,43 @@
+from squeez.encoder.chunking import chunk_output_lines
+
+
+class FakeTokenizer:
+    def __call__(
+        self,
+        text,
+        add_special_tokens=False,
+        truncation=False,
+        max_length=None,
+        return_attention_mask=False,
+        return_token_type_ids=False,
+        verbose=False,
+    ):
+        # Tokenize on whitespace for predictable chunk sizes in tests.
+        tokens = [len(part) for part in text.split() if part]
+        if truncation and max_length is not None:
+            tokens = tokens[:max_length]
+        return {"input_ids": tokens}
+
+
+def test_chunk_output_lines_splits_only_overlong_lines():
+    tokenizer = FakeTokenizer()
+    lines = [
+        "short line",
+        "a b c d e f g",
+        "tiny",
+    ]
+
+    chunks, chunk_to_line = chunk_output_lines(
+        tokenizer,
+        lines,
+        max_tokens_per_chunk=3,
+    )
+
+    assert chunks == [
+        [5, 4],  # short line
+        [1, 1, 1],  # first chunk of long line
+        [1, 1, 1],  # second chunk of long line
+        [1],  # third chunk of long line
+        [4],  # tiny
+    ]
+    assert chunk_to_line == [0, 1, 1, 1, 2]