py_xiaozhi/scripts/keyword_generator.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""唤醒词自动生成工具.

功能：
1. 输入中文自动转换为带声调拼音
2. 按字母分隔拼音（声母+韵母）
3. 验证token是否在tokens.txt中
4. 自动生成keywords.txt格式
"""

import sys
from pathlib import Path

try:
    from pypinyin import Style, lazy_pinyin
except ImportError:
    print("❌ 缺少依赖: pypinyin")
    print("请安装: pip install pypinyin")
    sys.exit(1)


class KeywordGenerator:
    def __init__(self, model_dir: Path):
        """初始化唤醒词生成器.

        Args:
            model_dir: 模型目录路径（包含tokens.txt和keywords.txt）
        """
        self.model_dir = Path(model_dir)
        self.tokens_file = self.model_dir / "tokens.txt"
        self.keywords_file = self.model_dir / "keywords.txt"

        # 加载已有的tokens
        self.available_tokens = self._load_tokens()

        # 声母表（需要分离的）
        self.initials = [
            "b",
            "p",
            "m",
            "f",
            "d",
            "t",
            "n",
            "l",
            "g",
            "k",
            "h",
            "j",
            "q",
            "x",
            "zh",
            "ch",
            "sh",
            "r",
            "z",
            "c",
            "s",
            "y",
            "w",
        ]

    def _load_tokens(self) -> set:
        """
        加载tokens.txt中的所有可用token.
        """
        if not self.tokens_file.exists():
            print(f"⚠️  警告: tokens文件不存在: {self.tokens_file}")
            return set()

        tokens = set()
        with open(self.tokens_file, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if line and not line.startswith("#"):
                    # 格式: "token id" 或 "token"
                    parts = line.split()
                    if parts:
                        tokens.add(parts[0])

        print(f"✅ 加载了 {len(tokens)} 个可用tokens")
        return tokens

    def _split_pinyin(self, pinyin: str) -> list:
        """将拼音按声母韵母分隔.

        例如: "xiǎo" -> ["x", "iǎo"]       "mǐ" -> ["m", "ǐ"]       "ài" -> ["ài"]  (零声母)
        """
        if not pinyin:
            return []

        # 按长度优先尝试匹配声母（zh, ch, sh优先）
        for initial in sorted(self.initials, key=len, reverse=True):
            if pinyin.startswith(initial):
                final = pinyin[len(initial) :]
                if final:
                    return [initial, final]
                else:
                    return [initial]

        # 没有声母（零声母）
        return [pinyin]

    def chinese_to_keyword_format(self, chinese_text: str) -> str:
        """将中文转换为keyword格式.

        Args:
            chinese_text: 中文文本，如"小米小米"

        Returns:
            keyword格式，如"x iǎo m ǐ x iǎo m ǐ @小米小米"
        """
        # 转换为带声调拼音
        pinyin_list = lazy_pinyin(chinese_text, style=Style.TONE)

        # 分割每个拼音
        split_parts = []
        missing_tokens = []

        for pinyin in pinyin_list:
            parts = self._split_pinyin(pinyin)

            # 验证每个part是否在tokens中
            for part in parts:
                if part not in self.available_tokens:
                    missing_tokens.append(part)
                split_parts.append(part)

        # 拼接结果
        pinyin_str = " ".join(split_parts)
        keyword_line = f"{pinyin_str} @{chinese_text}"

        # 如果有缺失的token，给出警告
        if missing_tokens:
            print(
                f"⚠️  警告: 以下token不在tokens.txt中: {', '.join(set(missing_tokens))}"
            )
            print(f"   生成的关键词可能无法正常工作")

        return keyword_line

    def add_keyword(self, chinese_text: str, append: bool = True) -> bool:
        """添加唤醒词到keywords.txt.

        Args:
            chinese_text: 中文唤醒词
            append: 是否追加（True）或覆盖（False）

        Returns:
            是否成功
        """
        try:
            # 生成keyword格式
            keyword_line = self.chinese_to_keyword_format(chinese_text)

            # 检查是否已存在
            if self.keywords_file.exists():
                with open(self.keywords_file, "r", encoding="utf-8") as f:
                    content = f.read()
                    if f"@{chinese_text}" in content:
                        print(f"⚠️  关键词 '{chinese_text}' 已存在")
                        return False

            # 写入文件
            mode = "a" if append else "w"
            with open(self.keywords_file, mode, encoding="utf-8") as f:
                f.write(keyword_line + "\n")

            print(f"✅ 成功添加: {keyword_line}")
            return True

        except Exception as e:
            print(f"❌ 添加失败: {e}")
            return False

    def batch_add_keywords(self, chinese_texts: list, overwrite: bool = False):
        """批量添加唤醒词.

        Args:
            chinese_texts: 中文列表
            overwrite: 是否覆盖原文件
        """
        if overwrite:
            print("⚠️  将覆盖现有keywords.txt")

        success_count = 0
        for text in chinese_texts:
            text = text.strip()
            if not text:
                continue

            if self.add_keyword(text, append=not overwrite):
                success_count += 1

            # 第一个后都追加
            overwrite = False

        print(f"\n📊 完成: 成功添加 {success_count}/{len(chinese_texts)} 个关键词")

    def list_keywords(self):
        """
        列出当前所有关键词.
        """
        if not self.keywords_file.exists():
            print("⚠️  keywords.txt 不存在")
            return

        print(f"\n📄 当前关键词列表 ({self.keywords_file}):")
        print("-" * 60)

        with open(self.keywords_file, "r", encoding="utf-8") as f:
            for i, line in enumerate(f, 1):
                line = line.strip()
                if line and not line.startswith("#"):
                    # 提取中文部分显示
                    if "@" in line:
                        pinyin_part, chinese_part = line.split("@", 1)
                        print(
                            f"{i}. {chinese_part.strip():15s} -> {pinyin_part.strip()}"
                        )
                    else:
                        print(f"{i}. {line}")

        print("-" * 60)


def main():
    """
    主函数.
    """
    import argparse

    parser = argparse.ArgumentParser(
        description="唤醒词自动生成工具",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例:
  # 添加单个关键词
  python keyword_generator.py -a "小米小米"

  # 批量添加关键词
  python keyword_generator.py -b "小米小米" "你好小智" "贾维斯"

  # 从文件批量导入（每行一个中文）
  python keyword_generator.py -f keywords_input.txt

  # 列出当前关键词
  python keyword_generator.py -l

  # 测试转换（不写入文件）
  python keyword_generator.py -t "小米小米"
        """,
    )

    parser.add_argument(
        "-m", "--model-dir", default="models", help="模型目录路径（默认: models）"
    )

    parser.add_argument("-a", "--add", help="添加单个关键词（中文）")

    parser.add_argument(
        "-b", "--batch", nargs="+", help="批量添加关键词（多个中文，空格分隔）"
    )

    parser.add_argument("-f", "--file", help="从文件批量导入（每行一个中文）")

    parser.add_argument("-l", "--list", action="store_true", help="列出当前所有关键词")

    parser.add_argument("-t", "--test", help="测试转换（不写入文件）")

    parser.add_argument(
        "--overwrite", action="store_true", help="覆盖模式（清空现有关键词）"
    )

    args = parser.parse_args()

    # 确定模型目录
    if Path(args.model_dir).is_absolute():
        model_dir = Path(args.model_dir)
    else:
        # 相对路径：相对于项目根目录
        script_dir = Path(__file__).parent
        project_root = script_dir.parent
        model_dir = project_root / args.model_dir

    if not model_dir.exists():
        print(f"❌ 模型目录不存在: {model_dir}")
        sys.exit(1)

    print(f"🔧 使用模型目录: {model_dir}")

    # 创建生成器
    generator = KeywordGenerator(model_dir)

    # 执行操作
    if args.test:
        # 测试模式
        print(f"\n🧪 测试转换:")
        keyword_line = generator.chinese_to_keyword_format(args.test)
        print(f"   输入: {args.test}")
        print(f"   输出: {keyword_line}")

    elif args.add:
        # 添加单个
        generator.add_keyword(args.add)

    elif args.batch:
        # 批量添加
        generator.batch_add_keywords(args.batch, overwrite=args.overwrite)

    elif args.file:
        # 从文件导入
        input_file = Path(args.file)
        if not input_file.exists():
            print(f"❌ 文件不存在: {input_file}")
            sys.exit(1)

        with open(input_file, "r", encoding="utf-8") as f:
            keywords = [line.strip() for line in f if line.strip()]

        print(f"📥 从文件导入 {len(keywords)} 个关键词")
        generator.batch_add_keywords(keywords, overwrite=args.overwrite)

    elif args.list:
        # 列出关键词
        generator.list_keywords()

    else:
        # 交互模式
        print("\n🎤 唤醒词生成工具（交互模式）")
        print("输入中文唤醒词，按 Ctrl+C 或输入 'q' 退出\n")

        try:
            while True:
                chinese = input("请输入中文唤醒词: ").strip()

                if not chinese or chinese.lower() == "q":
                    break

                generator.add_keyword(chinese)
                print()

        except KeyboardInterrupt:
            print("\n\n👋 已退出")

    # 最后列出所有关键词
    if not args.list and (args.add or args.batch or args.file):
        generator.list_keywords()


if __name__ == "__main__":
    main()