"""Batch analysis orchestrator for PsychiatryNLPKit.
Runs selected analyses on a pre-built ``TextData`` object and collects all
results into a single structured container. Optional model arguments
(``mask_filling_model``, ``vit_model``, ``image_paths``) are asserted at
construction time if the requested analyses require them.
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 collections.abc import Callable
from dataclasses import dataclass, field
from typing import Any, Literal
logger = logging.getLogger(__name__)
# -- Imports from sibling analysis modules ------------------------------------
from .Density import intrinsic_dimensionality_density, pca_density_metrics # noqa: E402
from .Lexicon import filler_words_count # noqa: E402
from .Perplexity import ( # noqa: E402
paragraph_level_perplexity,
paragraph_level_pseudo_perplexity,
sentence_level_perplexity,
sentence_level_pseudo_perplexity,
)
from .Similarity import ( # noqa: E402
sentence_level_cosine_similarity,
word_level_cosine_similarity,
)
from .Syntax import ( # noqa: E402
adjective_ratio,
adjective_sentence_length,
adverb_ratio,
clause_count,
coordinating_conjunction_ratio,
determiner_ratio,
modal_auxiliary_verb_ratio,
noun_group_count,
pronoun_ratio,
sentence_length,
stop_words_ratio,
syntax_depth,
unique_pos_tags,
)
# -- Data structures ----------------------------------------------------------
[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)
# -- Registry -----------------------------------------------------------------
# Each entry maps a function name to a tuple of (callable, data_source_key, kwargs_builder).
# callable – the analysis function to invoke
# data_source_key – attribute on TextData to access ("pos_tags", "syntax_trees", etc.)
# ``"data"`` means raw text dict; ``None`` means no data property
# kwargs_builder – callable(text_data, mask_filling_model) → dict of extra kwargs
# or None if the function only needs the data source as its first arg
_AnalysisEntry = tuple[Callable, str | None, Callable[[Any, Any], dict] | None]
def _build_lang_kwarg(text_data: Any, _mask_model: Any | None) -> dict:
return {"lang": text_data.lang}
def _build_language_kwarg(text_data: Any, _mask_model: Any | None) -> dict:
return {"language": text_data.lang}
def _build_generative_kwarg(text_data: Any, _mask_model: Any | None) -> dict:
return {"generative_model": text_data.generative_model, "manage_lifecycle": False}
def _build_mask_filling_kwarg(_text_data: Any, mask_model: Any | None) -> dict:
assert mask_model is not None, "mask_filling_model required for pseudo-perplexity"
return {"mask_filling_model": mask_model, "manage_lifecycle": False}
_ANALYSIS_REGISTRY: dict[str, _AnalysisEntry] = {
# Syntax — pos_tags based
"sentence_length": (sentence_length, "pos_tags", None),
"unique_pos_tags": (unique_pos_tags, "pos_tags", None),
"adverb_ratio": (adverb_ratio, "pos_tags", _build_lang_kwarg),
"coordinating_conjunction_ratio": (
coordinating_conjunction_ratio,
"pos_tags",
_build_lang_kwarg,
),
"adjective_ratio": (adjective_ratio, "pos_tags", _build_lang_kwarg),
"pronoun_ratio": (pronoun_ratio, "pos_tags", _build_lang_kwarg),
"determiner_ratio": (determiner_ratio, "pos_tags", _build_lang_kwarg),
"modal_auxiliary_verb_ratio": (
modal_auxiliary_verb_ratio,
"pos_tags",
_build_lang_kwarg,
),
"stop_words_ratio": (stop_words_ratio, "pos_tags", _build_lang_kwarg),
# Syntax — syntax_trees based
"syntax_depth": (syntax_depth, "syntax_trees", None),
"clause_count": (clause_count, "syntax_trees", None),
"noun_group_count": (noun_group_count, "syntax_trees", None),
"adjective_sentence_length": (
adjective_sentence_length,
"syntax_trees",
None,
),
# Similarity
"word_level_cosine_similarity": (
word_level_cosine_similarity,
"content_word_embedding_vectors",
None,
),
"sentence_level_cosine_similarity": (
sentence_level_cosine_similarity,
"sentence_embedding_vectors",
None,
),
# Perplexity — generative model
"paragraph_level_perplexity": (
paragraph_level_perplexity,
"paragraph_generative_tokens",
_build_generative_kwarg,
),
"sentence_level_perplexity": (
sentence_level_perplexity,
"sentence_generative_tokens",
_build_generative_kwarg,
),
# Perplexity — masked LM
"paragraph_level_pseudo_perplexity": (
paragraph_level_pseudo_perplexity,
"data",
_build_mask_filling_kwarg,
),
"sentence_level_pseudo_perplexity": (
sentence_level_pseudo_perplexity,
"data",
_build_mask_filling_kwarg,
),
# Density
"pca_density_metrics": (pca_density_metrics, "token_embedding_vectors", None),
"intrinsic_dimensionality_density": (
intrinsic_dimensionality_density,
"token_embedding_vectors",
None,
),
# Lexicon
"filler_words_count": (filler_words_count, "pos_tags", _build_language_kwarg),
}
# Canonical list of ALL analyses including special-dispatch entries.
_ALL_ANALYSES: list[str] = [
*_ANALYSIS_REGISTRY.keys(),
"structural_graph",
"image_text_similarity",
]
[docs]
class BatchAnalyzer:
"""Run selected analyses on a ``TextData`` object in batch.
The analyzer is an orchestrator only — it does not construct text data or
manage model lifecycles beyond what individual analysis functions require.
Args:
text_data: Pre-built ``TextData`` instance with sections and optional models.
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.
mask_filling_model: Required if any pseudo-perplexity function remains after
filtering.
vit_model: Required if ``"image_text_similarity"`` is in *analyses*.
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.
ValueError: If *analyses* contains an unknown function name (caught at
run time and recorded in :attr:`AnalysisResult.errors`).
Example:
.. code-block:: python
# Run all analyses (requires all models to be provided)
result = BatchAnalyzer(
text_data,
mask_filling_model=mask_lm,
vit_model=vit,
image_paths={"p1": "img1.jpg", "p2": "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,
mask_filling_model: Any | None = None,
vit_model: Any | None = None,
image_paths: dict[str, str] | None = None,
) -> None:
self.text_data = text_data
self.mask_filling_model = mask_filling_model
self.vit_model = vit_model
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 tier dependencies for requested analyses.
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 for requested analyses.
if (
requested
& {"paragraph_level_pseudo_perplexity", "sentence_level_pseudo_perplexity"}
) and mask_filling_model is None:
raise AssertionError(
"mask_filling_model required for pseudo-perplexity analysis"
)
if "image_text_similarity" in requested:
if 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()
elif name == "structural_graph":
sec_results = self._run_structural_graph()
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 function."""
func, data_key, kwarg_builder = _ANALYSIS_REGISTRY[name]
# Access the TextData property (triggers lazy computation).
data = getattr(self.text_data, data_key) if data_key else None
kwargs: dict = {}
if kwarg_builder is not None:
kwargs = kwarg_builder(self.text_data, self.mask_filling_model)
# Pre-load any models the function will need.
if "generative_model" in kwargs:
self._ensure_model_loaded(kwargs["generative_model"])
if "mask_filling_model" in kwargs:
self._ensure_model_loaded(kwargs["mask_filling_model"])
return func(data, **kwargs)
def _run_image_similarity(self) -> dict[str, dict[str, float]]:
"""Run image-text similarity per section."""
from .ImageSimilarity import image_text_similarity
assert self.vit_model is not None
assert self.image_paths is not None
self._ensure_model_loaded(self.vit_model)
merged: dict[str, dict[str, float]] = {}
for section in self.text_data.section_names:
sec_result = image_text_similarity(
image=self.image_paths[section],
text={section: self.text_data.data[section]},
vit_model=self.vit_model,
manage_lifecycle=False,
)
merged.update(sec_result)
return merged
def _run_structural_graph(self) -> dict[str, dict[str, float]]:
"""Run structural graph analysis."""
from .Graph import structural_graph
return structural_graph(self.text_data.content_words)
__all__: list[str] = [
"AnalysisResult",
"BatchAnalyzer",
]