knowledge_search/scripts/add_search_words.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
为关联特征生成检索词并去重

读取 associated_tags_results.json，为每个特征生成组合检索词，
并在同一结果项内去重。
"""

import json
import logging
from pathlib import Path
from typing import Dict, List, Any, Set
import argparse


# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger(__name__)


class SearchWordGenerator:
    """检索词生成器"""

    def __init__(self, input_path: str):
        """
        初始化生成器

        Args:
            input_path: 输入JSON文件路径
        """
        self.input_path = input_path
        self.data = self._load_json(input_path)
        self.stats = {
            '处理的结果项数': 0,
            '生成的总组合词数': 0,
            '唯一组合词数': 0,
            '重复过滤的词数': 0,
            '每项详情': []
        }

    def _load_json(self, file_path: str) -> List[Dict]:
        """加载JSON文件"""
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                return json.load(f)
        except Exception as e:
            logger.error(f"加载文件 {file_path} 失败: {e}")
            raise

    def generate_search_words(self) -> List[Dict[str, Any]]:
        """
        为所有结果项生成检索词

        Returns:
            增强后的数据列表
        """
        logger.info("=" * 60)
        logger.info("开始生成检索词")
        logger.info("=" * 60)

        enhanced_data = []

        for idx, result in enumerate(self.data, 1):
            logger.info(f"\n处理第 {idx}/{len(self.data)} 个结果项")

            # 获取基础词（人设特征名称）
            base_word = result.get('最高匹配信息', {}).get('人设特征名称', '')
            original_feature = result.get('原始特征名称', '')

            logger.info(f"  原始特征: {original_feature}")
            logger.info(f"  人设特征名称（基础词）: {base_word}")

            if not base_word:
                logger.warning(f"  警告：未找到人设特征名称，跳过")
                enhanced_data.append(result)
                continue

            # 用于去重的集合（在当前结果项范围内）
            seen_words: Set[str] = set()
            item_stats = {
                '原始特征': original_feature,
                '人设特征名称': base_word,
                '总特征数': 0,
                '唯一组合词数': 0,
                '重复词数': 0,
                '组合词列表': []
            }

            # 遍历所有关联
            associations = result.get('找到的关联', [])
            for assoc_idx, assoc in enumerate(associations):
                target_path = assoc.get('目标分类路径', '')
                features = assoc.get('特征列表', [])

                logger.info(f"  处理关联 {assoc_idx + 1}/{len(associations)}: {target_path}")
                logger.info(f"    特征数: {len(features)}")

                # 遍历特征列表
                for feature in features:
                    feature_name = feature.get('特征名称', '')
                    item_stats['总特征数'] += 1

                    if not feature_name:
                        feature['search_word'] = None
                        continue

                    # 生成组合词
                    search_word = f"{base_word} {feature_name}"

                    # 检查是否重复
                    if search_word not in seen_words:
                        # 首次出现，填充
                        feature['search_word'] = search_word
                        seen_words.add(search_word)
                        item_stats['唯一组合词数'] += 1
                        item_stats['组合词列表'].append(search_word)
                        logger.debug(f"      + 新增: {search_word}")
                    else:
                        # 重复，留空
                        feature['search_word'] = None
                        item_stats['重复词数'] += 1
                        logger.debug(f"      - 重复（留空）: {search_word}")

            # 记录统计
            logger.info(f"  完成：总特征 {item_stats['总特征数']} 个，"
                       f"唯一组合词 {item_stats['唯一组合词数']} 个，"
                       f"重复 {item_stats['重复词数']} 个")

            self.stats['处理的结果项数'] += 1
            self.stats['生成的总组合词数'] += item_stats['总特征数']
            self.stats['唯一组合词数'] += item_stats['唯一组合词数']
            self.stats['重复过滤的词数'] += item_stats['重复词数']
            self.stats['每项详情'].append(item_stats)

            enhanced_data.append(result)

        logger.info("\n" + "=" * 60)
        logger.info("生成完成")
        logger.info("=" * 60)
        logger.info(f"处理的结果项数: {self.stats['处理的结果项数']}")
        logger.info(f"生成的总组合词数: {self.stats['生成的总组合词数']}")
        logger.info(f"唯一组合词数: {self.stats['唯一组合词数']}")
        logger.info(f"重复过滤的词数: {self.stats['重复过滤的词数']}")

        return enhanced_data

    def save_results(self, enhanced_data: List[Dict[str, Any]], output_path: str):
        """保存增强后的数据"""
        try:
            with open(output_path, 'w', encoding='utf-8') as f:
                json.dump(enhanced_data, f, ensure_ascii=False, indent=2)
            logger.info(f"增强数据已保存到: {output_path}")
        except Exception as e:
            logger.error(f"保存结果失败: {e}")
            raise

    def save_stats(self, stats_path: str):
        """保存统计信息"""
        try:
            with open(stats_path, 'w', encoding='utf-8') as f:
                json.dump(self.stats, f, ensure_ascii=False, indent=2)
            logger.info(f"统计信息已保存到: {stats_path}")
        except Exception as e:
            logger.error(f"保存统计信息失败: {e}")
            raise


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='为关联特征生成检索词并去重')
    parser.add_argument(
        '--input',
        default='associated_tags_results.json',
        help='输入JSON文件路径（默认: associated_tags_results.json）'
    )
    parser.add_argument(
        '--output',
        default='associated_tags_results_with_search.json',
        help='输出JSON文件路径（默认: associated_tags_results_with_search.json）'
    )
    parser.add_argument(
        '--stats',
        default='search_words_stats.json',
        help='统计信息输出路径（默认: search_words_stats.json）'
    )
    parser.add_argument(
        '--debug',
        action='store_true',
        help='启用调试日志'
    )

    args = parser.parse_args()

    # 设置日志级别
    if args.debug:
        logger.setLevel(logging.DEBUG)

    # 创建生成器
    generator = SearchWordGenerator(input_path=args.input)

    # 生成检索词
    enhanced_data = generator.generate_search_words()

    # 保存结果
    generator.save_results(enhanced_data, args.output)
    generator.save_stats(args.stats)

    # 输出汇总
    logger.info("\n" + "=" * 60)
    logger.info("处理完成汇总")
    logger.info("=" * 60)
    logger.info(f"输入文件: {args.input}")
    logger.info(f"输出文件: {args.output}")
    logger.info(f"统计文件: {args.stats}")
    logger.info(f"")
    logger.info(f"处理结果:")
    logger.info(f"  - 结果项数: {generator.stats['处理的结果项数']}")
    logger.info(f"  - 总特征数: {generator.stats['生成的总组合词数']}")
    logger.info(f"  - 唯一组合词: {generator.stats['唯一组合词数']}")
    logger.info(f"  - 重复过滤: {generator.stats['重复过滤的词数']}")
    logger.info(f"  - 去重率: {generator.stats['重复过滤的词数'] / max(generator.stats['生成的总组合词数'], 1) * 100:.1f}%")


if __name__ == '__main__':
    main()