diff --git a/script/brev/README.md b/script/brev/README.md index 084df2f88..f437ddd1a 100644 --- a/script/brev/README.md +++ b/script/brev/README.md @@ -10,17 +10,16 @@ NeMo Safe Synthesizer without setting up CUDA, drivers, or Python locally. Nothing in this directory is executed by the repo or by CI. A Brev Launchable is configured in the Brev web console, and the setup script is pasted into a form field there. This directory exists so that configuration is versioned and reviewable rather -than living only in a browser. When you change `setup.sh`, you must also paste the new -contents into the console for the change to take effect. +than living only in a browser. When you change `setup.sh` or `welcome.md`, you must also +update it in the console for the change to take effect. ### Files - `setup.sh`: Pasted into the Launchable's Setup Script field. Installs the CUDA build of Safe Synthesizer into a dedicated venv, registers it as the default Jupyter kernel, and drops the tutorial notebooks in `$HOME`. -- `welcome.md`: Becomes the customer's `$HOME/README.md`. Fetched at provisioning - time from the same tarball as the tutorials, not baked into `setup.sh` -- it would - otherwise consume a tenth of the 16 KiB script budget. +- `welcome.md`: Added to the Launchable's Source files so it renders on the Launchable + webpage and appears as the customer's `$HOME/welcome.md`. ### Console configuration @@ -33,7 +32,7 @@ Launchable with these settings. | Software | Install Jupyter on the host | Enabled | | Software | Run a Setup Script | Enabled, contents of `setup.sh` | | Software | Image ID | Leave blank | -| Source | Code source | No code files (`setup.sh` downloads the tutorials itself) | +| Source | Code source | `welcome.md` | | Hardware | GPU | 1× 80 GiB VRAM, single GPU | | Hardware | Disk | 200 GiB or more -- not resizable after creation | | Network | Ports | 8888, named `jupyter` | @@ -67,8 +66,8 @@ there. Everything operational is a dotfile, which the browser hides by default. ```text $HOME/ tutorials/ the three tutorial notebooks and their datasets - README.md where to start, rendered on double-click - (SETUP-IN-PROGRESS.md until setup finishes) + welcome.md where to start, rendered on double-click + SETUP-IN-PROGRESS.md present only while setup is running or after failure .nss-venv/ cu129 venv, registered as the default kernel .cache/huggingface/ model cache (Hugging Face's default location) @@ -91,12 +90,12 @@ hard way on a real instance. them -- and they have to match the release being installed, not this repo's `main`. The script resolves the latest version from the PyPI JSON API, fetches that tag's `pyproject.toml`, and reads the CUDA index URLs out of it, then pins the install to - that exact version so the two cannot drift. Selection is keyed on the URL containing - `cu129`, not on the index name: the flashinfer entry was renamed - `flashinfer-jit-cache` → `flashinfer-jit-cache-cu129` between 0.1.8 and 0.1.9, so - names are not stable across releases. The parse runs inside a process substitution and - therefore cannot fail the script, so the count of discovered indexes is what validates - it. + that exact version so the two cannot drift. Selection uses the CUDA extra in each + index's name or URL and includes indexes referenced by `[tool.uv.sources]` for that + extra. The source lookup matters for variant-neutral indexes such as + `https://flashinfer.ai/whl/`, while the URL lookup handles names that changed between + releases. The parse runs inside a process substitution and therefore cannot fail the + script, so the count of discovered indexes is what validates it. - uv is installed from a checksum-verified tarball, not `curl | sh`. The `astral.sh/install.sh` path logs `no checksums to verify`, so nothing validated what it downloaded. The script fetches the pinned release tarball, compares it against the @@ -106,12 +105,11 @@ hard way on a real instance. accepts connections well before this script finishes, so a user who opens it early would otherwise see an empty or half-populated file browser and assume the Launchable is broken. `SETUP-IN-PROGRESS.md` is written before any slow work, rewritten by the - `ERR` trap if provisioning fails, and replaced by `README.md` on success. -- The welcome text lives in `welcome.md`, not a heredoc. It is pulled from the - same tarball as the tutorials, so the two always match, and it is staged as a dotfile - until the final step so it never appears while setup is still running. The fetch is - non-fatal: `script/brev/` exists in no released tag, so it resolves only from the - `main` fallback until a release includes it. + `ERR` trap if provisioning fails, and removed on success. +- The welcome text lives in the Launchable's Source configuration, not a heredoc or + release tarball. Brev renders it on the Launchable webpage and copies it to + `$HOME/welcome.md`; keeping the console copy synchronized with this directory is a + manual deployment step. - The setup script has a 16 KiB limit. Brev rejects anything larger, which is why the script carries short comments pointing here rather than full explanations. Check `wc -c script/brev/setup.sh` before pasting. diff --git a/script/brev/setup.sh b/script/brev/setup.sh index a123e5cf7..15d3350aa 100755 --- a/script/brev/setup.sh +++ b/script/brev/setup.sh @@ -21,11 +21,10 @@ readonly REPO_URL="https://github.com/NVIDIA-NeMo/Safe-Synthesizer" : "${HOME:?HOME is not set}" -# $HOME is the file browser root: only tutorials/ and README.md are visible. +# $HOME is the file browser root: only customer-facing files stay visible. readonly TUTORIALS_DIR="${HOME}/tutorials" -readonly README_FILE="${HOME}/README.md" +readonly WELCOME_FILE="${HOME}/welcome.md" readonly WAIT_FILE="${HOME}/SETUP-IN-PROGRESS.md" -readonly WELCOME_STAGED="${HOME}/.nss-welcome.md" readonly BIN_DIR="${HOME}/.local/bin" readonly VENV_DIR="${HOME}/.nss-venv" @@ -61,7 +60,7 @@ NeMo Safe Synthesizer is still installing -- roughly 5-10 minutes from when the instance started. Files appear as it progresses, so a partly-filled file browser is expected. Nothing here is ready to run yet. -When setup finishes, this file is replaced by README.md. Refresh to check. +When setup finishes, this file disappears. Open welcome.md to get started. EOF export PATH="${BIN_DIR}:${PATH}" @@ -123,8 +122,8 @@ else NSS_VERSION="$(curl -fsSL https://pypi.org/pypi/nemo-safe-synthesizer/json \ | "${VENV_DIR}/bin/python" -c 'import json, sys; print(json.load(sys.stdin)["info"]["version"])')" - # Indexes come from the installed release's pyproject. Match both generated - # names and URLs because static index names do not enforce the CUDA suffix. + # Indexes come from the installed release's pyproject. Match CUDA names and + # URLs plus source-mapped indexes whose names are variant-neutral. pyproject="$(mktemp)" curl -fsSL "${REPO_URL}/raw/v${NSS_VERSION}/pyproject.toml" -o "${pyproject}" index_args=() @@ -140,15 +139,29 @@ import sys import tomllib with open(sys.argv[1], "rb") as handle: - indexes = tomllib.load(handle)["tool"]["uv"]["index"] + uv_config = tomllib.load(handle)["tool"]["uv"] +indexes = uv_config["index"] cuda_extra = os.environ["CUDA_EXTRA"] + +# Some indexes carry no CUDA variant in their name or URL. Source config is +# not wheel metadata, so collect indexes mapped to packages for this extra. +source_indexes = { + entry["index"] + for value in uv_config.get("sources", {}).values() + for entry in (value if isinstance(value, list) else [value]) + if isinstance(entry, dict) + and entry.get("extra") == cuda_extra + and "index" in entry +} + print( "\n".join( index["url"] for index in indexes if index["name"].endswith(f"-{cuda_extra}") or f"/{cuda_extra}" in index["url"] + or index["name"] in source_indexes ) ) PY @@ -208,13 +221,6 @@ else # Written last; the guard keys on this, so partial runs are redone. : >"${TUTORIALS_DIR}/.fetched" log "tutorials extracted from ${ref}" - # Same tarball as the tutorials. Non-fatal -- see README. - if tar -xzf "${tarball}" -C "${tarball_dir}" --strip-components=3 \ - "${top}/script/brev/welcome.md" 2>/dev/null; then - mv "${tarball_dir}/welcome.md" "${WELCOME_STAGED}" - else - log "WARNING: welcome.md not present in ${ref}" - fi fetched=1 break fi @@ -343,6 +349,12 @@ if [[ "${registered}" -ne 1 ]]; then log "WARNING: kernel not registered; notebooks may open on the wrong Python" fi +# Pre-compile third-party packages that emit SyntaxWarnings on first import so +# the warnings go into the setup log rather than appearing in notebook output. +log "pre-compiling packages" +"${VENV_DIR}/bin/python" -W ignore::SyntaxWarning \ + -c "import torchao, range_regex" 2>/dev/null || true + # Smoke check -- fail provisioning loudly rather than handing over a broken VM. log "verifying install" @@ -350,12 +362,10 @@ log "verifying install" "${VENV_DIR}/bin/python" \ -c "import torch; print('cuda available:', torch.cuda.is_available())" -# Hand over: swap the "please wait" file for the welcome text. +# Hand over: the Source-provided welcome stays visible after setup completes. -if [[ -f "${WELCOME_STAGED}" ]]; then - mv "${WELCOME_STAGED}" "${README_FILE}" -else - log "WARNING: no welcome.md staged; skipping ${README_FILE}" +if [[ ! -f "${WELCOME_FILE}" ]]; then + log "WARNING: ${WELCOME_FILE} is missing; check the Launchable Source files" fi rm -f "${WAIT_FILE}" @@ -365,7 +375,7 @@ cat < QuantizationConfigMixin: ValueGTZero = ValueValidator(lambda p: range_validator(p, lambda v: v >= 0)) +def is_valid_warmup(value: float) -> bool: + """Whether a warmup setting is a usable ratio or step count. + + Mirrors how transformers interprets ``warmup_steps``: ``0`` disables warmup, + values below 1 are a ratio of total training steps, and values of 1 or more + are an absolute step count. Fractional values of 1 or more are rejected + because transformers truncates them (``1.5`` silently becomes ``1``), and + non-finite values are rejected because they raise ``OverflowError`` once + converted to an integer. + """ + return math.isfinite(value) and value >= 0 and (value < 1 or float(value).is_integer()) + + class TrainingHyperparams(Parameters): """Hyperparameters that control the training process behavior. @@ -210,15 +225,43 @@ class TrainingHyperparams(Parameters): ), ] = 0.01 - warmup_ratio: Annotated[ + warmup_steps: Annotated[ float, - ValueValidator(value_func=lambda v: v > 0), + ValueValidator(value_func=is_valid_warmup), Field( - title="warmup_ratio", - description="Ratio of total training steps used for a linear warmup from 0 to the learning rate. Must be > 0.", + title="warmup_steps", + description=( + "Linear warmup from 0 to the learning rate. " + "A whole number of 1 or more sets the exact number of warmup steps; " + "a float in (0, 1) is treated as a ratio of total training steps; " + "0 disables warmup. " + "Must be finite and >= 0, and cannot be fractional at or above 1." + ), ), ] = 0.05 + warmup_ratio: Annotated[ + float | None, + ValueValidator(value_func=lambda v: v is None or is_valid_warmup(v)), + Field( + title="warmup_ratio", + description="Deprecated. Use warmup_steps instead.", + exclude=True, + ), + ] = None + + @model_validator(mode="after") + def _migrate_warmup_ratio(self) -> TrainingHyperparams: + if self.warmup_ratio is not None: + warnings.warn( + "warmup_ratio is deprecated and will be removed in a future release. Use warmup_steps instead.", + DeprecationWarning, + stacklevel=2, + ) + if "warmup_steps" not in self.model_fields_set: + self.warmup_steps = self.warmup_ratio + return self + lr_scheduler: Annotated[ str, Field( diff --git a/src/nemo_safe_synthesizer/evaluation/components/text_semantic_similarity.py b/src/nemo_safe_synthesizer/evaluation/components/text_semantic_similarity.py index 57df93ee7..a7d877e24 100644 --- a/src/nemo_safe_synthesizer/evaluation/components/text_semantic_similarity.py +++ b/src/nemo_safe_synthesizer/evaluation/components/text_semantic_similarity.py @@ -4,6 +4,8 @@ from __future__ import annotations import logging +import warnings +from contextlib import contextmanager from functools import cached_property from typing import TYPE_CHECKING @@ -36,11 +38,32 @@ from . import multi_modal_figures as figures if TYPE_CHECKING: + from collections.abc import Iterator + from sentence_transformers import SentenceTransformer logger = get_logger(__name__) +@contextmanager +def _suppress_ks_exact_fallback() -> Iterator[None]: + """Silence SciPy's notice that ``ks_2samp`` fell back to the asymptotic method. + + ``method="auto"`` attempts the exact calculation and falls back to the + asymptotic approximation once the samples are large, which is the intended + behaviour here -- the resulting p-values are still valid. The notice is not + actionable, so keep it out of CLI, SDK, and notebook output rather than + suppressing it separately in each caller. + """ + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", + message="ks_2samp: Exact calculation unsuccessful", + category=RuntimeWarning, + ) + yield + + class TextSemanticSimilarityDatum(BaseModel): """Per-column text semantic similarity scores and PCA projections.""" @@ -401,12 +424,13 @@ def _get_text_semantic_similarity( # the minimum (most negative) difference between the empirical # distribution functions of the samples. The range of this statistic is # [0, 1], where 0 indicates no overfitting. - ks_test_overfitting = ks_2samp( - training_synth_similarity_matrix.max(axis=0), # F(x) - training_similarity_matrix.max(axis=0), # G(x) - alternative="less", - method="auto", - ) + with _suppress_ks_exact_fallback(): + ks_test_overfitting = ks_2samp( + training_synth_similarity_matrix.max(axis=0), # F(x) + training_similarity_matrix.max(axis=0), # G(x) + alternative="less", + method="auto", + ) # Underfitting is measured as the extent to which the synthetic # data is less similar to the test data than the test data is to @@ -417,12 +441,13 @@ def _get_text_semantic_similarity( # the minimum (most negative) difference between the empirical # distribution functions of the samples. The range of this statistic is # [0, 1], where 0 indicates no underfitting. - ks_test_underfitting = ks_2samp( - test_synth_similarity_matrix.max(axis=0), # F(x) - test_similarity_matrix.max(axis=0), # G(x) - alternative="greater", - method="auto", - ) + with _suppress_ks_exact_fallback(): + ks_test_underfitting = ks_2samp( + test_synth_similarity_matrix.max(axis=0), # F(x) + test_similarity_matrix.max(axis=0), # G(x) + alternative="greater", + method="auto", + ) # The overall semantic similarity score combines underfitting and overfitting # The range of this score is [0.37, 1], where 1 indicates perfect model and diff --git a/src/nemo_safe_synthesizer/pii_replacer/data_editor/detect.py b/src/nemo_safe_synthesizer/pii_replacer/data_editor/detect.py index ebe74ac27..0b13eae7f 100644 --- a/src/nemo_safe_synthesizer/pii_replacer/data_editor/detect.py +++ b/src/nemo_safe_synthesizer/pii_replacer/data_editor/detect.py @@ -5,6 +5,7 @@ import logging import os +import warnings from abc import ABC, abstractmethod from collections.abc import Callable, Iterable, Iterator from dataclasses import dataclass @@ -660,11 +661,15 @@ def get_entity_extractor( f"Loading NER model from filesystem to {map_location}", ) - extractor._model = GLiNER.from_pretrained( - clsfy_cfg.gliner_model, - map_location=map_location, - local_files_only=hf_offline_enabled(), - ) + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", message="The `resume_download` argument is deprecated", category=UserWarning + ) + extractor._model = GLiNER.from_pretrained( + clsfy_cfg.gliner_model, + map_location=map_location, + local_files_only=hf_offline_enabled(), + ) entity_types = DEFAULT_ENTITIES if clsfy_cfg.ner_entities: entity_types = clsfy_cfg.ner_entities @@ -691,17 +696,21 @@ def _predict_entities(self, text: str, entity_labels: list[str]) -> list[dict]: flat_ner=False, ) - return self._model.batch_predict_entities( - [text], - entity_labels, - threshold=self._ner_threshold, - flat_ner=False, - )[0] + inference = getattr(self._model, "inference", None) + if inference is not None: + return inference( + [text], + entity_labels, + threshold=self._ner_threshold, + flat_ner=False, + )[0] + + raise AttributeError("GLiNER model has neither inference nor predict_entities") def _batch_predict_entities(self, texts: list[str], entity_labels: list[str]) -> list[list[dict]]: - batch_predict_entities = getattr(self._model, "batch_predict_entities", None) - if batch_predict_entities is not None: - return batch_predict_entities( + inference = getattr(self._model, "inference", None) + if inference is not None: + return inference( texts, entity_labels, threshold=self._ner_threshold, @@ -720,7 +729,7 @@ def _batch_predict_entities(self, texts: list[str], entity_labels: list[str]) -> for text in texts ] - raise AttributeError("GLiNER model has neither batch_predict_entities nor predict_entities") + raise AttributeError("GLiNER model has neither inference nor predict_entities") def _detect_entities_chunked( self, diff --git a/src/nemo_safe_synthesizer/training/huggingface_backend.py b/src/nemo_safe_synthesizer/training/huggingface_backend.py index 6875c95a8..51da170e9 100644 --- a/src/nemo_safe_synthesizer/training/huggingface_backend.py +++ b/src/nemo_safe_synthesizer/training/huggingface_backend.py @@ -421,7 +421,7 @@ def _build_base_training_args(self) -> dict: learning_rate=self.params.training.learning_rate, eval_strategy=evaluation_strategy, weight_decay=self.params.training.weight_decay, - warmup_ratio=self.params.training.warmup_ratio, + warmup_steps=self.params.training.warmup_steps, eval_steps=EVAL_STEPS, do_eval=self.params.training.validation_ratio > 0, disable_tqdm=True, # The 🤗 progress bar doesn't play nice with our logging. diff --git a/tests/config/test_parameters.py b/tests/config/test_parameters.py index 20ad9b2f7..052f33c55 100644 --- a/tests/config/test_parameters.py +++ b/tests/config/test_parameters.py @@ -14,7 +14,7 @@ from nemo_safe_synthesizer.config.job import SafeSynthesizerJobConfig from nemo_safe_synthesizer.config.parameters import SafeSynthesizerParameters from nemo_safe_synthesizer.config.replace_pii import PiiReplacerConfig, StepDefinition -from nemo_safe_synthesizer.config.training import QuantizationScheme +from nemo_safe_synthesizer.config.training import QuantizationScheme, TrainingHyperparams from nemo_safe_synthesizer.configurator.parameter_paths import ( AmbiguousParameterName, ParameterFieldKind, @@ -885,3 +885,55 @@ def test_returned_config_is_independent_of_saved(self): assert saved.data.holdout != 0.42 assert saved.training.batch_size == 8 assert saved.generation.structured_generation.enabled is True + + +class TestWarmupSteps: + """`warmup_steps` mirrors the transformers contract: ratio below 1, whole steps at or above 1.""" + + @pytest.mark.parametrize("value", [0, 0.05, 0.5, 0.999, 1, 1.0, 10, 10.0, 500]) + def test_accepts_ratios_and_whole_step_counts(self, value): + """0 is legal and disables warmup, matching how transformers treats it.""" + assert TrainingHyperparams(warmup_steps=value).warmup_steps == value + + @pytest.mark.parametrize( + "value", + [ + 1.5, # transformers truncates to 1, so reject rather than silently change the schedule + 2.7, + float("inf"), # OverflowError once transformers converts it to an int + float("-inf"), + float("nan"), + -1, + ], + ) + def test_rejects_fractional_non_finite_and_negative(self, value): + with pytest.raises((ParameterError, ValidationError, ValueError)): + TrainingHyperparams(warmup_steps=value) + + def test_default_is_a_ratio(self): + assert TrainingHyperparams().warmup_steps == 0.05 + + +class TestWarmupRatioDeprecation: + """`warmup_ratio` stays accepted as a deprecated alias for `warmup_steps`.""" + + def test_migrates_to_warmup_steps_and_warns(self): + with pytest.warns(DeprecationWarning, match="warmup_ratio is deprecated"): + params = TrainingHyperparams(warmup_ratio=0.2) + assert params.warmup_steps == 0.2 + + def test_explicit_warmup_steps_wins_over_deprecated_alias(self): + with pytest.warns(DeprecationWarning): + params = TrainingHyperparams(warmup_steps=0.3, warmup_ratio=0.2) + assert params.warmup_steps == 0.3 + + @pytest.mark.parametrize("value", [1.5, float("inf")]) + def test_deprecated_alias_is_validated_too(self, value): + """The alias assigns to `warmup_steps` after field validation, so it needs its own guard.""" + with pytest.raises((ParameterError, ValidationError, ValueError)): + TrainingHyperparams(warmup_ratio=value) + + def test_not_serialized(self): + with pytest.warns(DeprecationWarning): + params = TrainingHyperparams(warmup_ratio=0.2) + assert "warmup_ratio" not in params.model_dump() diff --git a/tests/conftest.py b/tests/conftest.py index 5fd840d44..0ee52a342 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -166,7 +166,7 @@ def fixture_yaml_config_str() -> str: rope_scaling_factor: auto validation_ratio: 0.0 validation_steps: 15 - warmup_ratio: 0.05 + warmup_steps: 0.05 weight_decay: 0.01 """ diff --git a/tests/evaluation/components/test_text_similarity_context.py b/tests/evaluation/components/test_text_similarity_context.py index 090ed05e2..9bb5d51c8 100644 --- a/tests/evaluation/components/test_text_similarity_context.py +++ b/tests/evaluation/components/test_text_similarity_context.py @@ -3,6 +3,8 @@ from __future__ import annotations +import warnings + import pandas as pd import plotly.graph_objects as go import pytest @@ -11,6 +13,7 @@ from nemo_safe_synthesizer.evaluation.components.text_semantic_similarity import ( TextSemanticSimilarity, TextSemanticSimilarityDatum, + _suppress_ks_exact_fallback, ) from nemo_safe_synthesizer.evaluation.components.text_structure_similarity import ( TextDataSetStatistics, @@ -53,3 +56,16 @@ def test_text_structure_similarity_jinja_context_includes_column_heading( assert context["figures"][0]["title"] == "review" assert "plotly-graph-div" in context["figures"][0]["html"] + + +def test_suppress_ks_exact_fallback_is_scoped_to_the_scipy_notice(): + """Silences SciPy's asymptotic-fallback notice without swallowing other warnings.""" + # Verbatim from scipy/stats/_stats_py.py where ks_2samp abandons the exact method. + notice = "ks_2samp: Exact calculation unsuccessful. Switching to method=asymp." + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + with _suppress_ks_exact_fallback(): + warnings.warn(notice, RuntimeWarning) + warnings.warn("an unrelated problem", RuntimeWarning) + assert [str(w.message) for w in caught] == ["an unrelated problem"] diff --git a/tests/pii_replacer/test_detect.py b/tests/pii_replacer/test_detect.py index 79cd94d43..9144f408d 100644 --- a/tests/pii_replacer/test_detect.py +++ b/tests/pii_replacer/test_detect.py @@ -67,7 +67,7 @@ def test_gliner_batch_predict_config(): entity_extractor = EntityExtractorGliner.get_entity_extractor(cfg) entity_extractor.batch_update_cache(["abc"], None) assert entity_extractor._model is not None - entity_extractor._model.batch_predict_entities.assert_not_called() # ty: ignore[call-non-callable, unresolved-attribute] -- mock object + entity_extractor._model.inference.assert_not_called() # ty: ignore[call-non-callable, unresolved-attribute] -- mock object cfg = ClassifyConfig( valid_entities={"name"}, @@ -85,7 +85,7 @@ def test_gliner_batch_predict_config(): entity_extractor = EntityExtractorGliner.get_entity_extractor(cfg) entity_extractor.batch_update_cache(["abc"], None) assert entity_extractor._model is not None - entity_extractor._model.batch_predict_entities.assert_called() # ty: ignore[call-non-callable, unresolved-attribute] -- mock object + entity_extractor._model.inference.assert_called() # ty: ignore[call-non-callable, unresolved-attribute] -- mock object def test_gliner_entity_labels_are_indexable(): @@ -104,14 +104,14 @@ def test_gliner_entity_labels_are_indexable(): with patch("nemo_safe_synthesizer.pii_replacer.data_editor.detect.GLiNER") as mock_gliner: model = mock_gliner.from_pretrained.return_value model.predict_entities.return_value = [] - model.batch_predict_entities.return_value = [[]] + model.inference.return_value = [[]] entity_extractor = EntityExtractorGliner.get_entity_extractor(cfg) entity_extractor.extract_ner_predictions("abc", {"name", "email"}) entity_extractor.batch_update_cache(["abc"], None) assert model.predict_entities.call_args.args[1] == ["email", "name"] - assert model.batch_predict_entities.call_args.args[1] == ["email", "name"] + assert model.inference.call_args.args[1] == ["email", "name"] def test_gliner_batch_cache_falls_back_to_predict_entities_api(): @@ -146,8 +146,16 @@ def predict_entities(self, text, labels, **kwargs): assert model.calls[0][0] == "abc" assert model.calls[0][1] == ["email", "name"] + # A build exposing neither API must fail loudly, not return no entities. + with patch("nemo_safe_synthesizer.pii_replacer.data_editor.detect.GLiNER") as mock_gliner: + mock_gliner.from_pretrained.return_value = object() + entity_extractor = EntityExtractorGliner.get_entity_extractor(cfg) + + with pytest.raises(AttributeError, match="neither inference nor predict_entities"): + entity_extractor.batch_update_cache(["abc"], None) + -def test_gliner_single_prediction_falls_back_to_batch_api(): +def test_gliner_single_prediction_falls_back_to_inference_api(): cfg = ClassifyConfig( valid_entities={"email", "name"}, ner_threshold=0.8, @@ -164,7 +172,7 @@ class FakeGLiNER: def __init__(self): self.calls = [] - def batch_predict_entities(self, texts, labels, **kwargs): + def inference(self, texts, labels, **kwargs): self.calls.append((texts, labels, kwargs)) return [[]] @@ -178,6 +186,14 @@ def batch_predict_entities(self, texts, labels, **kwargs): assert model.calls[0][0] == ["abc"] assert model.calls[0][1] == ["email", "name"] + # A build exposing neither API must fail loudly, not return no entities. + with patch("nemo_safe_synthesizer.pii_replacer.data_editor.detect.GLiNER") as mock_gliner: + mock_gliner.from_pretrained.return_value = object() + entity_extractor = EntityExtractorGliner.get_entity_extractor(cfg) + + with pytest.raises(AttributeError, match="neither inference nor predict_entities"): + entity_extractor.extract_ner_predictions("abc", {"name", "email"}) + @pytest.mark.parametrize("offline_var", ["HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE"]) @pytest.mark.parametrize("env_value", ["1", "yes", "on"])