diff --git a/docs/indic-encoders.md b/docs/indic-encoders.md new file mode 100644 index 000000000..a15d48ae5 --- /dev/null +++ b/docs/indic-encoders.md @@ -0,0 +1,92 @@ +# Indic encoder backbones + +OpenMed can expose a user-supplied MuRIL or IndicBERT backbone to an Indic +entity adapter. This is an opt-in embedding contract, not a replacement for the +complete Hindi and Telugu PII predictors used by `extract_pii` and +`deidentify`. + +OpenMed does not bundle encoder weights and does not resolve a repository until +you configure one. Listing PII models never downloads weights. A missing +configuration, optional dependency, or checkpoint returns a skip reason so +offline and base-package imports remain deterministic. + +## Supported families + +| Family | License | Capability | +|---|---|---| +| MuRIL | Apache-2.0 | 17 Indian languages plus transliterated counterparts | +| IndicBERT | MIT | 12 languages | + +MuRIL was trained with transliterated segment pairs. It is therefore the +preferred backbone for Hinglish and other code-mixed paths where an Indian +language is written in Latin script. A backbone only produces aligned hidden +states; an application still needs an entity adapter or task head evaluated on +its own synthetic and governed data. + +## Configure and load + +Install the optional runtime without changing the base package footprint: + +```bash +pip install "openmed[hf]" torch +``` + +Then configure either a repository identifier that you have chosen or an +existing local directory: + +```python +from openmed.core.model_registry import ( + configure_indic_encoder, + get_pii_models_by_language, + load_configured_indic_encoder, +) + +configure_indic_encoder( + "google/muril-base-cased", + family="muril", + languages=("hi", "te"), + local_files_only=True, +) + +# Existing dedicated predictors remain present; the configured backbone is +# added only when its optional runtime dependencies are importable. +models = get_pii_models_by_language("hi") + +resolution = load_configured_indic_encoder() +if not resolution.available: + print(resolution.skip_reason) +else: + encoded = resolution.handle.encode("Patient Asha Verma ka record dekhiye.") + print(encoded.last_hidden_state.shape, encoded.offset_mapping) +``` + +Set `local_files_only=False` only when you intentionally want the explicitly +configured repository to be resolved through the model host. Private or gated +repositories can use the `token=` argument. Tokens are held only in the +process-local configuration and are excluded from its representation. + +The encoder output contains tokenizer tensors, an attention mask, character +offsets, final hidden states, and a SHA-256 input identifier. It contains no raw +input text. Encoding diagnostics log only that digest and aggregate lengths. + +Use `clear_indic_encoder_config()` to remove the process-local configuration. + +## Model-card provenance + +When publishing metadata for an adapter that consumes one of these backbones, +add an `encoder_provenance` object to the manifest row passed to the model-card +renderer: + +```json +{ + "family": "MuRIL", + "source": "google/muril-base-cased", + "license": "Apache-2.0", + "provenance": "user-supplied", + "weights": "user-supplied; not bundled", + "supports_transliterated_text": true +} +``` + +The rendered model card includes these values in a dedicated Encoder +Provenance section. diff --git a/mkdocs.yml b/mkdocs.yml index 4c88f4c67..3f7ae903d 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -33,6 +33,7 @@ nav: - ModelLoader & Pipelines: model-loader.md - Model Registry: model-registry.md - Model Manifest: model-manifest.md + - Indic Encoder Backbones: indic-encoders.md - Advanced NER & Output Formatting: output-formatting.md - Clinical Context & Extraction Depth: clinical/context-and-extraction.md - Medication Sig Parser: clinical/sig-parser.md diff --git a/openmed/core/model_card.py b/openmed/core/model_card.py index 2a08baea8..ab1da1441 100644 --- a/openmed/core/model_card.py +++ b/openmed/core/model_card.py @@ -215,6 +215,9 @@ def render_model_card(row: dict[str, Any]) -> str: training_provenance_lines = _training_provenance_block(row) if training_provenance_lines: lines.extend(["", "## Training Provenance", "", *training_provenance_lines]) + encoder_provenance_lines = _encoder_provenance_block(row) + if encoder_provenance_lines: + lines.extend(["", "## Encoder Provenance", "", *encoder_provenance_lines]) return "\n".join(lines) + "\n" @@ -769,6 +772,41 @@ def _training_provenance_block(row: dict[str, Any]) -> list[str]: return lines +def _encoder_provenance_block(row: dict[str, Any]) -> list[str]: + payload = row.get("encoder_provenance") + if not isinstance(payload, dict): + payload = row.get("indic_encoder") + if not isinstance(payload, dict): + return [] + + family = _string(payload.get("family"), "Not specified") + source = _string( + payload.get("source") or payload.get("repo_id"), + "Not specified", + ) + license_name = _string(payload.get("license"), "Not specified") + provenance = _string(payload.get("provenance"), "user-supplied") + weights = _string(payload.get("weights"), "user-supplied; not bundled") + transliterated = payload.get("supports_transliterated_text") + transliterated_value = ( + "Yes" + if transliterated is True + else "No" + if transliterated is False + else "Not specified" + ) + return [ + "| Field | Value |", + "|---|---|", + f"| Encoder family | {family} |", + f"| Source | `{source}` |", + f"| License | {license_name} |", + f"| Provenance | {provenance} |", + f"| Weights | {weights} |", + f"| Transliterated text | {transliterated_value} |", + ] + + def _gate_status(value: Any) -> str: if value is True: return "passed" diff --git a/openmed/core/model_registry.py b/openmed/core/model_registry.py index a3963c770..7d7d30cbc 100644 --- a/openmed/core/model_registry.py +++ b/openmed/core/model_registry.py @@ -2,12 +2,13 @@ from __future__ import annotations +import importlib.util import json import re from dataclasses import dataclass, field from itertools import chain from pathlib import Path -from typing import Any, Dict, Iterable, List, Optional, Tuple +from typing import Any, Dict, Iterable, List, Mapping, Optional, Tuple @dataclass @@ -38,6 +39,7 @@ class ModelInfo: license: Optional[str] = None reproducibility_hash: Optional[str] = None released: Optional[str] = None + provenance: Dict[str, Any] = field(default_factory=dict) @property def size_mb(self) -> Optional[int]: @@ -102,6 +104,88 @@ def permissive_license(self) -> bool: return self.license.lower() in PERMISSIVE_DRAFT_MODEL_LICENSES +@dataclass(frozen=True) +class IndicEncoderConfig: + """Explicit user configuration for an optional Indic encoder backbone.""" + + source: str + family: str + languages: tuple[str, ...] + cache_dir: Optional[str] = None + token: Optional[str] = field(default=None, repr=False, compare=False) + revision: Optional[str] = None + local_files_only: bool = False + device: Optional[str] = None + + +@dataclass(frozen=True) +class IndicEncoderSpec: + """Permissive-license metadata for a supported Indic encoder family.""" + + family: str + display_name: str + license_id: str + source_url: str + languages: frozenset[str] + supports_transliterated_text: bool = False + + +@dataclass(frozen=True) +class IndicEncoderLoadResult: + """Result of resolving optional encoder weights without hard failure.""" + + handle: Any | None + skip_reason: str | None + metadata: IndicEncoderSpec | None = None + + @property + def available(self) -> bool: + """Return whether a tokenizer/backbone handle was loaded.""" + + return self.handle is not None + + +INDIC_ENCODER_SPECS: Mapping[str, IndicEncoderSpec] = { + "muril": IndicEncoderSpec( + family="muril", + display_name="MuRIL", + license_id="Apache-2.0", + source_url="https://huggingface.co/google/muril-base-cased", + languages=frozenset( + { + "as", + "bn", + "en", + "gu", + "hi", + "kn", + "ks", + "ml", + "mr", + "ne", + "or", + "pa", + "sa", + "sd", + "ta", + "te", + "ur", + } + ), + supports_transliterated_text=True, + ), + "indicbert": IndicEncoderSpec( + family="indicbert", + display_name="IndicBERT", + license_id="MIT", + source_url="https://huggingface.co/ai4bharat/indic-bert", + languages=frozenset( + {"as", "bn", "en", "gu", "hi", "kn", "ml", "mr", "or", "pa", "ta", "te"} + ), + ), +} + + MANIFEST_PATH = Path(__file__).resolve().parents[2] / "models.jsonl" PERMISSIVE_DRAFT_MODEL_LICENSES = frozenset( @@ -128,6 +212,8 @@ def permissive_license(self) -> bool: ), } +_INDIC_ENCODER_CONFIG: Optional[IndicEncoderConfig] = None + _GENERATION_DRAFT_TARGET_ALIASES = { "laneformer-2b-it": "OpenMed/laneformer-2b-it-q4-mlx", "kogai/laneformer-2b-it": "OpenMed/laneformer-2b-it-q4-mlx", @@ -540,6 +626,9 @@ def _model_info_from_row(row: Dict[str, Any]) -> ModelInfo: license=row.get("license"), reproducibility_hash=row.get("reproducibility_hash"), released=row.get("released"), + provenance=dict(row.get("provenance") or {}) + if isinstance(row.get("provenance"), dict) + else {}, ) @@ -962,11 +1051,181 @@ def get_entity_types_by_category(category: str) -> List[str]: return sorted(entity_types) +def get_indic_encoder_spec( + family: str | None = None, + *, + source: str | None = None, +) -> IndicEncoderSpec: + """Return normalized family metadata, inferring from a known repo name.""" + + if family: + normalized = family.strip().casefold().replace("-", "").replace("_", "") + else: + normalized_source = (source or "").casefold().replace("-", "") + if "muril" in normalized_source: + normalized = "muril" + elif "indicbert" in normalized_source: + normalized = "indicbert" + else: + raise ValueError( + "family must be 'muril' or 'indicbert' for an unrecognized source" + ) + try: + return INDIC_ENCODER_SPECS[normalized] + except KeyError as exc: + allowed = ", ".join(sorted(INDIC_ENCODER_SPECS)) + raise ValueError( + f"unsupported Indic encoder family {family!r}; expected {allowed}" + ) from exc + + +def configure_indic_encoder( + source: str | Path, + *, + family: str | None = None, + languages: Iterable[str] | None = None, + cache_dir: str | Path | None = None, + token: str | None = None, + revision: str | None = None, + local_files_only: bool = False, + device: str | None = None, +) -> IndicEncoderConfig: + """Configure user-supplied MuRIL or IndicBERT weights for PII discovery. + + Configuration does not load or download weights. Loading happens only when + :func:`load_configured_indic_encoder` is called explicitly. + """ + + source_value = str(source).strip() + if not source_value: + raise ValueError("Indic encoder source must not be empty") + metadata = get_indic_encoder_spec(family, source=source_value) + + if languages is None: + from .pii_i18n import INDIC_ENCODER_PII_LANGUAGES + + normalized_languages = tuple(sorted(INDIC_ENCODER_PII_LANGUAGES)) + else: + normalized_languages = tuple( + sorted({str(language).strip().casefold() for language in languages}) + ) + if not normalized_languages or any( + not language for language in normalized_languages + ): + raise ValueError("at least one Indic encoder language must be configured") + unsupported = set(normalized_languages) - metadata.languages + if unsupported: + raise ValueError( + f"{metadata.display_name} does not support language codes: " + + ", ".join(sorted(unsupported)) + ) + + config = IndicEncoderConfig( + source=source_value, + family=metadata.family, + languages=normalized_languages, + cache_dir=str(cache_dir) if cache_dir is not None else None, + token=token, + revision=revision, + local_files_only=bool(local_files_only), + device=device, + ) + global _INDIC_ENCODER_CONFIG + _INDIC_ENCODER_CONFIG = config + return config + + +def get_indic_encoder_config() -> IndicEncoderConfig | None: + """Return the active Indic encoder configuration, if one was supplied.""" + + return _INDIC_ENCODER_CONFIG + + +def clear_indic_encoder_config() -> None: + """Remove the process-local Indic encoder configuration.""" + + global _INDIC_ENCODER_CONFIG + _INDIC_ENCODER_CONFIG = None + + +def load_configured_indic_encoder() -> IndicEncoderLoadResult: + """Load the configured encoder or return a result carrying a skip reason.""" + + config = get_indic_encoder_config() + if config is None: + return IndicEncoderLoadResult( + handle=None, + skip_reason="no Indic encoder weights are configured", + ) + + from openmed.ner.families.indic import load_indic_encoder + + return load_indic_encoder( + config.source, + family=config.family, + cache_dir=config.cache_dir, + token=config.token, + revision=config.revision, + local_files_only=config.local_files_only, + device=config.device, + ) + + +def _configured_indic_pii_models(lang: str) -> Dict[str, ModelInfo]: + config = get_indic_encoder_config() + if config is None or lang not in config.languages: + return {} + + if not _is_indic_encoder_runtime_available(): + return {} + metadata = get_indic_encoder_spec(config.family) + key = f"pii_{lang}_{metadata.family}_encoder" + transliteration = ( + " Includes transliterated-text support for Hinglish and code-mixed paths." + if metadata.supports_transliterated_text + else "" + ) + return { + key: ModelInfo( + model_id=config.source, + display_name=f"{metadata.display_name} Indic encoder adapter", + category="Privacy", + specialization="Indic encoder-backed PII", + description=( + "Optional backbone-only encoder configured from user-supplied weights." + + transliteration + ), + entity_types=list(_PII_ENTITY_TYPES), + size_category="Unknown", + family=metadata.display_name, + task="feature-extraction", + languages=list(config.languages), + architecture="Transformer encoder", + base_model=config.source, + formats=["transformers"], + license=metadata.license_id, + provenance={ + "source": config.source, + "weights": "user-supplied", + "bundled": False, + "revision": config.revision, + }, + ) + } + + +def _is_indic_encoder_runtime_available() -> bool: + return all( + importlib.util.find_spec(module_name) is not None + for module_name in ("transformers", "torch") + ) + + def get_pii_models_by_language(lang: str) -> Dict[str, ModelInfo]: - """Return all single-language PII models for a given language.""" + """Return PII models plus an explicitly configured Indic encoder adapter.""" if lang == "en": localized_prefixes = _LOCALIZED_PII_LANGUAGE_KEYS - return { + language_models = { key: info for key, info in OPENMED_MODELS.items() if key.startswith("pii_") @@ -974,33 +1233,35 @@ def get_pii_models_by_language(lang: str) -> Dict[str, ModelInfo]: and "en" in (info.languages or ["en"]) and not any(key.startswith(f"pii_{lc}_") for lc in localized_prefixes) } + else: + prefix = f"pii_{lang}_" + language_models = { + key: info + for key, info in OPENMED_MODELS.items() + if key.startswith(prefix) + and info.category == "Privacy" + and lang in (info.languages or []) + } - prefix = f"pii_{lang}_" - language_models = { - key: info - for key, info in OPENMED_MODELS.items() - if key.startswith(prefix) - and info.category == "Privacy" - and lang in (info.languages or []) - } - if language_models: - return language_models - - from .pii_i18n import DEFAULT_PII_MODELS - - default_model_id = DEFAULT_PII_MODELS.get(lang) - if not default_model_id: - return {} - # Internal fallback: DEFAULT_PII_MODELS is the validated source of truth, - # so language-pack callers can reuse multilingual privacy filters safely. - return { - key: info - for key, info in OPENMED_MODELS.items() - if key.startswith("pii_") - and info.category == "Privacy" - and info.model_id == default_model_id - and lang in (info.languages or []) - } + if not language_models: + from .pii_i18n import DEFAULT_PII_MODELS + + default_model_id = DEFAULT_PII_MODELS.get(lang) + if default_model_id: + # Internal fallback: DEFAULT_PII_MODELS is the validated source of + # truth, so language-pack callers can reuse multilingual privacy + # filters safely. + language_models = { + key: info + for key, info in OPENMED_MODELS.items() + if key.startswith("pii_") + and info.category == "Privacy" + and info.model_id == default_model_id + and lang in (info.languages or []) + } + + language_models.update(_configured_indic_pii_models(lang)) + return language_models def get_default_pii_model(lang: str) -> Optional[str]: diff --git a/openmed/core/pii_i18n.py b/openmed/core/pii_i18n.py index 5f6e29925..8371d29cf 100644 --- a/openmed/core/pii_i18n.py +++ b/openmed/core/pii_i18n.py @@ -135,6 +135,12 @@ "ro": "OpenMed/privacy-filter-multilingual", } +# Backbone-only Indic encoders are an opt-in supplement, not replacements for +# the complete Hindi and Telugu token-classification checkpoints above. Other +# family-supported language codes can be enabled explicitly through +# ``configure_indic_encoder(..., languages=...)`` in ``model_registry``. +INDIC_ENCODER_PII_LANGUAGES: frozenset[str] = frozenset({"hi", "te"}) + # --------------------------------------------------------------------------- # Financial Identifier Validators diff --git a/openmed/eval/datasets/__init__.py b/openmed/eval/datasets/__init__.py index 6aa0a2e02..9ef6cac2d 100644 --- a/openmed/eval/datasets/__init__.py +++ b/openmed/eval/datasets/__init__.py @@ -74,7 +74,14 @@ load_i2b2_deid, map_i2b2_phi_tag, ) -from .licenses import PUBLIC_DATASET_LICENSES, DatasetLicense, license_for +from .licenses import ( + PERMISSIVE_ENCODER_LICENSES, + PUBLIC_DATASET_LICENSES, + DatasetLicense, + EncoderLicense, + encoder_license_for, + license_for, +) from .multilingual_ner import ( CANTEMIST, CMEEE, @@ -122,6 +129,7 @@ "DUACorpusStub", "DUACredentialRequired", "DatasetLicense", + "EncoderLicense", "DatasetLoadResult", "DatasetUnavailable", "DRUGPROT", @@ -167,6 +175,7 @@ "MultilingualNerSpan", "PUBLIC_DATASETS", "PUBLIC_DATASET_LICENSES", + "PERMISSIVE_ENCODER_LICENSES", "PUBLIC_LABEL_MAPS", "PublicDatasetAdapter", "PublicDatasetRecord", @@ -182,6 +191,7 @@ "clinical_phi_manifest_hash", "i2b2_suite_metadata", "license_for", + "encoder_license_for", "load_clinical_phi_manifest", "load_biomedical_ner_corpus", "load_biomedical_ner_fixtures", diff --git a/openmed/eval/datasets/licenses.py b/openmed/eval/datasets/licenses.py index 0ea1d0f67..9bccb3bc4 100644 --- a/openmed/eval/datasets/licenses.py +++ b/openmed/eval/datasets/licenses.py @@ -24,6 +24,26 @@ def to_dict(self) -> dict[str, str]: } +@dataclass(frozen=True) +class EncoderLicense: + """License and provenance policy for an optional encoder backbone.""" + + family: str + license_id: str + source_url: str + redistribution: str + notes: str = "" + + def to_dict(self) -> dict[str, str]: + return { + "family": self.family, + "license_id": self.license_id, + "source_url": self.source_url, + "redistribution": self.redistribution, + "notes": self.notes, + } + + PUBLIC_DATASET_LICENSES: Mapping[str, DatasetLicense] = { "golden": DatasetLicense( dataset="golden", @@ -112,6 +132,27 @@ def to_dict(self) -> dict[str, str]: } +PERMISSIVE_ENCODER_LICENSES: Mapping[str, EncoderLicense] = { + "muril": EncoderLicense( + family="MuRIL", + license_id="Apache-2.0", + source_url="https://huggingface.co/google/muril-base-cased", + redistribution="user-supplied-reference-only", + notes=( + "OpenMed bundles no weights. MuRIL supports 17 Indian languages " + "and transliterated text, including Hinglish/code-mixed paths." + ), + ), + "indicbert": EncoderLicense( + family="IndicBERT", + license_id="MIT", + source_url="https://huggingface.co/ai4bharat/indic-bert", + redistribution="user-supplied-reference-only", + notes="OpenMed bundles no weights; users resolve an approved repo or local path.", + ), +} + + def license_for(dataset: str) -> DatasetLicense: try: return PUBLIC_DATASET_LICENSES[dataset] @@ -119,4 +160,21 @@ def license_for(dataset: str) -> DatasetLicense: raise ValueError(f"unknown public dataset: {dataset}") from exc -__all__ = ["DatasetLicense", "PUBLIC_DATASET_LICENSES", "license_for"] +def encoder_license_for(family: str) -> EncoderLicense: + """Return permissive license metadata for a supported Indic encoder.""" + + normalized = family.strip().casefold().replace("-", "").replace("_", "") + try: + return PERMISSIVE_ENCODER_LICENSES[normalized] + except KeyError as exc: + raise ValueError(f"unknown permissive encoder family: {family}") from exc + + +__all__ = [ + "DatasetLicense", + "EncoderLicense", + "PERMISSIVE_ENCODER_LICENSES", + "PUBLIC_DATASET_LICENSES", + "encoder_license_for", + "license_for", +] diff --git a/openmed/eval/suites/__init__.py b/openmed/eval/suites/__init__.py index 950f9f8e7..f85b6be6e 100644 --- a/openmed/eval/suites/__init__.py +++ b/openmed/eval/suites/__init__.py @@ -36,6 +36,12 @@ ) from openmed.eval.golden import load_benchmark_fixtures from openmed.eval.harness import BenchmarkFixture +from openmed.eval.suites.indic_encoder import ( + INDIC_ENCODER_RECALL_DELTA, + IndicRecallDeltaReport, + SyntheticIndicFixture, + run_indic_encoder_recall_delta, +) from openmed.eval.suites.multimodal_dicom import ( MULTIMODAL_DICOM, generate_synthetic_dicom_corpus, @@ -148,6 +154,7 @@ def suite_metadata(name: str, **kwargs: Any) -> dict[str, Any]: "BIOMEDICAL_NER", "MULTILINGUAL_NER", "MULTIMODAL_DICOM", + "INDIC_ENCODER_RECALL_DELTA", "RELATIONS", "RelationFixture", "RelationTrap", @@ -155,12 +162,15 @@ def suite_metadata(name: str, **kwargs: Any) -> dict[str, Any]: "ComparatorMatrixReport", "ComparatorMatrixRow", "ComparatorUnavailable", + "IndicRecallDeltaReport", + "SyntheticIndicFixture", "DEFAULT_SUITES", "validate_suite_name", "load_benchmark_fixtures", "load_suite_fixtures", "suite_metadata", "run_comparator_matrix", + "run_indic_encoder_recall_delta", "load_i2b2_deid", "i2b2_suite_metadata", "biomedical_ner_suite_metadata", diff --git a/openmed/eval/suites/indic_encoder.py b/openmed/eval/suites/indic_encoder.py new file mode 100644 index 000000000..92bc0c578 --- /dev/null +++ b/openmed/eval/suites/indic_encoder.py @@ -0,0 +1,210 @@ +"""Offline synthetic recall-delta check for Indic encoder-backed NER.""" + +from __future__ import annotations + +import re +from dataclasses import dataclass +from typing import Any, Protocol, Sequence + +from openmed.core.pii_i18n import get_patterns_for_language +from openmed.eval.metrics import EvalSpan, compute_leakage_rate, normalize_eval_spans +from openmed.ner.families.base import EncoderOutput, SupportsEncoding + +_PATTERN_LABELS = { + "date": "DATE", + "phone_number": "PHONE", + "national_id": "ID_NUM", + "street_address": "STREET_ADDRESS", + "postcode": "ZIPCODE", +} + + +@dataclass(frozen=True) +class SyntheticIndicFixture: + """One synthetic Indic or code-mixed note with exact gold spans.""" + + fixture_id: str + language: str + text: str + gold_spans: tuple[EvalSpan, ...] + + +@dataclass(frozen=True) +class IndicRecallDeltaReport: + """Entity recall and leakage comparison against the regex-only path.""" + + gold_entities: int + regex_matched_entities: int + encoder_matched_entities: int + regex_recall: float + encoder_recall: float + recall_delta: float + entity_leakage: int + encoder_leakage_rate: float + + def to_dict(self) -> dict[str, int | float]: + """Return a JSON-serializable report.""" + + return { + "gold_entities": self.gold_entities, + "regex_matched_entities": self.regex_matched_entities, + "encoder_matched_entities": self.encoder_matched_entities, + "regex_recall": self.regex_recall, + "encoder_recall": self.encoder_recall, + "recall_delta": self.recall_delta, + "entity_leakage": self.entity_leakage, + "encoder_leakage_rate": self.encoder_leakage_rate, + } + + +class IndicNerAdapter(Protocol): # pragma: no cover - interface only + """Adapter contract used by the deterministic Indic recall suite.""" + + encoder: SupportsEncoding + + def predict_entities( + self, + text: str, + *, + language: str, + encoding: EncoderOutput, + ) -> Sequence[Any]: ... + + +def _fixture( + fixture_id: str, + language: str, + text: str, + entities: Sequence[tuple[str, str]], +) -> SyntheticIndicFixture: + spans = tuple( + EvalSpan( + start=text.index(value), + end=text.index(value) + len(value), + label=label, + text=value, + language=language, + ) + for value, label in entities + ) + return SyntheticIndicFixture(fixture_id, language, text, spans) + + +SYNTHETIC_INDIC_GOLD: tuple[SyntheticIndicFixture, ...] = ( + _fixture( + "hi-devanagari", + "hi", + "रोगी आरव मेहता का फ़ोन 9876543210 है।", + (("आरव मेहता", "PERSON"), ("9876543210", "PHONE")), + ), + _fixture( + "hi-hinglish", + "hi", + "Patient Asha Verma ka phone 9123456780 hai.", + (("Asha Verma", "PERSON"), ("9123456780", "PHONE")), + ), + _fixture( + "te-native", + "te", + "రోగి అనన్య రెడ్డి ఫోన్ 9988776655.", + (("అనన్య రెడ్డి", "PERSON"), ("9988776655", "PHONE")), + ), +) + + +def run_indic_encoder_recall_delta( + adapter: IndicNerAdapter, + *, + fixtures: Sequence[SyntheticIndicFixture] = SYNTHETIC_INDIC_GOLD, +) -> IndicRecallDeltaReport: + """Compare an encoder-backed adapter with deterministic regex detection.""" + + gold_count = 0 + regex_matches = 0 + encoder_matches = 0 + encoder_leaked_chars = 0 + gold_chars = 0 + + for fixture in fixtures: + encoding = adapter.encoder.encode(fixture.text) + encoding.validate() + predicted = normalize_eval_spans( + adapter.predict_entities( + fixture.text, + language=fixture.language, + encoding=encoding, + ), + default_language=fixture.language, + source_text=fixture.text, + ) + regex_predicted = _regex_spans(fixture) + gold_count += len(fixture.gold_spans) + regex_matches += _matched_gold_count(fixture.gold_spans, regex_predicted) + encoder_matches += _matched_gold_count(fixture.gold_spans, predicted) + leakage = compute_leakage_rate( + fixture.gold_spans, + predicted, + default_language=fixture.language, + source_text=fixture.text, + ) + encoder_leaked_chars += leakage.leaked_chars + gold_chars += leakage.total_chars + + regex_recall = regex_matches / gold_count if gold_count else 1.0 + encoder_recall = encoder_matches / gold_count if gold_count else 1.0 + return IndicRecallDeltaReport( + gold_entities=gold_count, + regex_matched_entities=regex_matches, + encoder_matched_entities=encoder_matches, + regex_recall=regex_recall, + encoder_recall=encoder_recall, + recall_delta=encoder_recall - regex_recall, + entity_leakage=gold_count - encoder_matches, + encoder_leakage_rate=(encoder_leaked_chars / gold_chars if gold_chars else 0.0), + ) + + +def _regex_spans(fixture: SyntheticIndicFixture) -> list[EvalSpan]: + observed: dict[tuple[int, int, str], EvalSpan] = {} + for pattern in get_patterns_for_language(fixture.language): + label = _PATTERN_LABELS.get(pattern.entity_type) + if label is None: + continue + for match in re.finditer(pattern.pattern, fixture.text, pattern.flags): + value = match.group(0) + if pattern.validator is not None and not pattern.validator(value): + continue + key = (match.start(), match.end(), label) + observed[key] = EvalSpan( + start=match.start(), + end=match.end(), + label=label, + text=value, + language=fixture.language, + ) + return list(observed.values()) + + +def _matched_gold_count( + gold_spans: Sequence[EvalSpan], + predicted_spans: Sequence[EvalSpan], +) -> int: + predicted = { + (span.start, span.end, span.label.casefold()) for span in predicted_spans + } + return sum( + (span.start, span.end, span.label.casefold()) in predicted + for span in gold_spans + ) + + +__all__ = [ + "INDIC_ENCODER_RECALL_DELTA", + "IndicNerAdapter", + "IndicRecallDeltaReport", + "SYNTHETIC_INDIC_GOLD", + "SyntheticIndicFixture", + "run_indic_encoder_recall_delta", +] + +INDIC_ENCODER_RECALL_DELTA = "indic_encoder_recall_delta" diff --git a/openmed/ner/__init__.py b/openmed/ner/__init__.py index 822aad57c..ec426f8ef 100644 --- a/openmed/ner/__init__.py +++ b/openmed/ner/__init__.py @@ -9,11 +9,18 @@ from .adapter import TokenAnnotation, TokenClassificationResult, to_token_classification from .exceptions import MissingDependencyError from .families import ( + INDIC_ENCODER_SPECS, + EncoderOutput, + IndicEncoderLoadResult, + IndicEncoderSpec, ModelFamily, + SupportsEncoding, ensure_gliner2_available, ensure_gliner_available, is_gliner2_available, is_gliner_available, + is_indic_encoder_available, + load_indic_encoder, ) from .indexing import ( ModelIndex, @@ -39,6 +46,13 @@ "write_index", "load_index", "ModelFamily", + "EncoderOutput", + "SupportsEncoding", + "INDIC_ENCODER_SPECS", + "IndicEncoderSpec", + "IndicEncoderLoadResult", + "is_indic_encoder_available", + "load_indic_encoder", "MissingDependencyError", "ensure_gliner_available", "is_gliner_available", diff --git a/openmed/ner/families/__init__.py b/openmed/ner/families/__init__.py index 90c998f2d..eeba13a61 100644 --- a/openmed/ner/families/__init__.py +++ b/openmed/ner/families/__init__.py @@ -2,7 +2,7 @@ from __future__ import annotations -from .base import ModelFamily +from .base import EncoderOutput, ModelFamily, SupportsEncoding from .gliner import ensure_gliner_available, is_gliner_available from .gliner2 import ( clear_gliner2_cache, @@ -10,13 +10,31 @@ is_gliner2_available, load_gliner2_handle, ) +from .indic import ( + INDIC_ENCODER_SPECS, + IndicEncoderHandle, + IndicEncoderLoadResult, + IndicEncoderSpec, + get_indic_encoder_spec, + is_indic_encoder_available, + load_indic_encoder, +) __all__ = [ "ModelFamily", + "EncoderOutput", + "SupportsEncoding", "ensure_gliner_available", "is_gliner_available", "ensure_gliner2_available", "is_gliner2_available", "load_gliner2_handle", "clear_gliner2_cache", + "INDIC_ENCODER_SPECS", + "IndicEncoderHandle", + "IndicEncoderLoadResult", + "IndicEncoderSpec", + "get_indic_encoder_spec", + "is_indic_encoder_available", + "load_indic_encoder", ] diff --git a/openmed/ner/families/base.py b/openmed/ner/families/base.py index 06d4cd325..9f03a70b4 100644 --- a/openmed/ner/families/base.py +++ b/openmed/ner/families/base.py @@ -2,8 +2,9 @@ from __future__ import annotations +from dataclasses import dataclass from enum import Enum -from typing import Protocol +from typing import Any, Mapping, Protocol class ModelFamily(str, Enum): @@ -11,13 +12,107 @@ class ModelFamily(str, Enum): GLINER = "gliner" GLINER2 = "gliner2" + MURIL = "muril" + INDICBERT = "indicbert" OTHER = "other" +@dataclass(frozen=True) +class EncoderOutput: + """Stable output contract for backbone-only text encoders. + + The contract intentionally omits the source text. Consumers receive token + IDs, an attention mask, character offsets, hidden states, and a one-way + digest that can be used for diagnostics without persisting raw input. + """ + + tokenizer_outputs: Mapping[str, Any] + offset_mapping: tuple[tuple[int, int], ...] + last_hidden_state: Any + text_sha256: str + + @property + def input_ids(self) -> Any: + """Return the tokenizer input IDs.""" + + return self.tokenizer_outputs["input_ids"] + + @property + def attention_mask(self) -> Any: + """Return the tokenizer attention mask.""" + + return self.tokenizer_outputs["attention_mask"] + + @property + def token_count(self) -> int: + """Return the aligned sequence length, including special tokens.""" + + return len(self.offset_mapping) + + @property + def hidden_size(self) -> int: + """Return the final hidden-state width.""" + + return _shape(self.last_hidden_state)[-1] + + def validate(self) -> None: + """Raise ``ValueError`` when tensor shapes or offsets are misaligned.""" + + missing = {"input_ids", "attention_mask"} - set(self.tokenizer_outputs) + if missing: + raise ValueError( + "encoder tokenizer outputs are missing: " + ", ".join(sorted(missing)) + ) + + input_shape = _shape(self.input_ids) + mask_shape = _shape(self.attention_mask) + hidden_shape = _shape(self.last_hidden_state) + if len(input_shape) != 2 or input_shape[0] != 1: + raise ValueError("encoder input_ids must have shape [1, sequence]") + if mask_shape != input_shape: + raise ValueError("encoder attention_mask must align with input_ids") + if len(hidden_shape) != 3 or hidden_shape[:2] != input_shape: + raise ValueError( + "encoder last_hidden_state must have shape [1, sequence, hidden]" + ) + if len(self.offset_mapping) != input_shape[1]: + raise ValueError("encoder offsets must align with the token sequence") + if any(start < 0 or end < start for start, end in self.offset_mapping): + raise ValueError("encoder offsets must be ordered, non-negative spans") + if len(self.text_sha256) != 64 or any( + char not in "0123456789abcdef" for char in self.text_sha256 + ): + raise ValueError("encoder text_sha256 must be a lowercase SHA-256 digest") + + +def _shape(value: Any) -> tuple[int, ...]: + shape = getattr(value, "shape", None) + if shape is not None: + return tuple(int(dimension) for dimension in shape) + if isinstance(value, (list, tuple)): + if not value: + return (0,) + return (len(value), *_shape(value[0])) + return () + + class SupportsPrediction(Protocol): # pragma: no cover - interface only """Protocol for minimal predictor objects.""" def predict(self, text: str, *, labels: list[str] | None = None) -> object: ... -__all__ = ["ModelFamily", "SupportsPrediction"] +class SupportsEncoding(Protocol): # pragma: no cover - interface only + """Protocol implemented by reusable tokenizer/backbone handles.""" + + tokenizer: Any + + def encode(self, text: str, *, max_length: int = 512) -> EncoderOutput: ... + + +__all__ = [ + "EncoderOutput", + "ModelFamily", + "SupportsEncoding", + "SupportsPrediction", +] diff --git a/openmed/ner/families/indic.py b/openmed/ner/families/indic.py new file mode 100644 index 000000000..f9cee2298 --- /dev/null +++ b/openmed/ner/families/indic.py @@ -0,0 +1,258 @@ +"""Optional MuRIL and IndicBERT backbone-only encoder integration. + +Weights are never bundled with OpenMed. A caller must provide a local path or +repository identifier before this module imports the optional runtime stack or +attempts to resolve any model files. +""" + +from __future__ import annotations + +import hashlib +import importlib +import logging +from contextlib import nullcontext +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Mapping + +from openmed.core.model_registry import ( + INDIC_ENCODER_SPECS, + IndicEncoderLoadResult, + IndicEncoderSpec, + get_indic_encoder_spec, +) + +from .base import EncoderOutput + +logger = logging.getLogger(__name__) + +_INSTALL_HINT = ( + "Install the Hugging Face extra and a PyTorch runtime, for example " + "`pip install 'openmed[hf]' torch`." +) + + +@dataclass +class IndicEncoderHandle: + """Loaded tokenizer/backbone pair implementing the encoder contract.""" + + source: str + metadata: IndicEncoderSpec + tokenizer: Any + model: Any + torch_module: Any + device: str | None = None + + def encode(self, text: str, *, max_length: int = 512) -> EncoderOutput: + """Encode one string without logging or returning the raw input text.""" + + if not isinstance(text, str): + raise TypeError("encoder input must be a string") + if max_length <= 0: + raise ValueError("max_length must be positive") + + encoded = self.tokenizer( + text, + return_offsets_mapping=True, + return_tensors="pt", + truncation=True, + max_length=max_length, + ) + tokenizer_outputs = dict(encoded) + raw_offsets = tokenizer_outputs.pop("offset_mapping", None) + if raw_offsets is None: + raise ValueError( + "configured tokenizer does not expose character offset mappings" + ) + offsets = _normalize_offsets(raw_offsets) + model_inputs = _move_to_device(tokenizer_outputs, self.device) + + inference_mode = getattr(self.torch_module, "inference_mode", None) + context = inference_mode() if callable(inference_mode) else nullcontext() + with context: + model_output = self.model(**model_inputs, return_dict=True) + + hidden_states = _last_hidden_state(model_output) + digest = hashlib.sha256(text.encode("utf-8")).hexdigest() + output = EncoderOutput( + tokenizer_outputs=model_inputs, + offset_mapping=offsets, + last_hidden_state=hidden_states, + text_sha256=digest, + ) + output.validate() + logger.debug( + "Encoded Indic input sha256=%s characters=%d tokens=%d", + digest, + len(text), + output.token_count, + ) + return output + + +def is_indic_encoder_available() -> bool: + """Return whether the optional Transformers and PyTorch stack imports.""" + + return ( + _optional_module("transformers") is not None + and _optional_module("torch") is not None + ) + + +def load_indic_encoder( + source: str | None, + *, + family: str | None = None, + cache_dir: str | Path | None = None, + token: str | None = None, + revision: str | None = None, + local_files_only: bool = False, + device: str | None = None, +) -> IndicEncoderLoadResult: + """Resolve user-supplied MuRIL or IndicBERT weights. + + A missing source, dependency, or checkpoint returns a deterministic skip + reason. Repository identifiers may download only after the caller supplies + one explicitly; existing filesystem paths are always loaded locally. + """ + + if source is None or not str(source).strip(): + return IndicEncoderLoadResult( + handle=None, + skip_reason="no Indic encoder weights are configured", + ) + + source_value = str(source).strip() + metadata = get_indic_encoder_spec(family, source=source_value) + transformers_module = _optional_module("transformers") + torch_module = _optional_module("torch") + if transformers_module is None or torch_module is None: + missing = "transformers" if transformers_module is None else "torch" + return IndicEncoderLoadResult( + handle=None, + skip_reason=f"optional dependency {missing!r} is unavailable; {_INSTALL_HINT}", + metadata=metadata, + ) + + tokenizer_loader = getattr(transformers_module, "AutoTokenizer", None) + model_loader = getattr(transformers_module, "AutoModel", None) + if tokenizer_loader is None or model_loader is None: + return IndicEncoderLoadResult( + handle=None, + skip_reason="the installed transformers package lacks AutoTokenizer/AutoModel", + metadata=metadata, + ) + + load_kwargs: dict[str, Any] = { + "local_files_only": bool(local_files_only or _is_existing_path(source_value)), + "trust_remote_code": False, + } + if cache_dir is not None: + load_kwargs["cache_dir"] = str(cache_dir) + if token: + load_kwargs["token"] = token + if revision: + load_kwargs["revision"] = revision + + try: + tokenizer = tokenizer_loader.from_pretrained( + source_value, + use_fast=True, + **load_kwargs, + ) + model = model_loader.from_pretrained(source_value, **load_kwargs) + if hasattr(model, "eval"): + model.eval() + if device and hasattr(model, "to"): + model = model.to(device) + except (ImportError, OSError, RuntimeError, ValueError) as exc: + return IndicEncoderLoadResult( + handle=None, + skip_reason=( + f"configured {metadata.display_name} weights could not be loaded " + f"({type(exc).__name__})" + ), + metadata=metadata, + ) + + return IndicEncoderLoadResult( + handle=IndicEncoderHandle( + source=source_value, + metadata=metadata, + tokenizer=tokenizer, + model=model, + torch_module=torch_module, + device=device, + ), + skip_reason=None, + metadata=metadata, + ) + + +def _optional_module(name: str) -> Any | None: + try: + return importlib.import_module(name) + except (ImportError, OSError): + return None + + +def _is_existing_path(source: str) -> bool: + try: + return Path(source).expanduser().exists() + except OSError: + return False + + +def _normalize_offsets(value: Any) -> tuple[tuple[int, int], ...]: + if hasattr(value, "detach"): + value = value.detach() + if hasattr(value, "cpu"): + value = value.cpu() + if hasattr(value, "tolist"): + value = value.tolist() + if ( + isinstance(value, (list, tuple)) + and len(value) == 1 + and isinstance(value[0], (list, tuple)) + ): + value = value[0] + if not isinstance(value, (list, tuple)): + raise ValueError("encoder offsets must be a single batched sequence") + try: + return tuple((int(offset[0]), int(offset[1])) for offset in value) + except (IndexError, TypeError, ValueError) as exc: + raise ValueError("encoder offsets contain an invalid span") from exc + + +def _move_to_device( + tokenizer_outputs: Mapping[str, Any], + device: str | None, +) -> dict[str, Any]: + if not device: + return dict(tokenizer_outputs) + return { + key: value.to(device) if hasattr(value, "to") else value + for key, value in tokenizer_outputs.items() + } + + +def _last_hidden_state(model_output: Any) -> Any: + hidden_states = getattr(model_output, "last_hidden_state", None) + if hidden_states is not None: + return hidden_states + if isinstance(model_output, Mapping) and "last_hidden_state" in model_output: + return model_output["last_hidden_state"] + if isinstance(model_output, (list, tuple)) and model_output: + return model_output[0] + raise ValueError("configured encoder did not return last_hidden_state") + + +__all__ = [ + "INDIC_ENCODER_SPECS", + "IndicEncoderHandle", + "IndicEncoderLoadResult", + "IndicEncoderSpec", + "get_indic_encoder_spec", + "is_indic_encoder_available", + "load_indic_encoder", +] diff --git a/tests/unit/core/test_indic_encoder_metadata.py b/tests/unit/core/test_indic_encoder_metadata.py new file mode 100644 index 000000000..646f968e8 --- /dev/null +++ b/tests/unit/core/test_indic_encoder_metadata.py @@ -0,0 +1,49 @@ +from openmed.core.model_card import render_model_card +from openmed.eval.datasets.licenses import ( + PERMISSIVE_ENCODER_LICENSES, + encoder_license_for, +) +from openmed.ner.families.indic import INDIC_ENCODER_SPECS + + +def _row(): + return { + "repo_id": "Example/indic-pii-adapter", + "family": "PII", + "task": "token-classification", + "languages": ["hi", "te"], + "canonical_labels": ["PERSON", "PHONE"], + "license": "Apache-2.0", + "formats": ["safetensors"], + "encoder_provenance": { + "family": "MuRIL", + "source": "google/muril-base-cased", + "license": "Apache-2.0", + "provenance": "user-supplied", + "weights": "user-supplied; not bundled", + "supports_transliterated_text": True, + }, + } + + +def test_permissive_encoder_license_registry_matches_loader_metadata(): + assert set(PERMISSIVE_ENCODER_LICENSES) == {"muril", "indicbert"} + assert encoder_license_for("MuRIL").license_id == "Apache-2.0" + assert encoder_license_for("indic-bert").license_id == "MIT" + for family, metadata in INDIC_ENCODER_SPECS.items(): + registered = encoder_license_for(family) + assert registered.license_id == metadata.license_id + assert registered.source_url == metadata.source_url + assert registered.redistribution == "user-supplied-reference-only" + + +def test_model_card_surfaces_encoder_license_and_provenance(): + card = render_model_card(_row()) + + assert "## Encoder Provenance" in card + assert "| Encoder family | MuRIL |" in card + assert "| Source | `google/muril-base-cased` |" in card + assert "| License | Apache-2.0 |" in card + assert "| Provenance | user-supplied |" in card + assert "| Weights | user-supplied; not bundled |" in card + assert "| Transliterated text | Yes |" in card diff --git a/tests/unit/eval/test_indic_encoder_suite.py b/tests/unit/eval/test_indic_encoder_suite.py new file mode 100644 index 000000000..3c723daa9 --- /dev/null +++ b/tests/unit/eval/test_indic_encoder_suite.py @@ -0,0 +1,56 @@ +from __future__ import annotations + +import hashlib + +from openmed.eval.metrics import EvalSpan +from openmed.eval.suites.indic_encoder import ( + SYNTHETIC_INDIC_GOLD, + run_indic_encoder_recall_delta, +) +from openmed.ner.families.base import EncoderOutput + + +class _SyntheticEncoder: + tokenizer = object() + + def encode(self, text: str, *, max_length: int = 512) -> EncoderOutput: + del max_length + return EncoderOutput( + tokenizer_outputs={ + "input_ids": [[1]], + "attention_mask": [[1]], + }, + offset_mapping=((0, len(text)),), + last_hidden_state=[[[1.0, 0.0]]], + text_sha256=hashlib.sha256(text.encode()).hexdigest(), + ) + + +class _SyntheticPerfectAdapter: + encoder = _SyntheticEncoder() + + def predict_entities(self, text, *, language, encoding): + assert encoding.hidden_size == 2 + fixture = next(item for item in SYNTHETIC_INDIC_GOLD if item.text == text) + assert fixture.language == language + return [ + EvalSpan( + start=span.start, + end=span.end, + label=span.label, + text=span.text, + language=language, + ) + for span in fixture.gold_spans + ] + + +def test_encoder_backed_synthetic_recall_beats_regex_with_zero_leakage(): + report = run_indic_encoder_recall_delta(_SyntheticPerfectAdapter()) + + assert report.gold_entities == 6 + assert report.encoder_recall == 1.0 + assert report.encoder_recall > report.regex_recall + assert report.recall_delta > 0.0 + assert report.entity_leakage == 0 + assert report.encoder_leakage_rate == 0.0 diff --git a/tests/unit/ner/test_indic_encoder.py b/tests/unit/ner/test_indic_encoder.py new file mode 100644 index 000000000..1cd94b444 --- /dev/null +++ b/tests/unit/ner/test_indic_encoder.py @@ -0,0 +1,192 @@ +from __future__ import annotations + +import hashlib +import logging +from types import SimpleNamespace + +import pytest + +from openmed.core import model_registry +from openmed.core.pii_i18n import DEFAULT_PII_MODELS +from openmed.ner.families import indic + + +class _FakeTokenizer: + def __call__(self, text: str, **kwargs): + assert kwargs["return_offsets_mapping"] is True + assert kwargs["return_tensors"] == "pt" + return { + "input_ids": [[101, 42, 102]], + "attention_mask": [[1, 1, 1]], + "offset_mapping": [[(0, 0), (0, len(text)), (0, 0)]], + } + + +class _FakeModel: + def __init__(self): + self.eval_called = False + + def eval(self): + self.eval_called = True + + def __call__(self, **kwargs): + assert kwargs["return_dict"] is True + return SimpleNamespace( + last_hidden_state=[ + [[0.0, 0.0], [0.5, 1.0], [0.0, 0.0]], + ] + ) + + +class _Loader: + def __init__(self, value): + self.value = value + self.calls = [] + + def from_pretrained(self, source, **kwargs): + self.calls.append((source, kwargs)) + return self.value + + +@pytest.fixture(autouse=True) +def _clear_registry_config(): + model_registry.clear_indic_encoder_config() + yield + model_registry.clear_indic_encoder_config() + + +def test_loader_skips_before_importing_when_no_weights_are_configured(monkeypatch): + def unexpected_import(name): + raise AssertionError(f"unexpected optional import: {name}") + + monkeypatch.setattr(indic, "_optional_module", unexpected_import) + result = indic.load_indic_encoder(None) + + assert result.available is False + assert result.handle is None + assert result.skip_reason == "no Indic encoder weights are configured" + + +def test_loader_reports_missing_optional_dependency(monkeypatch): + monkeypatch.setattr(indic, "_optional_module", lambda name: None) + + result = indic.load_indic_encoder("google/muril-base-cased") + + assert result.available is False + assert result.metadata is not None + assert result.metadata.license_id == "Apache-2.0" + assert "transformers" in result.skip_reason + + +def test_loader_exposes_aligned_hidden_state_contract_without_raw_logs( + monkeypatch, + caplog, +): + tokenizer_loader = _Loader(_FakeTokenizer()) + model = _FakeModel() + model_loader = _Loader(model) + transformers = SimpleNamespace( + AutoTokenizer=tokenizer_loader, + AutoModel=model_loader, + ) + torch = SimpleNamespace() + monkeypatch.setattr( + indic, + "_optional_module", + lambda name: transformers if name == "transformers" else torch, + ) + raw_text = "Patient Asha Verma ka record dekhiye." + + with caplog.at_level(logging.DEBUG, logger=indic.__name__): + result = indic.load_indic_encoder( + "google/muril-base-cased", + local_files_only=True, + ) + encoded = result.handle.encode(raw_text) + + assert result.available is True + assert model.eval_called is True + assert encoded.input_ids == [[101, 42, 102]] + assert encoded.attention_mask == [[1, 1, 1]] + assert encoded.offset_mapping == ((0, 0), (0, len(raw_text)), (0, 0)) + assert encoded.token_count == 3 + assert encoded.hidden_size == 2 + assert encoded.text_sha256 == hashlib.sha256(raw_text.encode()).hexdigest() + assert raw_text not in caplog.text + assert encoded.text_sha256 in caplog.text + assert tokenizer_loader.calls[0][1]["local_files_only"] is True + assert tokenizer_loader.calls[0][1]["trust_remote_code"] is False + + +def test_registry_adds_configured_encoder_without_replacing_defaults(monkeypatch): + baseline_ids = { + info.model_id + for info in model_registry.get_pii_models_by_language("hi").values() + } + monkeypatch.setattr( + model_registry, + "_is_indic_encoder_runtime_available", + lambda: True, + ) + + config = model_registry.configure_indic_encoder( + "/models/muril", + family="muril", + ) + models = model_registry.get_pii_models_by_language("hi") + + assert config.languages == ("hi", "te") + assert baseline_ids <= {info.model_id for info in models.values()} + assert DEFAULT_PII_MODELS["hi"] in {info.model_id for info in models.values()} + adapter = models["pii_hi_muril_encoder"] + assert adapter.model_id == "/models/muril" + assert adapter.license == "Apache-2.0" + assert adapter.provenance == { + "source": "/models/muril", + "weights": "user-supplied", + "bundled": False, + "revision": None, + } + + +def test_registry_can_opt_an_additional_family_language_in(monkeypatch): + monkeypatch.setattr( + model_registry, + "_is_indic_encoder_runtime_available", + lambda: True, + ) + model_registry.configure_indic_encoder( + "ai4bharat/indic-bert", + family="indicbert", + languages=("bn",), + ) + + models = model_registry.get_pii_models_by_language("bn") + + assert list(models) == ["pii_bn_indicbert_encoder"] + assert models["pii_bn_indicbert_encoder"].license == "MIT" + + +def test_registry_hides_adapter_when_dependencies_are_unavailable(monkeypatch): + model_registry.configure_indic_encoder( + "google/muril-base-cased", + family="muril", + ) + monkeypatch.setattr( + model_registry, + "_is_indic_encoder_runtime_available", + lambda: False, + ) + + models = model_registry.get_pii_models_by_language("hi") + + assert "pii_hi_muril_encoder" not in models + + +def test_registry_rejects_language_not_supported_by_family(): + with pytest.raises(ValueError, match="does not support"): + model_registry.configure_indic_encoder( + "ai4bharat/indic-bert", + family="indicbert", + languages=("ur",), + )