feat(text): vary diacritics per document without inventing them

CHANGELOG.md

@@ -8,6 +8,19 @@ All notable changes to this project are documented here. The format follows
 
 ### Added
 
+- **Diacritics control** (`ocrsmith.text.diacritics`). Arabic OCR handles vocalisation
+  badly, and it is the first limitation AtlasOCR reports about itself. The cause is
+  distributional: real Arabic is *partially* diacritised, and the proportion varies by
+  genre. `DiacriticsPolicy` samples per document across four modes (`keep`, `strip`,
+  `partial`, `mixed`), and records the kept fraction in provenance so a diacritics
+  ablation is possible later. `DatasetStats` reports the corpus split across bare, partial
+  and fully marked pages.
+
+  Marks are only ever **removed**, never invented. Adding vocalisation to bare text needs
+  a diacritiser model and would make the label assert vowels nobody wrote — a fabricated
+  ground truth that looks entirely plausible. Point `text.source` at a diacritised corpus
+  and let the policy vary it downwards.
+
 - **`ocrsmith fetch-fonts`** — downloads open-licensed families from Google Fonts on
   demand. Font diversity is the highest-impact lever in synthetic text data, and a
   repository should not ship other people's typefaces. Only permissively licensed

src/ocrsmith/config/__init__.py

@@ -4,6 +4,7 @@
 from .schema import (
     BackgroundConfig,
     DegradationConfig,
+    DiacriticsConfig,
     FontConfig,
     GenerationConfig,
     NormalizationConfig,
@@ -19,6 +20,7 @@
     "DEFAULT_CONFIG_PATH",
     "BackgroundConfig",
     "DegradationConfig",
+    "DiacriticsConfig",
     "FontConfig",
     "GenerationConfig",
     "NormalizationConfig",

src/ocrsmith/config/default_config.yaml

@@ -30,6 +30,14 @@ text:
       - "يحتوي هذا التقرير على جداول وأرقام ومعلومات إضافية مفيدة للقارئ المهتم."
       - "في سنة 2024 أطلقت المجموعة نموذج OCR جديد يدعم اللغتين العربية والفرنسية."
       - "تعتمد الطريقة المقترحة على توليد بيانات اصطناعية متنوعة وواقعية قدر الإمكان."
+  diacritics:
+    # Marks are only ever removed, never invented: adding them would need a diacritiser
+    # model and would make the label assert vowels nobody wrote. Point this at a
+    # diacritised corpus and use "mixed" to reproduce how real Arabic actually varies -
+    # fully marked religious and pedagogical texts, lightly marked news, bare prose.
+    mode: keep          # keep | strip | partial | mixed
+    keep_range: [0.1, 0.6]
+    mixed_weights: [0.15, 0.25, 0.60]
   normalization:
     # Each of these rewrites the ground truth, so all default to off.
     strip_diacritics: false

src/ocrsmith/config/schema.py

@@ -21,6 +21,7 @@
     "FALLBACK_SENTENCES",
     "BackgroundConfig",
     "DegradationConfig",
+    "DiacriticsConfig",
     "FontConfig",
     "GenerationConfig",
     "OutputConfig",
@@ -93,6 +94,20 @@ def _path_required_for_file_sources(self) -> TextSourceConfig:
         return self
 
 
+class DiacriticsConfig(BaseModel):
+    """How vocalised the corpus is.
+
+    Real Arabic is *partially* diacritised and the proportion varies by genre, so a corpus
+    that is uniformly marked or uniformly bare teaches a model to expect that uniformity.
+    Marks are only ever removed, never invented: adding them would fabricate ground truth.
+    """
+
+    mode: Literal["keep", "strip", "partial", "mixed"] = "keep"
+    keep_range: tuple[float, float] = (0.1, 0.6)
+    #: Weights for "mixed": fully marked, partially marked, unmarked.
+    mixed_weights: tuple[float, float, float] = (0.15, 0.25, 0.60)
+
+
 class NormalizationConfig(BaseModel):
     """Label-affecting text transforms. Every one of these changes the ground truth."""
 
@@ -106,6 +121,7 @@ class NormalizationConfig(BaseModel):
 
 class TextConfig(BaseModel):
     source: TextSourceConfig = Field(default_factory=TextSourceConfig)
+    diacritics: DiacriticsConfig = Field(default_factory=DiacriticsConfig)
     normalization: NormalizationConfig = Field(default_factory=NormalizationConfig)
     #: "auto" reads the direction from the text itself.
     direction: Literal["auto", "rtl", "ltr"] = "auto"

src/ocrsmith/pipeline/factory.py

@@ -31,6 +31,7 @@
 )
 from ..core.fonts import FontPool
 from ..domain import Provenance, Sample
+from ..text.diacritics import DiacriticsMode, DiacriticsPolicy, apply_diacritics
 from ..text.script import Direction
 
 __all__ = ["SampleFactory"]
@@ -86,6 +87,7 @@ def create(self, index: int) -> Iterator[Sample]:
         direction = self._direction(rng)
         content = template.build(self.text, rng, direction=direction)
         content = self._with_footer(content, rng)
+        content, diacritics_kept = self._apply_diacritics(content, rng)
 
         spec = self._page_spec(self.config.page, direction, rng)
         typography = self._typography(content, rng)
@@ -121,6 +123,7 @@ def create(self, index: int) -> Iterator[Sample]:
                     degradations=tuple(record.to_dict() for record in records),
                     extra={
                         "index": index,
+                        "diacritics_kept": round(diacritics_kept, 3),
                         "page": rendered.number,
                         "preset": preset_name,
                         "paper": spec.width,
@@ -188,6 +191,50 @@ def _table_style(rng: random.Random) -> TableStyle:
             zebra_fill=(243, 245, 248) if shade > 0.75 else None,
         )
 
+    def _apply_diacritics(self, content, rng: random.Random):
+        """Vary how vocalised this document is, sampled once for the whole page.
+
+        Applied to the assembled content rather than at corpus load, so the fraction is a
+        property of the *document* - which is how real vocalisation works, a whole text
+        being marked or not - and so it can be recorded in provenance for an ablation.
+        """
+        settings = self.config.text.diacritics
+        if settings.mode == "keep":
+            return content, 1.0
+
+        policy = DiacriticsPolicy(
+            mode=DiacriticsMode(settings.mode),
+            keep_range=tuple(settings.keep_range),
+            mixed_weights=tuple(settings.mixed_weights),
+        )
+        kept = 1.0
+        blocks = []
+        for block in content.blocks:
+            text, kept = apply_diacritics(block.text, policy, rng) if block.text else (block.text, kept)
+            table = block.table
+            if table is not None:
+                table = type(table)(
+                    table.rows,
+                    table.cols,
+                    tuple(
+                        type(cell)(
+                            cell.row,
+                            cell.col,
+                            apply_diacritics(cell.text, policy, rng)[0],
+                            cell.bbox,
+                            cell.row_span,
+                            cell.col_span,
+                            cell.is_header,
+                            cell.lines,
+                        )
+                        for cell in table.cells
+                    ),
+                    table.has_header_row,
+                )
+            items = tuple(apply_diacritics(item, policy, rng)[0] for item in block.items)
+            blocks.append(type(block)(block.type, text, items, table, dict(block.attributes)))
+        return type(content)(tuple(blocks), content.direction, dict(content.metadata)), kept
+
     def _with_footer(self, content, rng: random.Random):
         if self.config.page.footer_probability <= 0 or rng.random() > self.config.page.footer_probability:
             return content

src/ocrsmith/quality/stats.py

@@ -22,6 +22,22 @@
 __all__ = ["DatasetStats", "scan_jsonl"]
 
 
+def _vocalisation_bucket(text: str) -> str:
+    """Coarse label for how diacritised a page is.
+
+    Bucketed rather than averaged because the interesting question is compositional - what
+    share of the corpus is bare, partial, or fully marked - not what the mean is.
+    """
+    from ..text.diacritics import diacritic_ratio
+
+    ratio = diacritic_ratio(text)
+    if ratio <= 0.01:
+        return "none"
+    if ratio < 0.35:
+        return "partial"
+    return "full"
+
+
 @dataclass
 class DatasetStats:
     """Streaming counters over a corpus."""
@@ -39,6 +55,8 @@ class DatasetStats:
     fonts: Counter = field(default_factory=Counter)
     characters_seen: Counter = field(default_factory=Counter)
     page_sizes: Counter = field(default_factory=Counter)
+    #: How vocalised each page was, bucketed. Arabic OCR lives or dies on this.
+    diacritics: Counter = field(default_factory=Counter)
     _line_heights: list = field(default_factory=list)
 
     def add(self, sample: Sample) -> DatasetStats:
@@ -59,6 +77,7 @@ def add(self, sample: Sample) -> DatasetStats:
         if provenance.font_path:
             self.fonts[Path(provenance.font_path).name] += 1
 
+        self.diacritics[_vocalisation_bucket(page.text)] += 1
         for region in page.regions:
             self.region_types[region.type.value] += 1
         for line in page.iter_lines():
@@ -89,6 +108,7 @@ def add_record(self, record: dict) -> DatasetStats:
         if provenance.get("font_path"):
             self.fonts[Path(provenance["font_path"]).name] += 1
 
+        self.diacritics[_vocalisation_bucket(record.get("text", ""))] += 1
         for region in page.get("regions", []):
             self.region_types[region.get("type", "unknown")] += 1
             for line in region.get("lines", []):
@@ -134,6 +154,7 @@ def to_dict(self) -> dict:
             "directions": dict(self.directions.most_common()),
             "region_types": dict(self.region_types.most_common()),
             "backgrounds": dict(self.backgrounds.most_common()),
+            "vocalisation": dict(self.diacritics.most_common()),
             "top_fonts": dict(self.fonts.most_common(10)),
             "top_page_sizes": dict(self.page_sizes.most_common(5)),
             "rarest_characters": [char for char, _ in self.characters_seen.most_common()[-15:]],
@@ -160,6 +181,7 @@ def to_markdown(self) -> str:
             ("Capture conditions", "degradation_presets"),
             ("Reading direction", "directions"),
             ("Region types", "region_types"),
+            ("Vocalisation", "vocalisation"),
         ):
             if data[key]:
                 lines.append("")

src/ocrsmith/text/__init__.py

@@ -10,6 +10,14 @@
     missing_glyphs,
     supports_text,
 )
+from .diacritics import (
+    DiacriticsMode,
+    DiacriticsPolicy,
+    apply_diacritics,
+    count_diacritics,
+    diacritic_ratio,
+    strip_partial,
+)
 from .normalization import (
     NormalizationPolicy,
     NumeralSystem,
@@ -49,6 +57,13 @@
     "TransparentShaper",
     "raqm_available",
     "resolve_shaper",
+    # diacritics
+    "DiacriticsMode",
+    "DiacriticsPolicy",
+    "apply_diacritics",
+    "count_diacritics",
+    "diacritic_ratio",
+    "strip_partial",
     # coverage
     "CoverageReport",
     "fonts_supporting",

src/ocrsmith/text/diacritics.py

@@ -0,0 +1,122 @@
+"""Arabic diacritics: measuring them, and varying how many survive.
+
+Diacritics (tashkeel) are the marks that fix short vowels and gemination. Arabic OCR
+systems handle them badly — it is called out as a systemic weakness in the Arabic
+benchmark literature, and it is the first limitation AtlasOCR reports about itself.
+
+The reason models struggle is distributional: real documents are *partially* diacritised.
+Religious and pedagogical texts are fully marked, newspapers carry a handful of
+disambiguating marks, and most prose has none. A corpus that is uniformly one or the other
+teaches a model to expect that uniformity.
+
+What this module deliberately does **not** do is invent diacritics. Adding marks to
+undiacritised text requires a diacritiser model and would fabricate ground truth — the
+label would assert vowels that no human wrote. Instead: start from a diacritised corpus
+and *remove* a sampled fraction. Removal is always truthful, because the remaining text is
+a form the source actually supports.
+"""
+
+from __future__ import annotations
+
+import random
+import re
+from dataclasses import dataclass
+from enum import Enum
+
+__all__ = [
+    "DiacriticsMode",
+    "DiacriticsPolicy",
+    "apply_diacritics",
+    "count_diacritics",
+    "diacritic_ratio",
+    "strip_partial",
+]
+
+#: Tashkeel, Quranic annotation marks and the superscript alef — the same set
+#: `normalization.strip_diacritics` removes, kept here as a character class for sampling.
+_MARK = re.compile("[ؐ-ًؚ-ٰٟۖ-ۜ۟-۪ۨ-ۭ࣓-ࣿ]")
+
+
+class DiacriticsMode(str, Enum):
+    """What happens to the marks a source text carries."""
+
+    #: Leave the source exactly as it is.
+    KEEP = "keep"
+    #: Remove every mark.
+    STRIP = "strip"
+    #: Remove a fixed fraction, sampled once per document.
+    PARTIAL = "partial"
+    #: Sample a mode per document, which is what a mixed real corpus looks like.
+    MIXED = "mixed"
+
+
+def count_diacritics(text: str) -> int:
+    """How many diacritic marks `text` carries."""
+    return len(_MARK.findall(text))
+
+
+def diacritic_ratio(text: str) -> float:
+    """Marks per non-mark character, a cheap measure of how vocalised a text is.
+
+    Fully vocalised Arabic sits near 0.5-0.7; newspaper prose is nearer 0.0-0.05. Useful
+    for reporting what a corpus actually contains rather than what it was assumed to.
+    """
+    marks = count_diacritics(text)
+    base = len(text) - marks
+    return marks / base if base else 0.0
+
+
+def strip_partial(text: str, keep_fraction: float, rng: random.Random) -> str:
+    """Keep `keep_fraction` of the marks, chosen uniformly at random.
+
+    Marks are dropped independently rather than by region, because partial vocalisation in
+    real documents is driven by ambiguity — a writer marks the word that would otherwise be
+    misread — and that is scattered, not clustered.
+    """
+    if keep_fraction >= 1.0:
+        return text
+    if keep_fraction <= 0.0:
+        return _MARK.sub("", text)
+    return _MARK.sub(lambda match: match.group(0) if rng.random() < keep_fraction else "", text)
+
+
+@dataclass(frozen=True, slots=True)
+class DiacriticsPolicy:
+    """How diacritics vary across a corpus."""
+
+    mode: DiacriticsMode = DiacriticsMode.KEEP
+    #: Range the kept fraction is sampled from, for PARTIAL and the partial branch of MIXED.
+    keep_range: tuple[float, float] = (0.1, 0.6)
+    #: Weights for MIXED: fully marked, partially marked, unmarked. Defaults approximate a
+    #: general-purpose corpus, where most prose is unmarked.
+    mixed_weights: tuple[float, float, float] = (0.15, 0.25, 0.60)
+
+    def apply(self, text: str, rng: random.Random) -> tuple[str, float]:
+        return apply_diacritics(text, self, rng)
+
+
+def apply_diacritics(text: str, policy: DiacriticsPolicy, rng: random.Random) -> tuple[str, float]:
+    """Apply `policy` to `text`.
+
+    Returns the text and the fraction of marks kept, so provenance can record how
+    vocalised each sample was — which is what makes a diacritics ablation possible later.
+    """
+    if not text or not count_diacritics(text):
+        return text, 1.0
+
+    mode = policy.mode
+    if mode is DiacriticsMode.MIXED:
+        mode = rng.choices(
+            [DiacriticsMode.KEEP, DiacriticsMode.PARTIAL, DiacriticsMode.STRIP],
+            weights=list(policy.mixed_weights),
+            k=1,
+        )[0]
+
+    if mode is DiacriticsMode.KEEP:
+        return text, 1.0
+    if mode is DiacriticsMode.STRIP:
+        return _MARK.sub("", text), 0.0
+
+    low, high = sorted(policy.keep_range)
+    keep = rng.uniform(low, high)
+    return strip_partial(text, keep, rng), keep