wuenlp_tools.run_pipeline

  1"""Run a configured wuenlp-tools pipeline from JSON or CLI flags."""
  2
  3from __future__ import annotations
  4
  5import argparse
  6import json
  7import sys
  8from pathlib import Path
  9from typing import Any
 10
 11from loguru import logger
 12from wuenlp import UIMADocument
 13from wuenlp.impl.UIMANLPStructs import UIMASystemScene
 14
 15from wuenlp_tools.pipeline import Pipeline
 16from wuenlp_tools.pipeline_catalog import (
 17    PIPELINE_PRESET_STEP_NAMES,
 18    build_steps,
 19    catalog_as_json,
 20    get_pipeline_presets,
 21)
 22
 23
 24def _load_input(path: Path) -> UIMADocument:
 25    suffix = path.name.lower()
 26    if suffix.endswith(".xmi.zip") or suffix.endswith(".xmi"):
 27        return UIMADocument.from_xmi(path)
 28    return UIMADocument.from_text(path.read_text(encoding="utf-8"))
 29
 30
 31def _wire_moment_embedder(pipeline: Pipeline, steps: list) -> None:
 32    moment_step = next((step for step in steps if step.name == "MomentEmbedder"), None)
 33    if moment_step is None:
 34        return
 35    moment_step.processor.set_include_features(
 36        pipeline.added_additional_features[UIMASystemScene]
 37    )
 38
 39
 40def run_pipeline_config(config: dict[str, Any]) -> Path:
 41    input_path = Path(config["input"]).resolve()
 42    output_path = Path(config["output"]).resolve()
 43    if not input_path.exists():
 44        raise FileNotFoundError(f"Input not found: {input_path}")
 45
 46    step_specs = config.get("steps", [])
 47    if isinstance(step_specs, str):
 48        preset = step_specs
 49        presets = get_pipeline_presets()
 50        if preset not in presets:
 51            raise KeyError(f"Unknown pipeline preset: {preset}")
 52        step_specs = presets[preset]
 53
 54    steps = build_steps(step_specs)
 55    if not steps:
 56        raise ValueError("Pipeline has no steps configured.")
 57
 58    logger.info("Pipeline steps: {}", " → ".join(step.name for step in steps))
 59
 60    intermediate_dir = config.get("intermediate_dir")
 61    intermediate_path = Path(intermediate_dir).resolve() if intermediate_dir else None
 62    overwrite = bool(config.get("overwrite", False))
 63
 64    logger.info("Loading {}", input_path)
 65    doc = _load_input(input_path)
 66
 67    pipeline = Pipeline(
 68        steps=steps,
 69        intermediate_file_dir=intermediate_path,
 70        overwrite=overwrite,
 71    )
 72    _wire_moment_embedder(pipeline, steps)
 73
 74    logger.info("Running pipeline with {} steps", len(steps))
 75    doc = pipeline(doc)
 76
 77    output_path.parent.mkdir(parents=True, exist_ok=True)
 78    doc.serialize(output_path)
 79    logger.info("Wrote {}", output_path)
 80    return output_path
 81
 82
 83def main() -> int:
 84    parser = argparse.ArgumentParser(description="Run a wuenlp-tools pipeline from configuration.")
 85    parser.add_argument("--config", type=Path, help="JSON config file")
 86    parser.add_argument("--input", type=Path, help="Input .txt / .xmi / .xmi.zip file")
 87    parser.add_argument("--output", type=Path, help="Output .xmi / .xmi.zip path")
 88    parser.add_argument(
 89        "--steps",
 90        type=str,
 91        help="Comma-separated step ids or a preset name (process_file, example, coref_only)",
 92    )
 93    parser.add_argument("--intermediate-dir", type=Path, default=None, help="Cache directory for step outputs")
 94    parser.add_argument("--overwrite", action="store_true", help="Overwrite cached intermediate results")
 95    parser.add_argument("--catalog", action="store_true", help="Print step catalog as JSON and exit")
 96    parser.add_argument("--debug", action="store_true", help="Enable debug logging")
 97    args = parser.parse_args()
 98
 99    if args.debug:
100        logger.remove()
101        logger.add(sink=sys.stdout, level="DEBUG")
102
103    if args.catalog:
104        logger.remove()
105        sys.stdout.write(catalog_as_json())
106        sys.stdout.write("\n")
107        return 0
108
109    if args.config:
110        config = json.loads(Path(args.config).read_text(encoding="utf-8"))
111    else:
112        if not args.input or not args.output or not args.steps:
113            parser.error("Provide --config or (--input, --output, and --steps)")
114        step_specs: str | list[str]
115        if args.steps in PIPELINE_PRESET_STEP_NAMES:
116            step_specs = args.steps
117        else:
118            step_specs = [part.strip() for part in args.steps.split(",") if part.strip()]
119        config = {
120            "input": str(args.input),
121            "output": str(args.output),
122            "steps": step_specs,
123            "intermediate_dir": str(args.intermediate_dir) if args.intermediate_dir else None,
124            "overwrite": args.overwrite,
125        }
126
127    try:
128        run_pipeline_config(config)
129    except Exception as exc:
130        logger.error("Pipeline failed: {}", exc)
131        return 1
132    return 0
133
134
135if __name__ == "__main__":
136    sys.exit(main())