mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-07-06 09:21: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
208
tests/enrich_hf_validation/enrichment_runner.py
Normal file
208
tests/enrich_hf_validation/enrichment_runner.py
Normal file
@@ -0,0 +1,208 @@
|
||||
"""Execute the ``enrich_hf_metadata`` skill serially over a list of models.
|
||||
|
||||
Design decisions (local Ollama, no rate limits):
|
||||
|
||||
- Sequential execution: one model at a time. 100 models at ~30-90 s/call
|
||||
→ roughly 1-2 h total.
|
||||
- Progress persisted to a JSON checkpoint file so the run can be resumed
|
||||
with ``--resume``.
|
||||
- Per-model timeout guards against a stuck Ollama inference.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_SKILL_NAME = "enrich_hf_metadata"
|
||||
|
||||
# How long to wait for a single LLM call before marking it timed-out.
|
||||
_PER_MODEL_TIMEOUT = 240 # seconds
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Progress checkpoint helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_PROGRESS_FILE = "progress.json"
|
||||
|
||||
|
||||
def _load_progress(output_dir: str) -> Dict[str, Any]:
|
||||
path = os.path.join(output_dir, _PROGRESS_FILE)
|
||||
if os.path.exists(path):
|
||||
with open(path, "r") as fh:
|
||||
return json.load(fh)
|
||||
return {"completed": [], "failed": [], "timed_out": []}
|
||||
|
||||
|
||||
def _save_progress(output_dir: str, progress: Dict[str, Any]) -> None:
|
||||
path = os.path.join(output_dir, _PROGRESS_FILE)
|
||||
with open(path, "w") as fh:
|
||||
json.dump(progress, fh, indent=2)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Core runner
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class EnrichmentRunner:
|
||||
"""Serial enrichment runner with checkpoint resume."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
output_dir: str,
|
||||
*,
|
||||
per_model_timeout: int = _PER_MODEL_TIMEOUT,
|
||||
) -> None:
|
||||
self._output_dir = output_dir
|
||||
self._per_model_timeout = per_model_timeout
|
||||
self._agent_service: Optional[Any] = None
|
||||
|
||||
async def _ensure_agent_service(self) -> Any:
|
||||
"""Lazy-init AgentService (expensive — needs LLMService init)."""
|
||||
if self._agent_service is not None:
|
||||
return self._agent_service
|
||||
from py.services.agent.agent_service import AgentService
|
||||
|
||||
self._agent_service = await AgentService.get_instance()
|
||||
return self._agent_service
|
||||
|
||||
async def run(
|
||||
self,
|
||||
model_paths: List[str],
|
||||
repos: List[str],
|
||||
) -> Dict[str, Any]:
|
||||
"""Run enrichment over *model_paths* (one-by-one).
|
||||
|
||||
Args:
|
||||
model_paths: model paths in the same order as *repos*.
|
||||
repos: HF repo IDs (for display / checkpoint labelling).
|
||||
|
||||
Returns:
|
||||
A dict with keys ``results``, ``progress``, ``durations``.
|
||||
"""
|
||||
assert len(model_paths) == len(repos)
|
||||
|
||||
progress = _load_progress(self._output_dir)
|
||||
completed_set = set(progress["completed"])
|
||||
failed_set = set(progress["failed"])
|
||||
timed_out_set = set(progress.get("timed_out", []))
|
||||
|
||||
agent = await self._ensure_agent_service()
|
||||
results: List[Dict[str, Any]] = []
|
||||
durations: Dict[str, float] = {}
|
||||
|
||||
total = len(model_paths)
|
||||
processed_before = len(completed_set | failed_set | timed_out_set)
|
||||
|
||||
logger.info(
|
||||
"Enrichment runner: %d models total, %d already processed",
|
||||
total,
|
||||
processed_before,
|
||||
)
|
||||
|
||||
for idx, (model_path, repo_id) in enumerate(zip(model_paths, repos)):
|
||||
if repo_id in completed_set:
|
||||
logger.info("[%d/%d] SKIP (already done): %s", idx + 1, total, repo_id)
|
||||
continue
|
||||
if repo_id in failed_set or repo_id in timed_out_set:
|
||||
logger.info(
|
||||
"[%d/%d] SKIP (previously failed/timeout): %s",
|
||||
idx + 1, total, repo_id,
|
||||
)
|
||||
continue
|
||||
|
||||
logger.info(
|
||||
"[%d/%d] Enriching %s ...", idx + 1, total, repo_id,
|
||||
)
|
||||
t0 = time.perf_counter()
|
||||
|
||||
try:
|
||||
result = await asyncio.wait_for(
|
||||
agent.execute_skill(
|
||||
skill_name=_SKILL_NAME,
|
||||
input_data={"model_paths": [model_path]},
|
||||
progress_callback=None,
|
||||
),
|
||||
timeout=self._per_model_timeout,
|
||||
)
|
||||
|
||||
elapsed = time.perf_counter() - t0
|
||||
durations[repo_id] = round(elapsed, 2)
|
||||
|
||||
if result.success:
|
||||
completed_set.add(repo_id)
|
||||
progress["completed"].append(repo_id)
|
||||
logger.info(
|
||||
" ✓ %s (%.1f s) — %s",
|
||||
repo_id, elapsed, result.summary,
|
||||
)
|
||||
else:
|
||||
failed_set.add(repo_id)
|
||||
progress["failed"].append(repo_id)
|
||||
logger.warning(
|
||||
" ✗ %s (%.1f s) — %s",
|
||||
repo_id, elapsed,
|
||||
"; ".join(result.errors) if result.errors else result.summary,
|
||||
)
|
||||
|
||||
results.append({
|
||||
"repo_id": repo_id,
|
||||
"model_path": model_path,
|
||||
"success": result.success,
|
||||
"updated_fields": result.updated_models,
|
||||
"errors": result.errors,
|
||||
"summary": result.summary,
|
||||
"duration_s": round(elapsed, 2),
|
||||
})
|
||||
|
||||
except asyncio.TimeoutError:
|
||||
elapsed = time.perf_counter() - t0
|
||||
durations[repo_id] = round(elapsed, 2)
|
||||
timed_out_set.add(repo_id)
|
||||
progress.setdefault("timed_out", []).append(repo_id)
|
||||
logger.warning(
|
||||
" ⏱ TIMEOUT %s (%.1f s, limit=%ds)",
|
||||
repo_id, elapsed, self._per_model_timeout,
|
||||
)
|
||||
results.append({
|
||||
"repo_id": repo_id,
|
||||
"model_path": model_path,
|
||||
"success": False,
|
||||
"errors": [f"Timeout after {self._per_model_timeout}s"],
|
||||
"summary": "LLM call timed out",
|
||||
"duration_s": round(elapsed, 2),
|
||||
})
|
||||
|
||||
except Exception as exc:
|
||||
elapsed = time.perf_counter() - t0
|
||||
durations[repo_id] = round(elapsed, 2)
|
||||
failed_set.add(repo_id)
|
||||
progress["failed"].append(repo_id)
|
||||
logger.error(
|
||||
" ✗ %s (%.1f s) — %s",
|
||||
repo_id, elapsed, exc,
|
||||
)
|
||||
results.append({
|
||||
"repo_id": repo_id,
|
||||
"model_path": model_path,
|
||||
"success": False,
|
||||
"errors": [str(exc)],
|
||||
"summary": f"Exception: {exc}",
|
||||
"duration_s": round(elapsed, 2),
|
||||
})
|
||||
|
||||
# Checkpoint after each model
|
||||
_save_progress(self._output_dir, progress)
|
||||
|
||||
return {
|
||||
"results": results,
|
||||
"progress": progress,
|
||||
"durations": durations,
|
||||
}
|
||||
Reference in New Issue
Block a user