Spaces:
Running on Zero
Running on Zero
File size: 4,218 Bytes
465b477 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 | """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")
|