hermes-benchmark 0.0.5__py3-none-any.whl

hermes_benchmark/capacity/__init__.py

@@ -0,0 +1,16 @@
+"""Wikipedia dump profiling and benchmark capacity estimation."""
+
+from .estimates import estimate_task_capacities
+from .models import CapacityReport, CapacityRange, CorpusProfile, TaskCapacity
+from .pilot import run_pilot
+from .profiler import profile_dump
+
+__all__ = [
+    "CapacityRange",
+    "CapacityReport",
+    "CorpusProfile",
+    "TaskCapacity",
+    "estimate_task_capacities",
+    "profile_dump",
+    "run_pilot",
+]

hermes_benchmark/capacity/estimates.py

@@ -0,0 +1,309 @@
+"""Task-capacity formulas compatible with the legacy 10K pipeline."""
+
+from typing import Dict, Iterable, List, Mapping, Optional, Tuple
+
+from hermes_benchmark.config import ServicesConfig, TaskConfig
+
+from .models import (
+    CapacityRange,
+    CapacityReport,
+    CorpusProfile,
+    PoolConstraint,
+    TaskCapacity,
+)
+
+
+_STATUS_PRIORITY = {
+    "disabled": 0,
+    "sufficient": 1,
+    "uncertain": 2,
+    "insufficient": 3,
+    "blocked": 4,
+}
+
+
+def _task_status(values: Iterable[CapacityRange]) -> str:
+    statuses = [value.status for value in values]
+    return max(statuses, key=lambda value: _STATUS_PRIORITY[value])
+
+
+def _missing_services(
+    services: ServicesConfig,
+    required: Iterable[str],
+) -> Tuple[str, ...]:
+    missing = []
+    for name in required:
+        if getattr(services, name) is None:
+            missing.append(name)
+    return tuple(missing)
+
+
+def _option(task: TaskConfig, name: str, default: Optional[int]) -> Optional[int]:
+    if name not in task.options:
+        return default
+    value = task.options[name]
+    if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
+        raise ValueError(f"Task option {name!r} must be a positive integer")
+    return value
+
+
+def _structural_range(
+    profile: CorpusProfile,
+    requested: int,
+    observed: int,
+    estimated: Optional[int],
+    upper: Optional[int],
+    basis: str,
+    missing: Tuple[str, ...] = (),
+) -> CapacityRange:
+    if missing:
+        status = "blocked"
+    elif observed >= requested:
+        status = "sufficient"
+    elif profile.complete:
+        status = "insufficient"
+    else:
+        status = "uncertain"
+    return CapacityRange(
+        requested=requested,
+        observed_candidates=observed,
+        lower_bound=observed,
+        estimated=estimated,
+        upper_bound=observed if profile.complete else upper,
+        confidence="exact" if profile.complete else "sample_projection",
+        status=status,
+        basis=basis,
+        dependencies=missing,
+    )
+
+
+def _proxy_range(
+    profile: CorpusProfile,
+    requested: int,
+    observed: int,
+    projected_upper: Optional[int],
+    basis: str,
+    missing: Tuple[str, ...],
+) -> CapacityRange:
+    upper = observed if profile.complete else projected_upper
+    if missing:
+        status = "blocked"
+    elif upper is not None and upper < requested and profile.complete:
+        status = "insufficient"
+    else:
+        status = "uncertain"
+    return CapacityRange(
+        requested=requested,
+        observed_candidates=observed,
+        lower_bound=0,
+        estimated=None,
+        upper_bound=upper,
+        confidence="proxy_upper_bound",
+        status=status,
+        basis=basis,
+        dependencies=missing,
+    )
+
+
+def _article_upper(profile: CorpusProfile, multiplier: int = 1) -> Optional[int]:
+    if profile.complete:
+        return profile.eligible_articles * multiplier
+    limits = [
+        value
+        for value in (profile.article_limit, profile.max_pages)
+        if value is not None
+    ]
+    return min(limits) * multiplier if limits else None
+
+
+def _disabled(name: str, task: TaskConfig) -> TaskCapacity:
+    return TaskCapacity(
+        name=name,
+        enabled=False,
+        requested=task.target,
+        status="disabled",
+        variants={},
+    )
+
+
+def estimate_task_capacities(
+    profile: CorpusProfile,
+    tasks: Mapping[str, TaskConfig],
+    services: ServicesConfig,
+) -> CapacityReport:
+    """Estimate each configured task without invoking external services."""
+
+    results: Dict[str, TaskCapacity] = {}
+    report_warnings: List[str] = list(profile.warnings)
+    embeddings = _missing_services(services, ("embeddings",))
+    knowledge_base = _missing_services(services, ("knowledge_base",))
+    llm_and_embeddings = _missing_services(services, ("llm", "embeddings"))
+
+    for name, task in tasks.items():
+        if not task.enabled:
+            results[name] = _disabled(name, task)
+            continue
+        variants: Dict[str, CapacityRange]
+        warnings: Tuple[str, ...] = ()
+        if name in {
+            "paragraph_coherence",
+            "paragraph_ordering",
+            "article_membership",
+            "next_paragraph",
+        }:
+            missing = () if name == "paragraph_ordering" else embeddings
+            value = _structural_range(
+                profile,
+                task.target,
+                profile.articles_four_plus,
+                profile.estimated_articles_four_plus,
+                _article_upper(profile),
+                "articles with at least four selected paragraphs",
+                missing,
+            )
+            variants = {"default": value}
+        elif name == "reranking":
+            candidate_count = _option(task, "candidate_count", 20)
+            query_per_title = _option(task, "query_per_title", 1)
+            assert candidate_count is not None
+            assert query_per_title is not None
+            observed = (
+                profile.reranking_articles * query_per_title
+                if profile.reranking_articles >= candidate_count
+                else 0
+            )
+            projected_articles = profile.estimated_reranking_articles
+            estimated = (
+                projected_articles * query_per_title
+                if projected_articles is not None
+                and projected_articles >= candidate_count
+                else (0 if projected_articles is not None else None)
+            )
+            variants = {
+                "type1": _structural_range(
+                    profile,
+                    task.target,
+                    observed,
+                    estimated,
+                    _article_upper(profile, query_per_title),
+                    (
+                        "reranking-ready articles; at least "
+                        f"{candidate_count} distinct titles required"
+                    ),
+                    embeddings,
+                )
+            }
+        elif name == "classification":
+            candidate_total = _option(task, "candidate_total", None)
+            observed = profile.category_articles
+            projected = profile.estimated_category_articles
+            if candidate_total is not None:
+                observed = min(observed, candidate_total)
+                projected = (
+                    min(projected, candidate_total) if projected is not None else None
+                )
+            value = _proxy_range(
+                profile,
+                task.target,
+                observed,
+                projected,
+                "category-bearing article proxy before clustering and LLM review",
+                llm_and_embeddings,
+            )
+            variants = {name: value for name in ("type1", "type2", "type3")}
+        elif name == "ner":
+            context = _proxy_range(
+                profile,
+                task.target,
+                profile.entity_contexts,
+                profile.estimated_entity_contexts,
+                "internal-link context proxy before Wikidata typing",
+                knowledge_base,
+            )
+            mention = _proxy_range(
+                profile,
+                task.target,
+                profile.entity_mentions,
+                profile.estimated_entity_mentions,
+                "internal-link mention proxy before Wikidata typing",
+                knowledge_base,
+            )
+            variants = {"type1": context, "type2": context, "type3": mention}
+        elif name == "extractive_qa":
+            qa_service_names = ["knowledge_base"]
+            if task.options.get("generation_method", "template") == "hybrid":
+                qa_service_names.append("llm")
+            qa_services = _missing_services(services, tuple(qa_service_names))
+            variants = {
+                "type1": _proxy_range(
+                    profile,
+                    task.target,
+                    profile.qa_capped_candidates,
+                    profile.estimated_qa_candidates,
+                    ("entity-bearing QA context proxy with per-label " "candidate cap"),
+                    qa_services,
+                )
+            }
+        else:
+            value = CapacityRange(
+                requested=task.target,
+                observed_candidates=0,
+                lower_bound=0,
+                estimated=None,
+                upper_bound=None,
+                confidence="unsupported",
+                status="blocked",
+                basis="unsupported task",
+            )
+            variants = {"default": value}
+            warnings = (f"No capacity model is registered for task {name!r}.",)
+            report_warnings.extend(warnings)
+
+        results[name] = TaskCapacity(
+            name=name,
+            enabled=True,
+            requested=task.target,
+            status=_task_status(variants.values()),
+            variants=variants,
+            warnings=warnings,
+        )
+
+    constraints = []
+    for constraint_name, pair in (
+        (
+            "multiple_choice_disjoint_articles",
+            ("paragraph_coherence", "paragraph_ordering"),
+        ),
+        (
+            "other_disjoint_articles",
+            ("article_membership", "next_paragraph"),
+        ),
+    ):
+        if not all(name in tasks and tasks[name].enabled for name in pair):
+            continue
+        requested = sum(tasks[name].target for name in pair)
+        if profile.articles_four_plus >= requested:
+            status = "sufficient"
+        elif profile.complete:
+            status = "insufficient"
+        else:
+            status = "uncertain"
+        constraints.append(
+            PoolConstraint(
+                name=constraint_name,
+                tasks=pair,
+                requested=requested,
+                observed_available=profile.articles_four_plus,
+                estimated_available=profile.estimated_articles_four_plus,
+                status=status,
+            )
+        )
+
+    return CapacityReport(
+        schema_version=1,
+        run_status="complete",
+        profile=profile,
+        tasks=results,
+        constraints=tuple(constraints),
+        warnings=tuple(dict.fromkeys(report_warnings)),
+    )

hermes_benchmark/capacity/models.py

@@ -0,0 +1,142 @@
+"""Public models for dump profiles and task-capacity reports."""
+
+from dataclasses import asdict, dataclass
+from typing import Any, Dict, Mapping, Optional, Tuple
+
+
+@dataclass(frozen=True)
+class CorpusProfile:
+    """Observed and projected properties of a configured corpus scan."""
+
+    mode: str
+    dump_path: str
+    complete: bool
+    completion_reason: str
+    pages_scanned: int
+    main_namespace_pages: int
+    redirects_skipped: int
+    empty_pages_skipped: int
+    title_suffixes_skipped: int
+    insufficient_paragraphs_skipped: int
+    eligible_articles: int
+    paragraphs_available: int
+    paragraphs_selected: int
+    articles_four_plus: int
+    category_articles: int
+    entity_contexts: int
+    entity_mentions: int
+    qa_candidate_contexts: int
+    qa_capped_candidates: int
+    reranking_articles: int
+    bytes_read: int
+    total_bytes: int
+    estimated_pages: Optional[int]
+    estimated_eligible_articles: Optional[int]
+    estimated_articles_four_plus: Optional[int]
+    estimated_category_articles: Optional[int]
+    estimated_entity_contexts: Optional[int]
+    estimated_entity_mentions: Optional[int]
+    estimated_qa_candidates: Optional[int]
+    estimated_reranking_articles: Optional[int]
+    article_limit: Optional[int]
+    max_pages: Optional[int]
+    warnings: Tuple[str, ...] = ()
+
+    @property
+    def fraction_read(self) -> Optional[float]:
+        if self.total_bytes <= 0:
+            return None
+        return min(1.0, self.bytes_read / self.total_bytes)
+
+    def to_dict(self) -> Dict[str, Any]:
+        value = asdict(self)
+        value["warnings"] = list(self.warnings)
+        value["fraction_read"] = self.fraction_read
+        return value
+
+
+@dataclass(frozen=True)
+class CapacityRange:
+    """Capacity information for one concrete output variant."""
+
+    requested: int
+    observed_candidates: int
+    lower_bound: int
+    estimated: Optional[int]
+    upper_bound: Optional[int]
+    confidence: str
+    status: str
+    basis: str
+    dependencies: Tuple[str, ...] = ()
+
+    def to_dict(self) -> Dict[str, Any]:
+        value = asdict(self)
+        value["dependencies"] = list(self.dependencies)
+        return value
+
+
+@dataclass(frozen=True)
+class TaskCapacity:
+    """Capacity result for one configured task and its variants."""
+
+    name: str
+    enabled: bool
+    requested: int
+    status: str
+    variants: Mapping[str, CapacityRange]
+    warnings: Tuple[str, ...] = ()
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "name": self.name,
+            "enabled": self.enabled,
+            "requested": self.requested,
+            "status": self.status,
+            "variants": {
+                name: value.to_dict() for name, value in self.variants.items()
+            },
+            "warnings": list(self.warnings),
+        }
+
+
+@dataclass(frozen=True)
+class PoolConstraint:
+    """A legacy disjoint-selection constraint shared by multiple tasks."""
+
+    name: str
+    tasks: Tuple[str, ...]
+    requested: int
+    observed_available: int
+    estimated_available: Optional[int]
+    status: str
+
+    def to_dict(self) -> Dict[str, Any]:
+        value = asdict(self)
+        value["tasks"] = list(self.tasks)
+        return value
+
+
+@dataclass(frozen=True)
+class CapacityReport:
+    """Complete, serializable capacity analysis result."""
+
+    schema_version: int
+    run_status: str
+    profile: CorpusProfile
+    tasks: Mapping[str, TaskCapacity]
+    constraints: Tuple[PoolConstraint, ...] = ()
+    warnings: Tuple[str, ...] = ()
+    output_dir: Optional[str] = None
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "schema_version": self.schema_version,
+            "run_status": self.run_status,
+            "profile": self.profile.to_dict(),
+            "tasks": {
+                name: value.to_dict() for name, value in self.tasks.items()
+            },
+            "constraints": [value.to_dict() for value in self.constraints],
+            "warnings": list(self.warnings),
+            "output_dir": self.output_dir,
+        }

hermes_benchmark/capacity/pilot.py

@@ -0,0 +1,133 @@
+"""Safe, isolated pilot corpus construction."""
+
+import json
+import os
+import shutil
+import tempfile
+import uuid
+from dataclasses import replace
+from pathlib import Path
+
+from hermes_benchmark.config import GenerationConfig
+from hermes_benchmark.progress import progress_bar
+from hermes_benchmark.wikipedia import build_corpus
+
+from .estimates import estimate_task_capacities
+from .models import CapacityReport
+from .profiler import profile_dump
+
+
+def _bounded(requested: int, configured: int) -> int:
+    return min(requested, configured)
+
+
+def run_pilot(
+    config: GenerationConfig,
+    *,
+    article_limit: int = 100,
+    max_pages: int = 10_000,
+) -> CapacityReport:
+    """Build and publish an isolated pilot corpus plus its capacity report."""
+
+    if article_limit <= 0:
+        raise ValueError("article_limit must be greater than zero")
+    if max_pages <= 0:
+        raise ValueError("max_pages must be greater than zero")
+
+    effective_articles = article_limit
+    if config.corpus.article_limit is not None:
+        effective_articles = _bounded(
+            effective_articles,
+            config.corpus.article_limit,
+        )
+    effective_pages = max_pages
+    if config.corpus.max_pages is not None:
+        effective_pages = _bounded(effective_pages, config.corpus.max_pages)
+
+    stage_progress = progress_bar(total=2, desc="Pilot", unit="stage")
+    try:
+        stage_progress.set_postfix_str("profiling", refresh=False)
+        profile = profile_dump(
+            config.dump.path,
+            config.corpus,
+            page_limit=effective_pages,
+            tasks=config.tasks,
+            mode="pilot",
+        )
+        base_report = estimate_task_capacities(
+            profile,
+            config.tasks,
+            config.services,
+        )
+        stage_progress.update(1)
+        stage_progress.set_postfix_str("building corpus", refresh=False)
+
+        output_root = Path(config.output_dir)
+        final_dir = output_root / "pilot"
+        output_root.mkdir(parents=True, exist_ok=True)
+        staging_dir = Path(
+            tempfile.mkdtemp(prefix=".pilot.", dir=str(output_root))
+        )
+        backup_dir = output_root / f".pilot.backup.{uuid.uuid4().hex}"
+        warnings = list(base_report.warnings)
+        pilot_corpus = config.corpus.model_copy(
+            update={
+                "article_limit": effective_articles,
+                "max_pages": effective_pages,
+            }
+        )
+        stats = build_corpus(
+            config.dump.path,
+            staging_dir,
+            pilot_corpus,
+        )
+        if stats.articles_written == 0:
+            run_status = "fail"
+            warnings.append("The pilot did not produce any eligible articles.")
+        elif stats.articles_written < effective_articles:
+            run_status = "pass_with_warnings"
+            warnings.append(
+                "The pilot reached the dump/page limit before the requested "
+                f"{effective_articles} articles; {stats.articles_written} were written."
+            )
+        else:
+            run_status = "pass"
+
+        report = replace(
+            base_report,
+            run_status=run_status,
+            warnings=tuple(dict.fromkeys(warnings)),
+            output_dir=str(final_dir),
+        )
+        report_path = staging_dir / "report.json"
+        with report_path.open("w", encoding="utf-8") as report_file:
+            json.dump(
+                report.to_dict(),
+                report_file,
+                ensure_ascii=False,
+                indent=2,
+                sort_keys=True,
+            )
+            report_file.write("\n")
+
+        moved_old = False
+        try:
+            if final_dir.exists():
+                os.replace(final_dir, backup_dir)
+                moved_old = True
+            os.replace(staging_dir, final_dir)
+        except BaseException:
+            if moved_old and backup_dir.exists() and not final_dir.exists():
+                os.replace(backup_dir, final_dir)
+            raise
+        if backup_dir.exists():
+            shutil.rmtree(backup_dir)
+        stage_progress.update(1)
+        stage_progress.set_postfix_str("complete", refresh=False)
+        return report
+    except BaseException:
+        if "staging_dir" in locals() and staging_dir.exists():
+            shutil.rmtree(staging_dir)
+        raise
+    finally:
+        stage_progress.close()