""" core/engine.py ────────────── DanbooruTagger 核心引擎 缓存格式（存于 cache_dir/ 目录）： embeddings.safetensors — 四路向量矩阵（FP16），行顺序与 metadata.parquet 完全对齐 metadata.parquet — DataFrame（name/cn_name/cn_core/wiki/nsfw/category/post_count） meta.json — 标量元数据（max_log_count、schema_version） """ from __future__ import annotations import asyncio import json import os import re import time from collections import OrderedDict from pathlib import Path from datetime import datetime from typing import Any, Optional import jieba import numpy as np import pandas as pd import torch from safetensors.torch import save_file as st_save from safetensors import safe_open from sentence_transformers import SentenceTransformer from .models import SearchRequest, SearchResponse, TagResult from platform_utils import ( PLATFORM, is_cloud, download_file, resolve_model_path, ) # 限制 PyTorch CPU 线程数，给 asyncio 事件循环留出至少一个核心。 torch.set_num_threads(max(1, (os.cpu_count() or 2) - 1)) # LRU 缓存 class LRUCache: def __init__(self, maxsize: int): self._cache: OrderedDict[Any, Any] = OrderedDict() self._maxsize = maxsize def get(self, key: Any) -> Any: if key not in self._cache: return None self._cache.move_to_end(key) return self._cache[key] def put(self, key: Any, value: Any) -> None: if key in self._cache: self._cache.move_to_end(key) else: if len(self._cache) >= self._maxsize: self._cache.popitem(last=False) self._cache[key] = value def __len__(self) -> int: return len(self._cache) # ────────────────────────────────────────────── # 常量 # ────────────────────────────────────────────── STOP_WORDS: frozenset[str] = frozenset({ ',', '.', ':', ';', '?', '!', '"', "'", '`', '(', ')', '[', ']', '{', '}', '<', '>', '-', '_', '=', '+', '/', '\\', '|', '@', '#', '$', '%', '^', '&', '*', '~', '，', '。', '：', '；', '？', '！', '\u201c', '\u201d', '\u2018', '\u2019', '（', '）', '【', '】', '《', '》', '、', '…', '—', '·', ' ', '\t', '\n', '\r', '的', '地', '得', '了', '着', '过', '是', '为', '被', '给', '把', '让', '由', '在', '从', '自', '向', '往', '对', '于', '和', '与', '及', '或', '且', '而', '但', '并', '即', '又', '也', '啊', '吗', '吧', '呢', '噢', '哦', '哈', '呀', '哇', '我', '你', '他', '她', '它', '我们', '你们', '他们', '这', '那', '此', '其', '谁', '啥', '某', '每', '这个', '那个', '这些', '那些', '这里', '那里', '个', '位', '只', '条', '张', '幅', '件', '套', '双', '对', '副', '种', '类', '群', '些', '点', '份', '部', '名', '很', '太', '更', '最', '挺', '特', '好', '真', '一', '一个', '一种', '一下', '一点', '一些', '有', '无', '非', '没', '不', '正在', '已经', '正', '刚', '开始', '继续', '一直', '不断', '穿着', '戴着', '穿', '戴', '带有', '具有', '拥有', '看起来', '看上去', '显得', '仿佛', '似乎', '十分', '非常', '特别', '比较', '图片', '画面', '图像', '位于', '处于', '许多', '大量', '各种', '所有', '其他', '其它', # ── 英文停用词 ── 'a', 'an', 'the', 'in', 'on', 'at', 'to', 'for', 'of', 'with', 'by', 'from', 'as', 'into', 'about', 'between', 'through', 'after', 'before', 'above', 'below', 'and', 'or', 'but', 'nor', 'so', 'yet', 'is', 'are', 'was', 'were', 'be', 'been', 'being', 'do', 'does', 'did', 'done', 'have', 'has', 'had', 'having', 'will', 'would', 'shall', 'should', 'can', 'could', 'may', 'might', 'must', 'not', 'no', 'very', 'too', 'also', 'just', 'only', 'even', 'still', 'i', 'me', 'my', 'we', 'our', 'you', 'your', 'he', 'him', 'his', 'she', 'her', 'it', 'its', 'they', 'them', 'their', 'this', 'that', 'these', 'those', 'which', 'who', 'whom', 'what', 'there', 'here', 'where', 'when', 'how', 'all', 'each', 'every', 'some', 'any', 'few', 'more', 'most', 'other', 'such', 'than', 'up', 'out', 'if', 'then', 'else', 'while', 'during', 'both', 'same', 'own', 'now', }) CAT_MAP: dict[str, str] = { '0': 'General', '1': 'Artist', '3': 'Copyright', '4': 'Character', '5': 'Meta', } SCHEMA_VERSION = 3 # 升级此值将自动触发全量重建，用于破坏性格式变更 # 用户显式分隔后，纯 CJK 片段超过此长度仍用 jieba 切分（避免长句被当作原子概念） _ATOMIC_CJK_MAX_LEN = 7 # 四路 embedding 层配置: (层名, tensor 属性名, DataFrame 列名) _LAYER_SPEC: list[tuple[str, str, str]] = [ ('英文', 'emb_en', 'name'), ('中文扩展词', 'emb_cn', 'cn_name'), ('释义', 'emb_wiki', 'wiki'), ('中文核心词', 'emb_cn_core', 'cn_core'), ] _ALL_LAYER_NAMES = [ln for ln, _, _ in _LAYER_SPEC] # ────────────────────────────────────────────── # 缓存路径助手 # ────────────────────────────────────────────── class _CachePaths: def __init__(self, cache_dir: str | Path): self.dir = Path(cache_dir) self.embeddings = self.dir / 'danbooru_multiview_embeddings.safetensors' self.metadata = self.dir / 'tags_metadata.parquet' self.meta_json = self.dir / 'version_data.json' def exists(self) -> bool: return ( self.embeddings.is_file() and self.metadata.is_file() and self.meta_json.is_file() ) def ensure_dir(self): self.dir.mkdir(parents=True, exist_ok=True) class DanbooruTagger: """核心搜索引擎（单例）""" _instance: Optional['DanbooruTagger'] = None _lock: Optional[asyncio.Lock] = None # 进程级 CPU 并发闸门：串行化所有 CPU 密集型操作（search / get_related / # get_group_candidates），避免并发抢占 CPU 而拖垮 asyncio 事件循环。 _cpu_sem: Optional[asyncio.Semaphore] = None @classmethod def is_ready(cls) -> bool: return cls._instance is not None and cls._instance.is_loaded @classmethod async def get_instance(cls, **kwargs) -> 'DanbooruTagger': if cls._lock is None: cls._lock = asyncio.Lock() async with cls._lock: if cls._instance is None: inst = cls(**kwargs) await asyncio.to_thread(inst.load) cls._instance = inst return cls._instance def __init__( self, model_path: Optional[str] = None, csv_file: str = 'origin_database/tags_enhanced.csv', cache_dir: str = 'tags_embedding', cooc_file: str = 'origin_database/cooccurrence_clean.csv', group_file: str = 'origin_database/tag_groups.json', tag_artist_file: str = 'origin_database/tag_artist_cooc.parquet', ): # 模型路径：优先使用显式传入，否则交由 platform_utils 解析 self.model_path = model_path or resolve_model_path() self.csv_path = csv_file self.device = 'cuda' if torch.cuda.is_available() else 'cpu' self.paths = _CachePaths(cache_dir) self.cooc_file = cooc_file self.group_file = group_file self.tag_artist_file = tag_artist_file self.model: Optional[SentenceTransformer] = None self.df: Optional[pd.DataFrame] = None self.emb_en: Optional[torch.Tensor] = None self.emb_cn: Optional[torch.Tensor] = None self.emb_wiki: Optional[torch.Tensor] = None self.emb_cn_core: Optional[torch.Tensor] = None self.max_log_count: float = 15.0 self.cooc: dict[str, list[tuple[str, int]]] = {} self._name_to_idx: dict[str, int] = {} self._tag_to_groups: dict[str, set[str]] = {} self._group_to_tags_idx: dict[str, np.ndarray] = {} self._group_cn_names: dict[str, str] = {} self._tag_artist_index: dict[str, list[tuple]] = {} self._artist_top_tags: dict[str, list[tuple]] = {} # artist → [(tag, npmi, cn_short), ...] self.is_loaded: bool = False # 预提取的列数组，避免热点路径上反复执行 df.iloc[idx] self._arr_name: Optional[np.ndarray] = None self._arr_cn_name: Optional[np.ndarray] = None self._arr_category: Optional[np.ndarray] = None self._arr_nsfw: Optional[np.ndarray] = None self._arr_wiki: Optional[np.ndarray] = None self._arr_post_count: Optional[np.ndarray] = None self._arr_pop_score: Optional[np.ndarray] = None # 三层 LRU 缓存（纯内存，重启后自动重热） # embedding 缓存：key=文本, value=归一化后的 1-D Tensor (D,)，约 40 MB self._emb_cache: LRUCache = LRUCache(maxsize=10_000) # 搜索结果缓存：key=请求参数 tuple, value=SearchResponse，约 100 MB self._search_cache: LRUCache = LRUCache(maxsize=5_000) # 关联推荐缓存：key=(seed_tuple, limit, show_nsfw), value=list[RelatedTag]，约 20 MB self._related_cache: LRUCache = LRUCache(maxsize=2_000) # ── 初始化 ──────────────────────────────────────────────────────────── def load(self) -> None: """同步加载，在线程池中调用。""" if self.is_loaded: return t0 = time.time() # ── 云端环境：从对应平台 Hub 拉取数据文件 ────────────────────────── if is_cloud(): self._pull_cloud_files() # ── 缓存校验与构建 ───────────────────────────────────────────────── if not self.paths.exists(): print('\n' + '=' * 50) print('[Engine] 未找到缓存，开始首次构建（约 1~3 分钟）...') print('=' * 50 + '\n') self._load_model() self._build_full() else: print(f'[Engine] 加载缓存 ({self.paths.dir}) ...') self._load_from_cache() if self._cached_schema_version() != SCHEMA_VERSION: print('[Engine] 缓存格式版本不符，触发全量重建...') self._load_model() self._build_full() elif os.path.exists(self.csv_path): self._load_model() self._smart_update() if self.model is None: self._load_model() self._setup_jieba_from_memory() self._load_cooc() self._load_tag_artist_cooc() self._name_to_idx = {n: i for i, n in enumerate(self.df['name'])} self._tag_names_set: set[str] = set(self._name_to_idx.keys()) self._rebuild_arrays_from_df() self._normalize_embeddings() self._load_groups() self.is_loaded = True print(f'[Engine] 初始化完成，耗时 {time.time() - t0:.2f}s') def _normalize_embeddings(self) -> None: """ 对四路 embedding 矩阵做 L2 归一化（in-place 替换 self.emb_*）。 归一化后 search 阶段可直接用矩阵乘法得到 cosine similarity， 无需每次调用 util.semantic_search 内部再做一次归一化。 """ for _, attr, _ in _LAYER_SPEC: t = getattr(self, attr) if t is None: continue setattr(self, attr, torch.nn.functional.normalize(t, p=2, dim=1)) def _rebuild_arrays_from_df(self) -> None: """ 将 DataFrame 中搜索热点路径需要的列预提取为 numpy 数组。 任何修改 self.df 行内容或行数的操作之后都必须调用此方法刷新。 """ if self.df is None: return self._arr_name = self.df['name'].to_numpy() self._arr_cn_name = self.df['cn_name'].to_numpy() self._arr_category = self.df['category'].astype(str).to_numpy() self._arr_nsfw = self.df['nsfw'].astype(str).to_numpy() self._arr_wiki = self.df['wiki'].astype(str).to_numpy() self._arr_post_count = self.df['post_count'].to_numpy() # 预算热度归一化分，避免 search 中每次 np.log1p max_log = self.max_log_count if self.max_log_count > 0 else 1.0 self._arr_pop_score = np.log1p(self._arr_post_count) / max_log def _pull_cloud_files(self) -> None: """ 从当前云平台拉取所有数据文件，并将路径写回实例属性。 HF / MS 的差异完全由 platform_utils.download_file() 屏蔽。 """ print(f'[Engine] 云端环境 ({PLATFORM})，开始拉取数据文件...') # ── HF 平台需要额外指定 repo_id（SPACE_ID）和 repo_type ────────── extra_hf_kwargs = {} if PLATFORM == 'hf': extra_hf_kwargs = { 'hf_repo_id': os.environ.get('SPACE_ID'), 'hf_repo_type': 'space', } def pull(filename: str) -> str: try: return download_file(filename, **extra_hf_kwargs) except Exception as e: print(f'[Engine] 拉取 {filename} 失败（非致命）: {e}') return filename # 回退到原始路径，让后续逻辑决定是否重建 self.csv_path = pull('origin_database/tags_enhanced.csv') self.cooc_file = pull('origin_database/cooccurrence_clean.parquet') self.group_file = pull('origin_database/tag_groups.json') self.tag_artist_file = pull('origin_database/tag_artist_cooc.parquet') meta_path = pull('tags_embedding/tags_metadata.parquet') emb_path = pull('tags_embedding/danbooru_multiview_embeddings.safetensors') json_path = pull('tags_embedding/version_data.json') # 只有三个缓存文件都成功拉取才覆盖路径，防止部分失败导致 exists() 误判 if all( Path(p).is_file() for p in (meta_path, emb_path, json_path) ): self.paths.metadata = Path(meta_path) self.paths.embeddings = Path(emb_path) self.paths.meta_json = Path(json_path) print('[Engine] 云端数据文件拉取完毕。') else: print('[Engine] 部分缓存文件拉取失败，将触发本地重建。') # ── 搜索 ────────────────────────────────────────────────────────────── def _encode_queries(self, queries: list[str]) -> tuple[torch.Tensor, list[bool]]: """批量编码查询词，命中 embedding 缓存的跳过 model.encode。 Returns: (q_emb, hit_mask): 编码后的张量 (Q, D) 以及每个 query 是否命中缓存。 """ cached_vecs: list[Optional[torch.Tensor]] = [self._emb_cache.get(q) for q in queries] uncached_idx = [i for i, v in enumerate(cached_vecs) if v is None] hit_mask = [v is not None for v in cached_vecs] if uncached_idx: uncached_texts = [queries[i] for i in uncached_idx] new_embs = self.model.encode( uncached_texts, convert_to_tensor=True, show_progress_bar=False, ).float() new_embs = torch.nn.functional.normalize(new_embs, p=2, dim=1) for j, i in enumerate(uncached_idx): emb = new_embs[j] self._emb_cache.put(queries[i], emb) cached_vecs[i] = emb return torch.stack(cached_vecs), hit_mask # type: ignore[arg-type] def search(self, request: SearchRequest) -> SearchResponse: if not self.is_loaded: self.load() cache_key = ( request.query, request.top_k, request.limit, request.popularity_weight, request.use_segmentation, tuple(sorted(request.target_layers)), tuple(sorted(request.target_categories)), request.group_mode, request.max_per_group, ) cached = self._search_cache.get(cache_key) if cached is not None: return cached if request.use_segmentation: raw_kw, raw_segments = self._smart_split(request.query) keywords = [w.strip() for w in raw_kw if w.strip() and w.strip() not in STOP_WORDS] # raw_segments: 分隔符切分后的原始片段（未经 jieba），作为从句级查询插入。 # 排除与完整 query 相同、以及已出现在 keywords 中的片段，避免重复编码和权重膨胀。 keywords_set = set(keywords) extra_segments = [s for s in raw_segments if s != request.query and s not in keywords_set] queries = [request.query] + extra_segments + keywords else: keywords = [] extra_segments = [] queries = [request.query] q_emb, hit_mask = self._encode_queries(queries) tl = request.target_layers k = request.top_k # 每个查询词单独做意图识别，避免长句意图污染短分词 query_weights = [self._detect_intent(q) for q in queries] active_layers = [ln for ln in _ALL_LAYER_NAMES if ln in tl] # 预算每个 query × 每个 layer 的 top_k 配额 # cur_pvk_per_q[i][ln] = 第 i 个 query 在 layer ln 的配额 cur_pvk_per_q: list[dict[str, int]] = [] for cur_weights in query_weights: if active_layers: aw = {l: cur_weights.get(l, 1.0) for l in active_layers} total_aw = sum(aw.values()) cur_pvk = {l: max(1, round(k * aw[l] / total_aw)) for l in active_layers} else: cur_pvk = {} cur_pvk_per_q.append(cur_pvk) target_cats = request.target_categories w_pop = request.popularity_weight final: dict[str, TagResult] = {} # 按 layer 批量做矩阵乘 + topk，合并 Q×L=20 次小调用为 L=4 次大调用 for ln, attr, _ in _LAYER_SPEC: if ln not in tl: continue emb_matrix = getattr(self, attr) # (N, D)，已归一化 if emb_matrix is None: continue # 该 layer 在所有 query 中的最大配额（少数 query 会算到多余的 hit，最后按各自配额截断） k_max = max((cur_pvk_per_q[i].get(ln, 1) for i in range(len(queries))), default=1) k_max = min(k_max, emb_matrix.shape[0]) scores = q_emb @ emb_matrix.T # (Q, N) top_v, top_i = scores.topk(k_max, dim=1) # (Q, k_max) top_v_list = top_v.tolist() top_i_list = top_i.tolist() for i, source_word in enumerate(queries): cur_weights = query_weights[i] kq = cur_pvk_per_q[i].get(ln, 1) layer_w = cur_weights.get(ln, 1.0) row_v = top_v_list[i] row_i = top_i_list[i] # 仅取该 query 自己的配额条数 for j in range(min(kq, len(row_v))): score = row_v[j] if score < 0.35: # topk 已按分数降序，后续都低于阈值，可提前结束 break idx = row_i[j] cat_text = CAT_MAP.get(self._arr_category[idx], 'Other') if cat_text not in target_cats: continue tag_name = self._arr_name[idx] count = self._arr_post_count[idx] pop_score = self._arr_pop_score[idx] final_score = score * layer_w * (1 - w_pop) + pop_score * w_pop if tag_name not in final or final_score > final[tag_name].final_score: final[tag_name] = TagResult( tag=tag_name, cn_name=self._arr_cn_name[idx], category=cat_text, nsfw=self._arr_nsfw[idx], final_score=round(float(final_score), 4), semantic_score=round(float(score), 4), count=int(count), source=source_word, layer=ln, wiki=self._arr_wiki[idx], ) # ── 全句语义一致性软重排 ────────────────────────────────────────── # 对每个候选标签，计算其与完整原始查询（而非分词片段）的语义相似度， # 将相似度作为软因子乘入 final_score，使仅由分词碎片匹配到的噪声 # 标签自然下沉，同时不硬过滤任何结果。 # 批量矩阵乘法替代逐条 torch.dot，O(R*L) 降为 O(L) + O(R)。 full_q = q_emb[0] # queries[0] 始终为完整原始查询 alpha = 0.3 if request.use_segmentation else 0 # 一致性调节强度（0=不调节, 1=完全按一致性重排），仅在启用分词时有意义 if alpha > 0 and final: tag_list = list(final.keys()) tag_indices = [self._name_to_idx[t] for t in tag_list] idx_tensor = torch.tensor(tag_indices, dtype=torch.long, device=full_q.device) max_co = torch.zeros(len(tag_indices), device=full_q.device) for ln, attr, _ in _LAYER_SPEC: if ln not in tl: continue emb_selected = getattr(self, attr)[idx_tensor] # (R, D) co = (full_q.unsqueeze(0) @ emb_selected.T).squeeze(0) # (R,) max_co = torch.maximum(max_co, co) for i, tag in enumerate(tag_list): r = final[tag] r.final_score = round(r.final_score * (1.0 - alpha + alpha * float(max_co[i])), 4) # Group expand 处理（在 guaranteed_tags 之前，因为会改分数） if request.group_mode == "expand" and self._tag_to_groups: self._apply_group_expand(final) # 收集每个查询源的 top-1 结果（高于阈值） guaranteed_tags: set[str] = set() for source_word in queries: best: TagResult | None = None for r in final.values(): if r.source == source_word and r.final_score > 0.45: if best is None or r.final_score > best.final_score: best = r if best is not None: guaranteed_tags.add(best.tag) # 对所有候选进行排序，然后在保留保证结果的同时截断至限制数量 sorted_results = sorted(final.values(), key=lambda r: r.final_score, reverse=True) valid: list[TagResult] = [] if request.group_mode == "diverse" and self._tag_to_groups: # diverse 模式：每个 group 最多保留 max_per_group 个标签 group_counter: dict[str, int] = {} max_per = request.max_per_group for r in sorted_results: if r.final_score <= 0.45: continue if r.tag in guaranteed_tags: # guaranteed_tags 豁免 group 上限 valid.append(r) continue groups = self._tag_to_groups.get(r.tag) if not groups: # 无 group 信息，不受限制 if len(valid) < request.limit: valid.append(r) continue # 检查是否有任一 group 达上限 if any(group_counter.get(g, 0) >= max_per for g in groups): continue if len(valid) < request.limit: valid.append(r) for g in groups: group_counter[g] = group_counter.get(g, 0) + 1 else: for r in sorted_results: if r.final_score <= 0.45: continue if len(valid) < request.limit or r.tag in guaranteed_tags: valid.append(r) tags_all = ', '.join(r.tag for r in valid) tags_sfw = ', '.join(r.tag for r in valid if r.nsfw != '1') cached_queries = [q for q, hit in zip(queries, hit_mask) if hit] response = SearchResponse( tags_all=tags_all, tags_sfw=tags_sfw, results=valid, keywords=keywords, segments=extra_segments, cached_queries=cached_queries, ) self._search_cache.put(cache_key, response) return response # ── CPU 并发闸门（类级信号量，所有 CPU 密集型操作共享）─────────────── @classmethod def _get_cpu_sem(cls) -> asyncio.Semaphore: if cls._cpu_sem is None: cls._cpu_sem = asyncio.Semaphore(1) return cls._cpu_sem async def search_async(self, request: SearchRequest) -> SearchResponse: """search() 的并发安全异步封装：共享闸门串行化 + 线程池执行。 所有异步入口（MCP / API / UI）都应改用本方法，而非各自 asyncio.to_thread(self.search)，以共享同一个 CPU 并发闸门。 包含 60 秒超时，防止异常卡死导致信号量永久泄漏。 """ async with self._get_cpu_sem(): return await asyncio.wait_for( asyncio.to_thread(self.search, request), timeout=120.0, ) async def get_related_async( self, seed_tags: list[str], exclude: set[str] | None = None, limit: int = 20, show_nsfw: bool = True, ) -> list: """get_related() 的并发安全异步封装，共享同一个 CPU 闸门。""" async with self._get_cpu_sem(): return await asyncio.to_thread( self.get_related, seed_tags, exclude, limit, show_nsfw, ) async def get_group_candidates_async( self, selected_tags: list[str], show_nsfw: bool = True, ) -> list[dict]: """get_group_candidates() 的并发安全异步封装，共享同一个 CPU 闸门。""" async with self._get_cpu_sem(): return await asyncio.to_thread( self.get_group_candidates, selected_tags, show_nsfw, ) async def search_artists_by_tags_async( self, tags: list[str], limit: int = 30, min_cooc: int = 5, ) -> list: """search_artists_by_tags() 的并发安全异步封装。""" async with self._get_cpu_sem(): return await asyncio.to_thread( self.search_artists_by_tags, tags, limit, min_cooc, ) async def search_artists_pipeline_async( self, query: str, limit: int = 30, min_cooc: int = 5, target_layers: list[str] | None = None, target_categories: list[str] | None = None, ) -> tuple: """search_artists_pipeline() 的并发安全异步封装。""" async with self._get_cpu_sem(): return await asyncio.wait_for( asyncio.to_thread( self.search_artists_pipeline, query, limit, min_cooc, target_layers, target_categories, ), timeout=120.0, ) def _apply_group_expand(self, final: dict[str, TagResult]) -> None: """expand 模式：提升同 group 标签的分数。""" BETA = 0.2 TOP_N = 20 sorted_items = sorted(final.values(), key=lambda r: r.final_score, reverse=True) top_n = min(TOP_N, len(sorted_items)) anchor_results = sorted_items[:top_n] # 收集锚点结果所属的所有 group active_groups: set[str] = set() for r in anchor_results: groups = self._tag_to_groups.get(r.tag) if groups: active_groups.update(groups) if not active_groups: return # 预计算每个 group 的锚点最大分 group_max_score: dict[str, float] = {} for g in active_groups: group_max_score[g] = max( (r.final_score for r in anchor_results if g in self._tag_to_groups.get(r.tag, set())), default=0.0, ) # 对所有候选应用 boost for r in final.values(): groups = self._tag_to_groups.get(r.tag) if not groups: continue overlap = groups & active_groups if not overlap: continue best_group_score = max(group_max_score[g] for g in overlap) boost = 1.0 + BETA * best_group_score r.final_score = round(r.final_score * boost, 4) # ── 全量构建 ────────────────────────────────────────────────────────── def _build_full(self) -> None: print(f'[Engine] 全量读取 {self.csv_path} ...') raw_df = self._read_csv_robust(self.csv_path) self.df = self._preprocess_raw_df(raw_df) self.max_log_count = float(np.log1p(self.df['post_count'].max())) self._encode_all_and_save() def _encode_all_and_save(self) -> None: print('[Engine] 全量编码...') for _, attr, col in _LAYER_SPEC: setattr(self, attr, self._encode_texts(self.df[col].tolist())) self._save_cache() # ── 增量更新 ────────────────────────────────────────────────────────── def _smart_update(self) -> None: print('[Engine] 检查增量变更...') t0 = time.time() raw_df = self._read_csv_robust(self.csv_path) new_df = self._preprocess_raw_df(raw_df) _SIG_COLS = ['cn_name', 'wiki', 'cn_core'] def _sig(df: pd.DataFrame, iloc_idx: int) -> tuple: row = df.iloc[iloc_idx] return tuple(str(row.get(c, '')) for c in _SIG_COLS) cached_idx: dict[str, int] = {n: i for i, n in enumerate(self.df['name'])} new_idx: dict[str, int] = {n: i for i, n in enumerate(new_df['name'])} added_names = [n for n in new_idx if n not in cached_idx] deleted_names = [n for n in cached_idx if n not in new_idx] changed_names = [ n for n in new_idx if n in cached_idx and _sig(new_df, new_idx[n]) != _sig(self.df, cached_idx[n]) ] if not added_names and not deleted_names and not changed_names: print('[Engine] 数据已是最新，无需更新。') return print(f'[Engine] 变更 → 新增: {len(added_names)} 修改: {len(changed_names)} 删除: {len(deleted_names)}') if deleted_names: keep_mask = ~self.df['name'].isin(set(deleted_names)) keep_pos = [i for i, v in enumerate(keep_mask) if v] self.df = self.df[keep_mask].reset_index(drop=True) for _, attr, _ in _LAYER_SPEC: setattr(self, attr, getattr(self, attr)[keep_pos]) cached_idx = {n: i for i, n in enumerate(self.df['name'])} if changed_names: changed_rows = new_df[new_df['name'].isin(set(changed_names))].reset_index(drop=True) _vecs = {attr: self._encode_texts(changed_rows[col].tolist()) for _, attr, col in _LAYER_SPEC} for j, name in enumerate(changed_rows['name']): ci = cached_idx[name] for _, attr, _ in _LAYER_SPEC: getattr(self, attr)[ci] = _vecs[attr][j] for col in changed_rows.columns: self.df.at[ci, col] = changed_rows.at[j, col] if added_names: added_rows = new_df[new_df['name'].isin(set(added_names))].reset_index(drop=True) for _, attr, col in _LAYER_SPEC: vecs = self._encode_texts(added_rows[col].tolist()) setattr(self, attr, torch.cat([getattr(self, attr), vecs], dim=0)) self.df = pd.concat([self.df, added_rows], ignore_index=True) self.max_log_count = float(np.log1p(self.df['post_count'].max())) self._name_to_idx = {n: i for i, n in enumerate(self.df['name'])} self._tag_names_set = set(self._name_to_idx.keys()) self._rebuild_arrays_from_df() self._load_groups() self._normalize_embeddings() self._save_cache() print(f'[Engine] 增量更新完成，耗时 {time.time() - t0:.2f}s（共 {len(self.df)} 条）') # ── 缓存 I/O ────────────────────────────────────────────────────────── def _save_cache(self) -> None: self.paths.ensure_dir() st_save( {attr: getattr(self, attr).half() for _, attr, _ in _LAYER_SPEC}, str(self.paths.embeddings), ) save_cols = ['name', 'cn_name', 'cn_core', 'wiki', 'nsfw', 'category', 'post_count'] self.df[save_cols].to_parquet(str(self.paths.metadata), index=False) current_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S') with open(self.paths.meta_json, 'w', encoding='utf-8') as f: json.dump({ 'schema_version': SCHEMA_VERSION, 'updated_at': current_time, }, f, ensure_ascii=False, indent=4) print(f'[Engine] 缓存保存完成（{len(self.df)} 条记录），生成时间: {current_time}') def _load_from_cache(self) -> None: t0 = time.time() emb_path = str(self.paths.embeddings) meta_path = str(self.paths.metadata) emb_size_mb = self.paths.embeddings.stat().st_size / 1024 / 1024 meta_size_mb = self.paths.metadata.stat().st_size / 1024 / 1024 # ── 步骤 1/3: embedding 文件（逐层加载以显示进度）───────────────── n_layers = len(_LAYER_SPEC) print(f' [1/3] 加载 embedding ({emb_size_mb:.0f} MB, {n_layers} 层) ...') with safe_open(emb_path, framework="pt", device=self.device) as f: for i, (name, attr, _) in enumerate(_LAYER_SPEC, 1): print(f' [{i}/{n_layers}] {attr:12s} ...', end=' ', flush=True) _t = time.time() tensor = f.get_tensor(attr) setattr(self, attr, tensor.float()) print(f'✓ {time.time() - _t:.1f}s') print(f' [1/3] ✓ {time.time() - t0:.1f}s') # ── 步骤 2/3: 元数据 ────────────────────────────────────────────── print(f' [2/3] 加载元数据 ({meta_size_mb:.0f} MB) ...', end=' ', flush=True) _t = time.time() self.df = pd.read_parquet(meta_path) print(f'✓ {time.time() - _t:.1f}s ({len(self.df):,} 条)') # ── 步骤 3/3: 统计信息 ──────────────────────────────────────────── print(f' [3/3] 计算统计信息 ...', end=' ', flush=True) _t = time.time() self.max_log_count = float(np.log1p(self.df['post_count'].max())) print(f'✓ {time.time() - _t:.1f}s') print(f'[Engine] 缓存加载完成 (总耗时 {time.time() - t0:.1f}s)') def _cached_schema_version(self) -> int: try: with open(self.paths.meta_json, 'r', encoding='utf-8') as f: return int(json.load(f).get('schema_version', 1)) except Exception: return 0 # ── 编码 & 预处理 ────────────────────────────────────────────────────── def _encode_texts(self, texts: list[str]) -> torch.Tensor: return self.model.encode( texts, batch_size=64, show_progress_bar=False, convert_to_tensor=True, ).float() def _load_model(self) -> None: if self.model is not None: return print(f'[Engine] 加载模型 (path={self.model_path}, device={self.device})...') try: self.model = SentenceTransformer(self.model_path, device=self.device) except Exception as e: print(f'[Engine] 指定路径加载失败，尝试重新解析: {e}') fallback = resolve_model_path() self.model = SentenceTransformer(fallback, device=self.device) def _read_csv_robust(self, path: str) -> pd.DataFrame: for enc in ['utf-8', 'gbk', 'gb18030']: try: return pd.read_csv(path, dtype=str, encoding=enc).fillna('') except UnicodeDecodeError: continue raise ValueError('CSV 读取失败，请检查编码') def _preprocess_raw_df(self, df: pd.DataFrame) -> pd.DataFrame: df = df.copy() df.dropna(subset=['name'], inplace=True) df = df[df['name'].str.strip() != ''] for col in ['cn_name', 'category', 'wiki', 'nsfw']: if col not in df.columns: df[col] = '' df['category'] = df['category'].fillna('0') df['nsfw'] = df['nsfw'].fillna('0') for char in ['，', '|', '、']: df['cn_name'] = df['cn_name'].str.replace(char, ',', regex=False) if 'post_count' not in df.columns: df['post_count'] = 0 df['post_count'] = pd.to_numeric(df['post_count'], errors='coerce').fillna(0) df['cn_name'] = df['cn_name'].fillna('') df['wiki'] = df['wiki'].fillna('') df['cn_core'] = df['cn_name'].str.split(',', n=1).str[0].str.strip().fillna('') df.drop_duplicates(subset=['name'], inplace=True) df.reset_index(drop=True, inplace=True) return df def _setup_jieba_from_memory(self) -> None: if self.df is None: return unique_words: set[str] = set() for text in self.df['cn_name'].dropna().astype(str): for part in text.replace(',', ' ').split(): part = part.strip() if len(part) > 1: unique_words.add(part) for word in unique_words: jieba.add_word(word, 2000) def _detect_intent(self, query: str) -> dict[str, float]: """ 根据查询词特征返回各视图的语义分系数。 系数 > 1 表示加权，< 1 表示降权。 """ stripped = query.replace(' ', '') cn_chars = sum(1 for c in stripped if '\u4e00' <= c <= '\u9fff') en_chars = sum(1 for c in stripped if c.isascii() and c.isalpha()) total = max(len(stripped), 1) is_long = len(query) > 8 is_cn = cn_chars / total > 0.5 is_en = en_chars / total > 0.5 if is_long and is_cn: return {'英文': 0.8, '中文核心词': 0.9, '中文扩展词': 1.1, '释义': 1.4} if is_long and is_en: return {'英文': 1.0, '中文核心词': 0.8, '中文扩展词': 0.9, '释义': 1.3} if is_cn: return {'英文': 0.8, '中文核心词': 1.3, '中文扩展词': 1.1, '释义': 0.6} if is_en: return {'英文': 1.3, '中文核心词': 1.0, '中文扩展词': 0.8, '释义': 0.6} return {'英文': 1.0, '中文核心词': 1.0, '中文扩展词': 1.0, '释义': 1.0} # ── 英文分词辅助 ────────────────────────────────────────────────────── _EN_MAX_COMPOUND = 4 # 复合标签最大单词数 def _tokenize_en_chunk(self, chunk: str) -> list[str]: """对一段英文文本做分词：清洗 → 按空格切分 → 过滤停用词/纯数 → 变体规范化 → 合并已知复合标签。 变体规范化指：未直接命中 tag_set 的 token 尝试 `连字符→下划线` / 复数还原（s/es/ies→y），仅在变体落在 tag_set 才采用。 """ cleaned = re.sub(r'[,()\[\]{}:]', ' ', chunk) raw = [p for p in cleaned.split() if p] tag_set = getattr(self, '_tag_names_set', None) tokens: list[str] = [] for part in raw: low = part.lower() # 已知标签直接保留（如用户输入了带下划线的 tag 名） if tag_set and low in tag_set: tokens.append(low) continue if low in STOP_WORDS: continue # 连字符/复数变体探测（仅当变体落在 tag_set 才采用） variant = self._resolve_tag_variant(low) if variant: tokens.append(variant) continue if part.isdigit(): # 仅过滤纯数字，保留 3d/2b 等含数字的词 continue tokens.append(low) if not tokens: return [] return self._merge_compound_english(tokens) def _resolve_tag_variant(self, low: str) -> str | None: """对未直接命中 tag_set 的英文 token 探测常见变体。 覆盖：连字符→下划线（cat-ears → cat_ears）、复数→单数 （cats → cat / dresses → dress / bunnies → bunny）。 仅在变体落在 _tag_names_set 中才返回，避免 'glass'→'glas' 之类误伤。 Returns: 命中的 tag 名；都不命中返回 None。 """ tag_set = getattr(self, '_tag_names_set', None) if not tag_set: return None # 连字符直接换成下划线若直接命中 tag 则优先返回 bases = [low] if '-' in low: hyphen_normalized = low.replace('-', '_') if hyphen_normalized in tag_set: return hyphen_normalized bases.append(hyphen_normalized) # 对每个基串尝试复数还原（按"剥离短→长"顺序，避免 houses→hous 误判） for base in bases: if base.endswith('s') and len(base) > 1: v = base[:-1] if v in tag_set: return v if base.endswith('es') and len(base) > 2: v = base[:-2] if v in tag_set: return v if base.endswith('ies') and len(base) > 3: v = base[:-3] + 'y' if v in tag_set: return v return None def _merge_compound_english(self, tokens: list[str]) -> list[str]: """将相邻英文单词合并为已知的 Danbooru 下划线复合标签。 贪心最长匹配：优先 4-gram，依次递减到 bigram，匹配到即消耗。 例: ['beam', 'rifle', 'scope'] → 如果 'beam_rifle' 是标签则合并， 否则保留原样。 """ tag_set = getattr(self, '_tag_names_set', None) if tag_set is None or len(tokens) < 2: return tokens result: list[str] = [] i = 0 max_w = min(self._EN_MAX_COMPOUND, len(tokens)) while i < len(tokens): merged = False for w in range(max_w, 1, -1): # 4, 3, 2 if i + w > len(tokens): continue candidate = '_'.join(tokens[i:i + w]) if candidate in tag_set: result.append(candidate) i += w merged = True break if not merged: result.append(tokens[i]) i += 1 return result # ── 查询切分 ────────────────────────────────────────────────────────── def _smart_split(self, text: str) -> tuple[list[str], list[str]]: """将查询文本拆分为关键词列表，同时返回从句级片段。 先把文本切成交替的 CN-region（含 CJK 字符）与 EN-region（无 CJK）； 中英文用各自的规则处理： 1. CN-region 按空格/CJK 标点切出"子句"（segments），中文自然语句里 这些符号是显式概念边界。子句的 token 切分策略由"整句是否含任何 分隔符"决定： - 整句无任何分隔符 → 视作自然句，jieba 切分； - 整句有分隔符 → 用户已标边界，每个短纯 CJK 子句原子保留， 超过 _ATOMIC_CJK_MAX_LEN 才走 jieba。 2. EN-region 仅走 _tokenize_en_chunk（停用词过滤 + 复合词合并）， 不产出 segments——英文里空格是词内分隔而非概念边界。 Returns: (tokens, segments): - tokens: 处理后的关键词列表 - segments: CN-region 切出的子句片段；纯英文查询为空列表 """ user_pieces = [s.strip() for s in re.split(r'[\s\n\r，、；。]+', text) if s.strip()] if not user_pieces: return [], [] has_boundary = len(user_pieces) > 1 # 整句是否含任何用户标注的概念边界 tokens: list[str] = [] segments: list[str] = [] # 把文本切成交替的 CN-region 与 EN-region。 # CN-region 允许内部以空格/CJK 标点连接相邻 CJK 块。 cjk_region = r'[一-龥]+(?:[\s\n\r，、；。]+[一-龥]+)*' parts = re.split(f'({cjk_region})', text) for part in parts: if not part.strip(): continue if re.search(r'[一-龥]', part): # CN region：产出子句 + tokens cn_segs = [s.strip() for s in re.split(r'[\s\n\r，、；。]+', part) if s.strip()] for seg in cn_segs: segments.append(seg) if has_boundary and re.match(r'^[一-龥]+$', seg) and len(seg) <= _ATOMIC_CJK_MAX_LEN: tokens.append(seg) # 短 → 原子概念 else: for chunk in re.split(r'([一-龥]+)', seg): if not chunk.strip(): continue if re.match(r'[一-龥]+', chunk): tokens.extend(jieba.cut(chunk)) else: tokens.extend(self._tokenize_en_chunk(chunk)) else: # EN region：仅 tokenize，不产出子句 tokens.extend(self._tokenize_en_chunk(part)) return tokens, segments # ── 关联推荐 ────────────────────────────────────────────────────────── def get_related( self, seed_tags: list[str], exclude: set[str] | None = None, limit: int = 20, show_nsfw: bool = True, ) -> list: from .models import RelatedTag import math if not self.cooc or not seed_tags: return [] exclude = exclude or set() related_key = (tuple(sorted(seed_tags)), tuple(sorted(exclude)), limit, show_nsfw) cached = self._related_cache.get(related_key) if cached is not None: return cached # 估算语料库总大小 N，取数据集中发帖量的最大值，并设置合理下限 total_posts = float(max(self.df['post_count'].max(), 7000000.0)) npmi_scores: dict[str, float] = {} total_cooc: dict[str, int] = {} tag_sources: dict[str, list[str]] = {} name_to_idx = self._name_to_idx arr_post_count = self._arr_post_count for seed in seed_tags: if seed not in name_to_idx: continue seed_count = float(arr_post_count[name_to_idx[seed]] or 1) for neighbor, cnt in self.cooc.get(seed, []): if neighbor in exclude or neighbor == seed: continue if neighbor not in name_to_idx: continue neighbor_count = float(arr_post_count[name_to_idx[neighbor]] or 1) cooc = min(float(cnt), seed_count, neighbor_count) if cooc <= 0: continue # 计算分子：(Cooc * N) / (Count(A) * Count(B)) numerator = (cooc * total_posts) / (seed_count * neighbor_count) # 忽略负相关或完全不相关的词条 if numerator <= 1.0: continue pmi = math.log(numerator) # 计算分母：-log(P(A, B)) p_a_b = cooc / total_posts if p_a_b >= 1.0: npmi = 1.0 else: npmi = pmi / -math.log(p_a_b) # 多词条搜索时累加 NPMI npmi_scores[neighbor] = npmi_scores.get(neighbor, 0.0) + npmi total_cooc[neighbor] = total_cooc.get(neighbor, 0) + cnt tag_sources.setdefault(neighbor, []).append(seed) if not npmi_scores: return [] # 归一化用于前端展示 max_score = max(npmi_scores.values()) sorted_candidates = sorted(npmi_scores.items(), key=lambda x: x[1], reverse=True) # ── 构建 NPMI 结果 ───────────────────────────────────────────── results: list = [] for tag_name, raw_score in sorted_candidates: if len(results) >= limit: break idx = name_to_idx[tag_name] nsfw = self._arr_nsfw[idx] if nsfw == '1' and not show_nsfw: continue cat = CAT_MAP.get(self._arr_category[idx], 'Other') results.append(RelatedTag( tag=tag_name, cn_name=str(self._arr_cn_name[idx]), category=cat, nsfw=nsfw, cooc_count=total_cooc.get(tag_name, 0), cooc_score=round(raw_score / max_score, 4), sources=tag_sources.get(tag_name, []), post_count=int(self._arr_post_count[idx]), wiki=str(self._arr_wiki[idx]) if self._arr_wiki is not None else '', )) self._related_cache.put(related_key, results) return results def get_group_candidates( self, selected_tags: list[str], show_nsfw: bool = True, ) -> list[dict]: """根据已选标签，返回候选 Group 及其成员标签。""" if not self._tag_to_groups or not selected_tags: return [] group_hit_count: dict[str, int] = {} for tag_name in selected_tags: groups = self._tag_to_groups.get(tag_name) if groups: for g in groups: group_hit_count[g] = group_hit_count.get(g, 0) + 1 if not group_hit_count: return [] selected_set = set(selected_tags) sorted_groups = sorted(group_hit_count.items(), key=lambda x: -x[1]) results = [] for group_name, hit_count in sorted_groups: member_idxs = self._group_to_tags_idx.get(group_name) if member_idxs is None: continue tags = [] for idx in member_idxs: tag_name = str(self._arr_name[idx]) if tag_name in selected_set: continue nsfw = self._arr_nsfw[idx] if nsfw == '1' and not show_nsfw: continue cat = CAT_MAP.get(self._arr_category[idx], 'Other') tags.append({ 'tag': tag_name, 'cn_name': str(self._arr_cn_name[idx]), 'category': cat, 'nsfw': nsfw, 'post_count': int(self._arr_post_count[idx]), 'wiki': str(self._arr_wiki[idx]) if self._arr_wiki is not None else '', }) tags.sort(key=lambda x: -x['post_count']) cn_name = self._group_cn_names.get(group_name, group_name) results.append({ 'group': group_name, 'group_cn_name': cn_name, 'hit_count': hit_count, 'tags': tags, }) return results # ── 画师查找 ────────────────────────────────────────────────────────── def search_artists_by_tags( self, tags: list[str], limit: int = 30, min_cooc: int = 5, ) -> list: """按标签查找画师：聚合多个标签的 NPMI 得分，返回排名靠前的画师。""" from .models import ArtistResult if not self._tag_artist_index or not tags: return [] artist_scores: dict[str, float] = {} artist_cooc: dict[str, int] = {} artist_post_count: dict[str, int] = {} artist_sources: dict[str, list[str]] = {} artist_hits: dict[str, int] = {} for tag in tags: tag = tag.strip().lower() if not tag or tag not in self._tag_artist_index: continue for artist, npmi, cooc, post_count in self._tag_artist_index[tag]: if cooc < min_cooc: continue artist_scores[artist] = artist_scores.get(artist, 0.0) + npmi artist_cooc[artist] = artist_cooc.get(artist, 0) + cooc artist_post_count[artist] = max(artist_post_count.get(artist, 0), post_count) artist_sources.setdefault(artist, []).append(tag) artist_hits[artist] = artist_hits.get(artist, 0) + 1 if not artist_scores: return [] scored = [] for artist, raw_score in artist_scores.items(): hit_bonus = 1.0 + 0.3 * (artist_hits[artist] - 1) final_score = raw_score * hit_bonus scored.append(( artist, final_score, artist_cooc[artist], artist_post_count[artist], artist_sources[artist], artist_hits[artist], )) scored.sort(key=lambda x: -x[1]) results = [] for artist, score, cooc, post_count, sources, hits in scored[:limit]: results.append(ArtistResult( artist=artist, score=round(score, 4), cooc_count=cooc, post_count=post_count, sources=sources, hit_count=hits, )) return results def get_artist_top_tags(self, artist_names: list[str], top_n: int = 10, show_nsfw: bool = True) -> dict[str, list[str]]: """获取画师的 top-N 共现标签（含中文简称）。 Returns: {artist: [f"{tag} ({cn_short})", ...]} """ result: dict[str, list[str]] = {} for artist in artist_names: entries = self._artist_top_tags.get(artist, []) items: list[str] = [] for tag, npmi, cooc in entries: if len(items) >= top_n: break # NSFW 过滤 if not show_nsfw and self._name_to_idx is not None and tag in self._name_to_idx: idx = self._name_to_idx[tag] if self._arr_nsfw is not None and self._arr_nsfw[idx] == '1': continue cn_short = "" if self._name_to_idx is not None and tag in self._name_to_idx: idx = self._name_to_idx[tag] cn_full = str(self._arr_cn_name[idx]) if self._arr_cn_name is not None else "" cn_short = cn_full.split(',')[0].strip() if cn_full else "" display = f"{tag}" if cn_short: display += f" ({cn_short})" items.append(display) result[artist] = items return result def search_artists_pipeline( self, query: str, limit: int = 30, min_cooc: int = 5, target_layers: list[str] | None = None, target_categories: list[str] | None = None, ) -> tuple[list, list[str], list[str], list[str]]: """画师查找完整管线：自然语言 → 标签搜索 → 提取最佳标签 → 画师查询。 Returns: (artist_results, seed_tags, found_tags, missing_tags) """ if target_layers is None: target_layers = ['英文', '中文扩展词', '释义', '中文核心词'] if target_categories is None: target_categories = ['General', 'Character', 'Copyright'] # Step 1: 自然语言 → 标签搜索 tag_request = SearchRequest( query=query, top_k=10, limit=80, popularity_weight=0.15, show_nsfw=True, use_segmentation=True, target_layers=target_layers, target_categories=target_categories, ) tag_response = self.search(tag_request) if not tag_response.results: return [], [], [], [] # Step 2: 按 source 分组，每组按分数降序取前 10 个候选，找到能命中画师的第一个 tag_artist = self._tag_artist_index source_candidates: dict[str, list[str]] = {} for r in tag_response.results: if r.final_score < 0.45: continue source_candidates.setdefault(r.source, []).append(r.tag) seed_tags: list[str] = [] seen = set() for candidates in source_candidates.values(): for tag in candidates[:10]: tag_lower = tag.strip().lower() if tag_lower in tag_artist and tag_lower not in seen: seed_tags.append(tag_lower) seen.add(tag_lower) break if not seed_tags: return [], [], [], [] # Step 3: 标签 → 画师查找 artist_results = self.search_artists_by_tags(seed_tags, limit, min_cooc) # 统计 found / missing（seed_tags 中哪些匹配到了画师） matched = {s for r in artist_results for s in r.sources} found_tags = [t for t in seed_tags if t in matched] missing_tags = [t for t in seed_tags if t not in matched] return artist_results, seed_tags, found_tags, missing_tags # ── 共现数据加载 ────────────────────────────────────────────────────── def _load_cooc(self) -> None: csv_path = Path(self.cooc_file) parquet_path = csv_path.with_suffix('.parquet') if parquet_path.is_file() and ( not csv_path.is_file() or parquet_path.stat().st_mtime >= csv_path.stat().st_mtime ): read_path = parquet_path is_parquet = True elif csv_path.is_file(): read_path = csv_path is_parquet = False else: print(f'[Engine] 未找到共现表 ({self.cooc_file})，关联推荐功能不可用。') return print(f'[Engine] 加载共现表 ({read_path.name})...') t0 = time.time() try: if is_parquet: df = pd.read_parquet(str(read_path)) else: df = self._read_csv_robust(str(read_path)) df['count'] = pd.to_numeric(df['count'], errors='coerce').fillna(0).astype(int) df.to_parquet(str(parquet_path), index=False) print(f'[Engine] 已将共现表缓存为 {parquet_path.name}，下次启动将直接加载。') tag_a = df['tag_a'].astype(str).to_numpy() tag_b = df['tag_b'].astype(str).to_numpy() counts = df['count'].astype(int).to_numpy() src = np.concatenate([tag_a, tag_b]) dst = np.concatenate([tag_b, tag_a]) cnt = np.concatenate([counts, counts]) sort_idx = np.lexsort((-cnt, src)) src = src[sort_idx] dst = dst[sort_idx] cnt = cnt[sort_idx] unique_srcs, first_pos = np.unique(src, return_index=True) end_pos = np.append(first_pos[1:], len(src)) cooc: dict[str, list[tuple[str, int]]] = {} for s, start, end in zip(unique_srcs, first_pos, end_pos): cooc[s] = list(zip(dst[start:end].tolist(), cnt[start:end].tolist())) self.cooc = cooc print( f'[Engine] 共现表加载完成，{len(cooc):,} 个 tag，' f'耗时 {time.time() - t0:.2f}s' ) except Exception as e: print(f'[Engine] 共现表加载失败: {e}') def _load_tag_artist_cooc(self) -> None: """加载标签-画师共现数据（tag_artist_cooc.parquet）。""" path = Path(self.tag_artist_file) if not path.is_file(): print(f'[Engine] 未找到标签-画师共现表 ({self.tag_artist_file})，画师查找功能不可用。') return print(f'[Engine] 加载标签-画师共现表 ({path.name})...') t0 = time.time() try: import math df = pd.read_parquet(str(path)) # 预计算语料库总大小和 tag → post_count 映射 total_posts = float(max(self.df['post_count'].max(), 7_000_000)) tag_post_count = dict(zip(self.df['name'], self.df['post_count'])) index: dict[str, list[tuple]] = {} skipped = 0 for _, row in df.iterrows(): tag = str(row["tag"]) artist = str(row["artist"]) cooc = int(row["cooc_count"]) artist_pc = int(row["artist_post_count"]) # 获取标签的 post_count tag_pc = tag_post_count.get(tag) if tag_pc is None or tag_pc <= 0: skipped += 1 continue # 即时计算 NPMI（与 get_related 相同算法） cooc_capped = min(float(cooc), float(tag_pc), float(artist_pc)) if cooc_capped <= 0: skipped += 1 continue numerator = (cooc_capped * total_posts) / (float(tag_pc) * float(artist_pc)) if numerator <= 1.0: skipped += 1 continue pmi = math.log(numerator) p_ab = cooc_capped / total_posts if p_ab >= 1.0: npmi_val = 1.0 else: npmi_val = pmi / -math.log(p_ab) if math.isnan(npmi_val): skipped += 1 continue index.setdefault(tag, []).append( (artist, npmi_val, cooc, artist_pc)) self._tag_artist_index = index # 构建反向索引：artist → top tags（按 NPMI 降序） artist_tags: dict[str, list[tuple]] = {} for tag, entries in index.items(): for artist, npmi, cooc, _ in entries: artist_tags.setdefault(artist, []).append((tag, npmi, cooc)) for artist, entries in artist_tags.items(): entries.sort(key=lambda x: -x[1]) self._artist_top_tags = artist_tags print( f'[Engine] 标签-画师共现表加载完成，{len(index):,} 个标签，' f'跳过 {skipped:,} 行，耗时 {time.time() - t0:.2f}s' ) except Exception as e: print(f'[Engine] 标签-画师共现表加载失败: {e}') def _load_groups(self) -> None: """加载 Tag Group 数据，构建 tag→group 和 group→idx 索引。""" if not Path(self.group_file).is_file(): print('[Engine] 未找到 Tag Group 数据，group 功能不可用。') return with open(self.group_file, 'r', encoding='utf-8') as f: data = json.load(f) raw_t2g = data.get('tag_to_groups', {}) name_to_idx = self._name_to_idx self._tag_to_groups = {} group_members: dict[str, list[int]] = {} for tag_name, groups in raw_t2g.items(): if tag_name not in name_to_idx: continue group_set = set(groups) self._tag_to_groups[tag_name] = group_set idx = name_to_idx[tag_name] for g in group_set: group_members.setdefault(g, []).append(idx) self._group_to_tags_idx = { g: np.array(idxs, dtype=np.int64) for g, idxs in group_members.items() } self._group_cn_names = data.get("group_cn_names", {}) print(f'[Engine] Tag Group loaded, {len(self._tag_to_groups)} tags, ' f'{len(self._group_to_tags_idx)} groups, ' f'{len(self._group_cn_names)} cn_names')