mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-07-06 01:11:17 -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
@@ -24,11 +24,16 @@ from typing import Any, Dict, List, Optional
|
||||
|
||||
import aiohttp
|
||||
|
||||
import os
|
||||
|
||||
from ..llm_service import LLMService
|
||||
from ..websocket_manager import ws_manager
|
||||
from .post_processor import PostProcessor
|
||||
from .skill_registry import SkillRegistry
|
||||
from .skills.enrich_hf_metadata.md_to_html import clean_readme_for_llm
|
||||
from .skills.enrich_hf_metadata.md_to_html import (
|
||||
clean_readme_for_llm,
|
||||
extract_relevant_section,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -344,6 +349,7 @@ class AgentService:
|
||||
|
||||
context: Dict[str, Any] = {
|
||||
"model_path": model_path,
|
||||
"model_basename": "",
|
||||
"hf_url": "",
|
||||
"repo": "",
|
||||
"readme_content": "",
|
||||
@@ -353,6 +359,11 @@ class AgentService:
|
||||
"priority_tags": "",
|
||||
}
|
||||
|
||||
# Extract model basename (filename without extension) for the LLM
|
||||
# to use when locating the matching section in collection repos.
|
||||
raw_basename = os.path.splitext(os.path.basename(model_path))[0]
|
||||
context["model_basename"] = raw_basename or ""
|
||||
|
||||
context["current_metadata"] = {
|
||||
"file_name": metadata.get("file_name", ""),
|
||||
"base_model": metadata.get("base_model", ""),
|
||||
@@ -369,7 +380,13 @@ class AgentService:
|
||||
context["repo"] = repo or ""
|
||||
if repo:
|
||||
readme = await self._fetch_readme(repo)
|
||||
cleaned = clean_readme_for_llm(readme) if readme else ""
|
||||
# Trim README to the section relevant to this model file
|
||||
# (collection repos often have multiple models in one README).
|
||||
if readme and raw_basename:
|
||||
trimmed = extract_relevant_section(readme, raw_basename)
|
||||
cleaned = clean_readme_for_llm(trimmed) if trimmed else ""
|
||||
else:
|
||||
cleaned = clean_readme_for_llm(readme) if readme else ""
|
||||
context["readme_content"] = cleaned if cleaned else "(README not available)"
|
||||
context["readme_content_full"] = readme or ""
|
||||
|
||||
|
||||
@@ -13,6 +13,7 @@ from __future__ import annotations
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
@@ -81,6 +82,8 @@ class PostProcessor:
|
||||
convert_readme_to_html,
|
||||
extract_gallery_images,
|
||||
extract_gallery_table_images,
|
||||
extract_simple_markdown_images,
|
||||
extract_html_img_tags,
|
||||
extract_repo_from_hf_url,
|
||||
)
|
||||
|
||||
@@ -101,9 +104,11 @@ class PostProcessor:
|
||||
|
||||
# trigger words → civitai.trainedWords
|
||||
new_triggers = llm_output.get("trigger_words", [])
|
||||
trigger_words_empty = True
|
||||
if isinstance(new_triggers, list):
|
||||
cleaned = [t.strip() for t in new_triggers if t.strip()]
|
||||
cleaned = [t for t in cleaned if t.lower() not in ("none", "null", "n/a")]
|
||||
trigger_words_empty = not cleaned
|
||||
current_civitai = metadata.get("civitai") or {}
|
||||
current_triggers = current_civitai.get("trainedWords") or []
|
||||
if self._should_overwrite_list(current_triggers, is_hf_model):
|
||||
@@ -152,8 +157,24 @@ class PostProcessor:
|
||||
existing_urls=existing_urls,
|
||||
default_width=rec_w, default_height=rec_h,
|
||||
)
|
||||
existing_urls.update(img["url"] for img in table_images if img.get("url"))
|
||||
|
||||
all_images = gallery + table_images
|
||||
# 3. Simple markdown images `` in the body
|
||||
simple_images = extract_simple_markdown_images(
|
||||
readme_content, repo,
|
||||
existing_urls=existing_urls,
|
||||
default_width=rec_w, default_height=rec_h,
|
||||
)
|
||||
existing_urls.update(img["url"] for img in simple_images if img.get("url"))
|
||||
|
||||
# 4. HTML `<img>` tags (used by many collection repos)
|
||||
html_images = extract_html_img_tags(
|
||||
readme_content, repo,
|
||||
existing_urls=existing_urls,
|
||||
default_width=rec_w, default_height=rec_h,
|
||||
)
|
||||
|
||||
all_images = gallery + table_images + simple_images + html_images
|
||||
if all_images:
|
||||
gallery_images = all_images
|
||||
current_civitai = metadata.get("civitai") or {}
|
||||
@@ -175,6 +196,23 @@ class PostProcessor:
|
||||
updates["metadata_source"] = "agent:enrich_hf_metadata"
|
||||
updates["llm_enriched_at"] = datetime.now(timezone.utc).isoformat()
|
||||
|
||||
# Store LLM confidence in metadata so it's accessible for evaluation
|
||||
raw_confidence = (llm_output.get("confidence") or "").strip()
|
||||
if raw_confidence:
|
||||
updates["_llm_confidence"] = raw_confidence
|
||||
|
||||
# Fallback: extract instance_prompt from YAML frontmatter when the LLM
|
||||
# returned empty trigger words but the README has instance_prompt.
|
||||
if trigger_words_empty:
|
||||
instance_prompt = _extract_yaml_instance_prompt(readme_content)
|
||||
if instance_prompt:
|
||||
current_civitai = metadata.get("civitai") or {}
|
||||
trig_civitai = dict(current_civitai)
|
||||
if "civitai" in updates and isinstance(updates["civitai"], dict):
|
||||
trig_civitai.update(updates["civitai"])
|
||||
trig_civitai["trainedWords"] = [instance_prompt]
|
||||
updates["civitai"] = trig_civitai
|
||||
|
||||
preview_remote_url = (llm_output.get("preview_url") or "").strip()
|
||||
# Fallback: if the LLM couldn't find a preview image in the cleaned
|
||||
# README, use the first gallery image extracted from the YAML widget
|
||||
@@ -249,3 +287,35 @@ class PostProcessor:
|
||||
merged.append(t)
|
||||
seen.add(t)
|
||||
return merged
|
||||
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Module-level helpers
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
|
||||
def _extract_yaml_instance_prompt(readme_content: str) -> str:
|
||||
"""Extract ``instance_prompt`` from the YAML frontmatter of a HF README.
|
||||
|
||||
Returns the prompt text, or empty string if not found. Handles
|
||||
``null`` / ``~`` YAML null values by returning empty string.
|
||||
"""
|
||||
if not readme_content or not readme_content.startswith("---"):
|
||||
return ""
|
||||
|
||||
# Find end of frontmatter
|
||||
end = readme_content.find("---", 3)
|
||||
if end == -1:
|
||||
return ""
|
||||
frontmatter = readme_content[3:end]
|
||||
|
||||
for line in frontmatter.split("\n"):
|
||||
line = line.strip()
|
||||
m = re.match(r"^instance_prompt:\s*(.*)", line)
|
||||
if m:
|
||||
val = m.group(1).strip().strip('"').strip("'")
|
||||
if val.lower() in ("null", "~", "none", ""):
|
||||
return ""
|
||||
return val
|
||||
|
||||
return ""
|
||||
|
||||
@@ -13,6 +13,7 @@ You are an expert assistant for AI image generation models. Your task is to extr
|
||||
|
||||
- **Repository**: {{hf_url}}
|
||||
- **Model file path**: {{model_path}}
|
||||
- **Model filename**: {{model_basename}}
|
||||
- **Repository ID**: {{repo}}
|
||||
|
||||
## Current Metadata (may be incomplete)
|
||||
@@ -36,6 +37,31 @@ These are the subjects, styles, and concepts the user considers useful for categ
|
||||
The following base models are currently valid in this system:
|
||||
{{base_models}}
|
||||
|
||||
## HuggingFace → CivitAI Base Model Mapping
|
||||
|
||||
HuggingFace repos often declare `base_model:` in YAML frontmatter using HuggingFace
|
||||
model names. Map them to CivitAI names using this reference:
|
||||
|
||||
| HuggingFace repo / model name | CivitAI base model |
|
||||
|-------------------------------------------------------|--------------------|
|
||||
| `runwayml/stable-diffusion-v1-5` | SD 1.5 |
|
||||
| `CompVis/stable-diffusion-v1-4` | SD 1.5 |
|
||||
| `stabilityai/stable-diffusion-2-1` / `-2-1-base` | SD 2.1 |
|
||||
| `stabilityai/stable-diffusion-xl-base-1.0` | SDXL 1.0 |
|
||||
| `stabilityai/stable-diffusion-3-5-large` | SD 3.5 Large |
|
||||
| `black-forest-labs/FLUX.1-dev` | Flux.1 D |
|
||||
| `black-forest-labs/FLUX.1-schnell` | Flux.1 S |
|
||||
| `black-forest-labs/FLUX.1-krea` | Flux.1 Krea |
|
||||
| `krea/Krea-2-Raw` or `krea/Krea-2-Turbo` | Krea 2 |
|
||||
| `Comfy-Org/z_image_turbo` | ZImageTurbo |
|
||||
| `alibaba-pai/Z-Image-Turbo-Fun-Controlnet-Union-2.1` | ZImageTurbo |
|
||||
| `CogVideo` | CogVideoX |
|
||||
| `Wan-AI/Wan2.1-T2V-14B` | Wan Video 2.2 T2V-A14B |
|
||||
| `stabilityai/stable-video-diffusion-img2vid` | SVD |
|
||||
|
||||
The model file name itself may also hint at the base model (e.g. "flux", "sdxl",
|
||||
"sd15", "krea2" in the filename).
|
||||
|
||||
## HuggingFace README Content
|
||||
|
||||
```
|
||||
@@ -49,24 +75,26 @@ Extract the following information from the README content above:
|
||||
### base_model
|
||||
The base model this model was trained on. Use EXACTLY one of the names from the **Available Base Models** list above. Do not invent new names or use aliases.
|
||||
|
||||
Check the YAML frontmatter (between --- markers) for `base_model:` first, then look at the description text and safetensors metadata. If you cannot determine it, return an empty string.
|
||||
Check the YAML frontmatter (between --- markers) for `base_model:` first, then look at the description text and safetensors metadata. If the YAML frontmatter uses a HuggingFace model name (e.g. `runwayml/stable-diffusion-v1-5`), use the mapping table above to find the correct CivitAI name. If you cannot determine it, return an empty string.
|
||||
|
||||
### trigger_words
|
||||
The trigger words or activation prompts needed to use this LoRA. Look for:
|
||||
- `instance_prompt:` in the YAML frontmatter
|
||||
- Phrases like "trigger word:", "trigger:", "use this prompt:", "activation prompt:"
|
||||
- In collection repos: the trigger section **specific to this model file** (look near matching download links or anchor IDs)
|
||||
- Example prompts at the start (usually the first word or phrase before any description)
|
||||
Return as an array of strings. If none found, return an empty array `[]`. **Never** return `["None"]` or any placeholder value — a truly empty list means no trigger words exist.
|
||||
|
||||
### short_description
|
||||
A concise 1-2 sentence summary of what this model does. Extract from the "Model description" section or the first paragraph. Return empty string if the README is too minimal.
|
||||
A concise 1-2 sentence summary of what this model does. Extract from the "Model description" section or the first paragraph. For collection repos, focus on the **specific model version** matching `{{model_basename}}`, not the repo as a whole. Return empty string if the README is too minimal.
|
||||
|
||||
### tags
|
||||
3-8 relevant tags for categorizing this model. **Quality over quantity.**
|
||||
|
||||
Sources to consider:
|
||||
- The YAML frontmatter `tags:` list
|
||||
- The YAML frontmatter `tags:` list (filter out technical ones — see below)
|
||||
- The subject, style, character, or concept the model represents
|
||||
- The model filename itself may give clues (e.g. "pokemon", "anime", "pixelart")
|
||||
|
||||
**Critical filtering rules — apply them strictly:**
|
||||
|
||||
@@ -82,10 +110,15 @@ Return empty array if no meaningful content tags remain after filtering.
|
||||
The recommended image generation resolution for this model, in pixels. Look for sections like "Best Dimensions", "Recommended size", "Suggested resolution", or similar phrasing in the README. Prefer the explicitly marked "Best" or default resolution. If the table/list has multiple entries (e.g. "768 x 1024 (Best)" and "1024 x 1024 (Default)"), use the one marked "Best". Return integers. If no resolution can be determined, return 0 for both.
|
||||
|
||||
### preview_url
|
||||
The URL of the most suitable preview image from the README. Look for image tags (e.g. ``) and the YAML frontmatter `widget:` section (which often has `output.url` fields). Choose the first image that appears to be a generation example (not a logo or diagram). Construct the absolute URL as `https://huggingface.co/{{repo}}/resolve/main/{filename}`. If no suitable image is found, return an empty string.
|
||||
The URL of the most suitable preview image from the README. Look for:
|
||||
- Image tags near the section matching the model filename (`{{model_basename}}`)
|
||||
- The YAML frontmatter `widget:` section (which often has `output.url` fields)
|
||||
- In collection repos: the sample images listed **under the section** for this specific model version
|
||||
- Generic `` in the body
|
||||
Choose the first image that appears to be a generation example (not a logo or diagram). Construct the absolute URL as `https://huggingface.co/{{repo}}/resolve/main/{filename}`. If no suitable image is found, return an empty string.
|
||||
|
||||
### notes
|
||||
A plain-text summary of the model card's key practical usage information. Combine trigger words, style modifiers, recommended parameters (steps, CFG, resolution, sampler), and any setup tips into a readable paragraph. Return empty string if the README has no useful usage info.
|
||||
A plain-text summary of the model card's key practical usage information. Combine trigger words, style modifiers, recommended parameters (steps, CFG, resolution, sampler), and any setup tips into a readable paragraph. For collection repos, focus on the **specific model version** matching `{{model_basename}}`. Return empty string if the README has no useful usage info.
|
||||
|
||||
### usage_tips
|
||||
A JSON string with structured usage recommendations. Extract from the README any explicit ranges or recommended values (e.g. "Set LoRA strength: **0.85 - 1.4**", "CLIP strength: 0.5"). Possible fields (include only those you can determine):
|
||||
@@ -109,6 +142,25 @@ Your confidence level in the extracted data:
|
||||
- "medium" — some fields were inferred from context
|
||||
- "low" — most fields are guesses based on limited information
|
||||
|
||||
## Important: Handling Collection Repos (multiple model files)
|
||||
|
||||
Many HuggingFace repos contain **multiple model files** in a single repository
|
||||
(e.g. a "LoRA collection" with different styles/characters in separate files).
|
||||
|
||||
The model file currently being enriched is: **`{{model_basename}}`**
|
||||
|
||||
To find the correct section in the README:
|
||||
|
||||
1. **Search for download links** containing the filename — the surrounding paragraph is your section.
|
||||
2. **Search for anchor IDs** (`<a id="...">`) or section headings whose text matches words from the filename.
|
||||
3. **Search for HTML headings** (`<h1>`, `<h2>`, `<span>`) containing parts of the filename.
|
||||
4. If no match is found, use the full README as usual — the model may be the only one in the repo.
|
||||
|
||||
When a matching section IS found, prefer metadata from that section.
|
||||
When no section matches (e.g. single-model repos or repos without per-file sections),
|
||||
extract metadata from the full README normally. Do not return empty data just
|
||||
because the filename doesn't appear in the README.
|
||||
|
||||
## Output Format
|
||||
|
||||
Return ONLY a JSON object with exactly these fields (no markdown fences, no extra text):
|
||||
|
||||
@@ -20,6 +20,155 @@ from typing import List, Tuple
|
||||
_REPO_URL_PATTERN = re.compile(r"https?://huggingface\.co/([^/]+/[^/]+)")
|
||||
|
||||
|
||||
def extract_simple_markdown_images(
|
||||
markdown_text: str,
|
||||
repo: str,
|
||||
existing_urls: set | None = None,
|
||||
default_width: int = 512,
|
||||
default_height: int = 512,
|
||||
) -> list[dict]:
|
||||
"""Extract standalone markdown images from the README body.
|
||||
|
||||
Matches ```` on lines that are NOT part of a markdown table
|
||||
and NOT inside fenced code blocks. These are common in DreamBooth
|
||||
training dumps where the user uploads example images with simple
|
||||
```` syntax.
|
||||
|
||||
Returns a list of dicts in the same ``civitai.images`` format as
|
||||
:func:`extract_gallery_images`.
|
||||
"""
|
||||
if not markdown_text or not repo:
|
||||
return []
|
||||
|
||||
base_url = f"https://huggingface.co/{repo}/resolve/main"
|
||||
images: list[dict] = []
|
||||
seen_urls: set = set(existing_urls) if existing_urls else set()
|
||||
|
||||
# Collect lines that are NOT inside fenced code blocks
|
||||
lines = markdown_text.split("\n")
|
||||
in_code_block = False
|
||||
n = len(lines)
|
||||
i = 0
|
||||
while i < n:
|
||||
line = lines[i]
|
||||
if line.strip().startswith("```"):
|
||||
in_code_block = not in_code_block
|
||||
i += 1
|
||||
continue
|
||||
if in_code_block:
|
||||
i += 1
|
||||
continue
|
||||
# Skip table rows
|
||||
if "|" in line and i + 1 < n and re.match(r"^\|[\s:-]+\|", lines[i + 1]):
|
||||
i += 1
|
||||
continue
|
||||
|
||||
m = re.match(r"^\s*!\[([^\]]*)\]\(([^)]+)\)\s*$", line)
|
||||
if m:
|
||||
raw_path = m.group(2).strip()
|
||||
if raw_path.startswith("http"):
|
||||
url = raw_path
|
||||
else:
|
||||
clean = raw_path.lstrip("./").lstrip("/")
|
||||
url = f"{base_url}/{clean}"
|
||||
if url not in seen_urls:
|
||||
seen_urls.add(url)
|
||||
images.append({
|
||||
"url": url,
|
||||
"type": "image",
|
||||
"nsfwLevel": 0,
|
||||
"width": default_width,
|
||||
"height": default_height,
|
||||
"meta": {"prompt": m.group(1), "negativePrompt": ""},
|
||||
"hasMeta": bool(m.group(1)),
|
||||
"hasPositivePrompt": bool(m.group(1)),
|
||||
})
|
||||
i += 1
|
||||
|
||||
return images
|
||||
|
||||
|
||||
def extract_html_img_tags(
|
||||
markdown_text: str,
|
||||
repo: str,
|
||||
existing_urls: set | None = None,
|
||||
default_width: int = 512,
|
||||
default_height: int = 512,
|
||||
) -> list[dict]:
|
||||
"""Extract image URLs from HTML ``<img src=\"...\">`` tags in the README.
|
||||
|
||||
Many HF collection repos (e.g. ``deadman44/Z-Image_LoRA``) use raw HTML
|
||||
``<img>`` tags exclusively for their sample images, with no markdown
|
||||
``![]()`` equivalents. This function finds those tags and constructs
|
||||
resolvable HF URLs.
|
||||
|
||||
Returns a list of dicts in the ``civitai.images`` format.
|
||||
"""
|
||||
if not markdown_text or not repo:
|
||||
return []
|
||||
|
||||
base_url = f"https://huggingface.co/{repo}/resolve/main"
|
||||
images: list[dict] = []
|
||||
seen_urls: set = set(existing_urls) if existing_urls else set()
|
||||
|
||||
for m in re.finditer(
|
||||
r'<img\s[^>]*src=\"([^\"]+)\"',
|
||||
markdown_text,
|
||||
re.IGNORECASE,
|
||||
):
|
||||
raw_path = m.group(1).strip()
|
||||
if not raw_path:
|
||||
continue
|
||||
|
||||
if raw_path.startswith("http"):
|
||||
url = raw_path
|
||||
else:
|
||||
clean = raw_path.lstrip("./").lstrip("/")
|
||||
url = f"{base_url}/{clean}"
|
||||
|
||||
if url and url not in seen_urls:
|
||||
seen_urls.add(url)
|
||||
images.append({
|
||||
"url": url,
|
||||
"type": "image",
|
||||
"nsfwLevel": 0,
|
||||
"width": default_width,
|
||||
"height": default_height,
|
||||
"meta": {"prompt": "", "negativePrompt": ""},
|
||||
"hasMeta": False,
|
||||
"hasPositivePrompt": False,
|
||||
})
|
||||
|
||||
# Also try single-quoted src attributes
|
||||
for m in re.finditer(
|
||||
r"<img\s[^>]*src='([^']+)'",
|
||||
markdown_text,
|
||||
re.IGNORECASE,
|
||||
):
|
||||
raw_path = m.group(1).strip()
|
||||
if not raw_path:
|
||||
continue
|
||||
if raw_path.startswith("http"):
|
||||
url = raw_path
|
||||
else:
|
||||
clean = raw_path.lstrip("./").lstrip("/")
|
||||
url = f"{base_url}/{clean}"
|
||||
if url and url not in seen_urls:
|
||||
seen_urls.add(url)
|
||||
images.append({
|
||||
"url": url,
|
||||
"type": "image",
|
||||
"nsfwLevel": 0,
|
||||
"width": default_width,
|
||||
"height": default_height,
|
||||
"meta": {"prompt": "", "negativePrompt": ""},
|
||||
"hasMeta": False,
|
||||
"hasPositivePrompt": False,
|
||||
})
|
||||
|
||||
return images
|
||||
|
||||
|
||||
def extract_repo_from_hf_url(hf_url: str) -> str:
|
||||
"""Extract ``user/repo`` from a HuggingFace URL."""
|
||||
m = _REPO_URL_PATTERN.match(hf_url)
|
||||
@@ -62,43 +211,20 @@ def extract_gallery_images(
|
||||
if not widget_match:
|
||||
return images
|
||||
|
||||
# Split entries starting with `- text:`
|
||||
entries = re.split(r"\n- text:", frontmatter[widget_match.end():])
|
||||
# Split entries by YAML list marker `\n- `. Each entry is one widget list
|
||||
# item with `output:` and optionally `text:` sub-keys.
|
||||
entries_raw = frontmatter[widget_match.end():]
|
||||
entries = re.split(r"\n- ", entries_raw) if "\n- " in entries_raw else [entries_raw]
|
||||
|
||||
for entry in entries:
|
||||
if not entry.strip():
|
||||
entry = entry.strip()
|
||||
if not entry:
|
||||
continue
|
||||
|
||||
entry = entry.strip()
|
||||
|
||||
# Extract text (prompt) — handles three YAML scalar styles:
|
||||
# 1. "quoted inline"
|
||||
# 2. >- folded block
|
||||
# 3. plain (unquoted) multi-line
|
||||
text = ""
|
||||
qm = re.match(r'^"((?:[^"\\]|\\.)*)"', entry)
|
||||
if qm:
|
||||
text = qm.group(1)
|
||||
else:
|
||||
mm = re.match(r"^>(?:-\s*)?\n((?:.+(?:\n|$))+)", entry, re.MULTILINE)
|
||||
if mm:
|
||||
raw = mm.group(1)
|
||||
else:
|
||||
# Plain YAML scalar — take lines until a YAML key
|
||||
raw = entry
|
||||
if raw:
|
||||
text_lines: list[str] = []
|
||||
for line in raw.split("\n"):
|
||||
if re.match(r"^\s*\w+:", line):
|
||||
break
|
||||
stripped = line.strip()
|
||||
if stripped:
|
||||
text_lines.append(stripped)
|
||||
text = " ".join(text_lines)
|
||||
|
||||
# Extract output.url
|
||||
# Extract output.url — handle both `output:` and `- output:` (dash prefix)
|
||||
url = ""
|
||||
url_match = re.search(
|
||||
r"^\s*output:\s*\n\s+url:\s*(.+?)\s*$", entry, re.MULTILINE
|
||||
r"^-?\s*output:\s*\n\s+url:\s*(.+?)\s*$", entry, re.MULTILINE
|
||||
)
|
||||
if url_match:
|
||||
raw_path = url_match.group(1).strip().strip("'\"")
|
||||
@@ -107,6 +233,16 @@ def extract_gallery_images(
|
||||
elif raw_path.startswith("http"):
|
||||
url = raw_path
|
||||
|
||||
# Extract text (prompt) — from YAML `text:` sub-key
|
||||
text = ""
|
||||
text_match = re.search(
|
||||
r"^-?\s*text:\s*(.+?)\s*$", entry, re.MULTILINE
|
||||
)
|
||||
if text_match:
|
||||
raw_text = text_match.group(1).strip().strip("'\"")
|
||||
if raw_text and raw_text != "-":
|
||||
text = raw_text
|
||||
|
||||
if url:
|
||||
image: dict = {
|
||||
"url": url,
|
||||
@@ -338,20 +474,23 @@ def _strip_fenced_code_blocks(text: str) -> str:
|
||||
|
||||
|
||||
def _strip_standalone_images(text: str) -> str:
|
||||
"""Strip image embeds that occupy their own line.
|
||||
"""Strip/compress image embeds for LLM-prompt injection.
|
||||
|
||||
Preserves the alt text from markdown images (```` → ``alt``)
|
||||
since it often describes what the model generates, which is useful signal
|
||||
for tag/description extraction.
|
||||
Markdown images (````) are **kept intact** so the LLM can
|
||||
extract a ``preview_url`` from them. Only the alt text is needed for
|
||||
content signal; the URL is needed for image extraction.
|
||||
|
||||
HTML ``<img>`` tags on their own line are replaced by their alt text
|
||||
(if any) or removed, since the LLM has difficulty extracting URLs from
|
||||
raw HTML attributes.
|
||||
"""
|
||||
# Markdown: ```` on its own line → keep alt text
|
||||
# HTML: ``<img src="..." alt="..." ...>`` on its own line → keep alt text
|
||||
text = re.sub(
|
||||
r"^\s*!\[([^\]]*)\]\([^)]+\)\s*$",
|
||||
r'^\s*<img\s[^>]*alt="([^"]*)"[^>]*/?>(?:</img>)?\s*$',
|
||||
r"\1",
|
||||
text,
|
||||
flags=re.MULTILINE,
|
||||
flags=re.MULTILINE | re.IGNORECASE,
|
||||
)
|
||||
# HTML: ``<img src="..." ...>`` on its own line → remove entirely
|
||||
text = re.sub(
|
||||
r'^\s*<img\s[^>]+/?>(?:</img>)?\s*$',
|
||||
"",
|
||||
@@ -479,6 +618,155 @@ def _strip_massive_lists(text: str) -> str:
|
||||
return "\n".join(out)
|
||||
|
||||
|
||||
def extract_relevant_section(
|
||||
readme_content: str,
|
||||
model_basename: str,
|
||||
*,
|
||||
context_lines: int = 30,
|
||||
) -> str:
|
||||
"""Find the section of a HuggingFace README relevant to a specific model file.
|
||||
|
||||
Many HF repos bundle multiple model files (LoRA collections) in a single
|
||||
repo. This function trims the README to only the portion that references
|
||||
the given *model_basename* (e.g. ``lora_zimage_myjs_turbo_beta01``),
|
||||
dramatically reducing prompt length and focusing the LLM on the correct
|
||||
metadata.
|
||||
|
||||
Matching strategies, tried in order:
|
||||
|
||||
1. **Download link** — a markdown or HTML link whose URL contains the
|
||||
basename. Returns the surrounding block (previous heading → next
|
||||
heading).
|
||||
2. **Anchor ID** — an ``<a id="...">`` whose ``id`` matches a token from
|
||||
the basename (split on ``_``/``-``).
|
||||
3. **Section heading** — an ``<h1>``-``<h4>`` or markdown heading whose
|
||||
text overlaps with tokens from the basename.
|
||||
4. **Fallback** — the full README unchanged.
|
||||
|
||||
Args:
|
||||
readme_content: Raw README.md content.
|
||||
model_basename: Model file basename *without* extension (e.g.
|
||||
``lora_zimage_myjs_turbo_beta01``).
|
||||
context_lines: Number of lines of context before/after a matched
|
||||
download link to include (default 30).
|
||||
|
||||
Returns:
|
||||
Trimmed README text, or the original when no matching section is
|
||||
found.
|
||||
"""
|
||||
if not readme_content or not model_basename:
|
||||
return readme_content or ""
|
||||
|
||||
lines = readme_content.split("\n")
|
||||
n = len(lines)
|
||||
basename_lower = model_basename.lower()
|
||||
# Tokens from the basename split on common separators
|
||||
tokens = {t for t in re.split(r"[_\-.\s]+", basename_lower) if len(t) > 2}
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Strategy 1: Find a download link containing the basename
|
||||
# ------------------------------------------------------------------
|
||||
for idx, line in enumerate(lines):
|
||||
# Match markdown links: [text](url) and HTML links: <a href="url">
|
||||
# whose URL contains the basename.
|
||||
if basename_lower in line.lower() and _looks_like_download_link(line):
|
||||
return _extract_section(lines, idx, context_lines)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Strategy 2: Find an anchor ID matching a token
|
||||
# ------------------------------------------------------------------
|
||||
for idx, line in enumerate(lines):
|
||||
m = re.search(r'<a\s+id="([^"]+)"', line, re.IGNORECASE)
|
||||
if m:
|
||||
aid = m.group(1).lower()
|
||||
if any(token in aid for token in tokens):
|
||||
return _extract_section(lines, idx, context_lines)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Strategy 3: Find an HTML or markdown heading with overlapping tokens
|
||||
# ------------------------------------------------------------------
|
||||
for idx, line in enumerate(lines):
|
||||
# HTML heading: <h1...>, <h2...>, <h3...>, <h4...>
|
||||
hm = re.search(r'<h[1-4][^>]*>(.+?)</h[1-4]>', line, re.IGNORECASE)
|
||||
heading_text = ""
|
||||
if hm:
|
||||
heading_text = hm.group(1)
|
||||
else:
|
||||
# Markdown heading: ## text
|
||||
mm = re.match(r"^#{1,4}\s+(.+?)\s*#*$", line)
|
||||
if mm:
|
||||
heading_text = mm.group(1)
|
||||
if heading_text:
|
||||
heading_lower = heading_text.lower()
|
||||
# Check if any token appears in the heading
|
||||
if any(token in heading_lower for token in tokens):
|
||||
return _extract_section(lines, idx, context_lines)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Fallback: return FULL readme
|
||||
# ------------------------------------------------------------------
|
||||
return readme_content
|
||||
|
||||
|
||||
def _looks_like_download_link(line: str) -> bool:
|
||||
"""Heuristic: does *line* look like it contains a model-file download link?"""
|
||||
# Markdown: [text](url)
|
||||
if re.search(r'\[([^\]]*)\]\(([^)]*\.safetensors[^)]*)\)', line):
|
||||
return True
|
||||
# Markdown: [text](url) containing "download" or "resolve"
|
||||
if re.search(r'\[([^\]]*)\]\(([^)]*/(resolve|download)/[^)]*)\)', line):
|
||||
return True
|
||||
# HTML <a href="...safetensors">
|
||||
if re.search(r'<a\s[^>]*href="[^"]*\.safetensors', line, re.IGNORECASE):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _extract_section(
|
||||
lines: list[str], match_idx: int, context_lines: int,
|
||||
) -> str:
|
||||
"""Return a window of *lines* around *match_idx* bounded by headings.
|
||||
|
||||
When *match_idx* is itself a heading line, the section starts *at*
|
||||
that heading (no backward walk), avoiding pulling in content from
|
||||
earlier sibling sections.
|
||||
"""
|
||||
n = len(lines)
|
||||
|
||||
# Determine start — if match is a heading, start right there
|
||||
if _is_heading(lines[match_idx]):
|
||||
start = match_idx
|
||||
else:
|
||||
# Walk backward to find the nearest heading
|
||||
start = max(0, match_idx - context_lines)
|
||||
for i in range(match_idx - 1, max(-1, match_idx - context_lines * 3), -1):
|
||||
if i < 0:
|
||||
start = 0
|
||||
break
|
||||
if _is_heading(lines[i]):
|
||||
start = i
|
||||
break
|
||||
|
||||
# Walk forward to find the next heading at same or higher level
|
||||
end = min(n, match_idx + context_lines)
|
||||
for i in range(match_idx + 1, min(n, match_idx + context_lines * 3)):
|
||||
if _is_heading(lines[i]):
|
||||
end = i
|
||||
break
|
||||
|
||||
return "\n".join(lines[start:end])
|
||||
|
||||
|
||||
def _is_heading(line: str) -> bool:
|
||||
"""Return True if *line* is a markdown or HTML heading (not a closing tag)."""
|
||||
stripped = line.strip()
|
||||
if re.match(r"^#{1,4}\s", stripped):
|
||||
return True
|
||||
if re.match(r"^<h[1-4](?:\s|>)", stripped, re.IGNORECASE):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _compress_blank_lines(text: str) -> str:
|
||||
"""Collapse runs of 3+ blank lines down to 2."""
|
||||
return re.sub(r"\n{3,}", "\n\n", text)
|
||||
|
||||
@@ -551,7 +551,103 @@ class LLMService:
|
||||
try:
|
||||
return json.loads(content)
|
||||
except (json.JSONDecodeError, TypeError) as parse_err:
|
||||
# Last resort: attempt to salvage partial JSON (closing unclosed
|
||||
# brackets/braces, truncating incomplete strings, etc.)
|
||||
salvaged = _try_salvage_json(content)
|
||||
if salvaged is not None:
|
||||
logger.warning(
|
||||
"LLM JSON salvaged from partial content (%d chars raw)",
|
||||
len(content),
|
||||
)
|
||||
return salvaged
|
||||
raise LLMResponseError(
|
||||
f"LLM response could not be parsed as JSON after retry: {parse_err}\n"
|
||||
f"Raw content: {content[:500]}"
|
||||
) from parse_err
|
||||
|
||||
|
||||
def _try_salvage_json(raw: str) -> Dict[str, Any] | None:
|
||||
"""Attempt to repair and parse a truncated JSON string.
|
||||
|
||||
Handles common truncation patterns:
|
||||
|
||||
* Incomplete string value at the end (``"foo`` → ``"foo"``)
|
||||
* Missing closing ``}`` or ``]`` (respecting nesting order)
|
||||
* Trailing comma before closing bracket
|
||||
* Extra text after the JSON object (e.g. markdown fences)
|
||||
|
||||
Returns the parsed dict on success, ``None`` if repair is impossible.
|
||||
"""
|
||||
if not raw:
|
||||
return None
|
||||
|
||||
text = raw.strip()
|
||||
|
||||
# Strip markdown fences if the LLM wrapped the JSON
|
||||
if text.startswith("```"):
|
||||
end = text.find("\n")
|
||||
text = text[end + 1:] if end != -1 else text[3:]
|
||||
if text.endswith("```"):
|
||||
text = text[:-3].rstrip()
|
||||
|
||||
# Find the first '{' and strip everything before it
|
||||
start = text.find("{")
|
||||
if start == -1:
|
||||
return None
|
||||
text = text[start:]
|
||||
|
||||
# Try to close an incomplete string at the end (e.g. ``"https://huggingf``)
|
||||
# Pattern: ends mid-string (last quote is open)
|
||||
if text.count('"') % 2 == 1:
|
||||
text += '"'
|
||||
|
||||
# Ensure trailing commas before closing braces work
|
||||
text = _strip_trailing_commas(text)
|
||||
|
||||
# Walk through the text character by character to find unclosed
|
||||
# brackets and close them in the correct (LIFO) order.
|
||||
# We ignore brackets inside quoted strings.
|
||||
stack: list[str] = []
|
||||
in_string = False
|
||||
escape = False
|
||||
for ch in text:
|
||||
if escape:
|
||||
escape = False
|
||||
continue
|
||||
if ch == "\\":
|
||||
escape = True
|
||||
continue
|
||||
if ch == '"':
|
||||
in_string = not in_string
|
||||
continue
|
||||
if in_string:
|
||||
continue
|
||||
if ch in ("{", "["):
|
||||
stack.append(ch)
|
||||
elif ch == "}":
|
||||
if stack and stack[-1] == "{":
|
||||
stack.pop()
|
||||
else:
|
||||
return None # Unmatched closer — unrecoverable
|
||||
elif ch == "]":
|
||||
if stack and stack[-1] == "[":
|
||||
stack.pop()
|
||||
else:
|
||||
return None
|
||||
|
||||
# Close remaining open brackets in reverse order
|
||||
for opener in reversed(stack):
|
||||
text += "}" if opener == "{" else "]"
|
||||
|
||||
try:
|
||||
return json.loads(text)
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
return None
|
||||
|
||||
|
||||
def _strip_trailing_commas(text: str) -> str:
|
||||
"""Remove commas that appear before a closing brace/bracket."""
|
||||
import re as _re
|
||||
text = _re.sub(r",\s*}", "}", text)
|
||||
text = _re.sub(r",\s*]", "]", text)
|
||||
return text
|
||||
|
||||
1
tests/__init__.py
Normal file
1
tests/__init__.py
Normal file
@@ -0,0 +1 @@
|
||||
# Test suite package.
|
||||
1
tests/enrich_hf_validation/__init__.py
Normal file
1
tests/enrich_hf_validation/__init__.py
Normal file
@@ -0,0 +1 @@
|
||||
# HF Metadata Enrichment validation suite.
|
||||
97
tests/enrich_hf_validation/config.py
Normal file
97
tests/enrich_hf_validation/config.py
Normal file
@@ -0,0 +1,97 @@
|
||||
"""Configuration for the HF metadata enrichment validation suite.
|
||||
|
||||
Loads user settings, defines paths, and pulls constants from the main
|
||||
codebase (``py.utils.constants``).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from typing import Any, Dict, List
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Default paths
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_DEFAULT_MODELS_FILE = os.path.expanduser(
|
||||
"~/Documents/hf_lora_models.txt"
|
||||
)
|
||||
_DEFAULT_SETTINGS_PATH = os.path.expanduser(
|
||||
"~/.config/ComfyUI-LoRA-Manager/settings.json"
|
||||
)
|
||||
_DEFAULT_OUTPUT_DIR = "/tmp/hf_enrich_validation"
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Constants from the main codebase (copied at import time)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# Priority tags used in the LLM prompt for tag selection guidance.
|
||||
CIVITAI_MODEL_TAGS: List[str] = [
|
||||
"character", "concept", "clothing", "realistic", "anime", "toon",
|
||||
"furry", "style", "poses", "background", "tool", "vehicle",
|
||||
"buildings", "objects", "assets", "animal", "action",
|
||||
]
|
||||
|
||||
# Base models recognised as valid values.
|
||||
SUPPORTED_BASE_MODELS: List[str] = [
|
||||
"SD 1.4", "SD 1.5", "SD 1.5 LCM", "SD 1.5 Hyper",
|
||||
"SD 2.0", "SD 2.1",
|
||||
"SD 3", "SD 3.5", "SD 3.5 Medium", "SD 3.5 Large", "SD 3.5 Large Turbo",
|
||||
"SDXL 1.0", "SDXL Lightning", "SDXL Hyper",
|
||||
"Flux.1 D", "Flux.1 S", "Flux.1 Krea", "Flux.1 Kontext",
|
||||
"Flux.2 D", "Flux.2 Klein 9B", "Flux.2 Klein 9B-base",
|
||||
"Flux.2 Klein 4B", "Flux.2 Klein 4B-base",
|
||||
"AuraFlow", "Chroma", "PixArt a", "PixArt E",
|
||||
"Hunyuan 1", "Lumina", "Kolors",
|
||||
"NoobAI", "Illustrious", "Pony", "Pony V7",
|
||||
"HiDream", "Qwen", "ZImageTurbo", "ZImageBase",
|
||||
"SVD", "LTXV", "LTXV2", "LTXV 2.3",
|
||||
"CogVideoX", "Mochi",
|
||||
"Wan Video", "Wan Video 1.3B t2v", "Wan Video 14B t2v",
|
||||
"Wan Video 14B i2v 480p", "Wan Video 14B i2v 720p",
|
||||
"Wan Video 2.2 TI2V-5B", "Wan Video 2.2 T2V-A14B",
|
||||
"Wan Video 2.2 I2V-A14B",
|
||||
"Wan Video 2.5 T2V", "Wan Video 2.5 I2V",
|
||||
"Hunyuan Video", "Anima", "Ernie", "Ernie Turbo",
|
||||
"Nucleus", "Krea 2",
|
||||
]
|
||||
|
||||
# Placeholder values the LLM sometimes emits that should count as "empty".
|
||||
PLACEHOLDER_VALUES = frozenset({
|
||||
"none", "null", "n/a", "unknown", "not available",
|
||||
"not specified", "no trigger words", "no trigger word",
|
||||
})
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# User settings loader
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def load_settings(settings_path: str) -> Dict[str, Any]:
|
||||
"""Load LoRA Manager settings from *settings_path*.
|
||||
|
||||
Returns a flat dict with the LLM configuration fields that the
|
||||
enrichment pipeline depends on.
|
||||
"""
|
||||
path = os.path.expanduser(settings_path)
|
||||
if not os.path.exists(path):
|
||||
raise FileNotFoundError(
|
||||
f"Settings file not found: {path}\n"
|
||||
"Please provide a valid --settings path."
|
||||
)
|
||||
|
||||
with open(path, "r", encoding="utf-8") as fh:
|
||||
raw: Dict[str, Any] = json.load(fh)
|
||||
|
||||
# Extract LLM-relevant config
|
||||
return {
|
||||
"llm_provider": raw.get("llm_provider", "ollama"),
|
||||
"llm_model": raw.get("llm_model", "qwen3.5:9b"),
|
||||
"llm_api_base": raw.get("llm_api_base", "http://localhost:11434/v1"),
|
||||
"llm_api_key": raw.get("llm_api_key", ""),
|
||||
"settings_path": path,
|
||||
}
|
||||
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,
|
||||
}
|
||||
352
tests/enrich_hf_validation/evaluation_engine.py
Normal file
352
tests/enrich_hf_validation/evaluation_engine.py
Normal file
@@ -0,0 +1,352 @@
|
||||
"""Evaluate enriched ``.metadata.json`` quality across multiple dimensions.
|
||||
|
||||
Scoring rubric (per field):
|
||||
|
||||
- **Completeness**: Is the field populated with meaningful content?
|
||||
- **Validity**: Does the value conform to expected constraints (controlled
|
||||
vocab, non-placeholder, parsable JSON)?
|
||||
- **Accuracy**: (sub-sample only — requires manual verification against
|
||||
the HF README).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from typing import Any, Dict, List, Optional, Set
|
||||
|
||||
from .config import (
|
||||
CIVITAI_MODEL_TAGS,
|
||||
PLACEHOLDER_VALUES,
|
||||
SUPPORTED_BASE_MODELS,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Scoring helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_MIN_TAGS = 1
|
||||
_MAX_TAGS = 8
|
||||
_MIN_DESC_LENGTH = 20
|
||||
_MIN_NOTES_LENGTH = 30
|
||||
|
||||
# Tags that the LLM sometimes emits but which are not meaningful content tags.
|
||||
_TECH_TAGS = frozenset({
|
||||
"lora", "dreambooth", "text-to-image", "diffusers", "flux",
|
||||
"sdxl", "checkpoint", "pytorch", "safetensors", "fine-tuning",
|
||||
"stable-diffusion", "training", "stablediffusion",
|
||||
})
|
||||
|
||||
|
||||
def _is_placeholder(val: str) -> bool:
|
||||
return val.strip().lower() in PLACEHOLDER_VALUES
|
||||
|
||||
|
||||
def _is_valid_trigger_words(words: List[str]) -> bool:
|
||||
"""Return True if *words* is a non-empty list of real trigger words."""
|
||||
if not words:
|
||||
return False
|
||||
cleaned = [w.strip() for w in words if w.strip()]
|
||||
if not cleaned:
|
||||
return False
|
||||
# Reject if ALL entries are placeholders
|
||||
non_placeholder = [w for w in cleaned if not _is_placeholder(w)]
|
||||
return len(non_placeholder) > 0
|
||||
|
||||
|
||||
def _is_valid_tags(tags: List[str]) -> bool:
|
||||
"""Return True if *tags* is a reasonable list of content tags."""
|
||||
if not tags:
|
||||
return False
|
||||
cleaned = [t.strip().lower() for t in tags if t.strip()]
|
||||
if not cleaned:
|
||||
return False
|
||||
# At least one tag that isn't a technical keyword
|
||||
meaningful = [t for t in cleaned if t not in _TECH_TAGS]
|
||||
return len(meaningful) >= _MIN_TAGS
|
||||
|
||||
|
||||
def _tag_priority_coverage(tags: List[str]) -> float:
|
||||
"""Fraction of tags that align with the user's priority tag vocabulary."""
|
||||
if not tags:
|
||||
return 0.0
|
||||
priority_lower = {t.lower() for t in CIVITAI_MODEL_TAGS}
|
||||
matched = sum(1 for t in tags if t.strip().lower() in priority_lower)
|
||||
return matched / len(tags)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Per-model evaluation
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# Type alias for a score record
|
||||
ScoreRecord = Dict[str, Any]
|
||||
|
||||
|
||||
def evaluate_model(
|
||||
metadata: Dict[str, Any],
|
||||
model_path: str,
|
||||
repo_id: str,
|
||||
*,
|
||||
enrichment_success: bool,
|
||||
enrichment_errors: List[str],
|
||||
) -> ScoreRecord:
|
||||
"""Score a single enriched model's metadata.
|
||||
|
||||
Returns a dict with per-field scores, a total score, and a list of
|
||||
flagged issues.
|
||||
"""
|
||||
civitai = metadata.get("civitai") or {}
|
||||
trained_words: List[str] = civitai.get("trainedWords") or metadata.get("trainedWords") or []
|
||||
short_desc: str = civitai.get("description") or ""
|
||||
tags: List[str] = metadata.get("tags") or []
|
||||
notes: str = metadata.get("notes") or ""
|
||||
usage_tips_raw: str = metadata.get("usage_tips") or "{}"
|
||||
model_description: str = metadata.get("modelDescription") or ""
|
||||
base_model: str = metadata.get("base_model") or ""
|
||||
preview_url: str = metadata.get("preview_url") or ""
|
||||
confidence: str = metadata.get("_llm_confidence") or ""
|
||||
|
||||
# --- base_model ---
|
||||
base_model_valid = base_model in SUPPORTED_BASE_MODELS
|
||||
base_model_filled = bool(base_model) and base_model != "Unknown"
|
||||
|
||||
# --- trigger_words (trainedWords) ---
|
||||
triggers_valid = _is_valid_trigger_words(trained_words)
|
||||
|
||||
# --- short_description (civitai.description) ---
|
||||
desc_filled = len(short_desc.strip()) >= _MIN_DESC_LENGTH
|
||||
|
||||
# --- tags ---
|
||||
tags_valid = _is_valid_tags(tags)
|
||||
tags_priority_coverage = _tag_priority_coverage(tags)
|
||||
tags_no_technical = (
|
||||
sum(1 for t in tags if t.strip().lower() not in _TECH_TAGS) >= _MIN_TAGS
|
||||
if tags else False
|
||||
)
|
||||
|
||||
# --- notes ---
|
||||
notes_filled = len(notes.strip()) >= _MIN_NOTES_LENGTH
|
||||
|
||||
# --- usage_tips ---
|
||||
usage_tips_valid = False
|
||||
if usage_tips_raw.strip() and usage_tips_raw.strip() != "{}":
|
||||
try:
|
||||
parsed = json.loads(usage_tips_raw)
|
||||
if isinstance(parsed, dict) and len(parsed) > 0:
|
||||
usage_tips_valid = True
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
pass
|
||||
|
||||
# --- modelDescription (README → HTML) ---
|
||||
desc_html_filled = len(model_description.strip()) > 100
|
||||
|
||||
# --- preview_url ---
|
||||
preview_filled = bool(preview_url) and os.path.exists(preview_url)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Composite score (0-100)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
field_scores = {
|
||||
"base_model": _score_bool(base_model_filled and base_model_valid, weight=15),
|
||||
"trigger_words": _score_bool(triggers_valid, weight=15),
|
||||
"short_description": _score_bool(desc_filled, weight=10),
|
||||
"tags": _score_bool(tags_valid, weight=15),
|
||||
"tags_priority_coverage": _score_continuous(tags_priority_coverage, weight=5),
|
||||
"notes": _score_bool(notes_filled, weight=5),
|
||||
"usage_tips": _score_bool(usage_tips_valid, weight=5),
|
||||
"modelDescription_html": _score_bool(desc_html_filled, weight=10),
|
||||
"preview_downloaded": _score_bool(preview_filled, weight=10),
|
||||
}
|
||||
|
||||
# Deduct points for enrichment-level failures
|
||||
penalty = 0
|
||||
if enrichment_errors:
|
||||
penalty += 10
|
||||
if not enrichment_success:
|
||||
penalty += 20
|
||||
|
||||
total_raw = sum(field_scores.values())
|
||||
total = max(0, min(100, total_raw - penalty))
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Flagged issues
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
issues: List[str] = []
|
||||
if not base_model_filled:
|
||||
issues.append("base_model is empty or 'Unknown'")
|
||||
elif not base_model_valid:
|
||||
issues.append(f"base_model '{base_model}' not in SUPPORTED_BASE_MODELS")
|
||||
if not triggers_valid:
|
||||
issues.append("trigger_words are missing or contain only placeholders")
|
||||
if not desc_filled:
|
||||
issues.append("short_description is too short or empty")
|
||||
if not tags_valid:
|
||||
issues.append("tags are missing, too few, or purely technical")
|
||||
if tags_valid and tags_priority_coverage < 0.5:
|
||||
issues.append("tags have low overlap with priority_tags (< 50%)")
|
||||
if not notes_filled:
|
||||
issues.append("notes are too short or empty")
|
||||
if not usage_tips_valid:
|
||||
issues.append("usage_tips is empty or invalid JSON")
|
||||
if not desc_html_filled:
|
||||
issues.append("modelDescription is too short (README may not have been converted)")
|
||||
if not preview_filled:
|
||||
issues.append("preview image not downloaded (URL missing or download failed)")
|
||||
|
||||
return {
|
||||
"repo_id": repo_id,
|
||||
"model_path": model_path,
|
||||
"enrichment_success": enrichment_success,
|
||||
"total_score": total,
|
||||
"field_scores": field_scores,
|
||||
"issues": issues,
|
||||
"confidence_from_llm": confidence,
|
||||
"raw_values": {
|
||||
"base_model": base_model,
|
||||
"trigger_words": trained_words,
|
||||
"short_description": short_desc,
|
||||
"tags": tags,
|
||||
"notes": notes,
|
||||
"usage_tips": usage_tips_raw,
|
||||
"preview_url": preview_url,
|
||||
"has_modelDescription": len(model_description) > 0,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _score_bool(condition: bool, weight: int = 10) -> int:
|
||||
return weight if condition else 0
|
||||
|
||||
|
||||
def _score_continuous(value: float, weight: int = 10) -> int:
|
||||
"""Linear interpolation: value 0.0 → 0, value 1.0 → *weight*."""
|
||||
return int(round(value * weight))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Batch evaluation
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def evaluate_batch(
|
||||
enriched: List[Dict[str, Any]],
|
||||
) -> List[ScoreRecord]:
|
||||
"""Evaluate a list of enrichment results.
|
||||
|
||||
Each entry in *enriched* should have keys:
|
||||
``repo_id``, ``model_path``, ``metadata`` (the enriched dict),
|
||||
``success``, ``errors``.
|
||||
"""
|
||||
scores: List[ScoreRecord] = []
|
||||
for entry in enriched:
|
||||
record = evaluate_model(
|
||||
metadata=entry.get("metadata", {}),
|
||||
model_path=entry.get("model_path", ""),
|
||||
repo_id=entry.get("repo_id", ""),
|
||||
enrichment_success=entry.get("success", False),
|
||||
enrichment_errors=entry.get("errors", []),
|
||||
)
|
||||
scores.append(record)
|
||||
return scores
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Aggregate statistics
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def aggregate_scores(scores: List[ScoreRecord]) -> Dict[str, Any]:
|
||||
"""Compute aggregate stats across all scored models."""
|
||||
n = len(scores)
|
||||
if n == 0:
|
||||
return {"error": "no scores to aggregate"}
|
||||
|
||||
field_names = [
|
||||
"base_model", "trigger_words", "short_description", "tags",
|
||||
"tags_priority_coverage", "notes", "usage_tips",
|
||||
"modelDescription_html", "preview_downloaded",
|
||||
]
|
||||
possible = {f: 15 if f == "base_model" or f == "trigger_words" or f == "tags" else
|
||||
10 if f == "short_description" or f == "modelDescription_html" or f == "preview_downloaded" else
|
||||
5
|
||||
for f in field_names}
|
||||
|
||||
# Per-field aggregate
|
||||
field_agg: Dict[str, Any] = {}
|
||||
for fn in field_names:
|
||||
vals = [s["field_scores"].get(fn, 0) for s in scores]
|
||||
max_per_field = possible[fn]
|
||||
field_agg[fn] = {
|
||||
"mean": round(sum(vals) / n, 1) if n else 0,
|
||||
"fill_rate_pct": round(
|
||||
sum(1 for v in vals if v >= max_per_field) / n * 100, 1
|
||||
) if n else 0.0,
|
||||
"partial_rate_pct": round(
|
||||
sum(1 for v in vals if 0 < v < max_per_field) / n * 100, 1
|
||||
) if n else 0.0,
|
||||
"empty_rate_pct": round(
|
||||
sum(1 for v in vals if v == 0) / n * 100, 1
|
||||
) if n else 0.0,
|
||||
}
|
||||
|
||||
# Total score distribution
|
||||
total_scores = [s["total_score"] for s in scores]
|
||||
total_agg = {
|
||||
"mean": round(sum(total_scores) / n, 1) if n else 0,
|
||||
"median": _median(total_scores),
|
||||
"min": min(total_scores) if total_scores else 0,
|
||||
"max": max(total_scores) if total_scores else 0,
|
||||
"bins": {
|
||||
"excellent_80+": sum(1 for s in total_scores if s >= 80),
|
||||
"good_60_79": sum(1 for s in total_scores if 60 <= s < 80),
|
||||
"fair_40_59": sum(1 for s in total_scores if 40 <= s < 60),
|
||||
"poor_20_39": sum(1 for s in total_scores if 20 <= s < 40),
|
||||
"bad_0_19": sum(1 for s in total_scores if s < 20),
|
||||
},
|
||||
}
|
||||
|
||||
# Issue frequency
|
||||
issue_counter: Dict[str, int] = {}
|
||||
for s in scores:
|
||||
for issue in s["issues"]:
|
||||
issue_counter[issue] = issue_counter.get(issue, 0) + 1
|
||||
top_issues = sorted(issue_counter.items(), key=lambda x: -x[1])
|
||||
|
||||
# Confidence distribution
|
||||
conf_counter: Dict[str, int] = {"high": 0, "medium": 0, "low": 0, "": 0}
|
||||
for s in scores:
|
||||
c = (s.get("confidence_from_llm") or "").strip().lower()
|
||||
if c in conf_counter:
|
||||
conf_counter[c] += 1
|
||||
else:
|
||||
conf_counter[""] += 1
|
||||
|
||||
# Success / timeout / failure stats
|
||||
success_count = sum(1 for s in scores if s["enrichment_success"])
|
||||
fail_count = n - success_count
|
||||
|
||||
return {
|
||||
"model_count": n,
|
||||
"success_count": success_count,
|
||||
"fail_count": fail_count,
|
||||
"total_score": total_agg,
|
||||
"field_aggregates": field_agg,
|
||||
"top_issues": top_issues[:15],
|
||||
"confidence_distribution": conf_counter,
|
||||
}
|
||||
|
||||
|
||||
def _median(values: List[float]) -> float:
|
||||
if not values:
|
||||
return 0.0
|
||||
sorted_v = sorted(values)
|
||||
m = len(sorted_v) // 2
|
||||
if len(sorted_v) % 2 == 0:
|
||||
return round((sorted_v[m - 1] + sorted_v[m]) / 2, 1)
|
||||
return round(sorted_v[m], 1)
|
||||
148
tests/enrich_hf_validation/metadata_constructor.py
Normal file
148
tests/enrich_hf_validation/metadata_constructor.py
Normal file
@@ -0,0 +1,148 @@
|
||||
"""Construct initial ``.metadata.json`` sidecars for HF model repos.
|
||||
|
||||
Each HF repo ID gets a minimal metadata file — no real model file is needed.
|
||||
The enrichment pipeline reads only the sidecar.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from typing import Dict, List
|
||||
|
||||
from .config import CIVITAI_MODEL_TAGS
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def load_repo_ids(path: str, max_models: int | None = None) -> List[str]:
|
||||
"""Read HF repo IDs from *path* (one per line, ignoring blanks/comments)."""
|
||||
path = os.path.expanduser(path)
|
||||
if not os.path.exists(path):
|
||||
raise FileNotFoundError(f"Models file not found: {path}")
|
||||
|
||||
repos: List[str] = []
|
||||
with open(path, "r", encoding="utf-8") as fh:
|
||||
for line in fh:
|
||||
line = line.strip()
|
||||
if not line or line.startswith("#"):
|
||||
continue
|
||||
repos.append(line)
|
||||
|
||||
if max_models is not None and max_models > 0:
|
||||
repos = repos[:max_models]
|
||||
|
||||
logger.info("Loaded %d HF repo IDs from %s", len(repos), path)
|
||||
return repos
|
||||
|
||||
|
||||
def sanitize_repo_id(repo_id: str) -> str:
|
||||
"""Turn ``user/repo-name`` into a safe directory name."""
|
||||
return repo_id.replace("/", "__").replace(".", "_")
|
||||
|
||||
|
||||
def build_model_dir(output_dir: str, repo_id: str) -> str:
|
||||
"""Return the per-model working directory."""
|
||||
return os.path.join(output_dir, "models", sanitize_repo_id(repo_id))
|
||||
|
||||
|
||||
def build_model_path(model_dir: str) -> str:
|
||||
"""Return a synthetic model file path (no real file will exist)."""
|
||||
return os.path.join(model_dir, "model.safetensors")
|
||||
|
||||
|
||||
def build_metadata_path(model_path: str) -> str:
|
||||
"""Return the sidecar path for a model file.
|
||||
|
||||
This MUST match the convention used by ``MetadataManager`` /
|
||||
``apply_metadata_updates``, which derives the sidecar path via
|
||||
``os.path.splitext(model_path)[0] + '.metadata.json'``.
|
||||
For a model file ``model.safetensors`` the sidecar is
|
||||
``model.metadata.json`` — *not* ``model.safetensors.metadata.json``.
|
||||
"""
|
||||
return f"{os.path.splitext(model_path)[0]}.metadata.json"
|
||||
|
||||
|
||||
def create_initial_metadata(
|
||||
output_dir: str,
|
||||
repo_id: str,
|
||||
) -> str:
|
||||
"""Write a minimal ``.metadata.json`` for *repo_id*.
|
||||
|
||||
Returns the **model path** (the ``.safetensors`` path whose sidecar was
|
||||
written). The caller passes this path to ``AgentService.execute_skill``.
|
||||
"""
|
||||
model_dir = build_model_dir(output_dir, repo_id)
|
||||
os.makedirs(model_dir, exist_ok=True)
|
||||
model_path = build_model_path(model_dir)
|
||||
metadata_path = build_metadata_path(model_path)
|
||||
|
||||
hf_url = f"https://huggingface.co/{repo_id}"
|
||||
file_name = repo_id.split("/")[-1]
|
||||
|
||||
metadata: Dict = {
|
||||
"file_name": file_name,
|
||||
"model_name": file_name,
|
||||
"file_path": model_path.replace(os.sep, "/"),
|
||||
"size": 0,
|
||||
"modified": 0,
|
||||
"sha256": "",
|
||||
"base_model": "Unknown",
|
||||
"preview_url": "",
|
||||
"preview_nsfw_level": 0,
|
||||
"notes": "",
|
||||
"from_civitai": False,
|
||||
"civitai": {},
|
||||
"tags": [],
|
||||
"modelDescription": "",
|
||||
"civitai_deleted": False,
|
||||
"favorite": False,
|
||||
"exclude": False,
|
||||
"db_checked": False,
|
||||
"skip_metadata_refresh": False,
|
||||
"metadata_source": "",
|
||||
"last_checked_at": 0,
|
||||
"hash_status": "completed",
|
||||
"trainedWords": [],
|
||||
"hf_url": hf_url,
|
||||
"usage_tips": "{}",
|
||||
}
|
||||
|
||||
with open(metadata_path, "w", encoding="utf-8") as fh:
|
||||
json.dump(metadata, fh, indent=2, ensure_ascii=False)
|
||||
|
||||
logger.debug("Created initial metadata for %s -> %s", repo_id, metadata_path)
|
||||
return model_path
|
||||
|
||||
|
||||
def create_all_initial_metadata(
|
||||
repos: List[str],
|
||||
output_dir: str,
|
||||
*,
|
||||
skip_existing: bool = True,
|
||||
) -> List[str]:
|
||||
"""Create initial metadata for every repo in *repos*.
|
||||
|
||||
Returns a list of model paths in the same order as *repos*.
|
||||
``skip_existing=True`` skips repos whose metadata already exists,
|
||||
allowing safe re-run.
|
||||
"""
|
||||
model_paths: List[str] = []
|
||||
for repo_id in repos:
|
||||
model_dir = build_model_dir(output_dir, repo_id)
|
||||
model_path = build_model_path(model_dir)
|
||||
metadata_path = build_metadata_path(model_path)
|
||||
|
||||
if skip_existing and os.path.exists(metadata_path):
|
||||
model_paths.append(model_path)
|
||||
continue
|
||||
|
||||
model_paths.append(create_initial_metadata(output_dir, repo_id))
|
||||
|
||||
logger.info(
|
||||
"Constructed initial metadata for %d/%d repos",
|
||||
len(model_paths),
|
||||
len(repos),
|
||||
)
|
||||
return model_paths
|
||||
326
tests/enrich_hf_validation/report_generator.py
Normal file
326
tests/enrich_hf_validation/report_generator.py
Normal file
@@ -0,0 +1,326 @@
|
||||
"""Generate structured reports from evaluation results.
|
||||
|
||||
Produces:
|
||||
|
||||
1. A JSON data dump (``report.json``) with all scores and aggregations.
|
||||
2. A human-readable Markdown report (``report.md``) with summary stats,
|
||||
issue patterns, and actionable optimisation suggestions.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from datetime import datetime
|
||||
from typing import Any, Dict, List
|
||||
|
||||
from .evaluation_engine import ScoreRecord
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Markdown report
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _fmt_pct(value: float) -> str:
|
||||
return f"{value:.1f}%"
|
||||
|
||||
|
||||
def _bar(value: float, width: int = 20) -> str:
|
||||
filled = int(round(value / 100 * width))
|
||||
return "█" * filled + "░" * (width - filled)
|
||||
|
||||
|
||||
def generate_optimisation_suggestions(
|
||||
agg: Dict[str, Any],
|
||||
scores: List[ScoreRecord],
|
||||
) -> List[str]:
|
||||
"""Analyse evaluation results and produce concrete suggestions."""
|
||||
suggestions: List[str] = []
|
||||
fa = agg.get("field_aggregates", {})
|
||||
|
||||
# --- base_model ---
|
||||
bm = fa.get("base_model", {})
|
||||
if bm and bm.get("empty_rate_pct", 0) > 30:
|
||||
suggestions.append(
|
||||
"- **base_model 空置率高 ({:.0f}%)**: 多数 HF 模型卡片未在 YAML frontmatter 中声明 "
|
||||
"`base_model:` 字段，LLM 无法推断。可考虑在 prompt 中增加 \"look at the model file name "
|
||||
"for clues\" 的引导，或在后处理中增加基于文件名规则的 fallback 猜测。".format(
|
||||
bm.get("empty_rate_pct", 0)
|
||||
)
|
||||
)
|
||||
bm_invalid = sum(
|
||||
1
|
||||
for s in scores
|
||||
if s["raw_values"]["base_model"]
|
||||
and s["raw_values"]["base_model"] != "Unknown"
|
||||
and s["raw_values"]["base_model"] not in {
|
||||
"SD 1.4", "SD 1.5", "SD 1.5 LCM", "SD 1.5 Hyper",
|
||||
"SD 2.0", "SD 2.1", "SD 3", "SD 3.5", "SD 3.5 Medium",
|
||||
"SD 3.5 Large", "SD 3.5 Large Turbo",
|
||||
"SDXL 1.0", "SDXL Lightning", "SDXL Hyper",
|
||||
"Flux.1 D", "Flux.1 S", "Flux.1 Krea", "Flux.1 Kontext",
|
||||
"Flux.2 D", "Flux.2 Klein 9B", "Flux.2 Klein 9B-base",
|
||||
"Flux.2 Klein 4B", "Flux.2 Klein 4B-base",
|
||||
"AuraFlow", "Chroma", "PixArt a", "PixArt E",
|
||||
"Hunyuan 1", "Lumina", "Kolors",
|
||||
"NoobAI", "Illustrious", "Pony", "Pony V7",
|
||||
"HiDream", "Qwen", "ZImageTurbo", "ZImageBase",
|
||||
"SVD", "LTXV", "LTXV2", "LTXV 2.3",
|
||||
"CogVideoX", "Mochi",
|
||||
"Wan Video", "Wan Video 1.3B t2v", "Wan Video 14B t2v",
|
||||
"Wan Video 14B i2v 480p", "Wan Video 14B i2v 720p",
|
||||
"Wan Video 2.2 TI2V-5B", "Wan Video 2.2 T2V-A14B",
|
||||
"Wan Video 2.2 I2V-A14B",
|
||||
"Wan Video 2.5 T2V", "Wan Video 2.5 I2V",
|
||||
"Hunyuan Video", "Anima", "Ernie", "Ernie Turbo",
|
||||
"Nucleus", "Krea 2",
|
||||
}
|
||||
)
|
||||
if bm_invalid > 5:
|
||||
suggestions.append(
|
||||
"- **base_model 含非标准值 ({} 个)**: LLM 输出了未在 `SUPPORTED_DOWNLOAD_SKIP_BASE_MODELS` "
|
||||
"中的 base model 名称。建议在 prompt 中强调 \"Use EXACTLY one name from the list\" 并在 "
|
||||
"`PostProcessor` 中加一层验证过滤，非标准值直接丢弃。".format(bm_invalid)
|
||||
)
|
||||
|
||||
# --- trigger_words ---
|
||||
tw = fa.get("trigger_words", {})
|
||||
if tw and tw.get("empty_rate_pct", 0) > 40:
|
||||
suggestions.append(
|
||||
"- **trigger_words 空置率高 ({:.0f}%)**: 大量 HF 模型卡没有明确的 "
|
||||
"`instance_prompt:` 或 trigger word 说明。当前 prompt 已覆盖常见模式。若确认这些模型确实"
|
||||
"没有 trigger words（例如 style lora），空数组是正确结果，不需优化。".format(
|
||||
tw.get("empty_rate_pct", 0)
|
||||
)
|
||||
)
|
||||
|
||||
# --- tags ---
|
||||
tag = fa.get("tags", {})
|
||||
if tag and tag.get("empty_rate_pct", 0) > 30:
|
||||
suggestions.append(
|
||||
"- **tags 空置率高 ({:.0f}%)**: 当前 prompt 要求 tags 必须与 "
|
||||
"`priority_tags`（CIVITAI_MODEL_TAGS）对齐。HF 模型的标签体系与 Civitai 不同，"
|
||||
"很多 model card 使用细粒度标签（如 `pokemon`、`watercolor`）而不在 priority list 中。"
|
||||
"建议: 扩大 priority_tags 范围，或允许 LLM 自由生成 tags 后只做去重不做严格过滤。".format(
|
||||
tag.get("empty_rate_pct", 0)
|
||||
)
|
||||
)
|
||||
|
||||
# --- tags priority coverage ---
|
||||
low_coverage = sum(
|
||||
1
|
||||
for s in scores
|
||||
if s["field_scores"].get("tags_priority_coverage", 5) < 3 # < 60% of max
|
||||
and s["field_scores"].get("tags", 0) > 0
|
||||
)
|
||||
if low_coverage > 10:
|
||||
suggestions.append(
|
||||
"- **{} 个模型的 tags 与 priority_tags 匹配度低于 60%**: "
|
||||
"LLM 生成了有意义但不属于 CIVITAI_MODEL_TAGS 的标签。这说明 priority_tags "
|
||||
"的覆盖范围对 HF 模型不足，建议按 HF 模型的实际分布补充新类别。".format(low_coverage)
|
||||
)
|
||||
|
||||
# --- preview ---
|
||||
prev = fa.get("preview_downloaded", {})
|
||||
if prev and prev.get("empty_rate_pct", 0) > 50:
|
||||
suggestions.append(
|
||||
"- **预览图下载成功率低 ({:.0f}%)**: 很多 HF 模型卡没有 embed 图片（仅使用 YAML widget "
|
||||
"或 external link）。当前 `md_to_html.py` 的 `extract_gallery_images` 和 "
|
||||
"`extract_gallery_table_images` 已覆盖了多数场景。若预览图不重要，可降低此字段权重。".format(
|
||||
prev.get("empty_rate_pct", 0)
|
||||
)
|
||||
)
|
||||
|
||||
# --- usage_tips ---
|
||||
ut = fa.get("usage_tips", {})
|
||||
if ut and ut.get("empty_rate_pct", 0) > 70:
|
||||
suggestions.append(
|
||||
"- **usage_tips 空置率极高 ({:.0f}%)**: 这是预期行为。HF 模型卡通常不包含 LoRA "
|
||||
"强度/CLIP skip 等结构化参数。当前提取策略已合理。若需要可用数据，" "可以考虑使用模型类型的通用默认值。".format(
|
||||
ut.get("empty_rate_pct", 0)
|
||||
)
|
||||
)
|
||||
|
||||
# --- short_description ---
|
||||
sd = fa.get("short_description", {})
|
||||
if sd and sd.get("empty_rate_pct", 0) > 40:
|
||||
suggestions.append(
|
||||
"- **short_description 空置率 ({:.0f}%)**: 部分 HF 模型卡 README 内容极少（仅含标签和训练参数）。".format(
|
||||
sd.get("empty_rate_pct", 0)
|
||||
)
|
||||
)
|
||||
|
||||
if not suggestions:
|
||||
suggestions.append("- 未发现明显问题模式，各字段填充率均在可接受范围。")
|
||||
|
||||
return suggestions
|
||||
|
||||
|
||||
def generate_markdown_report(
|
||||
agg: Dict[str, Any],
|
||||
scores: List[ScoreRecord],
|
||||
output_dir: str,
|
||||
duration_summary: Dict[str, Any] | None = None,
|
||||
) -> str:
|
||||
"""Write ``report.md`` and return its content."""
|
||||
lines: List[str] = []
|
||||
def wl(text: str = "") -> None:
|
||||
lines.append(text)
|
||||
|
||||
wl("# HF Metadata Enrichment Validation Report")
|
||||
wl()
|
||||
wl(f"Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
|
||||
wl(f"Models evaluated: **{agg.get('model_count', 0)}**")
|
||||
wl(f"Successful enrichments: **{agg.get('success_count', 0)}**")
|
||||
wl(f"Failures: **{agg.get('fail_count', 0)}**")
|
||||
wl()
|
||||
|
||||
# ---- Duration ----
|
||||
if duration_summary:
|
||||
wl("## Timing")
|
||||
wl()
|
||||
wl(f"- Total wall time: **{duration_summary.get('total_wall_s', 0):.0f} s** ")
|
||||
wl(f" ({duration_summary.get('total_wall_s', 0) / 60:.1f} min)")
|
||||
wl(f"- Mean per model: **{duration_summary.get('mean_s', 0):.1f} s**")
|
||||
wl(f"- Median per model: **{duration_summary.get('median_s', 0):.1f} s**")
|
||||
wl(f"- Fastest: **{duration_summary.get('min_s', 0):.1f} s**")
|
||||
wl(f"- Slowest: **{duration_summary.get('max_s', 0):.1f} s**")
|
||||
wl()
|
||||
|
||||
# ---- Overall score ----
|
||||
ts = agg.get("total_score", {})
|
||||
wl("## Overall Score Distribution (0–100)")
|
||||
wl()
|
||||
wl(f"| Metric | Value |")
|
||||
wl(f"|--------|-------|")
|
||||
wl(f"| Mean | {ts.get('mean', 'N/A')} |")
|
||||
wl(f"| Median | {ts.get('median', 'N/A')} |")
|
||||
wl(f"| Min | {ts.get('min', 'N/A')} |")
|
||||
wl(f"| Max | {ts.get('max', 'N/A')} |")
|
||||
wl()
|
||||
for label, key in [
|
||||
("Excellent (≥80)", "excellent_80+"),
|
||||
("Good (60–79)", "good_60_79"),
|
||||
("Fair (40–59)", "fair_40_59"),
|
||||
("Poor (20–39)", "poor_20_39"),
|
||||
("Bad (<20)", "bad_0_19"),
|
||||
]:
|
||||
count = ts.get("bins", {}).get(key, 0)
|
||||
pct = count / agg["model_count"] * 100 if agg["model_count"] else 0
|
||||
wl(f"- **{label}**: {count} models ({_fmt_pct(pct)})")
|
||||
wl()
|
||||
|
||||
# ---- Per-field aggregates ----
|
||||
wl("## Per-Field Completeness")
|
||||
wl()
|
||||
wl("| Field | Mean Score | Fill Rate | Empty Rate |")
|
||||
wl("|-------|-----------:|----------:|-----------:|")
|
||||
fa = agg.get("field_aggregates", {})
|
||||
for fn in [
|
||||
"base_model", "trigger_words", "short_description", "tags",
|
||||
"tags_priority_coverage", "notes", "usage_tips",
|
||||
"modelDescription_html", "preview_downloaded",
|
||||
]:
|
||||
f = fa.get(fn, {})
|
||||
if not f:
|
||||
continue
|
||||
wl(
|
||||
f"| {fn} "
|
||||
f"| {f.get('mean', 'N/A')} "
|
||||
f"| {_fmt_pct(f.get('fill_rate_pct', 0))} "
|
||||
f"| {_fmt_pct(f.get('empty_rate_pct', 0))} |"
|
||||
)
|
||||
wl()
|
||||
|
||||
# ---- Confidence distribution ----
|
||||
wl("## LLM Confidence Distribution")
|
||||
wl()
|
||||
cd = agg.get("confidence_distribution", {})
|
||||
total_conf = sum(cd.values()) or 1
|
||||
for level in ["high", "medium", "low", ""]:
|
||||
count = cd.get(level, 0)
|
||||
label = level if level else "(not reported)"
|
||||
pct = count / total_conf * 100
|
||||
bar = _bar(pct)
|
||||
wl(f"- **{label}**: {count} {bar} {_fmt_pct(pct)}")
|
||||
wl()
|
||||
|
||||
# ---- Top issues ----
|
||||
wl("## Most Frequent Issues")
|
||||
wl()
|
||||
for issue, count in agg.get("top_issues", []):
|
||||
pct = count / agg["model_count"] * 100 if agg["model_count"] else 0
|
||||
wl(f"- **{issue}** — {count}/{agg['model_count']} ({_fmt_pct(pct)})")
|
||||
wl()
|
||||
|
||||
# ---- Optimisation suggestions ----
|
||||
wl("## Optimisation Suggestions")
|
||||
wl()
|
||||
suggestions = generate_optimisation_suggestions(agg, scores)
|
||||
for s in suggestions:
|
||||
wl(s)
|
||||
wl()
|
||||
|
||||
# ---- Per-model detail ----
|
||||
wl("## Per-Model Detail")
|
||||
wl()
|
||||
wl("<details>")
|
||||
wl("<summary>Click to expand</summary>")
|
||||
wl()
|
||||
wl("| # | Repo ID | Score | Issues | Confidence |")
|
||||
wl("|---|---------|------:|--------|------------|")
|
||||
for i, s in enumerate(scores, 1):
|
||||
issue_count = len(s["issues"])
|
||||
issue_str = (
|
||||
f"{issue_count} issue(s)" if issue_count else "✓ ok"
|
||||
)
|
||||
wl(
|
||||
f"| {i} "
|
||||
f"| {s['repo_id']} "
|
||||
f"| {s['total_score']} "
|
||||
f"| {issue_str} "
|
||||
f"| {s.get('confidence_from_llm', '') or '-'} |"
|
||||
)
|
||||
wl()
|
||||
wl("</details>")
|
||||
wl()
|
||||
|
||||
content = "\n".join(lines)
|
||||
report_path = os.path.join(output_dir, "report.md")
|
||||
with open(report_path, "w", encoding="utf-8") as fh:
|
||||
fh.write(content)
|
||||
logger.info("Markdown report written to %s", report_path)
|
||||
return content
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# JSON dump
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def save_json_report(
|
||||
agg: Dict[str, Any],
|
||||
scores: List[ScoreRecord],
|
||||
enrichment_results: List[Dict[str, Any]],
|
||||
output_dir: str,
|
||||
duration_summary: Dict[str, Any] | None = None,
|
||||
) -> str:
|
||||
"""Write ``report.json`` and return the path."""
|
||||
report: Dict[str, Any] = {
|
||||
"metadata": {
|
||||
"generated_at": datetime.now().isoformat(),
|
||||
"model_count": agg.get("model_count", 0),
|
||||
},
|
||||
"aggregate": agg,
|
||||
"timing": duration_summary or {},
|
||||
"per_model_scores": scores,
|
||||
"enrichment_results": enrichment_results,
|
||||
}
|
||||
path = os.path.join(output_dir, "report.json")
|
||||
with open(path, "w", encoding="utf-8") as fh:
|
||||
json.dump(report, fh, indent=2, ensure_ascii=False)
|
||||
logger.info("JSON report written to %s", path)
|
||||
return path
|
||||
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