mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-07-06 17:31:16 -03:00
feat(agent): enrich_hf_metadata — filename-aware section matching, preview extraction for markdown/HTML/widget, JSON salvage, instance_prompt fallback, and validation suite
- extract_relevant_section(): trim README to model-filename-matching section for collection repos (download link, anchor ID, heading strategies) - _strip_standalone_images(): preserve markdown image URLs so LLM can extract preview_url; strip only HTML <img> tags - extract_simple_markdown_images(): extract civitai.images from ![]() body - extract_html_img_tags(): extract from <img src="..."> (deadman44-style) - extract_gallery_images(): fix widget parser for YAML - output: dash prefix - _is_heading: exclude </hN> closing tags from boundary detection - _extract_section: start at matching heading when match IS a heading line - _try_salvage_json(): recover truncated JSON (close braces/brackets in LIFO order, close unterminated strings, strip trailing commas) - PostProcessor: store _llm_confidence, add instance_prompt YAML fallback - agent_service: pass model_basename to prompt, trim README via extract_relevant_section before clean_readme_for_llm - Add tests/enrich_hf_validation/ suite: 100-model pipeline with progress checkpoint/resume, per-field scoring, markdown+JSON reporting - Fix evaluation_engine: read _llm_confidence (not _llm_response)
This commit is contained in:
Will Miao
294
tests/enrich_hf_validation/run_validation.py
Normal file
294
tests/enrich_hf_validation/run_validation.py
Normal file
@@ -0,0 +1,294 @@
|
||||
#!/usr/bin/env python3
|
||||
"""CLI entry point for the HF metadata enrichment validation suite.
|
||||
|
||||
Usage::
|
||||
|
||||
# Full run (100 models, serial, ~1-2 h)
|
||||
python -m tests.enrich_hf_validation.run_validation \\
|
||||
--output /tmp/hf_enrich_validation
|
||||
|
||||
# Quick smoke test with 2 models
|
||||
python -m tests.enrich_hf_validation.run_validation --sample 2
|
||||
|
||||
# Resume from a previous partial run
|
||||
python -m tests.enrich_hf_validation.run_validation --resume
|
||||
|
||||
# Custom settings file
|
||||
python -m tests.enrich_hf_validation.run_validation \\
|
||||
--settings /custom/path/settings.json
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from typing import Any, Dict, List
|
||||
|
||||
# Ensure the project root is on sys.path so that ``from py import ...`` works.
|
||||
_PROJECT_ROOT = os.path.normpath(
|
||||
os.path.join(os.path.dirname(__file__), "..", "..")
|
||||
)
|
||||
if _PROJECT_ROOT not in sys.path:
|
||||
sys.path.insert(0, _PROJECT_ROOT)
|
||||
|
||||
from tests.enrich_hf_validation.config import load_settings
|
||||
from tests.enrich_hf_validation.metadata_constructor import (
|
||||
create_all_initial_metadata,
|
||||
load_repo_ids,
|
||||
)
|
||||
from tests.enrich_hf_validation.enrichment_runner import EnrichmentRunner
|
||||
from tests.enrich_hf_validation.evaluation_engine import (
|
||||
aggregate_scores,
|
||||
evaluate_batch,
|
||||
)
|
||||
from tests.enrich_hf_validation.report_generator import (
|
||||
generate_markdown_report,
|
||||
save_json_report,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _setup_logging(verbose: bool) -> None:
|
||||
level = logging.DEBUG if verbose else logging.INFO
|
||||
fmt = "%(asctime)s [%(levelname)s] %(name)s: %(message)s"
|
||||
logging.basicConfig(level=level, format=fmt, stream=sys.stderr)
|
||||
|
||||
# Quiet noisy third-party loggers
|
||||
for name in ("aiohttp", "asyncio", "urllib3"):
|
||||
logging.getLogger(name).setLevel(logging.WARNING)
|
||||
|
||||
|
||||
def _parse_args(argv: List[str]) -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Validate and optimise HF metadata enrichment via LLM.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--models",
|
||||
default="~/Documents/hf_lora_models.txt",
|
||||
help="Path to the HF repo ID list (one per line)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--settings",
|
||||
default="~/.config/ComfyUI-LoRA-Manager/settings.json",
|
||||
help="Path to LoRA Manager settings.json",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
default="/tmp/hf_enrich_validation",
|
||||
help="Output directory for reports and intermediate data",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--sample",
|
||||
type=int,
|
||||
default=0,
|
||||
help="Process only the first N models (for quick smoke tests)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--resume",
|
||||
action="store_true",
|
||||
help="Resume from previous partial run (uses progress.json)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-enrich",
|
||||
action="store_true",
|
||||
help="Skip enrichment phase (evaluate existing metadata only)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--timeout",
|
||||
type=int,
|
||||
default=240,
|
||||
help="Per-model LLM timeout in seconds (default: 240)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"-v", "--verbose",
|
||||
action="store_true",
|
||||
help="Enable debug logging",
|
||||
)
|
||||
return parser.parse_args(argv)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Phase helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _phase_header(label: str) -> None:
|
||||
sep = "=" * 60
|
||||
print(f"\n{sep}", file=sys.stderr)
|
||||
print(f" PHASE: {label}", file=sys.stderr)
|
||||
print(sep, file=sys.stderr)
|
||||
|
||||
|
||||
async def _run_enrichment(
|
||||
model_paths: List[str],
|
||||
repos: List[str],
|
||||
output_dir: str,
|
||||
timeout: int,
|
||||
verbose: bool,
|
||||
) -> Dict[str, Any]:
|
||||
"""Execute the enrichment phase."""
|
||||
runner = EnrichmentRunner(
|
||||
output_dir=output_dir,
|
||||
per_model_timeout=timeout,
|
||||
)
|
||||
result = await runner.run(model_paths, repos)
|
||||
|
||||
# Print quick summary
|
||||
progress = result["progress"]
|
||||
total_done = (
|
||||
len(progress.get("completed", []))
|
||||
+ len(progress.get("failed", []))
|
||||
+ len(progress.get("timed_out", []))
|
||||
)
|
||||
print(
|
||||
f"\n Enrichment complete: {total_done} processed "
|
||||
f"({len(progress.get('completed', []))} ok, "
|
||||
f"{len(progress.get('failed', []))} failed, "
|
||||
f"{len(progress.get('timed_out', []))} timed out)",
|
||||
file=sys.stderr,
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
def _collect_enriched_metadata(
|
||||
model_paths: List[str],
|
||||
repos: List[str],
|
||||
results: List[Dict[str, Any]],
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""Read enriched .metadata.json for each model.
|
||||
|
||||
Uses the same path convention as the rest of the codebase:
|
||||
``os.path.splitext(model_path)[0] + '.metadata.json'``.
|
||||
|
||||
Returns a list of dicts with keys: repo_id, model_path, success,
|
||||
errors, metadata.
|
||||
"""
|
||||
enriched: List[Dict[str, Any]] = []
|
||||
# Build a lookup from repo_id → enrichment result
|
||||
result_lookup: Dict[str, Dict[str, Any]] = {}
|
||||
for r in results:
|
||||
result_lookup[r["repo_id"]] = r
|
||||
|
||||
for model_path, repo_id in zip(model_paths, repos):
|
||||
res = result_lookup.get(repo_id, {})
|
||||
metadata_path = f"{os.path.splitext(model_path)[0]}.metadata.json"
|
||||
metadata: Dict[str, Any] = {}
|
||||
if os.path.exists(metadata_path):
|
||||
try:
|
||||
with open(metadata_path, "r", encoding="utf-8") as fh:
|
||||
metadata = json.load(fh)
|
||||
except (json.JSONDecodeError, OSError) as exc:
|
||||
logger.warning("Failed to read %s: %s", metadata_path, exc)
|
||||
else:
|
||||
logger.warning(
|
||||
"Metadata file not found for %s (expected: %s)",
|
||||
repo_id, metadata_path,
|
||||
)
|
||||
|
||||
enriched.append({
|
||||
"repo_id": repo_id,
|
||||
"model_path": model_path,
|
||||
"success": res.get("success", False),
|
||||
"errors": res.get("errors", []),
|
||||
"metadata": metadata,
|
||||
})
|
||||
|
||||
return enriched
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Main
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def main(argv: List[str]) -> int:
|
||||
args = _parse_args(argv)
|
||||
_setup_logging(args.verbose)
|
||||
|
||||
output_dir = os.path.abspath(os.path.expanduser(args.output))
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
settings = load_settings(args.settings)
|
||||
logger.info(
|
||||
"LLM config: provider=%s model=%s api_base=%s",
|
||||
settings["llm_provider"],
|
||||
settings["llm_model"],
|
||||
settings["llm_api_base"],
|
||||
)
|
||||
|
||||
# ---- Phase 1: Load repo IDs & construct initial metadata ----
|
||||
_phase_header("Load repo IDs & construct initial metadata")
|
||||
repos = load_repo_ids(args.models, max_models=args.sample if args.sample > 0 else None)
|
||||
model_paths = create_all_initial_metadata(
|
||||
repos, output_dir, skip_existing=True,
|
||||
)
|
||||
print(f" {len(model_paths)} repos ready", file=sys.stderr)
|
||||
|
||||
# ---- Phase 2: Enrichment ----
|
||||
enrichment_results: List[Dict[str, Any]] = []
|
||||
t_start = time.perf_counter()
|
||||
if not args.no_enrich:
|
||||
_phase_header("Enrich metadata via LLM")
|
||||
enrichment_out = await _run_enrichment(
|
||||
model_paths, repos, output_dir, args.timeout, args.verbose,
|
||||
)
|
||||
enrichment_results = enrichment_out["results"]
|
||||
else:
|
||||
print(" Enrichment skipped (--no-enrich)", file=sys.stderr)
|
||||
|
||||
t_enrich = time.perf_counter()
|
||||
|
||||
# ---- Phase 3: Evaluation ----
|
||||
_phase_header("Evaluate enriched metadata")
|
||||
enriched = _collect_enriched_metadata(model_paths, repos, enrichment_results)
|
||||
scores = evaluate_batch(enriched)
|
||||
agg = aggregate_scores(scores)
|
||||
print(
|
||||
f" Mean total score: {agg.get('total_score', {}).get('mean', 'N/A')} / 100",
|
||||
file=sys.stderr,
|
||||
)
|
||||
print(
|
||||
f" Models scored: {agg.get('model_count', 0)}",
|
||||
file=sys.stderr,
|
||||
)
|
||||
|
||||
# ---- Phase 4: Report generation ----
|
||||
_phase_header("Generate reports")
|
||||
duration_summary: Dict[str, Any] | None = None
|
||||
if enrichment_results:
|
||||
durations = [r.get("duration_s", 0) for r in enrichment_results if r.get("duration_s")]
|
||||
if durations:
|
||||
sorted_d = sorted(durations)
|
||||
m = len(sorted_d) // 2
|
||||
duration_summary = {
|
||||
"total_wall_s": round(t_enrich - t_start, 1),
|
||||
"mean_s": round(sum(durations) / len(durations), 1),
|
||||
"median_s": round(sorted_d[m] if len(sorted_d) % 2 else (sorted_d[m - 1] + sorted_d[m]) / 2, 1),
|
||||
"min_s": round(min(durations), 1),
|
||||
"max_s": round(max(durations), 1),
|
||||
}
|
||||
|
||||
save_json_report(agg, scores, enrichment_results, output_dir, duration_summary)
|
||||
generate_markdown_report(agg, scores, output_dir, duration_summary)
|
||||
|
||||
# ---- Final summary ----
|
||||
total_wall = time.perf_counter() - t_start
|
||||
print(f"\n Done in {total_wall:.0f}s ({total_wall / 60:.1f} min)", file=sys.stderr)
|
||||
print(f" Reports: {output_dir}/report.md, {output_dir}/report.json", file=sys.stderr)
|
||||
print(file=sys.stderr)
|
||||
|
||||
return 0 if agg.get("success_count", 0) > 0 else 1
|
||||
|
||||
|
||||
def entry_point() -> int:
|
||||
return asyncio.run(main(sys.argv[1:]))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(entry_point())
|
||||
Reference in New Issue
Block a user