# 引入 sentencepiece 库的 SentencePieceProcessor 模块，用于进行分词操作
from sentencepiece import SentencePieceProcessor
# 引入 logging 库的 getLogger 模块，用于生成日志
from logging import getLogger
# 引入 typing 库的 List 模块，用于注释函数参数或返回值的类型
from typing import List
# 引入 os 库，提供了大量与操作系统进行交互的接口
import os

# 创建一个日志记录器
logger = getLogger()

# 定义一个 Tokenizer 类
class Tokenizer:
    # 初始化函数，参数为 SentencePiece 模型的路径
    def __init__(self, model_path: str):
        # 判断指定的模型文件是否存在
        assert os.path.isfile(model_path), model_path
        # 加载 SentencePiece 模型
        self.sp_model = SentencePieceProcessor(model_file=model_path)
        # 记录日志，提示模型加载成功
        logger.info(f"Reloaded SentencePiece model from {model_path}")

        # 获取模型的词汇量、开始标记 ID、结束标记 ID、填充标记 ID
        self.n_words: int = self.sp_model.vocab_size()
        self.bos_id: int = self.sp_model.bos_id()
        self.eos_id: int = self.sp_model.eos_id()
        self.pad_id: int = self.sp_model.pad_id()
        # 记录日志，显示获取的信息
        logger.info(
            f"#words: {self.n_words} - BOS ID: {self.bos_id} - EOS ID: {self.eos_id}"
        )
        # 确保模型的词汇量与词片段大小一致
        assert self.sp_model.vocab_size() == self.sp_model.get_piece_size()

    # 编码函数，将输入的字符串编码为 token id 列表
    def encode(self, s: str, bos: bool, eos: bool) -> List[int]:
        # 检查输入的是否是字符串
        assert type(s) is str
        # 使用 SentencePiece 模型将字符串编码为 token id 列表
        t = self.sp_model.encode(s)
        # 如果需要在开头添加开始标记，就将开始标记 id 添加到列表的开头
        if bos:
            t = [self.bos_id] + t
        # 如果需要在结尾添加结束标记，就将结束标记 id 添加到列表的结尾
        if eos:
            t = t + [self.eos_id]
        # 返回 token id 列表
        return t

    # 解码函数，将 token id 列表解码为字符串
    def decode(self, t: List[int]) -> str:
        # 使用 SentencePiece 模型将 token id 列表解码为字符串
        return self.sp_model.decode(t)

**class** **RMSNorm**(torch**.**nn**.**Module):
    **def** __init__(self, dim: int, eps: float **=** 1e-6):
        super()**.**__init__()
        self**.**eps **=** eps
        self**.**weight **=** nn**.**Parameter(torch**.**ones(dim))

    **def** **_norm**(self, x):
            # torch.rsqrt是开平方并取倒数
        **return** x ***** torch**.**rsqrt(x**.**pow(2)**.**mean(**-**1, keepdim**=**True) **+** self**.**eps)

    **def** **forward**(self, x):
        output **=** self**.**_norm(x**.**float())**.**type_as(x)
        **return** output ***** self**.**weight

**class** **Attention**(nn**.**Module):
    **def** __init__(self, args: ModelArgs):
        super()**.**__init__()

        self**.**n_local_heads **=** args**.**n_heads **//** fs_init**.**get_model_parallel_world_size()
        self**.**head_dim **=** args**.**dim **//** args**.**n_heads

        self**.**wq **=** ColumnParallelLinear(
            args**.**dim,
            args**.**n_heads ***** self**.**head_dim,
            bias**=**False,
            gather_output**=**False,
            init_method**=lambda** x: x,
        )
        self**.**wk **=** ColumnParallelLinear(
            args**.**dim,
            args**.**n_heads ***** self**.**head_dim,
            bias**=**False,
            gather_output**=**False,
            init_method**=lambda** x: x,
        )
        self**.**wv **=** ColumnParallelLinear(
            args**.**dim,
            args**.**n_heads ***** self**.**head_dim,
            bias**=**False,
            gather_output**=**False,
            init_method**=lambda** x: x,
        )
        self**.**wo **=** RowParallelLinear(
            args**.**n_heads ***** self**.**head_dim,
            args**.**dim,
            bias**=**False,
            input_is_parallel**=**True,
            init_method**=lambda** x: x,
        )

        self**.**cache_k **=** torch**.**zeros(
            (args**.**max_batch_size, args**.**max_seq_len, self**.**n_local_heads, self**.**head_dim)
        )**.**cuda()
        self**.**cache_v **=** torch**.**zeros(
            (args**.**max_batch_size, args**.**max_seq_len, self**.**n_local_heads, self**.**head_dim)
        )**.**cuda()

    **def** **forward**(self, x: torch**.**Tensor, start_pos: int, freqs_cis: torch**.**Tensor, mask: Optional[torch**.**Tensor]):
        bsz, seqlen, _ **=** x**.**shape
        xq, xk, xv **=** self**.**wq(x), self**.**wk(x), self**.**wv(x)

        xq **=** xq**.**view(bsz, seqlen, self**.**n_local_heads, self**.**head_dim)
        xk **=** xk**.**view(bsz, seqlen, self**.**n_local_heads, self**.**head_dim)
        xv **=** xv**.**view(bsz, seqlen, self**.**n_local_heads, self**.**head_dim)

        xq, xk **=** apply_rotary_emb(xq, xk, freqs_cis**=**freqs_cis)

        self**.**cache_k **=** self**.**cache_k**.**to(xq)
        self**.**cache_v **=** self**.**cache_v**.**to(xq)

        self**.**cache_k[:bsz, start_pos : start_pos **+** seqlen] **=** xk
        self**.**cache_v[:bsz, start_pos : start_pos **+** seqlen] **=** xv

        keys **=** self**.**cache_k[:bsz, : start_pos **+** seqlen]
        values **=** self**.**cache_v[:bsz, : start_pos **+** seqlen]

        xq **=** xq**.**transpose(1, 2)
        keys **=** keys**.**transpose(1, 2)
        values **=** values**.**transpose(1, 2)
        scores **=** torch**.**matmul(xq, keys**.**transpose(2, 3)) **/** math**.**sqrt(self**.**head_dim)
        **if** mask **is** **not** None:
            scores **=** scores **+** mask  *# (bs, n_local_heads, slen, cache_len + slen)*scores **=** F**.**softmax(scores**.**float(), dim**=-**1)**.**type_as(xq)
        output **=** torch**.**matmul(scores, values)  *# (bs, n_local_heads, slen, head_dim)*output **=** output**.**transpose(
            1, 2
        )**.**contiguous()**.**view(bsz, seqlen, **-**1)

        **return** self**.**wo(output)

self.cache_k = torch.zeros((args.max_batch_size, args.max_seq_len, self.n_local_heads, self.head_dim)).cuda()
self.cache_v = torch.zeros((args.max_batch_size, args.max_seq_len, self.n_local_heads, self.head_dim)).cuda()

self.cache_k[:bsz, start_pos : start_pos + seqlen] = xk
self.cache_v[:bsz, start_pos : start_pos + seqlen] = xv
keys = self.cache_k[:bsz, : start_pos + seqlen]
values = self.cache_v[:bsz, : start_pos + seqlen]