Spaces:
Running on Zero
Running on Zero
| """YAML configuration, model registry and run manifests.""" | |
| from __future__ import annotations | |
| import json | |
| import platform | |
| import subprocess | |
| from datetime import UTC, datetime | |
| from importlib import metadata | |
| from pathlib import Path | |
| from typing import Any | |
| import yaml | |
| from pydantic import BaseModel, Field | |
| from .bench import BenchConfig | |
| from .chunking import ChunkingConfig | |
| from .llm import CachedClient, ModelSpec, build_client | |
| from .llm.limit import ConcurrencyLimited | |
| from .methods import MethodConfig | |
| from .prompts import PROMPT_VERSION | |
| from .tokens import tokenizer_name | |
| class Paths(BaseModel): | |
| raw: str = "data/raw" | |
| processed: str = "data/processed" | |
| graph: str = "data/graph" | |
| bench: str = "data/bench" | |
| runs: str = "runs" | |
| llm_cache: str = "cache/llm" | |
| class CorpusConfig(BaseModel): | |
| splits: list[str] = ["test", "validation"] | |
| n_papers: int = 120 | |
| min_mentions: int = 2 | |
| class ExtractionConfig(BaseModel): | |
| model: str = "extractor" | |
| workers: int = 4 | |
| nli_confidence: bool = True | |
| nli_model: str = "cross-encoder/nli-deberta-v3-small" | |
| curated_aliases: str = "configs/aliases.yaml" | |
| class ProjectConfig(BaseModel): | |
| paths: Paths = Field(default_factory=Paths) | |
| corpus: CorpusConfig = Field(default_factory=CorpusConfig) | |
| chunking: ChunkingConfig = Field(default_factory=ChunkingConfig) | |
| extraction: ExtractionConfig = Field(default_factory=ExtractionConfig) | |
| bench: BenchConfig = Field(default_factory=BenchConfig) | |
| paraphrase_model: str | None = None | |
| models_file: str = "configs/models.yaml" | |
| class ExperimentConfig(BaseModel): | |
| name: str | |
| bench_file: str | |
| allow_unreviewed: bool = False # only for pipeline smoke tests; never for reported results | |
| limit: int | None = None | |
| base_model: str | |
| judge_model: str | None = None | |
| nli_verifier: bool = True | |
| nli_model: str = "cross-encoder/nli-deberta-v3-small" | |
| min_support: float | None = 0.5 | |
| workers: int = 4 | |
| methods: list[MethodConfig] | |
| reference_method: str = "text_rag" | |
| def load_yaml(path: str | Path) -> dict[str, Any]: | |
| return yaml.safe_load(Path(path).read_text(encoding="utf-8")) or {} | |
| def load_project(path: str | Path = "configs/default.yaml") -> ProjectConfig: | |
| p = Path(path) | |
| return ProjectConfig(**load_yaml(p)) if p.exists() else ProjectConfig() | |
| def load_models(path: str | Path) -> dict[str, ModelSpec]: | |
| return {name: ModelSpec(**spec) for name, spec in load_yaml(path).get("models", {}).items()} | |
| def make_client(name: str, project: ProjectConfig, offline: bool = False) -> CachedClient: | |
| models = load_models(project.models_file) | |
| if name not in models: | |
| raise KeyError(f"model '{name}' is not defined in {project.models_file}; known: {sorted(models)}") | |
| spec = models[name] | |
| # Cache outside the limiter: hits never wait; misses respect the model's max_concurrency. | |
| return CachedClient(ConcurrencyLimited(build_client(spec), spec.max_concurrency), project.paths.llm_cache, offline=offline) | |
| def _git_commit() -> str | None: | |
| try: | |
| return subprocess.run(["git", "rev-parse", "HEAD"], capture_output=True, text=True, check=True, timeout=10).stdout.strip() | |
| except Exception: # noqa: BLE001 - git may be missing or the tree may not be a repo; the commit is optional metadata | |
| return None | |
| def manifest(extra: dict[str, Any]) -> dict[str, Any]: | |
| pkgs = {} | |
| for name in ("graphtrace", "networkx", "rank-bm25", "rapidfuzz", "tiktoken", "pydantic", "anthropic", "httpx", "sentence-transformers", "torch"): | |
| try: | |
| pkgs[name] = metadata.version(name) | |
| except metadata.PackageNotFoundError: | |
| pass | |
| return { | |
| "created_utc": datetime.now(UTC).isoformat(), | |
| "prompt_version": PROMPT_VERSION, | |
| "tokenizer": tokenizer_name(), | |
| "git_commit": _git_commit(), | |
| "python": platform.python_version(), | |
| "platform": platform.platform(), | |
| "packages": pkgs, | |
| **extra, | |
| } | |
| def write_json(path: str | Path, obj: Any) -> None: | |
| p = Path(path) | |
| p.parent.mkdir(parents=True, exist_ok=True) | |
| p.write_text(json.dumps(obj, indent=2, ensure_ascii=False, default=str), encoding="utf-8") | |