classifyre-cli 0.4.2__py3-none-any.whl

src/detectors/custom/runners/_regex.py

@@ -0,0 +1,147 @@
+"""Regex pipeline runner."""
+
+from __future__ import annotations
+
+import logging
+import re
+import time
+from datetime import UTC, datetime
+from typing import Any
+
+from ....models.generated_detectors import (
+    PipelineResult,
+    RegexPatternDefinition,
+    RegexPipelineSchema,
+)
+from ...dependencies import MissingDependencyError, require_module
+from ._base import BaseRunner
+
+logger = logging.getLogger(__name__)
+
+
+def _load_regex_engine() -> tuple[Any, bool]:
+    """Try to load google-re2, fall back to stdlib re."""
+    try:
+        re2_module = require_module("re2", "regex", ["regex"])
+        logger.info("Using google-re2 engine for regex patterns")
+        return re2_module, True
+    except MissingDependencyError:
+        logger.info(
+            "google-re2 not available, using stdlib re (install with: uv sync --group regex)"
+        )
+        return re, False
+
+
+class RegexRunner(BaseRunner):
+    """Regex pipeline — uses google-re2 when available, falls back to stdlib re."""
+
+    def __init__(
+        self, schema: RegexPipelineSchema, detector_key: str = "", detector_name: str = ""
+    ) -> None:
+        self._schema = schema
+        self._detector_key = detector_key
+        self._detector_name = detector_name
+        self._engine, self._using_re2 = _load_regex_engine()
+        self._compiled: dict[str, tuple[re.Pattern[str], RegexPatternDefinition]] = {}
+        self._compile_patterns()
+
+    def _compile_patterns(self) -> None:
+        patterns = self._schema.patterns or {}
+        for name, defn in patterns.items():
+            try:
+                compiled = self._compile_one(defn)
+                self._compiled[name] = (compiled, defn)
+            except Exception as exc:
+                logger.warning(
+                    "Invalid regex pattern '%s' in detector '%s': %s",
+                    name,
+                    self._detector_key,
+                    exc,
+                )
+
+    def _compile_one(self, defn: RegexPatternDefinition) -> re.Pattern[str]:
+        case_sensitive = defn.case_sensitive if defn.case_sensitive is not None else True
+        dot_nl = defn.dot_nl or False
+        literal = defn.literal or False
+        longest_match = defn.longest_match or False
+        max_mem = defn.max_mem
+
+        legacy_flags = defn.flags or 0
+        if isinstance(legacy_flags, int) and legacy_flags & re.IGNORECASE:
+            case_sensitive = False
+        if isinstance(legacy_flags, int) and legacy_flags & re.DOTALL:
+            dot_nl = True
+
+        if self._using_re2:
+            options = self._engine.Options()
+            options.case_sensitive = case_sensitive
+            options.dot_nl = dot_nl
+            options.literal = literal
+            options.longest_match = longest_match
+            if max_mem is not None:
+                options.max_mem = max_mem
+            return self._engine.compile(defn.pattern, options=options)
+
+        flags = legacy_flags
+        if not case_sensitive:
+            flags |= re.IGNORECASE
+        if dot_nl:
+            flags |= re.DOTALL
+        if literal:
+            return re.compile(re.escape(defn.pattern), flags)
+        if longest_match:
+            logger.debug("longest_match is a RE2-only feature, ignored with stdlib re")
+        if max_mem is not None:
+            logger.debug("max_mem is a RE2-only feature, ignored with stdlib re")
+        return re.compile(defn.pattern, flags)
+
+    def run(self, text: str) -> PipelineResult:
+        start_ms = time.monotonic()
+        entities: dict[str, list[dict[str, object]]] = {}
+
+        for name, (rx, defn) in self._compiled.items():
+            group_idx = defn.group or 0
+            spans: list[dict[str, object]] = []
+            for match in rx.finditer(text):
+                actual_group = group_idx
+                try:
+                    value = match.group(group_idx)
+                except IndexError:
+                    value = match.group(0)
+                    actual_group = 0
+                    logger.warning(
+                        "Capture group %d does not exist in pattern '%s', using group 0",
+                        group_idx,
+                        name,
+                    )
+
+                start = match.start(actual_group)
+                end = match.end(actual_group)
+
+                span: dict[str, object] = {
+                    "value": value or "",
+                    "confidence": 1.0,
+                    "start": start,
+                    "end": end,
+                }
+                if defn.severity is not None:
+                    span["severity"] = str(defn.severity)
+                if match.lastindex:
+                    span["groups"] = match.groups()
+
+                spans.append(span)
+            if spans:
+                entities[name] = spans
+
+        latency_ms = round((time.monotonic() - start_ms) * 1000)
+        engine_tag = "RE2" if self._using_re2 else "stdlib-re"
+        return PipelineResult(
+            entities=entities,
+            classification={},
+            metadata={
+                "runner": "REGEX",
+                "engine": engine_tag,
+                "latency_ms": latency_ms,
+                "timestamp": datetime.now(UTC).isoformat(),
+            },
+        )

src/detectors/custom/runners/_text_classification.py

@@ -0,0 +1,109 @@
+"""Text classification pipeline runner."""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from ....models.generated_detectors import Severity, TextClassificationPipelineSchema
+from ....models.generated_single_asset_scan_results import DetectionResult
+from ...dependencies import ensure_torch, require_module
+from ._base import _TEXT_CONTENT_TYPES, BaseRunner, _resolve_pipeline_severity
+
+logger = logging.getLogger(__name__)
+
+
+def _chunk_text(text: str, chunk_size: int | None, chunk_overlap: int) -> list[str]:
+    """Split text into chunks. Returns [text] when chunk_size is not set."""
+    if not chunk_size:
+        return [text]
+    step = max(1, chunk_size - chunk_overlap)
+    return [text[i : i + chunk_size] for i in range(0, len(text), step)]
+
+
+class TextClassificationRunner(BaseRunner):
+    """Text classification via a single HuggingFace text-classification pipeline."""
+
+    def __init__(
+        self,
+        schema: TextClassificationPipelineSchema,
+        detector_key: str = "",
+        detector_name: str = "",
+    ) -> None:
+        self._schema = schema
+        self._detector_key = detector_key
+        self._detector_name = detector_name
+        ensure_torch("text_classification", ["custom", "detectors"])
+        transformers = require_module(
+            "transformers", "text_classification", ["custom", "detectors"]
+        )
+        pipeline_kwargs: dict[str, Any] = {
+            "model": schema.model,
+            "device": schema.device or "cpu",
+        }
+        if schema.model_revision:
+            pipeline_kwargs["revision"] = schema.model_revision
+        if schema.top_k is not None:
+            pipeline_kwargs["top_k"] = schema.top_k
+        if schema.function_to_apply is not None:
+            pipeline_kwargs["function_to_apply"] = str(schema.function_to_apply)
+        self._pipe: Any = transformers.pipeline("text-classification", **pipeline_kwargs)
+
+    def run(self, text: str) -> None:  # type: ignore[override]  # pragma: no cover
+        raise NotImplementedError("TextClassificationRunner uses detect() directly")
+
+    def detect(self, content: str | bytes, content_type: str) -> list[DetectionResult]:
+        if isinstance(content, bytes):
+            return []
+        if content_type not in _TEXT_CONTENT_TYPES:
+            return []
+        text = content.strip()
+        if not text:
+            return []
+
+        schema = self._schema
+        chunk_size: int | None = getattr(schema.chunk_size, "root", schema.chunk_size)
+        chunk_overlap: int = getattr(schema.chunk_overlap, "root", schema.chunk_overlap) or 0
+        max_length: int | None = getattr(schema.max_length, "root", schema.max_length)
+        threshold = schema.confidence_threshold if schema.confidence_threshold is not None else 0.7
+        default_severity = schema.severity or Severity.info
+
+        best_scores: dict[str, float] = {}
+        try:
+            for chunk in _chunk_text(text, chunk_size, chunk_overlap):
+                call_kwargs: dict[str, Any] = {"truncation": True}
+                if max_length is not None:
+                    call_kwargs["max_length"] = max_length
+                raw = self._pipe(chunk, **call_kwargs) or []
+                preds: list[dict[str, Any]] = raw[0] if raw and isinstance(raw[0], list) else raw
+                for pred in preds:
+                    label: str = pred.get("label", "unknown")
+                    score: float = float(pred.get("score", 0.0))
+                    if score > best_scores.get(label, 0.0):
+                        best_scores[label] = score
+        except Exception as exc:
+            logger.error(
+                "text_classification error (model=%s): %s", schema.model, exc, exc_info=True
+            )
+
+        results: list[DetectionResult] = []
+        for label, score in best_scores.items():
+            if score < threshold:
+                continue
+            severity = _resolve_pipeline_severity(label, schema.severity_map, default_severity)
+            results.append(
+                self._make_result(
+                    finding_type=f"classification:{label}",
+                    category="CONTENT",
+                    severity=severity,
+                    confidence=score,
+                    matched_content=text[:512],
+                    location=None,
+                    metadata={"model": schema.model, "predicted_label": label, "score": score},
+                )
+            )
+        results.sort(key=lambda r: r.confidence, reverse=True)
+        return results
+
+    def get_supported_content_types(self) -> list[str]:
+        return list(_TEXT_CONTENT_TYPES)

src/detectors/custom/trainer.py

@@ -0,0 +1,293 @@
+"""GLiNER2 pipeline fine-tuning trainer.
+
+Handles two training modes in a single pass:
+  - NER fine-tuning via the base GLiNER model (entity examples with span values)
+  - Zero-shot classification fine-tuning via SetFit (text + label pairs)
+
+Artifacts are written to an output directory structured as:
+  <output_dir>/
+    gliner2/       -- fine-tuned GLiNER2 model weights (HuggingFace format)
+    setfit/<task>/ -- one SetFit model per classification task
+    manifest.json  -- training metadata for the runner
+
+The trainer produces a JSON result dict to stdout, which the API reads to update
+the training run record.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import time
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+# Minimum examples needed before we attempt fine-tuning (not just annotation storage)
+_MIN_NER_EXAMPLES = 5
+_MIN_SETFIT_PER_CLASS = 2
+
+
+@dataclass
+class TrainingExample:
+    label: str
+    text: str
+    value: str | None = None  # specific entity span text (NER only)
+    accepted: bool = True
+    source: str | None = None
+
+
+@dataclass
+class TrainingResult:
+    status: str
+    trained_examples: int
+    positive_examples: int
+    negative_examples: int
+    model_artifact_path: str
+    metrics: dict[str, Any] = field(default_factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "status": self.status,
+            "trained_examples": self.trained_examples,
+            "positive_examples": self.positive_examples,
+            "negative_examples": self.negative_examples,
+            "model_artifact_path": self.model_artifact_path,
+            "metrics": self.metrics,
+        }
+
+
+class GLiNER2Trainer:
+    """Orchestrates NER + classification fine-tuning for a GLiNER2 pipeline."""
+
+    def __init__(
+        self,
+        pipeline_schema: dict[str, Any],
+        examples_raw: list[dict[str, Any]],
+        output_dir: Path,
+    ) -> None:
+        self._schema = pipeline_schema
+        self._output_dir = output_dir
+        self._examples: list[TrainingExample] = [
+            TrainingExample(
+                label=str(ex.get("label", "")),
+                text=str(ex.get("text", "")),
+                value=ex.get("value") or None,
+                accepted=bool(ex.get("accepted", True)),
+                source=ex.get("source") or None,
+            )
+            for ex in examples_raw
+            if ex.get("label") and ex.get("text")
+        ]
+
+    def train(self) -> TrainingResult:
+        t0 = time.monotonic()
+        self._output_dir.mkdir(parents=True, exist_ok=True)
+
+        positive = [e for e in self._examples if e.accepted]
+        negative = [e for e in self._examples if not e.accepted]
+        metrics: dict[str, Any] = {}
+
+        entities: dict[str, Any] = self._schema.get("entities") or {}
+        classification: dict[str, Any] = self._schema.get("classification") or {}
+        base_model: str = (self._schema.get("model") or {}).get("name") or "fastino/gliner2-base-v1"
+
+        # ── NER fine-tuning ────────────────────────────────────────────────────
+        entity_labels = set(entities.keys())
+        ner_examples = [e for e in positive if e.label in entity_labels and e.value]
+        if ner_examples:
+            metrics["ner"] = self._train_ner(ner_examples, base_model)
+        else:
+            metrics["ner"] = {"skipped": True, "reason": "No span-annotated NER examples"}
+
+        # ── Classification fine-tuning (SetFit) ────────────────────────────────
+        if classification:
+            metrics["classification"] = self._train_classification(positive, classification)
+        else:
+            metrics["classification"] = {
+                "skipped": True,
+                "reason": "No classification tasks defined",
+            }
+
+        # Write manifest so the runner knows what's available
+        manifest: dict[str, Any] = {
+            "schema_type": self._schema.get("type", "GLINER2"),
+            "base_model": base_model,
+            "trained_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
+            "metrics": metrics,
+        }
+        (self._output_dir / "manifest.json").write_text(json.dumps(manifest, indent=2))
+
+        metrics["duration_s"] = round(time.monotonic() - t0, 2)
+
+        return TrainingResult(
+            status="SUCCEEDED",
+            trained_examples=len(positive),
+            positive_examples=len(positive),
+            negative_examples=len(negative),
+            model_artifact_path=str(self._output_dir),
+            metrics=metrics,
+        )
+
+    # ── NER ────────────────────────────────────────────────────────────────────
+
+    def _train_ner(self, examples: list[TrainingExample], base_model: str) -> dict[str, Any]:
+        if len(examples) < _MIN_NER_EXAMPLES:
+            return {
+                "skipped": True,
+                "reason": f"Need ≥{_MIN_NER_EXAMPLES} span-annotated examples (got {len(examples)})",
+            }
+
+        # Build GLiNER-format span annotations
+        train_data: list[dict[str, Any]] = []
+        for ex in examples:
+            if not ex.value:
+                continue
+            start = ex.text.find(ex.value)
+            if start < 0:
+                continue
+            train_data.append(
+                {
+                    "text": ex.text,
+                    "ner": [{"start": start, "end": start + len(ex.value), "label": ex.label}],
+                }
+            )
+
+        if len(train_data) < _MIN_NER_EXAMPLES:
+            return {
+                "skipped": True,
+                "reason": f"Too few locatable spans after search (got {len(train_data)})",
+            }
+
+        gliner_out = self._output_dir / "gliner2"
+        gliner_out.mkdir(parents=True, exist_ok=True)
+
+        try:
+            from gliner import GLiNER  # type: ignore[import-untyped]
+
+            model = GLiNER.from_pretrained(base_model)
+
+            # Attempt to use the trainer API if available
+            try:
+                from gliner.training import Trainer as GLiNERTrainer  # type: ignore[import-untyped]
+                from gliner.training import (
+                    TrainingArguments as GLiNERArgs,  # type: ignore[import-untyped]
+                )
+
+                args = GLiNERArgs(
+                    output_dir=str(gliner_out),
+                    num_train_epochs=3,
+                    per_device_train_batch_size=min(4, len(train_data)),
+                    warmup_ratio=0.1,
+                    save_steps=0,
+                )
+                trainer = GLiNERTrainer(model=model, args=args, train_dataset=train_data)
+                trainer.train()
+            except ImportError:
+                # Fallback: manual training loop if the Trainer API isn't available
+                logger.warning("gliner.training not available — using manual fine-tuning loop")
+                self._manual_ner_train(model, train_data)
+
+            model.save_pretrained(str(gliner_out))
+            logger.info("NER model saved to %s", gliner_out)
+            return {"examples": len(train_data), "epochs": 3, "saved_to": "gliner2/"}
+
+        except ImportError as e:
+            return {"skipped": True, "reason": f"gliner not installed: {e}"}
+        except Exception as e:
+            logger.warning("NER fine-tuning failed: %s", e, exc_info=True)
+            return {"skipped": True, "reason": str(e)}
+
+    def _manual_ner_train(self, model: Any, train_data: list[dict[str, Any]]) -> None:
+        """Simple SGD loop when Trainer API is unavailable."""
+        import torch  # type: ignore[import-untyped]
+
+        optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5)
+        model.train()
+        for _epoch in range(3):
+            for item in train_data:
+                optimizer.zero_grad()
+                loss = model.compute_loss(item)
+                if loss is not None:
+                    loss.backward()
+                    optimizer.step()
+
+    # ── Classification (SetFit) ────────────────────────────────────────────────
+
+    def _train_classification(
+        self,
+        positive: list[TrainingExample],
+        classification: dict[str, Any],
+    ) -> dict[str, Any]:
+        task_results: dict[str, Any] = {}
+        for task_name, task_defn in classification.items():
+            labels: list[str] = task_defn.get("labels") or []
+            task_examples = [e for e in positive if e.label in labels]
+            if len(task_examples) < _MIN_SETFIT_PER_CLASS * max(len(labels), 1):
+                task_results[task_name] = {
+                    "skipped": True,
+                    "reason": (
+                        f"Need ≥{_MIN_SETFIT_PER_CLASS} examples per label "
+                        f"(got {len(task_examples)} across {len(labels)} labels)"
+                    ),
+                }
+                continue
+            try:
+                task_results[task_name] = self._train_setfit_task(task_name, task_examples, labels)
+            except Exception as e:
+                logger.warning("SetFit training failed for task '%s': %s", task_name, e)
+                task_results[task_name] = {"skipped": True, "reason": str(e)}
+        return task_results
+
+    def _train_setfit_task(
+        self,
+        task_name: str,
+        examples: list[TrainingExample],
+        labels: list[str],
+    ) -> dict[str, Any]:
+        try:
+            from datasets import Dataset  # type: ignore[import-untyped]
+            from setfit import (  # type: ignore[import-untyped]
+                SetFitModel,
+                Trainer,
+                TrainingArguments,
+            )
+        except ImportError as e:
+            return {"skipped": True, "reason": f"setfit/datasets not installed: {e}"}
+
+        label2id = {label: i for i, label in enumerate(labels)}
+        dataset = Dataset.from_dict(
+            {
+                "text": [ex.text for ex in examples],
+                "label": [label2id.get(ex.label, 0) for ex in examples],
+            }
+        )
+
+        model = SetFitModel.from_pretrained(
+            "sentence-transformers/paraphrase-MiniLM-L6-v2",
+            labels=labels,
+        )
+
+        save_path = self._output_dir / "setfit" / task_name
+        save_path.mkdir(parents=True, exist_ok=True)
+
+        args = TrainingArguments(
+            output_dir=str(save_path),
+            num_epochs=1,
+            batch_size=min(8, len(examples)),
+        )
+        trainer = Trainer(model=model, args=args, train_dataset=dataset)
+        trainer.train()
+        model.save_pretrained(str(save_path))
+
+        # Persist label ordering so the runner can decode predictions
+        (save_path / "labels.json").write_text(json.dumps(labels))
+
+        logger.info("SetFit model for task '%s' saved to %s", task_name, save_path)
+        return {
+            "examples": len(examples),
+            "labels": labels,
+            "saved_to": f"setfit/{task_name}/",
+        }

src/detectors/dependencies.py

@@ -0,0 +1,109 @@
+"""Helpers for optional detector dependencies."""
+
+from __future__ import annotations
+
+import importlib
+import logging
+from types import ModuleType
+
+from src.utils.uv_sync import auto_install_enabled, sync_group
+
+logger = logging.getLogger(__name__)
+
+
+class MissingDependencyError(RuntimeError):
+    """Raised when an optional detector dependency is unavailable or invalid."""
+
+    def __init__(
+        self,
+        detector_name: str,
+        dependencies: list[str],
+        uv_groups: list[str],
+        detail: str | None = None,
+    ) -> None:
+        self.detector_name = detector_name
+        self.dependencies = dependencies
+        self.uv_groups = uv_groups
+        self.detail = detail
+
+        deps = ", ".join(dependencies)
+        group_hint = " or ".join(f"`uv sync --group {group}`" for group in uv_groups)
+        message = (
+            f"{detector_name} detector requires optional dependencies ({deps}). "
+            f"Install with {group_hint}."
+        )
+        if detail:
+            message = f"{message} {detail}"
+
+        super().__init__(message)
+
+
+def _ordered_groups(groups: list[str]) -> list[str]:
+    unique = list(dict.fromkeys(groups))
+    return sorted(unique, key=lambda group: (group == "detectors", group))
+
+
+def require_module(
+    module_name: str,
+    detector_name: str,
+    uv_groups: list[str],
+    detail: str | None = None,
+) -> ModuleType:
+    """Import a module or raise a MissingDependencyError with uv guidance."""
+    try:
+        return importlib.import_module(module_name)
+    except Exception as exc:  # pragma: no cover - exercised indirectly in integration setups
+        detail_messages: list[str] = [f"Original error: {exc}"]
+
+        if auto_install_enabled() and uv_groups:
+            for group in _ordered_groups(uv_groups):
+                success, install_detail = sync_group(group)
+                if install_detail:
+                    detail_messages.append(install_detail)
+                if not success:
+                    continue
+
+                try:
+                    importlib.invalidate_caches()
+                    return importlib.import_module(module_name)
+                except Exception as retry_exc:  # pragma: no cover
+                    detail_messages.append(
+                        f"Module '{module_name}' still unavailable after installing '{group}': {retry_exc}"
+                    )
+
+        base_detail = detail or "Optional dependency import failed"
+        error_detail = (
+            f"{base_detail}. {'; '.join(detail_messages)}" if detail_messages else base_detail
+        )
+        raise MissingDependencyError(
+            detector_name=detector_name,
+            dependencies=[module_name.split(".", maxsplit=1)[0]],
+            uv_groups=uv_groups,
+            detail=error_detail,
+        ) from exc
+
+
+def ensure_torch(detector_name: str, uv_groups: list[str]) -> ModuleType:
+    """Verify PyTorch is importable and looks like a valid install."""
+    torch_module = require_module("torch", detector_name, uv_groups)
+    if not hasattr(torch_module, "no_grad"):
+        raise MissingDependencyError(
+            detector_name=detector_name,
+            dependencies=["torch"],
+            uv_groups=uv_groups,
+            detail=(
+                "Detected a module named 'torch' but it is missing `no_grad`. "
+                "Ensure PyTorch is installed via uv and no local `torch.py` shadows it."
+            ),
+        )
+    if not hasattr(torch_module, "_utils"):
+        raise MissingDependencyError(
+            detector_name=detector_name,
+            dependencies=["torch"],
+            uv_groups=uv_groups,
+            detail=(
+                "Detected an incomplete/broken PyTorch install (`torch._utils` missing). "
+                "Reinstall torch via uv for this environment."
+            ),
+        )
+    return torch_module