Source code for PsychiatryNLPKit.analysis.batch

"""Batch analysis orchestrator for PsychiatryNLPKit.

``BatchAnalyzer`` is a thin orchestrator over ``TextData.compute``: it runs
the requested analyses, keeps any required models loaded for the duration of
the run, and merges per-section metrics into a single :class:`AnalysisResult`.

Required models live on the ``TextData`` instance (``data.generative_model``,
``data.mask_filling_model``, ``data.vit_model``) rather than on the analyzer.

Typical usage::

    from PsychiatryNLPKit.analysis import BatchAnalyzer, AnalysisResult

    analyzer = BatchAnalyzer(
        text_data, included_analyses=["sentence_length", "adverb_ratio"]
    )
    result: AnalysisResult = analyzer.run()

    # Run every analysis
    analyzer_all = BatchAnalyzer(text_data)  # included_analyses="all" by default
    result_all = analyzer_all.run()
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field
from typing import Any, Literal

from .registry import ALL_ANALYSES, get_analysis_spec

logger = logging.getLogger(__name__)


[docs] @dataclass class AnalysisResult: """Container for batch analysis output. Attributes: results: Merged per-section metrics keyed by section name, then metric name. Matches the ``{section → {metric: value}}`` return shape of individual analysis functions. sections: Section names in order (from ``TextData.section_names``). analyses_run: Analysis function names that executed successfully. errors: Per-function failures mapped as ``function_name → error message``. """ results: dict[str, dict[str, float]] = field(default_factory=dict) sections: list[str] = field(default_factory=list) analyses_run: list[str] = field(default_factory=list) errors: dict[str, str] = field(default_factory=dict)
[docs] class BatchAnalyzer: """Run selected analyses on a ``TextData`` object in batch. The analyzer is an orchestrator only — every analysis is executed through ``TextData.compute``. Models required by the requested analyses must be attached to the ``TextData`` instance; the analyzer keeps them loaded for the whole run and unloads them afterwards. Args: text_data: Pre-built ``TextData`` instance with sections and any models required by the requested analyses. included_analyses: ``"all"`` runs every registered analysis. Pass an explicit list of function names to run a subset. excluded_analyses: Function names to remove from *included_analyses*. Every name here must already be in the resolved inclusion list; otherwise an :class:`AssertionError` is raised. image_paths: Mapping of section name → image file path. Required if ``"image_text_similarity"`` is in *analyses*; must cover every section in ``text_data.section_names``. Raises: AssertionError: If a requested analysis requires a model or data argument that was not provided. Example: .. code-block:: python # Run all analyses (requires all models attached to text_data) result = BatchAnalyzer( text_data, image_paths={"Paragraph 1": "img1.jpg", "Paragraph 2": "img2.jpg"} ).run() # Selective analyses with language-dependent metrics result = BatchAnalyzer( text_data, included_analyses=["sentence_length", "adverb_ratio", "filler_words_count"], ).run() """
[docs] def __init__( self, text_data: Any, included_analyses: Literal["all"] | list[str] = "all", excluded_analyses: list[str] | None = None, image_paths: dict[str, str] | None = None, ) -> None: self.text_data = text_data self.image_paths = image_paths # Resolve inclusion list. if included_analyses == "all": self._analyses: list[str] = list(ALL_ANALYSES) else: self._analyses = list(included_analyses) # Apply exclusions. if excluded_analyses is not None: for name in excluded_analyses: if name not in self._analyses: raise AssertionError( f"Cannot exclude '{name}': not in included analyses" ) self._analyses = [ n for n in self._analyses if n not in set(excluded_analyses) ] # Validate optional-dependency availability. requested = set(self._analyses) if "structural_graph" in requested: self._check_graph_available() if "image_text_similarity" in requested: self._check_image_available() # Validate required models attached to text_data. if ( requested & {"paragraph_level_pseudo_perplexity", "sentence_level_pseudo_perplexity"} ) and text_data.mask_filling_model is None: raise AssertionError( "mask_filling_model required for pseudo-perplexity analysis" ) if "image_text_similarity" in requested: if text_data.vit_model is None: raise AssertionError( "vit_model required for image-text similarity analysis" ) if image_paths is None: raise AssertionError( "image_paths must map every section name to an image path" ) missing = set(text_data.section_names) - set(image_paths.keys()) if missing: raise AssertionError( f"image_paths missing entries for sections: {', '.join(sorted(missing))}" ) # Track which models need to be loaded/unloaded for this run. self._models_to_load: list[Any] = [] self._models_loaded: set[int] = set()
@staticmethod def _check_graph_available() -> None: """Raise ImportError with install hint if networkx is unavailable.""" try: import networkx # noqa: F401 except ImportError: raise ImportError( "structural_graph requires the 'graph' extra. " "Install with: pip install PsychiatryNLPKit[graph]" ) from None @staticmethod def _check_image_available() -> None: """Raise ImportError with install hint if pillow is unavailable.""" try: import PIL # noqa: F401 except ImportError: raise ImportError( "image_text_similarity requires the 'image' extra. " "Install with: pip install PsychiatryNLPKit[image]" ) from None def _ensure_model_loaded(self, model: Any) -> None: """Load model if not already loaded, track for later unload.""" if model is None: return model_id = id(model) if model_id not in self._models_loaded: model.load() self._models_loaded.add(model_id) self._models_to_load.append(model) def _unload_tracked_models(self) -> None: """Unload all models that were loaded during this run.""" for model in self._models_to_load: try: model.unload() except Exception as exc: logger.warning("Failed to unload model %s: %s", model.name, exc) self._models_to_load.clear() self._models_loaded.clear()
[docs] def run(self) -> AnalysisResult: """Execute all requested analyses and return merged results.""" result = AnalysisResult(sections=list(self.text_data.section_names)) try: for name in self._analyses: try: if name == "image_text_similarity": sec_results = self._run_image_similarity() else: sec_results = self._run_registered(name) # Merge into result. for section, metrics in sec_results.items(): if section not in result.results: result.results[section] = {} result.results[section].update(metrics) result.analyses_run.append(name) except Exception as exc: logger.warning("Analysis %s failed: %s", name, exc) result.errors[name] = str(exc) finally: # Always unload models we loaded, even if an analysis failed. self._unload_tracked_models() return result
def _run_registered(self, name: str) -> dict[str, dict[str, float]]: """Dispatch a registered analysis through ``TextData.compute``.""" spec = get_analysis_spec(name) for attr in spec.requires: self._ensure_model_loaded(getattr(self.text_data, attr, None)) return self.text_data.compute(name, manage_lifecycle=False) def _run_image_similarity(self) -> dict[str, dict[str, float]]: """Run image-text similarity per section through ``TextData.compute``.""" assert self.image_paths is not None spec = get_analysis_spec("image_text_similarity") for attr in spec.requires: self._ensure_model_loaded(getattr(self.text_data, attr, None)) merged: dict[str, dict[str, float]] = {} for section in self.text_data.section_names: sec_result = self.text_data.compute( "image_text_similarity", image=self.image_paths[section], sections=[section], manage_lifecycle=False, ) merged.update(sec_result) return merged
__all__: list[str] = [ "AnalysisResult", "BatchAnalyzer", ]