本文收集封装了几种常用的中文分词器：

• Unicode，基于 Unicode 编码，合并高频编码值；
• Jieba，基于 Jieba 库的中文分词，创建词汇库；
• Pinyin，基于拼音库（带声调）；
• Bert，基于 bert-base-chinese 模型；
• GPT2，基于 GPT2 模型。

"""
@Title: 分词器
@Time: 2025/2/6
@Author: Michael Jie
"""

import operator
import os
from abc import ABC, abstractmethod
from typing import Union, Optional, Sequence

import jieba
from bidict import bidict
from transformers import BertTokenizer, GPT2Tokenizer

from ztools.utils import MjUtil, MjText


# 分词器
class Tokenizer(ABC):
    # 创建词汇库
    def create_vocab(self, text):
        """
        处理语料：分词，过滤停词，编码等，以创建词汇映射关系
        :param text: 语料
        :return: 词汇库 dict | bidict {词汇: idx}
        """
        ...

    # 扩充词汇库
    def add_token(self, token):
        """
        向词汇库中添加新的词汇映射关系
        :param token: 词汇
        :return: 编码 int
        """
        ...

    # text -> tokens -> ids
    @abstractmethod
    def text2token(self, text):
        """
        根据词汇库将文本转换成 token 序列，一般用于数据集预处理
        :param text: 文本
        :return: ids list[int] | tuple[int]
        """
        ...

    # ids -> tokens -> text
    @abstractmethod
    def token2text(self, ids):
        """
        text2token 函数的逆操作，根据词汇库将 token 序列转换成可视化文本
        :param ids: 编码
        :return: text str
        """
        ...


# unicode 编码，默认为 utf-8
class TokenizerUnicode(Tokenizer):
    def __init__(self, vocab: Optional[dict] = None):
        if vocab is None:  # 初始化
            vocab = {}
        self.vocab = bidict(vocab)  # 词汇 <-> token 映射

    # 创建、扩充词汇库
    def create_vocab(self, text: str, num: Optional[int] = 10, idx: Optional[int] = 256) -> None:
        unicode = TokenizerUnicode.text2unicode(text)  # unicode 编码值
        # 统计，根据出现频率倒序
        counts = TokenizerUnicode.stats(unicode)
        counts = dict(sorted(counts.items(), key=operator.itemgetter(1), reverse=True))
        # 取高频词汇
        for inx, pair in enumerate(counts):
            if inx >= num:  # 默认取前 10 个
                break
            self.vocab[pair] = idx  # 默认 0-255 以外的正整数
            idx += 1

    # text -> tokens -> ids
    def text2token(self, text: str) -> list[int]:
        unicode = TokenizerUnicode.text2unicode(text)  # unicode 编码值
        # 根据词汇库合并高频词汇
        for pair in self.vocab:
            if not set(pair).issubset(set(unicode)):
                continue
            unicode = TokenizerUnicode.merge(unicode, pair, self.vocab.get(pair))

        return unicode

    # ids -> tokens -> text
    def token2text(self, ids: list[int]) -> str:
        tokens = []
        for item in ids:
            if item in self.vocab.inverse:
                # 根据词汇库还原新增的 token
                tokens.extend(self.vocab.inverse.get(item))
            else:
                tokens.append(item)
        text = TokenizerUnicode.unicode2text(tokens)

        return text

    # 文本转 unicode
    @staticmethod
    def text2unicode(text: str,
                     encoding: Optional[str] = "utf-8",
                     to_num: Optional[bool] = True) -> Union[list, bytes]:
        """
        :param text: 文本
        :param encoding: 编码类型，默认为 utf-8
        :param to_num: 是否转换成整数，默认为 True
        :return: list | bytes
        :example: "欢迎" -> [230, 172, 162, 232, 191, 142] | b'\xe6\xac\xa2\xe8\xbf\x8e'
        """
        unicode = text.encode(encoding)
        if to_num:
            # 字节转整数
            unicode = list(unicode)

        return unicode

    # unicode 转文本
    @staticmethod
    def unicode2text(unicode: Union[list, bytes], encoding: Optional[str] = "utf-8") -> str:
        """
        :param unicode: unicode 值
        :param encoding: 编码类型，默认为 utf-8
        :return: str
        :example: [230, 172, 162, 232, 191, 142] | b'\xe6\xac\xa2\xe8\xbf\x8e' -> "欢迎"
        """
        text = b""
        if type(unicode) is list:
            # 整数转字节
            for item in unicode:
                text += item.to_bytes(length=1, byteorder="big")
        text = text.decode(encoding)

        return text

    # 统计
    @staticmethod
    def stats(ids: list[int], counts: Optional[dict] = None) -> dict:
        """
        :param ids: token 列表
        :param counts: token 对频率字典，默认为 None
        :return: dict
        :example: [1, 2, 3, 1, 2] -> {(1, 2): 2, (2, 3): 1, (3, 1): 1}
        """
        if counts is None:
            counts = {}
        for pair in zip(ids, ids[1:]):
            counts[pair] = counts.get(pair, 0) + 1

        return counts

    # 合并
    @staticmethod
    def merge(ids: list[int], pair: tuple[int, int], idx: int) -> list:
        """
        :param ids: token 列表
        :param pair: token 对
        :param idx: 新 token
        :return: list
        :example: ids=[1, 2, 3, 1, 2], pair=(1, 2), idx=4 -> [4, 3, 4]
        """
        new_ids = []
        i = 0
        while i < len(ids):
            # 替换
            if ids[i] == pair[0] and i + 1 < len(ids) and ids[i + 1] == pair[1]:
                new_ids.append(idx)
                i += 2
            else:
                new_ids.append(ids[i])
                i += 1

        return new_ids


# Jieba 分词
class TokenizerJieba(Tokenizer):
    def __init__(self, vocab: Optional[Union[list, dict]] = None):
        if vocab is None:  # 初始化
            vocab = {"<pad>": 0, "<bos>": 1, "<eos>": 2, "<sep>": 3, "<num>": 4, "<unk>": 5}
        if type(vocab) is list:
            vocab = {token: inx for inx, token in enumerate(vocab)}
        self.vocab = bidict(vocab)  # 词汇 <-> token 映射

    # 创建词汇库
    def create_vocab(self, text: str) -> None:
        # 分词并转换为 Token
        tokens = TokenizerJieba.cut(text)  # 词汇
        # 扩充词汇库
        for token in tokens:
            self.add_token(token)

    # 扩充词汇库
    def add_token(self, token: str) -> None:
        if token not in self.vocab:
            self.vocab[token] = len(self.vocab)

    # text -> tokens -> ids
    def text2token(self, text: str, add_end: Optional[bool] = False) -> list[int]:
        tokens = TokenizerJieba.cut(text)
        ids = []
        for token in tokens:
            if token in self.vocab:
                ids.append(self.vocab.get(token))
            else:  # 未知词汇
                ids.append(self.vocab.get("<unk>"))
        if add_end:  # 起止符
            ids.insert(0, self.vocab.get("<bos>"))
            ids.append(self.vocab.get("<eos>"))

        return ids

    # ids -> tokens -> text
    def token2text(self, ids: list[int]) -> str:
        tokens = [self.vocab.inverse.get(s) for s in ids]

        return " ".join(tokens)

    # 分词
    @staticmethod
    def cut(text: str, stop: Optional[Sequence] = None) -> list:
        """
        :param text: 中文预料
        :param stop: 停词
        :return: list
        :example: "欢迎光临！今天的“天气”真美好！" ->
        ['欢迎光临', '今天', '的', '天气', '真', '美好']
        """
        lst = jieba.lcut(text)  # 精确模式
        words = []  # 过滤后的词汇
        for s in lst:
            if MjUtil.is_none(s):  # 空字符
                continue
            elif stop is not None and s in stop:  # 停词
                continue
            elif MjUtil.is_or_not(s, "symbol"):  # 符号
                words.append("<sep>")
            elif MjUtil.is_number(s):  # 数字
                words.append("<num>")
            else:
                words.append(s)

        return words


# 拼音分词
class TokenizerPinyin(Tokenizer):
    def __init__(self, vocab: Optional[dict] = None):
        if vocab is None:  # 加载默认词汇库
            root = os.path.dirname(os.path.abspath(__file__))
            path = os.path.join(root, "pinyin", "vocab.txt")
            lst = MjText.load_text_file(path)
            vocab = {token: inx for inx, token in enumerate(lst)}
        self.vocab = bidict(vocab)  # 词汇 <-> token 映射

    # text -> tokens -> ids
    def text2token(self,
                   text: Union[str, list],
                   sep: Optional[str] = " ",
                   add_end: Optional[bool] = False) -> list[int]:
        if type(text) is str:
            text = text.split(sep)
        ids = []
        for s in text:
            if s in self.vocab:
                ids.append(self.vocab.get(s))
            else:  # 未知词汇
                ids.append(self.vocab.get("<unk>"))
        if add_end:  # 起止符
            ids.insert(0, self.vocab.get("<bos>"))
            ids.append(self.vocab.get("<eos>"))

        return ids

    # ids -> tokens -> text
    def token2text(self, ids: list[int]) -> str:
        text = [self.vocab.inverse.get(s) for s in ids]

        return " ".join(text)


# Bert 分词
class TokenizerBert(Tokenizer):
    def __init__(self, path: Optional[str] = None):
        if path is None:  # 默认包路径
            # config.json, pytorch_model.bin, vocab.json
            root = os.path.dirname(os.path.abspath(__file__))
            path = os.path.join(root, "bert-base-chinese/")
        # 加载 Bert 分词器
        self.tokenizer = BertTokenizer.from_pretrained(path)

    # text -> tokens -> ids
    def text2token(self, text: str, add_end: Optional[bool] = False) -> list[int]:
        # 分词并转换为 Token
        tokens = self.tokenizer.tokenize(text)
        if add_end:  # 起止符
            tokens.insert(0, "[CLS]")
            tokens.append("[SEP]")
        # 转换为 Token ID
        ids = self.tokenizer.convert_tokens_to_ids(tokens)

        return ids

    # ids -> tokens -> text
    def token2text(self, ids: list[int]) -> str:
        tokens = self.tokenizer.convert_ids_to_tokens(ids)
        text = self.tokenizer.convert_tokens_to_string(tokens)

        return text


# GPT2 分词
class TokenizerGPT2(Tokenizer):
    def __init__(self, path: Optional[str] = None):
        if path is None:  # 默认包路径
            # config.json, merges.txt, pytorch_model.bin, vocab.json
            root = os.path.dirname(os.path.abspath(__file__))
            path = os.path.join(root, "gpt2/")
        # 加载 GPT2 分词器
        self.tokenizer = GPT2Tokenizer.from_pretrained(path)

    # text -> tokens -> ids
    def text2token(self, text: str, add_end: Optional[bool] = False) -> list[int]:
        # 分词并转换为 Token
        tokens = self.tokenizer.tokenize(text)
        if add_end:  # 起止符
            tokens.insert(0, "<|startoftext|>")
            tokens.append("<|endoftext|>")
        # 转换为 Token ID
        ids = self.tokenizer.convert_tokens_to_ids(tokens)

        return ids

    # ids -> tokens -> text
    def token2text(self, ids: list[int]) -> str:
        tokens = self.tokenizer.convert_ids_to_tokens(ids)
        text = self.tokenizer.convert_tokens_to_string(tokens)

        return text


if __name__ == '__main__':
    t1 = TokenizerUnicode()
    print(t1.text2token("hi，吃饭了吗？"))
    print(t1.token2text([104, 105, 239, 188, 140, 229, 144, 131, 233, 165, 173, 228, 186, 134, 229, 144, 151, 239, 188, 159]))
    print("--------------------")

    t2 = TokenizerJieba()
    t2.create_vocab("hi，吃饭了吗？")
    print(t2.text2token("hi，吃饭了吗？"))
    print(t2.token2text([6, 3, 7, 8, 9, 3]))

    print("--------------------")

    t3 = TokenizerPinyin()
    print(t3.text2token(["chi1", "fan4", "le5", "ma5"]))
    print(t3.token2text([137, 296, 575, 652]))

    print("--------------------")

    t4 = TokenizerBert()
    print(t4.text2token("hi，吃饭了吗？", True))
    print(t4.token2text([8913, 8024, 1391, 7649, 749, 1408, 8043]))

    print("--------------------")

    t5 = TokenizerGPT2()
    print(t5.text2token("hi，吃饭了吗？"))
    print(t5.token2text([5303, 171, 120, 234, 28938, 225, 165, 98, 255, 12859, 228, 28938, 245, 171, 120, 253]))