Update app.py

app.py

@@ -40,6 +40,9 @@ from rule_detector import DEFAULT_RULEBASE_ALLOWED_ENTITY_TYPES, detect_by_rules
 DEFAULT_MODEL_ID = os.getenv("MODEL_ID", "AITeamUIT/gliner2-multi-v1-e3-25-6")
 DEFAULT_THRESHOLD = float(os.getenv("THRESHOLD", "0.5"))
 DEFAULT_CHUNK_CHARS = int(os.getenv("CHUNK_CHARS", "1000"))
+DEFAULT_CHUNK_TOKENS = int(os.getenv("CHUNK_TOKENS", "768"))
+CHUNK_SAFETY_MARGIN = int(os.getenv("CHUNK_SAFETY_MARGIN", "20"))
+MAX_CHARS_PER_CHUNK = int(os.getenv("MAX_CHARS_PER_CHUNK", "8000"))
 MAX_WIDTH = int(os.getenv("MAX_WIDTH", "30"))
 
 PII_LABEL_DESCRIPTIONS: Dict[str, str] = {
@@ -235,7 +238,76 @@ def manual_chunks(text: str, max_chars: int) -> List[Tuple[str, int]]:
     return chunks
 
 
-def find_all_occurrences(text: str, value: str) -> List[Tuple[int, int]]:
+def word_token_count(text: str) -> int:
+    """Count word tokens the way the (patched) model tokenizer splits them."""
+    try:
+        from gliner2.processor import WhitespaceTokenSplitter
+
+        splitter = WhitespaceTokenSplitter()
+        return sum(1 for _ in splitter(text, lower=False))
+    except Exception:
+        return len(text.split())
+
+
+class RecursiveWordChunker:
+    """Label-aware recursive chunker built on chonkie.RecursiveChunker.
+
+    The chunk-size budget is measured in the model's own word-token space and
+    already subtracts the entity-schema label tokens, so each emitted chunk plus
+    the labels fits the encoder. Recursion falls back paragraph -> sentence ->
+    whitespace. ``chunk()`` returns ``(chunk_text, start_index)`` tuples.
+    """
+
+    def __init__(self, budget_words: int, max_chars: int = MAX_CHARS_PER_CHUNK, min_chars_per_chunk: int = 24):
+        # chonkie is a pip dependency (not a cross-folder import).
+        from chonkie.chunker.recursive import RecursiveChunker
+        from chonkie.types import RecursiveLevel, RecursiveRules
+
+        self.budget = budget_words
+        self.max_chars = max_chars
+        self._chunker = RecursiveChunker(
+            tokenizer=word_token_count,
+            chunk_size=budget_words,
+            rules=RecursiveRules(
+                levels=[
+                    RecursiveLevel(delimiters=["\n\n", "\r\n", "\n", "\r"]),
+                    RecursiveLevel(delimiters=[". ", "! ", "? ", ".\n", "!\n", "?\n"]),
+                    RecursiveLevel(whitespace=True),
+                ]
+            ),
+            min_characters_per_chunk=min_chars_per_chunk,
+        )
+
+    def chunk(self, text: str) -> List[Tuple[str, int]]:
+        return [(c.text, int(c.start_index)) for c in self._chunker.chunk(text)]
+
+
+@lru_cache(maxsize=4)
+def build_word_chunker(
+    chunk_tokens: int = DEFAULT_CHUNK_TOKENS,
+    safety_margin: int = CHUNK_SAFETY_MARGIN,
+    max_chars: int = MAX_CHARS_PER_CHUNK,
+) -> RecursiveWordChunker:
+    """Build a :class:`RecursiveWordChunker` with a label-aware token budget.
+
+    The label tokens (``"<LABEL> <description>"`` plus a +2 separator budget each)
+    are subtracted from ``chunk_tokens`` so the encoder never overflows.
+    """
+    labels_words = [f"{k} {v}" for k, v in PII_LABEL_DESCRIPTIONS.items()]
+    label_token_count = sum(word_token_count(lbl) + 2 for lbl in labels_words)
+    budget = max(64, chunk_tokens - label_token_count - safety_margin)
+    return RecursiveWordChunker(budget_words=budget, max_chars=max_chars)
+
+
+def chunk_text(text: str) -> List[Tuple[str, int]]:
+    """Split text into (chunk_text, start_index) using the recursive word chunker.
+
+    Falls back to a plain whitespace chunker only if chonkie is unavailable.
+    """
+    try:
+        return build_word_chunker(DEFAULT_CHUNK_TOKENS).chunk(text)
+    except Exception:
+        return manual_chunks(text, DEFAULT_CHUNK_CHARS)
     spans = []
     if not value:
         return spans
@@ -417,24 +489,24 @@ def run_inference(text: str, model_id: str, threshold: float, use_rules: bool, u
     schema = build_schema(model, use_descriptions)
     all_predictions: List[Dict] = []
 
-    chunks = manual_chunks(text, DEFAULT_CHUNK_CHARS)
-    for chunk_idx, (chunk_text, chunk_start) in enumerate(chunks):
+    chunks = chunk_text(text)
+    for chunk_idx, (chunk_text_value, chunk_start) in enumerate(chunks):
         # Store chunk info
-        debug_info["chunks"].append({"index": chunk_idx, "start": chunk_start, "end": chunk_start + len(chunk_text), "text": chunk_text})
+        debug_info["chunks"].append({"index": chunk_idx, "start": chunk_start, "end": chunk_start + len(chunk_text_value), "text": chunk_text_value})
 
         # Tokenize chunk (show what WhitespaceTokenSplitter produces)
-        tokens = tokenize_chunk(chunk_text)
+        tokens = tokenize_chunk(chunk_text_value)
         debug_info["tokens_per_chunk"].append({"chunk_index": chunk_idx, "tokens": [(t, s, e) for t, s, e in tokens]})
 
         raw = model.extract(
-            chunk_text,
+            chunk_text_value,
             schema,
             threshold=threshold,
             include_spans=not emit_duplicates,
             include_confidence=not emit_duplicates,
             format_results=emit_duplicates,
         )
-        predictions = normalize_predictions(raw, chunk_text, emit_duplicates=emit_duplicates)
+        predictions = normalize_predictions(raw, chunk_text_value, emit_duplicates=emit_duplicates)
 
         # Store per-chunk raw predictions (before offset adjustment)
         debug_info["raw_preds_per_chunk"].append({"chunk_index": chunk_idx, "predictions": [dict(p) for p in predictions]})