CS336-1-基础部分

class BytePairEncoder:
    """完整的BPE编码器实现"""

    @staticmethod
    def char_to_bytes(text):
        """将文本转换为UTF-8字节列表"""
        return list(text.encode('utf-8'))

    @staticmethod
    def bytes_to_text(byte_list):
        """将字节列表转换回文本"""
        return bytes(byte_list).decode('utf-8', errors='replace')

    def __init__(self, vocab_size=1000):
        """
        初始化BPE编码器

        Args:
            vocab_size: 目标词汇表大小（包括基础的256字节）
        """
        # 基础词汇表：256个字节
        self.base_vocab = {bytes([i]): i for i in range(256)}
        self.reverse_base = {i: bytes([i]) for i in range(256)}

        # BPE合并规则
        self.merges = []  # 存储(byte1, byte2)合并规则
        self.vocab = self.base_vocab.copy()  # 完整词汇表
        self.reverse_vocab = self.reverse_base.copy()

        self.vocab_size = vocab_size
        self.next_token_id = 256  # 新token从256开始

    def train(self, corpus, verbose=False):
        """
        在语料上训练BPE

        Args:
            corpus: 文本列表
            verbose: 是否显示训练过程
        """
        if verbose:
            print(f"开始BPE训练，目标词汇表大小: {self.vocab_size}")
            print(f"基础词汇表大小: 256")
            print(f"需要学习 {self.vocab_size - 256} 个合并规则")

        # 将语料转换为字节序列 - 使用类方法
        tokenized_corpus = []
        for text in corpus:
            byte_list = self.char_to_bytes(text)  # 改为self.char_to_bytes
            # 转换为字节对象列表
            tokens = [bytes([b]) for b in byte_list]
            tokenized_corpus.append(tokens)

        # BPE训练：迭代合并
        merge_count = 0
        while len(self.vocab) < self.vocab_size:
            # 统计所有相邻pair的频率
            pair_freq = {}

            for tokens in tokenized_corpus:
                for i in range(len(tokens) - 1):
                    pair = (tokens[i], tokens[i+1])
                    pair_freq[pair] = pair_freq.get(pair, 0) + 1

            if not pair_freq:
                if verbose:
                    print(f"没有更多可合并的pair，提前停止")
                break

            # 找到最常见的pair
            best_pair = max(pair_freq.items(), key=lambda x: x[1])[0]
            self.merges.append(best_pair)

            # 创建新token
            new_token = best_pair[0] + best_pair[1]
            token_id = self.next_token_id
            self.vocab[new_token] = token_id
            self.reverse_vocab[token_id] = new_token
            self.next_token_id += 1

            if verbose and merge_count % 10 == 0:
                print(f"合并 #{merge_count+1}: {best_pair} -> ID:{token_id}")
                print(f"  新token: {new_token} (长度:{len(new_token)}字节)")

            # 在语料中应用这个合并
            new_tokenized_corpus = []
            for tokens in tokenized_corpus:
                new_tokens = []
                i = 0
                while i < len(tokens):
                    if (i < len(tokens) - 1 and
                        tokens[i] == best_pair[0] and
                        tokens[i+1] == best_pair[1]):
                        # 合并这一对
                        new_tokens.append(new_token)
                        i += 2  # 跳过已合并的第二个元素
                    else:
                        new_tokens.append(tokens[i])
                        i += 1
                new_tokenized_corpus.append(new_tokens)

            tokenized_corpus = new_tokenized_corpus
            merge_count += 1

        if verbose:
            print(f"训练完成！共学习 {len(self.merges)} 个合并规则")
            print(f"最终词汇表大小: {len(self.vocab)}")

        return self.merges

    def encode(self, text):
        """
        使用训练好的BPE编码文本

        Args:
            text: 输入文本

        Returns:
            token_ids: token ID列表
            tokens: 对应的字节对象列表
        """
        # 1. 转换为字节列表 - 使用类方法
        byte_list = self.char_to_bytes(text)  # 改为self.char_to_bytes

        # 2. 初始化为单个字节的token
        tokens = [bytes([b]) for b in byte_list]

        # 3. 按顺序应用所有合并规则
        for pair in self.merges:
            new_tokens = []
            i = 0
            while i < len(tokens):
                if (i < len(tokens) - 1 and
                    tokens[i] == pair[0] and
                    tokens[i+1] == pair[1]):
                    # 合并这一对
                    new_token = pair[0] + pair[1]
                    new_tokens.append(new_token)
                    i += 2
                else:
                    new_tokens.append(tokens[i])
                    i += 1
            tokens = new_tokens

        # 4. 转换为token ID
        token_ids = [self.vocab[token] for token in tokens]

        return token_ids, tokens

    def decode(self, token_ids):
        """
        将token IDs解码回文本

        Args:
            token_ids: token ID列表

        Returns:
            text: 解码后的文本
        """
        # 1. 将IDs转换回字节
        byte_sequence = b''
        for token_id in token_ids:
            if token_id in self.reverse_vocab:
                byte_sequence += self.reverse_vocab[token_id]
            else:
                # 如果ID不在词汇表中，使用基础字节
                byte_sequence += bytes([token_id % 256])

        # 2. 解码为文本
        try:
            text = byte_sequence.decode('utf-8')
        except UnicodeDecodeError:
            # 如果有无效序列，使用错误处理
            text = byte_sequence.decode('utf-8', errors='replace')

        return text

    def get_vocab_size(self):
        """获取当前词汇表大小"""
        return len(self.vocab)

    def print_vocab_sample(self, n=20):
        """打印词汇表示例"""
        print("\n词汇表示例:")
        print("-" * 40)
        print(f"{'Token ID':<10} {'字节长度':<10} {'内容(可打印部分)'}")
        print("-" * 40)

        # 按ID排序
        sorted_items = sorted(self.reverse_vocab.items(), key=lambda x: x[0])

        for i, (token_id, token_bytes) in enumerate(sorted_items):
            if i >= n:
                print(f"... 还有 {len(self.vocab) - n} 个token")
                break

            # 尝试解码为可打印字符
            try:
                content = token_bytes.decode('utf-8')
                display = repr(content)
            except:
                display = repr(token_bytes)

            print(f"{token_id:<10} {len(token_bytes):<10} {display}")
        print("-" * 40)