polystring 0.1.0__py3-none-any.whl

polystring/__init__.py

@@ -0,0 +1,29 @@
+"""polystring — span-level language detection for mixed-language text."""
+from __future__ import annotations
+
+from polystring._analyzer import analyze
+from polystring._exceptions import (
+    InputTooShortError,
+    PolyStringError,
+    UnsupportedLanguageError,
+)
+from polystring._models import PolyStringResult, Span, Token
+
+__version__ = "0.1.0"
+__all__ = [
+    "__version__",
+    "analyze",
+    "supported_languages",
+    "Span",
+    "Token",
+    "PolyStringResult",
+    "PolyStringError",
+    "UnsupportedLanguageError",
+    "InputTooShortError",
+]
+
+
+def supported_languages() -> list[str]:
+    """Return sorted list of ISO 639-1 codes supported for detection."""
+    from polystring._analyzer import _LINGUA_SUPPORTED
+    return sorted(_LINGUA_SUPPORTED)

polystring/_analyzer.py

@@ -0,0 +1,133 @@
+from __future__ import annotations
+
+from typing import Literal
+
+from polystring._exceptions import InputTooShortError, UnsupportedLanguageError
+from polystring._models import PolyStringResult
+from polystring._pipeline import (
+    stage1_preprocess,
+    stage2_script,
+    stage3_classify,
+    stage4_context,
+    stage5_merge,
+)
+from polystring.lexicons import add_custom_lexicon
+
+# Languages supported by lingua's ISO 639-1 codes (subset used for validation)
+# We rely on lingua raising its own error if a code is truly unknown;
+# this set is used only for fast pre-validation of the hint list.
+_LINGUA_SUPPORTED: frozenset[str] = frozenset({
+    "af", "sq", "ar", "hy", "az", "eu", "be", "bn", "bs", "bg", "ca", "zh",
+    "hr", "cs", "da", "nl", "en", "eo", "et", "fi", "fr", "lg", "ka", "de",
+    "el", "gu", "he", "hi", "hu", "is", "id", "ga", "it", "ja", "kn", "kk",
+    "ko", "la", "lv", "lt", "mk", "ms", "mi", "mr", "mn", "ne", "nb", "nn",
+    "fa", "pl", "pt", "pa", "ro", "ru", "sr", "sn", "sk", "sl", "so", "st",
+    "es", "sw", "sv", "tl", "ta", "te", "th", "ts", "tn", "tr", "uk", "ur",
+    "vi", "cy", "xh", "yo", "zu",
+})
+
+
+def analyze(
+    text: str,
+    *,
+    languages: list[str] | None = None,
+    granularity: Literal["span", "token"] = "span",
+    min_confidence: float = 0.70,
+    low_accuracy_mode: bool = False,
+    normalize: bool = True,
+    custom_lexicon: dict[str, list[str]] | None = None,
+) -> PolyStringResult:
+    """Detect languages of each span in mixed-language text.
+
+    Parameters
+    ----------
+    text:
+        Input text to analyse.
+    languages:
+        Restrict detection to these ISO 639-1 codes. Speeds up detection and
+        reduces false positives on known language sets.
+    granularity:
+        "span" (default) merges adjacent same-language tokens into spans.
+        "token" also populates result.tokens with per-token data.
+    min_confidence:
+        Tokens below this threshold are tagged "und". Default 0.70.
+    low_accuracy_mode:
+        Skip the lingua model entirely; use only lexicons and script detection.
+        Much faster but lower recall.
+    normalize:
+        Run NFC normalisation before analysis. Set False to skip.
+    custom_lexicon:
+        Additional {lang_code: [word, ...]} entries merged into the lexicons
+        before analysis.
+    """
+    if not isinstance(text, str):
+        raise TypeError(f"text must be str, got {type(text).__name__}")
+
+    if languages:
+        for code in languages:
+            if code not in _LINGUA_SUPPORTED:
+                raise UnsupportedLanguageError(code)
+
+    if custom_lexicon:
+        add_custom_lexicon(custom_lexicon)
+
+    languages_key = frozenset(languages) if languages else None
+
+    stage1 = stage1_preprocess.run(
+        text,
+        normalize=normalize,
+    )
+
+    if len(stage1.linguistic_tokens) < 2:
+        raise InputTooShortError(
+            "Input has fewer than 2 tokens after special token removal. "
+            "Cannot perform reliable language detection."
+        )
+
+    script_tokens, latin_tokens = stage2_script.run(stage1.linguistic_tokens)
+
+    if low_accuracy_mode:
+        from polystring._models import Token
+        from polystring.lexicons import lexicon_lookup
+        latin_classified: list[Token] = []
+        for rt in latin_tokens:
+            lex = lexicon_lookup(rt.text)
+            if lex:
+                lang, conf = lex
+                latin_classified.append(Token(
+                    text=rt.text,
+                    language=lang if lang != "amb" else "und",
+                    token_type="text",
+                    confidence=conf if lang != "amb" else 0.0,
+                    start=rt.start,
+                    end=rt.end,
+                ))
+            else:
+                latin_classified.append(Token(
+                    text=rt.text,
+                    language="und",
+                    token_type="text",
+                    confidence=0.0,
+                    start=rt.start,
+                    end=rt.end,
+                ))
+    else:
+        latin_classified = stage3_classify.run(
+            latin_tokens,
+            languages_hint=languages_key,
+            min_confidence=min_confidence,
+        )
+
+    all_tokens = sorted(
+        script_tokens + latin_classified,
+        key=lambda t: t.start,
+    )
+
+    all_tokens = stage4_context.run(all_tokens)
+
+    return stage5_merge.run(
+        all_tokens,
+        stage1.special_tokens,
+        stage1.normalized_text,
+        granularity=granularity,
+    )

polystring/_detector.py

@@ -0,0 +1,74 @@
+from __future__ import annotations
+
+import functools
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from lingua import LanguageDetector
+
+_detector: LanguageDetector | None = None
+_detector_languages: frozenset[str] | None = None
+
+
+def _build_detector(languages: list[str] | None = None) -> LanguageDetector:
+    from lingua import Language, LanguageDetectorBuilder
+
+    if languages:
+        from lingua import IsoCode639_1
+        lang_objs = []
+        for code in languages:
+            try:
+                iso = IsoCode639_1[code.upper()]
+                lang = Language.from_iso_code_639_1(iso)
+                lang_objs.append(lang)
+            except (KeyError, Exception):
+                pass
+        if not lang_objs:
+            builder = LanguageDetectorBuilder.from_all_languages()
+        else:
+            builder = LanguageDetectorBuilder.from_languages(*lang_objs)
+    else:
+        builder = LanguageDetectorBuilder.from_all_languages()
+
+    return builder.with_preloaded_language_models().build()
+
+
+def get_detector(languages: list[str] | None = None) -> LanguageDetector:
+    global _detector, _detector_languages
+
+    key = frozenset(languages) if languages else None
+    if _detector is None or _detector_languages != key:
+        _detector = _build_detector(languages)
+        _detector_languages = key
+        lingua_top2.cache_clear()
+    return _detector
+
+
+@functools.lru_cache(maxsize=4096)
+def lingua_top2(
+    text: str, languages_key: frozenset[str] | None = None
+) -> list[tuple[str, float]]:
+    detector = get_detector(list(languages_key) if languages_key else None)
+    confidence_values = detector.compute_language_confidence_values(text)
+    out: list[tuple[str, float]] = []
+    for cv in confidence_values[:2]:
+        code = cv.language.iso_code_639_1.name.lower()
+        out.append((code, cv.value))
+    return out
+
+
+def lingua_confidence_for(text: str, lang_code: str) -> float:
+    """Return lingua's confidence that `text` is in `lang_code`."""
+    from lingua import IsoCode639_1, Language
+
+    detector = get_detector()
+    try:
+        iso = IsoCode639_1[lang_code.upper()]
+        lang = Language.from_iso_code_639_1(iso)
+    except (KeyError, Exception):
+        return 0.0
+
+    for cv in detector.compute_language_confidence_values(text):
+        if cv.language == lang:
+            return cv.value
+    return 0.0

polystring/_exceptions.py

@@ -0,0 +1,17 @@
+class PolyStringError(Exception): ...
+
+
+class UnsupportedLanguageError(PolyStringError):
+    def __init__(self, code: str) -> None:
+        super().__init__(
+            f"'{code}' is not a supported language code. "
+            f"Call polystring.supported_languages() for the full list."
+        )
+
+
+class InputTooShortError(PolyStringError):
+    def __init__(
+        self,
+        message: str = "Input too short: need at least 2 tokens.",
+    ) -> None:
+        super().__init__(message)

polystring/_models.py

@@ -0,0 +1,106 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    import pandas as pd
+
+
+@dataclass
+class Token:
+    text: str
+    language: str
+    token_type: str
+    confidence: float
+    start: int
+    end: int
+    ambiguous_candidates: list[str] = field(default_factory=list)
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "text": self.text,
+            "language": self.language,
+            "token_type": self.token_type,
+            "confidence": self.confidence,
+            "start": self.start,
+            "end": self.end,
+            "ambiguous_candidates": self.ambiguous_candidates,
+        }
+
+
+@dataclass
+class Span:
+    text: str
+    language: str
+    token_type: str
+    confidence: float
+    start: int
+    end: int
+    is_foreign: bool = False
+    ambiguous_candidates: list[str] = field(default_factory=list)
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "text": self.text,
+            "language": self.language,
+            "token_type": self.token_type,
+            "confidence": self.confidence,
+            "start": self.start,
+            "end": self.end,
+            "is_foreign": self.is_foreign,
+            "ambiguous_candidates": self.ambiguous_candidates,
+        }
+
+
+@dataclass
+class PolyStringResult:
+    text: str
+    spans: list[Span]
+    tokens: list[Token] | None
+    languages: set[str]
+    dominant_language: str
+    is_mixed: bool
+    confidence: float
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "text": self.text,
+            "spans": [s.to_dict() for s in self.spans],
+            "languages": list(self.languages),
+            "dominant_language": self.dominant_language,
+            "is_mixed": self.is_mixed,
+            "confidence": self.confidence,
+        }
+
+    def to_dataframe(self) -> pd.DataFrame:
+        try:
+            import pandas as pd
+        except ImportError as e:
+            raise ImportError(
+                "pandas is required: pip install polystring[pandas]"
+            ) from e
+        return pd.DataFrame([s.to_dict() for s in self.spans])
+
+    def highlight(self) -> str:
+        # ANSI colour codes per language (cycles through a palette)
+        _PALETTE = [
+            "\033[91m", "\033[92m", "\033[93m", "\033[94m",
+            "\033[95m", "\033[96m", "\033[97m",
+        ]
+        _RESET = "\033[0m"
+        lang_colour: dict[str, str] = {}
+        colour_idx = 0
+        parts: list[str] = []
+        for span in self.spans:
+            if span.language not in lang_colour:
+                lang_colour[span.language] = _PALETTE[colour_idx % len(_PALETTE)]
+                colour_idx += 1
+            parts.append(
+                f"{lang_colour[span.language]}[{span.language}]{span.text}{_RESET}"
+            )
+        return " ".join(parts)
+
+    def linguistic_spans(self) -> list[Span]:
+        _NON_LINGUISTIC = {"url", "mention", "hashtag", "emoji", "num"}
+        return [s for s in self.spans if s.token_type not in _NON_LINGUISTIC]

polystring/_ngram.py

@@ -0,0 +1,144 @@
+"""Character n-gram language scorer for low-resource romanised languages.
+
+Loaded once at import time from pre-built JSON profiles in polystring/data/.
+Used in stage 3 between lexicon lookup and lingua for languages where lingua
+has insufficient training data (ur-Latn, tl, sw).
+
+Architecture: discriminative hit-count scoring.
+
+Each profile contains only *discriminative* n-grams: n-grams that appear
+significantly more often in that language than in all competitor languages
+(other target languages + English background), as determined at build time
+by a log-prob margin threshold.
+
+At inference time we count how many of a token's n-grams match each
+language's discriminative profile.  The language with the most hits wins,
+provided it leads the runner-up by at least _MIN_GAP_HITS and has at least
+_MIN_HITS total.  Ties are broken by the average log-prob of matched n-grams.
+
+This avoids the cross-contamination problem that affects plain LLR scoring:
+because the profiles are pre-filtered to exclude n-grams shared across
+ur-Latn/tl/sw, an Urdu word cannot "accidentally" accumulate Tagalog hits.
+"""
+from __future__ import annotations
+
+import json
+import re
+from pathlib import Path
+
+_DATA_DIR = Path(__file__).parent / "data"
+
+NGRAM_LANGUAGES: frozenset[str] = frozenset({"ur-Latn", "tl", "sw"})
+
+# Scoring thresholds (tuned empirically on build corpora test words)
+_MIN_HITS = 2           # winner must have at least this many discriminative n-gram hits
+_MIN_GAP_HITS = 1       # winner must lead runner-up by at least this many hits
+_MIN_TOKEN_LEN = 4      # tokens shorter than this are not scored (too noisy)
+
+_CLEAN = re.compile(r"[^a-z'\-]")
+
+_MODELS: dict[str, dict[str, dict[str, float]]] = {}
+_NGRAM_SIZES: dict[str, list[int]] = {}
+_LOADED = False
+
+
+def _load() -> None:
+    global _LOADED
+    if _LOADED:
+        return
+
+    for lang in NGRAM_LANGUAGES:
+        fname = _DATA_DIR / f"{lang.replace('-', '_')}_ngram.json"
+        if not fname.exists():
+            continue
+        payload = json.loads(fname.read_text(encoding="utf-8"))
+        _MODELS[lang] = payload["profile"]
+        _NGRAM_SIZES[lang] = payload["ngram_sizes"]
+
+    _LOADED = True
+
+
+def _hit_score(
+    cleaned: str, profile: dict[str, dict[str, float]], sizes: list[int]
+) -> tuple[int, float]:
+    """Count discriminative n-gram hits and sum their log-probs.
+
+    Returns (hit_count, avg_log_prob_of_hits) where avg is 0 when hit_count=0.
+    """
+    hit_count = 0
+    lp_sum = 0.0
+
+    for n in sizes:
+        table = profile[str(n)]
+        padded = f"{'_' * (n - 1)}{cleaned}{'_' * (n - 1)}"
+        for i in range(len(padded) - n + 1):
+            ng = padded[i:i + n]
+            v = table.get(ng)
+            if v is not None:
+                hit_count += 1
+                lp_sum += v
+
+    avg_lp = lp_sum / hit_count if hit_count > 0 else 0.0
+    return hit_count, avg_lp
+
+
+def score(
+    token: str, candidates: frozenset[str] | None = None
+) -> tuple[str, float] | None:
+    """Score token using discriminative character n-gram hit counts.
+
+    Each language's profile contains only n-grams exclusive to that language
+    (built with a log-prob margin vs. all competitor languages + English).
+    The winner is the language that matches the most of the token's n-grams.
+
+    Parameters
+    ----------
+    token:
+        Raw token text; cleaned internally.
+    candidates:
+        Restrict scoring to languages in this set that also have n-gram models.
+
+    Returns
+    -------
+    (lang, confidence) with confidence ∈ [0.60, 0.95], or None if no model
+    wins convincingly.
+    """
+    _load()
+
+    if not _MODELS:
+        return None
+
+    cleaned = _CLEAN.sub("", token.lower()).strip("-'")
+    if len(cleaned) < _MIN_TOKEN_LEN:
+        return None
+
+    langs_to_score = set(_MODELS.keys())
+    if candidates is not None:
+        langs_to_score &= candidates
+    if not langs_to_score:
+        return None
+
+    results: list[tuple[int, float, str]] = []
+    for lang in langs_to_score:
+        hits, avg_lp = _hit_score(cleaned, _MODELS[lang], _NGRAM_SIZES[lang])
+        results.append((hits, avg_lp, lang))
+
+    # Sort: primary by hit count (desc), secondary by avg log-prob (desc)
+    results.sort(key=lambda x: (x[0], x[1]), reverse=True)
+    best_hits, best_avg, best_lang = results[0]
+
+    if best_hits < _MIN_HITS:
+        return None
+
+    if len(results) > 1 and (best_hits - results[1][0]) < _MIN_GAP_HITS:
+        return None
+
+    # Map hit count to confidence: 2 hits → 0.65, 10+ hits → 0.90
+    confidence = max(0.60, min(0.95, 0.60 + best_hits * 0.03))
+    return best_lang, confidence
+
+
+def available_languages() -> frozenset[str]:
+    """Return languages for which a model file is present."""
+    _load()
+    return frozenset(_MODELS.keys())

polystring/_pipeline/stage1_preprocess.py

@@ -0,0 +1,134 @@
+from __future__ import annotations
+
+import re
+import unicodedata
+from dataclasses import dataclass
+
+import regex as _regex
+
+_SPECIAL_PATTERNS: list[tuple[str, re.Pattern]] = [
+    ("url",     re.compile(r"https?://\S+|www\.\S+")),
+    ("mention", re.compile(r"@\w+")),
+    ("hashtag", re.compile(r"#\w+")),
+    ("emoji",   _regex.compile(
+        "[\U0001F600-\U0001F64F"
+        "\U0001F300-\U0001F5FF"
+        "\U0001F680-\U0001F6FF"
+        "\U0001F1E0-\U0001F1FF"
+        "\U00002702-\U000027B0"
+        "\U0001F900-\U0001F9FF"
+        "☀-⛿"
+        "✀-➿]+"
+    )),
+    ("num",     re.compile(r"\b\d+[a-zA-Z]*\b|\b[a-zA-Z]*\d+\b")),
+]
+
+
+@dataclass
+class SpecialToken:
+    text: str
+    token_type: str   # url | mention | hashtag | emoji | num
+    start: int
+    end: int
+    hashtag_lang: str | None = None
+    hashtag_confidence: float = 0.0
+
+
+@dataclass
+class RawToken:
+    text: str
+    start: int
+    end: int
+    is_ne_candidate: bool = False
+
+
+@dataclass
+class Stage1Result:
+    linguistic_tokens: list[RawToken]
+    special_tokens: list[SpecialToken]
+    normalized_text: str    # NFC-normalized original (offsets valid against this)
+
+
+def _nfc(text: str) -> str:
+    return unicodedata.normalize("NFC", text)
+
+
+def _extract_special_tokens(text: str) -> tuple[list[SpecialToken], str]:
+    """Extract special tokens and replace them with whitespace-width placeholders.
+
+    Returns (special_tokens, masked_text) where masked_text has the same byte
+    offsets but non-linguistic tokens replaced with spaces so downstream
+    tokenisation still splits correctly.
+    """
+    specials: list[SpecialToken] = []
+    chars = list(text)
+
+    for tok_type, pattern in _SPECIAL_PATTERNS:
+        for m in pattern.finditer(text):
+            already = any(s.start <= m.start() < s.end for s in specials)
+            if already:
+                continue
+            specials.append(SpecialToken(
+                text=m.group(),
+                token_type=tok_type,
+                start=m.start(),
+                end=m.end(),
+            ))
+            for i in range(m.start(), m.end()):
+                chars[i] = " "
+
+    specials.sort(key=lambda s: s.start)
+    masked = "".join(chars)
+    return specials, masked
+
+
+def _tokenize(masked_text: str) -> list[tuple[str, int, int]]:
+    """Split masked text into (token, start, end) by whitespace and punctuation."""
+    tokens: list[tuple[str, int, int]] = []
+    for m in re.finditer(r"\S+", masked_text):
+        token_text = m.group()
+        stripped = token_text.strip(".,!?;:\"'()[]{}")
+        if not stripped:
+            continue
+        offset = token_text.index(stripped[0]) if stripped else 0
+        end = m.start() + offset + len(stripped)
+        tokens.append((stripped, m.start() + offset, end))
+    return tokens
+
+
+def _is_ne_candidate(token: str, idx: int, tokens: list[tuple[str, int, int]]) -> bool:
+    """True if a mid-sentence capitalised token that may be a named entity."""
+    if idx == 0:
+        return False
+    if not token[0].isupper():
+        return False
+    if idx > 0:
+        prev = tokens[idx - 1][0]
+        if prev.endswith((".", "!", "?")):
+            return False
+    return True
+
+
+def run(
+    text: str,
+    normalize: bool = True,
+) -> Stage1Result:
+    """Stage 1: extract special tokens, NFC normalise, tokenise, tag NE candidates."""
+    normalized = _nfc(text) if normalize else text
+    specials, masked = _extract_special_tokens(normalized)
+
+    raw_tokens_raw = _tokenize(masked)
+    linguistic_tokens: list[RawToken] = []
+    for idx, (tok, start, end) in enumerate(raw_tokens_raw):
+        if not tok.strip():
+            continue
+        is_ne = _is_ne_candidate(tok, idx, raw_tokens_raw)
+        linguistic_tokens.append(
+            RawToken(text=tok, start=start, end=end, is_ne_candidate=is_ne)
+        )
+
+    return Stage1Result(
+        linguistic_tokens=linguistic_tokens,
+        special_tokens=specials,
+        normalized_text=normalized,
+    )