mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-07-06 17:31:16 -03:00
feat(agent): fix extract_relevant_section false positives, add validation pipeline audit
- extract_relevant_section: raise token threshold >3, verify anchor sections contain basename, require 2+ heading token overlaps, skip TOC-style headings (markdown links), verify heading section size - metadata_constructor: parse repo_id,model_name.safetensors format so model_path basename matches real filename - config: replace hardcoded SUPPORTED_BASE_MODELS with dynamic init_supported_base_models() using production list_base_models() - preprocessing_auditor: new Phase 1.5 audit module — fetches each README, runs extract_relevant_section + clean_readme_for_llm, records stats and flags, saves raw READMEs for cross-reference - run_validation: integrate audit phase, add --audit-only mode, add LLM config consistency check, add ComfyUI root to sys.path - report_generator: add Preprocessing Audit and Config Warnings sections to both markdown and JSON reports
This commit is contained in:
467
tests/enrich_hf_validation/preprocessing_auditor.py
Normal file
467
tests/enrich_hf_validation/preprocessing_auditor.py
Normal file
@@ -0,0 +1,467 @@
|
||||
"""Preprocessing audit for the HF metadata enrichment validation pipeline.
|
||||
|
||||
Phase 1.5 — runs between Phase 1 (metadata creation) and Phase 2 (enrichment).
|
||||
|
||||
Audits the README preprocessing pipeline (section extraction + cleaning)
|
||||
for each repo in the dataset, capturing intermediate outputs so we can
|
||||
distinguish between:
|
||||
|
||||
(A) Preprocessing failed → LLM never saw the right content
|
||||
(B) Preprocessing succeeded → LLM/prompt needs improvement
|
||||
|
||||
This prevents wasted effort optimizing prompts when the actual problem is
|
||||
that ``extract_relevant_section`` or ``clean_readme_for_llm`` removed or
|
||||
misaligned the content the LLM needed.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
from dataclasses import dataclass, field, asdict
|
||||
from typing import Any, Dict, List, Tuple
|
||||
|
||||
import aiohttp
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Audit record
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@dataclass
|
||||
class AuditRecord:
|
||||
"""Preprocessing audit for a single repo entry."""
|
||||
|
||||
# Identity
|
||||
repo_id: str
|
||||
safetensors_name: str
|
||||
basename: str # filename without .safetensors
|
||||
|
||||
# Raw README stats
|
||||
raw_readme_length: int
|
||||
raw_readme_line_count: int
|
||||
has_yaml_frontmatter: bool
|
||||
yaml_has_base_model: bool
|
||||
yaml_has_tags: bool
|
||||
|
||||
# Section extraction
|
||||
section_extraction_activated: bool # output < 95% of input length
|
||||
section_length: int
|
||||
section_line_count: int
|
||||
basename_in_section: bool # basename appears in extracted section text
|
||||
|
||||
# Cleaning
|
||||
cleaned_length: int
|
||||
cleaned_line_count: int
|
||||
compression_pct: float # (1 - cleaned/raw) * 100
|
||||
|
||||
# Widget section (stripped by _strip_widget_section)
|
||||
widget_section_found: bool
|
||||
widget_section_length: int
|
||||
|
||||
# Flags (list of anomaly descriptions)
|
||||
flags: List[str] = field(default_factory=list)
|
||||
|
||||
# Local file path to the saved raw README (for cross-reference)
|
||||
readme_file: str = ""
|
||||
|
||||
# Staged intermediate output for report detail
|
||||
raw_readme_preview: str = "" # first 200 chars
|
||||
section_preview: str = "" # first 300 chars
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Constants
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_HF_RAW_URL = "https://huggingface.co/{repo_id}/raw/main/README.md"
|
||||
|
||||
# Thresholds for flagging
|
||||
_SECTION_ACTIVATION_RATIO = 0.95
|
||||
_MIN_CLEANED_LENGTH = 100
|
||||
_MAX_COMPRESSION_PCT = 99.0
|
||||
_MIN_SECTION_LINES = 3
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Module loader — bypasses parent-package __init__ that imports ComfyUI
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_readme_processor_module = None
|
||||
|
||||
|
||||
def _load_readme_processor():
|
||||
"""Import ``readme_processor`` without triggering ``folder_paths`` import.
|
||||
|
||||
The normal import path (``py.services.agent.skills.enrich_hf_metadata.
|
||||
readme_processor``) triggers ``py.services.agent.__init__`` which
|
||||
imports ``agent_service.py`` → ``py/config.py`` → ComfyUI's
|
||||
``folder_paths``, which is not available in standalone mode.
|
||||
"""
|
||||
global _readme_processor_module
|
||||
if _readme_processor_module is not None:
|
||||
return _readme_processor_module
|
||||
|
||||
import importlib.util
|
||||
|
||||
_RP_PATH = os.path.join(
|
||||
os.path.dirname(__file__), # tests/enrich_hf_validation/
|
||||
"..", "..",
|
||||
"py", "services", "agent", "skills", "enrich_hf_metadata",
|
||||
"readme_processor.py",
|
||||
)
|
||||
rp_path = os.path.normpath(_RP_PATH)
|
||||
if not os.path.exists(rp_path):
|
||||
logger.error("readme_processor.py not found at %s", rp_path)
|
||||
return None
|
||||
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
"readme_processor", rp_path,
|
||||
)
|
||||
if spec is None or spec.loader is None:
|
||||
logger.error("Could not create spec for readme_processor.py")
|
||||
return None
|
||||
|
||||
mod = importlib.util.module_from_spec(spec)
|
||||
try:
|
||||
spec.loader.exec_module(mod)
|
||||
except Exception as exc:
|
||||
logger.error("Failed to load readme_processor.py: %s", exc)
|
||||
return None
|
||||
|
||||
_readme_processor_module = mod
|
||||
return mod
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# HF README fetcher
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def _fetch_readme(repo_id: str, session: aiohttp.ClientSession) -> str:
|
||||
"""Fetch the raw README.md from HuggingFace."""
|
||||
url = _HF_RAW_URL.format(repo_id=repo_id)
|
||||
try:
|
||||
async with session.get(url, timeout=aiohttp.ClientTimeout(total=30)) as resp:
|
||||
if resp.status == 200:
|
||||
return await resp.text()
|
||||
logger.warning("Failed to fetch README for %s: HTTP %d", repo_id, resp.status)
|
||||
return ""
|
||||
except (asyncio.TimeoutError, aiohttp.ClientError) as exc:
|
||||
logger.warning("Failed to fetch README for %s: %s", repo_id, exc)
|
||||
return ""
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Analysis helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _has_yaml_frontmatter(text: str) -> bool:
|
||||
return bool(text.strip().startswith("---"))
|
||||
|
||||
|
||||
def _extract_yaml_field(text: str, field: str) -> bool:
|
||||
"""Check if the given YAML field exists in the frontmatter."""
|
||||
lines = text.split("\n")
|
||||
if not lines or not lines[0].strip().startswith("---"):
|
||||
return False
|
||||
end = 1
|
||||
while end < len(lines):
|
||||
if lines[end].strip().startswith("---"):
|
||||
break
|
||||
end += 1
|
||||
if end >= len(lines):
|
||||
return False
|
||||
frontmatter = "\n".join(lines[1:end])
|
||||
pattern = rf"^{field}:"
|
||||
return bool(re.search(pattern, frontmatter, re.MULTILINE))
|
||||
|
||||
|
||||
def _find_widget_section_length(text: str) -> int:
|
||||
"""Find the ``widget:`` YAML section and return its length (0 if none)."""
|
||||
if not _has_yaml_frontmatter(text):
|
||||
return 0
|
||||
frontmatter_end = text.find("---", 3)
|
||||
if frontmatter_end == -1:
|
||||
return 0
|
||||
frontmatter = text[3:frontmatter_end]
|
||||
|
||||
# Match widget: through to the next top-level key or frontmatter end
|
||||
m = re.search(r"\nwidget:", frontmatter)
|
||||
if not m:
|
||||
return 0
|
||||
# Length from widget: to end of frontmatter (the next \n\w+: or \n---)
|
||||
return len(frontmatter[m.start():])
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Core auditor
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def run_audit(
|
||||
entries: List[Tuple[str, str]],
|
||||
*,
|
||||
concurrency: int = 10,
|
||||
readmes_dir: str | None = None,
|
||||
) -> Tuple[List[AuditRecord], Dict[str, Any]]:
|
||||
"""Run the preprocessing audit over all repo entries.
|
||||
|
||||
Args:
|
||||
entries: List of ``(repo_id, safetensors_name)``.
|
||||
concurrency: Max parallel fetches to HuggingFace.
|
||||
readmes_dir: If set, saves each fetched README as
|
||||
``{sanitized_repo_id}.md`` in this directory for offline
|
||||
cross-reference against audit results.
|
||||
|
||||
Returns:
|
||||
Tuple of ``(records, summary)`` where *summary* is a dict with
|
||||
aggregate statistics.
|
||||
"""
|
||||
semaphore = asyncio.Semaphore(concurrency)
|
||||
records: List[AuditRecord] = []
|
||||
flag_counter: Dict[str, int] = {}
|
||||
|
||||
if readmes_dir:
|
||||
os.makedirs(readmes_dir, exist_ok=True)
|
||||
|
||||
connector = aiohttp.TCPConnector(limit=concurrency)
|
||||
async with aiohttp.ClientSession(connector=connector) as session:
|
||||
tasks = [_audit_one(entry, session, semaphore, readmes_dir=readmes_dir) for entry in entries]
|
||||
gathered = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
|
||||
for entry, result in zip(entries, gathered):
|
||||
if isinstance(result, Exception):
|
||||
logger.error("Audit failed for %s: %s", entry[0], result)
|
||||
records.append(
|
||||
AuditRecord(
|
||||
repo_id=entry[0],
|
||||
safetensors_name=entry[1],
|
||||
basename=os.path.splitext(entry[1])[0],
|
||||
raw_readme_length=0,
|
||||
raw_readme_line_count=0,
|
||||
has_yaml_frontmatter=False,
|
||||
yaml_has_base_model=False,
|
||||
yaml_has_tags=False,
|
||||
section_extraction_activated=False,
|
||||
section_length=0,
|
||||
section_line_count=0,
|
||||
basename_in_section=False,
|
||||
cleaned_length=0,
|
||||
cleaned_line_count=0,
|
||||
compression_pct=0.0,
|
||||
widget_section_found=False,
|
||||
widget_section_length=0,
|
||||
readme_file="",
|
||||
flags=[f"Audit exception: {result}"],
|
||||
)
|
||||
)
|
||||
continue
|
||||
|
||||
# The continue above ensures result is AuditRecord here
|
||||
assert isinstance(result, AuditRecord)
|
||||
records.append(result)
|
||||
for flag in result.flags:
|
||||
flag_counter[flag] = flag_counter.get(flag, 0) + 1
|
||||
|
||||
summary = _build_summary(records, flag_counter)
|
||||
return records, summary
|
||||
|
||||
|
||||
def _sanitize_repo_id(repo_id: str) -> str:
|
||||
"""Turn ``user/repo-name`` into a safe filename."""
|
||||
return repo_id.replace("/", "__").replace(".", "_")
|
||||
|
||||
|
||||
async def _audit_one(
|
||||
entry: Tuple[str, str],
|
||||
session: aiohttp.ClientSession,
|
||||
semaphore: asyncio.Semaphore,
|
||||
*,
|
||||
readmes_dir: str | None = None,
|
||||
) -> AuditRecord:
|
||||
"""Audit a single repo entry."""
|
||||
repo_id, safetensors_name = entry
|
||||
basename = os.path.splitext(safetensors_name)[0]
|
||||
|
||||
async with semaphore:
|
||||
# Import production preprocessing functions.
|
||||
# Use importlib to bypass py.services.agent.__init__ which triggers
|
||||
# ComfyUI's folder_paths module (not available in standalone mode).
|
||||
_rp = _load_readme_processor()
|
||||
if _rp is None:
|
||||
return AuditRecord(
|
||||
repo_id=repo_id,
|
||||
safetensors_name=safetensors_name,
|
||||
basename=basename,
|
||||
raw_readme_length=0, raw_readme_line_count=0,
|
||||
has_yaml_frontmatter=False, yaml_has_base_model=False, yaml_has_tags=False,
|
||||
readme_file="",
|
||||
section_extraction_activated=False, section_length=0, section_line_count=0,
|
||||
basename_in_section=False, cleaned_length=0, cleaned_line_count=0,
|
||||
compression_pct=0.0, widget_section_found=False, widget_section_length=0,
|
||||
flags=["IMPORT_FAILED"],
|
||||
)
|
||||
clean_readme_for_llm = _rp.clean_readme_for_llm
|
||||
extract_relevant_section = _rp.extract_relevant_section
|
||||
|
||||
# Step 1: Fetch the raw README
|
||||
raw_text = await _fetch_readme(repo_id, session)
|
||||
if not raw_text:
|
||||
return AuditRecord(
|
||||
repo_id=repo_id,
|
||||
safetensors_name=safetensors_name,
|
||||
basename=basename,
|
||||
raw_readme_length=0,
|
||||
raw_readme_line_count=0,
|
||||
has_yaml_frontmatter=False,
|
||||
yaml_has_base_model=False,
|
||||
yaml_has_tags=False,
|
||||
section_extraction_activated=False,
|
||||
section_length=0,
|
||||
section_line_count=0,
|
||||
basename_in_section=False,
|
||||
readme_file="",
|
||||
cleaned_length=0,
|
||||
cleaned_line_count=0,
|
||||
compression_pct=0.0,
|
||||
widget_section_found=False,
|
||||
widget_section_length=0,
|
||||
flags=["README_FETCH_FAILED"],
|
||||
)
|
||||
|
||||
# Save the raw README to disk for offline cross-reference
|
||||
readme_path = ""
|
||||
if readmes_dir:
|
||||
safe_name = _sanitize_repo_id(repo_id)
|
||||
readme_path = os.path.join(readmes_dir, f"{safe_name}.md")
|
||||
try:
|
||||
with open(readme_path, "w", encoding="utf-8") as fh:
|
||||
fh.write(raw_text)
|
||||
except OSError as exc:
|
||||
logger.warning("Failed to save README for %s: %s", repo_id, exc)
|
||||
readme_path = ""
|
||||
|
||||
raw_lines = raw_text.split("\n")
|
||||
raw_len = len(raw_text)
|
||||
raw_line_count = len(raw_lines)
|
||||
|
||||
# Step 2: Analyze raw README
|
||||
yaml_fm = _has_yaml_frontmatter(raw_text)
|
||||
yaml_has_bm = _extract_yaml_field(raw_text, "base_model") if yaml_fm else False
|
||||
yaml_has_tg = _extract_yaml_field(raw_text, "tags") if yaml_fm else False
|
||||
widget_len = _find_widget_section_length(raw_text)
|
||||
|
||||
# Step 3: Section extraction
|
||||
section = extract_relevant_section(raw_text, basename)
|
||||
section_len = len(section)
|
||||
section_line_count = len(section.split("\n"))
|
||||
section_activated = section_len < raw_len * _SECTION_ACTIVATION_RATIO
|
||||
basename_in_sec = basename.lower() in section.lower()
|
||||
|
||||
# Step 4: Cleaning for LLM
|
||||
cleaned = clean_readme_for_llm(section)
|
||||
cleaned_len = len(cleaned)
|
||||
cleaned_line_count = len(cleaned.split("\n"))
|
||||
compression_pct = round((1 - cleaned_len / raw_len) * 100, 1) if raw_len else 0.0
|
||||
|
||||
# Step 5: Flag anomalies
|
||||
flags: List[str] = []
|
||||
if not raw_text.strip():
|
||||
flags.append("README_EMPTY")
|
||||
if not yaml_fm:
|
||||
flags.append("NO_YAML_FRONTMATTER")
|
||||
if not section_activated:
|
||||
# Check if basename is extremely short/generic (likely synthetic)
|
||||
if len(basename) <= 5:
|
||||
flags.append("BASENAME_TOO_SHORT_SECTION_NOT_EXPECTED")
|
||||
else:
|
||||
flags.append("SECTION_EXTRACTION_NOT_ACTIVATED")
|
||||
elif not basename_in_sec:
|
||||
flags.append("BASENAME_NOT_IN_EXTRACTED_SECTION")
|
||||
if widget_len == 0:
|
||||
# Not necessarily a problem — many repos lack a widget section
|
||||
pass
|
||||
if cleaned_len < _MIN_CLEANED_LENGTH:
|
||||
flags.append("CLEANED_README_TOO_SHORT")
|
||||
if compression_pct > _MAX_COMPRESSION_PCT:
|
||||
flags.append("EXTREME_COMPRESSION")
|
||||
if section_activated and section_line_count < _MIN_SECTION_LINES:
|
||||
flags.append("SECTION_TOO_SMALL")
|
||||
|
||||
return AuditRecord(
|
||||
repo_id=repo_id,
|
||||
safetensors_name=safetensors_name,
|
||||
basename=basename,
|
||||
raw_readme_length=raw_len,
|
||||
raw_readme_line_count=raw_line_count,
|
||||
has_yaml_frontmatter=yaml_fm,
|
||||
yaml_has_base_model=yaml_has_bm,
|
||||
yaml_has_tags=yaml_has_tg,
|
||||
section_extraction_activated=section_activated,
|
||||
section_length=section_len,
|
||||
section_line_count=section_line_count,
|
||||
basename_in_section=basename_in_sec,
|
||||
cleaned_length=cleaned_len,
|
||||
cleaned_line_count=cleaned_line_count,
|
||||
compression_pct=compression_pct,
|
||||
widget_section_found=widget_len > 0,
|
||||
widget_section_length=widget_len,
|
||||
readme_file=readme_path,
|
||||
flags=flags,
|
||||
raw_readme_preview=raw_text[:200],
|
||||
section_preview=section[:300],
|
||||
)
|
||||
|
||||
|
||||
def _build_summary(
|
||||
records: List[AuditRecord],
|
||||
flag_counter: Dict[str, int],
|
||||
) -> Dict[str, Any]:
|
||||
"""Aggregate audit statistics."""
|
||||
n = len(records)
|
||||
if n == 0:
|
||||
return {"error": "no records", "model_count": 0}
|
||||
|
||||
activated = sum(1 for r in records if r.section_extraction_activated)
|
||||
basename_hit = sum(1 for r in records if r.basename_in_section)
|
||||
with_yaml = sum(1 for r in records if r.has_yaml_frontmatter)
|
||||
with_widget = sum(1 for r in records if r.widget_section_found)
|
||||
fetch_failed = sum(1 for r in records if "README_FETCH_FAILED" in r.flags)
|
||||
|
||||
avg_compression = round(
|
||||
sum(r.compression_pct for r in records if r.raw_readme_length > 0) / max(n - fetch_failed, 1),
|
||||
1,
|
||||
)
|
||||
avg_cleaned = round(
|
||||
sum(r.cleaned_length for r in records if r.raw_readme_length > 0) / max(n - fetch_failed, 1),
|
||||
)
|
||||
|
||||
top_flags = sorted(flag_counter.items(), key=lambda x: -x[1])[:10]
|
||||
|
||||
return {
|
||||
"model_count": n,
|
||||
"fetch_failed_count": fetch_failed,
|
||||
"section_extraction_activated": activated,
|
||||
"section_extraction_pct": round(activated / max(n - fetch_failed, 1) * 100, 1),
|
||||
"basename_in_section": basename_hit,
|
||||
"basename_in_section_pct": round(basename_hit / max(n - fetch_failed, 1) * 100, 1),
|
||||
"with_yaml_frontmatter": with_yaml,
|
||||
"with_yaml_frontmatter_pct": round(with_yaml / max(n - fetch_failed, 1) * 100, 1),
|
||||
"with_widget_section": with_widget,
|
||||
"avg_compression_pct": avg_compression,
|
||||
"avg_cleaned_length": avg_cleaned,
|
||||
"top_flags": top_flags,
|
||||
}
|
||||
|
||||
|
||||
def audit_records_to_serializable(records: List[AuditRecord]) -> List[Dict[str, Any]]:
|
||||
"""Convert AuditRecord dataclasses to plain dicts for JSON serialization."""
|
||||
return [asdict(r) for r in records]
|
||||
Reference in New Issue
Block a user