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
2026-07-04 12:00:15 +08:00
parent a1fd4e150b
commit 170c8068c5
13 changed files with 1998 additions and 48 deletions

View 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