修复多用户kv_cacke共享的Bug,优化交互逻辑,新增垃圾回收

Encoder.py

@@ -1,76 +1,76 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from Affine import Affine
-
-#借来一用，简单改改
-class Qwen2RMSNorm(nn.Module):
-    def __init__(self, embedding_dim, eps=1e-6):
-        """
-        Qwen2RMSNorm is equivalent to T5LayerNorm
-        """
-        super().__init__()
-        self.weight = nn.Parameter(torch.ones(embedding_dim))
-        self.variance_epsilon = eps
-
-    def forward(self, hidden_states):
-        # input_dtype = hidden_states.dtype
-        # hidden_states = hidden_states.to(torch.float32)
-        variance = hidden_states.pow(2).mean(-1, keepdim=True)
-        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
-        return self.weight * hidden_states#.to(input_dtype)
-
-    def extra_repr(self):
-        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
-        
-#针对每个词嵌入的前馈网络
-class PositionWiseFeedForward(nn.Module):
-    def __init__(self,embedding_dim,feed_forward_dim,enable_affine):
-        super(PositionWiseFeedForward, self).__init__()
-        self.w1 = nn.Linear(embedding_dim, feed_forward_dim, bias=False)
-        self.w2 = nn.Linear(feed_forward_dim, embedding_dim, bias=False)
-        self.enable_affine = enable_affine
-        if enable_affine:
-            self.a1 = Affine(1.0)
-            self.a2 = Affine(1.0)
-        
-    def forward(self, x):
-        if self.enable_affine:
-            x = F.relu(self.w1(self.a1(x)))
-            return F.relu(self.w2(self.a2(x)))
-        else:
-            x = F.relu(self.w1(x))
-            return F.relu(self.w2(x))
-
-#编码器层
-class EncoderLayer(nn.Module):
-    def __init__(self,multi_head_attention,mask_future,position_wise_feed_forward,enable_layer_norm,dropout_rate):
-        super(EncoderLayer,self).__init__()
-        self.multi_head_attention = multi_head_attention
-        self.position_wise_feed_forward = position_wise_feed_forward
-        self.mask_future = mask_future
-        if enable_layer_norm == True:
-            self.layer_norm = Qwen2RMSNorm(multi_head_attention.embedding_dim)
-        else:
-            self.layer_norm = None
-
-        self.dropout_layer = nn.Dropout(p=dropout_rate)
-
-    def forward(self,query,q_mask):
-        #绝对不能用+=，那是原地修改，没法算梯度
-        query = query + self.dropout_layer(self.multi_head_attention(query,q_mask,query,self.mask_future))
-        query = query + self.dropout_layer(self.position_wise_feed_forward(query))
-        if self.layer_norm is not None:
-            query = self.layer_norm(query)
-        return query
-
-#编码器
-class Encoder(nn.Module):
-    def __init__(self, encoder_layers):
-        super(Encoder, self).__init__()
-        self.encoder_layers = encoder_layers
-        
-    def forward(self, query, q_mask):
-        for encoder_layer in self.encoder_layers:
-            query = encoder_layer(query,q_mask)
-        return query
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from Affine import Affine
+
+#借来一用，简单改改
+class Qwen2RMSNorm(nn.Module):
+    def __init__(self, embedding_dim, eps=1e-6):
+        """
+        Qwen2RMSNorm is equivalent to T5LayerNorm
+        """
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(embedding_dim))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        # input_dtype = hidden_states.dtype
+        # hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states#.to(input_dtype)
+
+    def extra_repr(self):
+        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+        
+#针对每个词嵌入的前馈网络
+class PositionWiseFeedForward(nn.Module):
+    def __init__(self,embedding_dim,feed_forward_dim,enable_affine):
+        super(PositionWiseFeedForward, self).__init__()
+        self.w1 = nn.Linear(embedding_dim, feed_forward_dim, bias=False)
+        self.w2 = nn.Linear(feed_forward_dim, embedding_dim, bias=False)
+        self.enable_affine = enable_affine
+        if enable_affine:
+            self.a1 = Affine(1.0)
+            self.a2 = Affine(1.0)
+        
+    def forward(self, x):
+        if self.enable_affine:
+            x = F.relu(self.w1(self.a1(x)))
+            return F.relu(self.w2(self.a2(x)))
+        else:
+            x = F.relu(self.w1(x))
+            return F.relu(self.w2(x))
+
+#编码器层
+class EncoderLayer(nn.Module):
+    def __init__(self,multi_head_attention,mask_future,position_wise_feed_forward,enable_layer_norm,dropout_rate):
+        super(EncoderLayer,self).__init__()
+        self.multi_head_attention = multi_head_attention
+        self.position_wise_feed_forward = position_wise_feed_forward
+        self.mask_future = mask_future
+        if enable_layer_norm == True:
+            self.layer_norm = Qwen2RMSNorm(multi_head_attention.embedding_dim)
+        else:
+            self.layer_norm = None
+
+        self.dropout_layer = nn.Dropout(p=dropout_rate)
+
+    def forward(self,query,q_mask,session_id):
+        #绝对不能用+=，那是原地修改，没法算梯度
+        query = query + self.dropout_layer(self.multi_head_attention(query,q_mask,query,self.mask_future,session_id))
+        query = query + self.dropout_layer(self.position_wise_feed_forward(query))
+        if self.layer_norm is not None:
+            query = self.layer_norm(query)
+        return query
+
+#编码器
+class Encoder(nn.Module):
+    def __init__(self, encoder_layers):
+        super(Encoder, self).__init__()
+        self.encoder_layers = encoder_layers
+        
+    def forward(self, query, q_mask,session_id):
+        for encoder_layer in self.encoder_layers:
+            query = encoder_layer(query,q_mask,session_id)
+        return query

LazyCache.py

@@ -0,0 +1,93 @@
+import time
+import threading
+from collections import defaultdict
+
+class ExpiringDict(dict):
+    """带过期时间的字典"""
+    def __init__(self, ttl=600, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.ttl = ttl  # 秒
+        self._timestamps = {}
+        self._lock = threading.Lock()
+
+    def __setitem__(self, key, value):
+        with self._lock:
+            super().__setitem__(key, value)
+            self._timestamps[key] = time.time()
+
+    def __getitem__(self, key):
+        with self._lock:
+            if key in self._timestamps and (time.time() - self._timestamps[key] > self.ttl):
+                super().__delitem__(key)
+                del self._timestamps[key]
+                raise KeyError(f"{key} 已过期")
+            # 访问时更新活跃时间
+            self._timestamps[key] = time.time()
+            return super().__getitem__(key)
+
+    def get(self, key, default=None):
+        try:
+            return self.__getitem__(key)
+        except KeyError:
+            return default
+
+    def cleanup(self):
+        with self._lock:
+            now = time.time()
+            expired = [k for k, t in self._timestamps.items() if now - t > self.ttl]
+            for k in expired:
+                super().__delitem__(k)
+                del self._timestamps[k]
+
+    def start_auto_cleanup(self, interval=1):
+        def loop():
+            while True:
+                time.sleep(interval)
+                self.cleanup()
+        threading.Thread(target=loop, daemon=True).start()
+
+
+class ExpiringDefaultDict(defaultdict):
+    """带过期时间的 defaultdict"""
+    def __init__(self, default_factory=None, ttl=600, *args, **kwargs):
+        super().__init__(default_factory, *args, **kwargs)
+        self.ttl = ttl
+        self._timestamps = {}
+        self._lock = threading.Lock()
+
+    def __setitem__(self, key, value):
+        with self._lock:
+            super().__setitem__(key, value)
+            self._timestamps[key] = time.time()
+
+    def __getitem__(self, key):
+        with self._lock:
+            if key in self._timestamps and (time.time() - self._timestamps[key] > self.ttl):
+                super().__delitem__(key)
+                del self._timestamps[key]
+                raise KeyError(f"{key} 已过期")
+            # 如果 key 不存在，则会调用 default_factory
+            val = super().__getitem__(key)
+            self._timestamps[key] = time.time()
+            return val
+
+    def get(self, key, default=None):
+        try:
+            return self.__getitem__(key)
+        except KeyError:
+            return default
+
+    def cleanup(self):
+        with self._lock:
+            now = time.time()
+            expired = [k for k, t in self._timestamps.items() if now - t > self.ttl]
+            for k in expired:
+                super().__delitem__(k)
+                del self._timestamps[k]
+
+    def start_auto_cleanup(self, interval=1):
+        def loop():
+            while True:
+                time.sleep(interval)
+                self.cleanup()
+        threading.Thread(target=loop, daemon=True).start()

MultiHeadAttention.py

@@ -1,397 +1,406 @@
-import math
-import numpy as np
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from Affine import Affine
-
-#获取相对位置矩阵
-def get_relative_mat(height,width,k=0):
-    posi_i = np.arange(k,height+k) #列的范围
-    posi_j = np.arange(0,width) #行的范围
-    posi_grid = np.meshgrid(posi_i, posi_j, indexing='ij')
-    return abs(posi_grid[0]-posi_grid[1])
-    
-#用于添加绝对位置信息的掩码
-def get_relative_dist(i,j,block_size,i_end,j_end):
-    if block_size == 0:
-        assert i==0 and j==0 ,"i!=0 or j!=0"
-        return get_relative_mat(i_end,j_end,k=0)
-    #i,j:当前分块的起始位置
-    #block_size:分块大小
-    #i_end,j_end:序列的长度
-    height = block_size #高度，也就是第一个序列中截取的长度，与分块大小相等
-    width  = block_size * 3 #宽度，也就是第二个序列中截取的长度，为了更长的上下文，还需要考虑上一个分块和下一个分块
-    #创建用来遮挡未来信息的标准掩码
-    #i越大，可见的部分越多，j相反，+block_size是因为上一个分块可见。
-    rela_dist = get_relative_mat(height,width,k=block_size+i-j)
-    #边界超出处理
-    #下超出
-    down_out = max(0,i+height-i_end)
-    #左超出
-    left_out = max(0,block_size-j)
-    #右超出
-    right_out = max(0,j+block_size*2-j_end)
-    #边界内截取
-    rela_dist = rela_dist[:height-down_out,left_out:width-right_out]
-    return rela_dist.astype(np.float32)
-
-#用于添加绝对位置信息的掩码
-def get_absolute_mask(i,j,block_size,i_end,j_end):
-    if block_size == 0:
-        assert i==0 and j==0 ,"i!=0 or j!=0"
-        return np.triu(np.ones((i_end,j_end),dtype='bool'), k=0)
-    #i,j:当前分块的起始位置
-    #block_size:分块大小
-    #i_end,j_end:序列的长度
-    height = block_size #高度，也就是第一个序列中截取的长度，与分块大小相等
-    width  = block_size * 3 #宽度，也就是第二个序列中截取的长度，为了更长的上下文，还需要考虑上一个分块和下一个分块
-    #创建用来遮挡未来信息的标准掩码
-    #i越大，可见的部分越多，j相反，+block_size是因为上一个分块可见。
-    abs_mask = np.triu(np.ones((height,width),dtype='bool'), k=block_size+i-j)
-    #边界超出处理
-    #下超出
-    down_out = max(0,i+height-i_end)
-    #左超出
-    left_out = max(0,block_size-j)
-    #右超出
-    right_out = max(0,j+block_size*2-j_end)
-    #边界内截取
-    abs_mask = abs_mask[:height-down_out,left_out:width-right_out]
-    return abs_mask
-
-#用于遮挡未来信息的标准掩码
-def get_std_mask(i,j,block_size,i_end,j_end):
-    if block_size == 0:
-        assert i==0 and j==0 ,"i!=0 or j!=0"
-        return np.triu(np.ones((i_end,j_end),dtype='bool'), k=1) == False
-    #i,j:当前分块的起始位置
-    #block_size:分块大小
-    #i_end,j_end:序列的长度
-    height = block_size #高度，也就是第一个序列中截取的长度，与分块大小相等
-    width  = block_size * 3 #宽度，也就是第二个序列中截取的长度，为了更长的上下文，还需要考虑上一个分块和下一个分块
-    #创建用来遮挡未来信息的标准掩码
-    #i越大，可见的部分越多，j相反，+block_size是因为上一个分块可见。
-    std_mask = np.triu(np.ones((height,width),dtype='bool'), k=1+block_size+i-j)
-    #边界超出处理
-    #下超出
-    down_out = max(0,i+height-i_end)
-    #左超出
-    left_out = max(0,block_size-j)
-    #右超出
-    right_out = max(0,j+block_size*2-j_end)
-    #边界内截取
-    std_mask = std_mask[:height-down_out,left_out:width-right_out]
-    return std_mask == False
-
-#标记一个需要多次使用的tensor
-def ident(p_list):
-    i,j,block_size,i_end,j_end = p_list[1:]
-    ret = [p_list[0]]
-    if p_list[0]=='r' or p_list[0]=='a':
-        if block_size == 0:
-            ret += [i_end,j_end,0]
-        else:
-            height = block_size
-            width  = block_size * 3
-            ret += [height,width,block_size+i-j]
-            down_out = max(0,i+height-i_end)
-            left_out = max(0,block_size-j)
-            right_out = max(0,j+block_size*2-j_end)
-            ret += [height-down_out,left_out,width-right_out]
-    else:
-        if block_size == 0:
-            ret += [i_end,j_end,1]
-        else:
-            height = block_size
-            width  = block_size * 3
-            ret += [height,width,1+block_size+i-j]
-            down_out = max(0,i+height-i_end)
-            left_out = max(0,block_size-j)
-            right_out = max(0,j+block_size*2-j_end)
-            ret += [height-down_out,left_out,width-right_out]
-    return str(ret)
-
-#缓存字典与定时器
-reg_dict = dict()
-reg_timer = dict()
-
-#查看是否未注册
-def un_reg(p):
-    return not p in reg_dict
-
-#注册需要重复使用的tensor
-def reg(p,v):
-    #找缓冲中用的最少的
-    keys = [k for k in reg_dict]
-    time_min = 0
-    if len(keys) != 0:
-        key_min = keys[0]
-        time_min = reg_timer[key_min]
-        for k in keys:
-            if reg_timer[k]<time_min:
-                key_min = k
-                time_min = reg_timer[key_min]
-    #计数
-    if not p in reg_timer:
-        reg_timer[p] = 1
-    else:
-        reg_timer[p] += 1
-    #缓冲满了就删掉最少用的
-    if len(keys) > 12:
-        del reg_dict[key_min]
-    #比最小的值大就保留
-    if reg_timer[p] > time_min or len(keys) < 12:
-        reg_dict[p] = v
-
-#从缓冲区中获取可重复使用的张量
-def get_reg(p):
-    reg_timer[p] += 1
-    return reg_dict[p]
-
-
-#多头注意力
-class MultiHeadAttention(nn.Module):
-    def __init__(self,embedding_dim,key_dim,head_number,position_information_type,enable_affine,enable_talking_head, \
-                 self_attention_block_size,dropout_rate,enable_el_cache):
-        super(MultiHeadAttention, self).__init__()
-        self.embedding_dim              = embedding_dim
-        self.key_dim                    = key_dim
-        self.head_number                = head_number
-        self.position_information_type  = position_information_type
-        self.enable_talking_head        = enable_talking_head
-        self.self_attention_block_size  = self_attention_block_size
-        self.dropout_layer              = nn.Dropout(p=dropout_rate)
-        self.enable_affine              = enable_affine
-
-        self.query_w = nn.Linear(embedding_dim,key_dim*head_number,bias=False)
-        self.key_w   = nn.Linear(embedding_dim,key_dim*head_number,bias=False)
-        self.value_w = nn.Linear(embedding_dim,key_dim*head_number,bias=False)
-        self.out_w   = nn.Linear(key_dim*head_number,embedding_dim,bias=False)
-
-        self.enable_el_cache = enable_el_cache
-        self.kv_cache = None
-        self.temp = None
-        self.cnt  = 0
-
-        if enable_affine == True:
-            self.query_a = Affine(1.0)
-            self.key_a   = Affine(1.0)
-            self.value_a = Affine(1.0)
-            self.out_a   = Affine(1.0)
-
-        if enable_talking_head == True:
-            self.talking_before_softmax = nn.Linear(head_number,head_number,bias=False)
-            self.talking_after_softmax  = nn.Linear(head_number,head_number,bias=False)
-        else:
-            self.talking_before_softmax = None
-            self.talking_after_softmax  = None
-
-        if position_information_type == "mask":
-            self.absolute_affine = Affine(1.0,grad_factor=1.0)
-            self.relative_affine = Affine(0.1,grad_factor=1.0)
-        else:
-            self.absolute_affine = None
-            self.relative_affine = None
-
-    #注意力运算
-    def attention(self, query, q_mask, key_value, mask_future):
-        #为了使用EL-Attention需要修改参数传递方式
-        absolute_affine = self.absolute_affine
-        relative_affine = self.relative_affine
-        talking_before_softmax = self.talking_before_softmax
-        talking_after_softmax = self.talking_after_softmax
-        block_size = self.self_attention_block_size
-        #提前调整q_mask的形状，方便广播
-        #query:[batch,head,query_len,emb_dim]
-        #q_mask:[batch,query_len]
-        #q_mask:[batch,query_len]->[batch,1,query_len]
-        #q_mask:[batch,1,query_len]->[batch,head,query_len]
-        q_mask = q_mask.unsqueeze(1).expand(*(query.size()[:-1]))
-        #判断是否需要分块运算
-        if block_size == 0:
-            #不进行分块
-            #计算scores
-            scores = torch.matmul(query,key_value.transpose(-1,-2))
-            if self.enable_affine == True:
-                scores = scores+self.temp
-            scores = scores/math.sqrt(self.key_dim)
-            #尝试添加相对位置信息
-            if relative_affine is not None:
-                if self.enable_el_cache and query.size(-2) == 1:
-                    p = ident(['er',0,0,0,query.size(-2),key_value.size(-2)])
-                    if un_reg(p):
-                        rela_dist = np.arange(self.cnt,-1,-1).reshape(1,-1)
-                        rela_dist = torch.from_numpy(rela_dist).detach().to(query.device)
-                        reg(p,rela_dist)
-                    else:
-                        rela_dist = get_reg(p)
-                else:
-                    p = ident(['r',0,0,0,query.size(-2),key_value.size(-2)])
-                    if un_reg(p):
-                        rela_dist = get_relative_dist(0,0,0,query.size(-2),key_value.size(-2))
-                        #直接广播更高效
-                        rela_dist = torch.from_numpy(rela_dist).detach().to(query.device)
-                        reg(p,rela_dist)
-                    else:
-                        rela_dist = get_reg(p)
-                dist_decay= rela_dist.mul(relative_affine(1.0)).add(1.0).reciprocal()
-                scores    = scores.mul(dist_decay)
-            #尝试添加绝对位置信息
-            if absolute_affine is not None:
-                if self.enable_el_cache and query.size(-2) == 1:
-                    p = ident(['ea',0,0,0,query.size(-2),key_value.size(-2)])
-                    if un_reg(p):
-                        abs_mask = np.array([[False]*(self.cnt)+[True]])
-                        abs_mask = torch.from_numpy(abs_mask).unsqueeze_(0).unsqueeze_(0).detach().to(query.device)
-                        reg(p,abs_mask)
-                    else:
-                        abs_mask  = get_reg(p)
-                else:
-                    p = ident(['a',0,0,0,query.size(-2),key_value.size(-2)])
-                    if un_reg(p):
-                        abs_mask = get_absolute_mask(0,0,0,query.size(-2),key_value.size(-2))
-                        #mask:[query_len,key_len]->[batch,head,query_len,key_len]
-                        abs_mask = torch.from_numpy(abs_mask).unsqueeze_(0).unsqueeze_(0).detach().to(query.device)
-                        reg(p,abs_mask)
-                    else:
-                        abs_mask  = get_reg(p)
-                abs_mask  = abs_mask.expand(*(scores.size()))
-                value_to_sub = absolute_affine(1.0)
-                scores = torch.where(abs_mask == 0, scores - value_to_sub, scores)
-            #遮挡信息之前先talk，这样数值稳定
-            if talking_before_softmax is not None:
-                scores = talking_before_softmax(scores.transpose(-1,-3)).transpose(-1,-3)
-            #是否需要遮挡未来信息
-            if mask_future == True:
-                p = ident(['f',0,0,0,query.size(-2),key_value.size(-2)])
-                if un_reg(p):
-                    #创建遮挡未来信息的掩码
-                    #mask:[query_len,key_len]->[batch,head,query_len,key_len]
-                    std_mask = get_std_mask(0,0,0,query.size(-2),key_value.size(-2))
-                    std_mask = torch.from_numpy(std_mask).unsqueeze_(0).unsqueeze_(0).detach().to(query.device)
-                    reg(p,std_mask)
-                else:
-                    std_mask = get_reg(p)
-                std_mask = std_mask.expand(*(scores.size()))
-                #q_mask:[batch,head,query_len]->[batch,head,query_len,key_len]
-                std_mask = q_mask.unsqueeze_(-1).expand(*(std_mask.size())) & std_mask
-                scores.masked_fill_(std_mask == 0.0,-1e3)
-            #计算概率权重
-            p_attn = F.softmax(scores, dim = -1)
-            #权重talk
-            if talking_after_softmax is not None:
-                p_attn = talking_after_softmax(p_attn.transpose(-1,-3)).transpose(-1,-3)
-            if self.enable_affine:
-                temp = p_attn.sum(dim=-1,keepdim=True)*self.value_a.bias*self.value_a.grad_factor
-            #计算加权求和的结果
-            ret = torch.matmul(p_attn, key_value)
-        else:
-            #分块时需要一个空间存放最终计算结果
-            ret = torch.zeros_like(query)
-            temp = torch.zeros_like(query[...,:1])
-            #分块操作
-            for i in range(0,query.size(-2),block_size):
-                #进行分块
-                query_block  =  query[...,i:i+block_size,:]
-                q_mask_block = q_mask[...,i:i+block_size]
-                key_value_block  =  key_value[...,max(0,i-block_size):i+block_size*2,:]
-                #计算scores
-                scores = torch.matmul(query_block,key_value_block.transpose(-1,-2))
-                if self.enable_affine == True:
-                    scores = scores+self.temp[:,:,i:i+block_size]
-                scores = scores/math.sqrt(self.key_dim)
-                #尝试添加相对位置信息
-                if relative_affine is not None: 
-                    p = ident(['r',i,i,block_size,query.size(-2),key_value.size(-2)])
-                    if un_reg(p):
-                        rela_dist = get_relative_dist(i,i,block_size,query.size(-2),key_value.size(-2))
-                        rela_dist = torch.from_numpy(rela_dist).detach().to(query.device)
-                        reg(p,rela_dist)
-                    else:
-                        rela_dist = get_reg(p)
-                    # dist_decay= 1.0 / (1 + rela_dist*relative_affine(1.0))
-                    dist_decay= rela_dist.mul(relative_affine(1.0)).add(1.0).reciprocal()
-                    scores    = scores.mul(dist_decay)
-                    
-                #尝试添加绝对位置信息
-                if absolute_affine is not None:
-                    p = ident(['a',i,i,block_size,query.size(-2),key_value.size(-2)])
-                    if un_reg(p):
-                        abs_mask  = get_absolute_mask(i,i,block_size,query.size(-2),key_value.size(-2))
-                        abs_mask  = torch.from_numpy(abs_mask).unsqueeze_(0).unsqueeze_(0).detach().to(query.device)
-                        reg(p,abs_mask)
-                    else:
-                        abs_mask = get_reg(p)
-                    abs_mask  = abs_mask.expand(*(scores.size()))
-                    value_to_sub = absolute_affine(1.0)
-                    scores = torch.where(abs_mask == 0, scores - value_to_sub, scores)
-        
-                #遮挡信息之前先talk，这样数值稳定
-                if talking_before_softmax is not None:
-                    scores = talking_before_softmax(scores.transpose(-1,-3)).transpose(-1,-3)
-                    
-                #是否需要遮挡未来信息
-                if mask_future == True:
-                    p = ident(['f',i,i,block_size,query.size(-2),key_value.size(-2)])
-                    if un_reg(p):
-                        #创建遮挡未来信息的掩码，因为是批次操作，需要进行升维
-                        std_mask = get_std_mask(i,i,block_size,query.size(-2),key_value.size(-2))
-                        std_mask = torch.from_numpy(std_mask).unsqueeze_(0).unsqueeze_(0).detach().to(query.device)
-                        reg(p,std_mask)
-                    else:
-                        std_mask = get_reg(p)
-                    std_mask = std_mask.expand(*(scores.size()))
-                    std_mask = q_mask_block.unsqueeze(-1).expand(*(std_mask.size())) & std_mask
-                    scores.masked_fill_(std_mask == 0.0,-1e3)
-        
-                #计算概率权重
-                p_attn = F.softmax(scores, dim = -1)
-        
-                #权重talk
-                if talking_after_softmax is not None:
-                    p_attn = talking_after_softmax(p_attn.transpose(-1,-3)).transpose(-1,-3)
-                if self.enable_affine:
-                    temp[...,i:i+block_size,:] = p_attn.sum(dim=-1,keepdim=True)*self.value_a.bias*self.value_a.grad_factor
-                #计算加权求和的结果
-                ret[...,i:i+block_size,:] = torch.matmul(p_attn, key_value_block)
-        if self.enable_affine:
-            ret = ret * self.value_a.value * self.value_a.grad_factor
-        ret = torch.matmul(ret,self.value_w.weight.view(self.head_number,self.key_dim,self.embedding_dim).transpose(1,2)) + temp
-        return ret
-        
-    def forward(self, query, q_mask, key_value, mask_future):
-        #采用EL-Attention方案
-        if self.enable_el_cache:
-            if query.size(-2) > 1:
-                self.cnt = query.size(-2) - 1
-                self.kv_cache = key_value
-            else:
-                self.cnt += 1
-                self.kv_cache = torch.cat((self.kv_cache,key_value),1)
-                key_value = self.kv_cache
-                mask_future = False
-        #经过线性变换得到真正的QKV
-        query = self.query_w(query)
-        batch_size = query.size(0)
-        query = query.view(batch_size, -1, self.head_number, self.key_dim).transpose(1,2)
-        #进行仿射变换，加快训练速度
-        if self.enable_affine == True:
-            query = self.query_a(query)
-            self.temp = query.sum(dim=-1,keepdim=True)*self.key_a.bias*self.key_a.grad_factor
-            query = query*self.key_a.value*self.key_a.grad_factor
-        #划分注意力头
-        query = torch.matmul(query,self.key_w.weight.view(self.head_number, self.key_dim, self.embedding_dim))
-        key_value = key_value.view(batch_size,-1,1,self.embedding_dim).transpose(1,2)
-        #query:[batch,head,seq_len,emd_dim]
-        #key_value:[batch,1,seq_len,emd_dim]
-        #计算多头注意力
-        out = self.attention(query, q_mask, key_value, mask_future)
-        self.temp = None
-        #将计算完注意力的结果拼接回去
-        out = out.transpose(1,2).contiguous().view(batch_size, -1, self.head_number * self.key_dim)
-        if self.enable_affine:
-            return self.dropout_layer(self.out_a(self.out_w(out)))
-        else:
-            return self.dropout_layer(self.out_w(out))
+import math
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from LazyCache import ExpiringDict, ExpiringDefaultDict
+from Affine import Affine
+
+#获取相对位置矩阵
+def get_relative_mat(height,width,k=0):
+    posi_i = np.arange(k,height+k) #列的范围
+    posi_j = np.arange(0,width) #行的范围
+    posi_grid = np.meshgrid(posi_i, posi_j, indexing='ij')
+    return abs(posi_grid[0]-posi_grid[1])
+    
+#用于添加绝对位置信息的掩码
+def get_relative_dist(i,j,block_size,i_end,j_end):
+    if block_size == 0:
+        assert i==0 and j==0 ,"i!=0 or j!=0"
+        return get_relative_mat(i_end,j_end,k=0)
+    #i,j:当前分块的起始位置
+    #block_size:分块大小
+    #i_end,j_end:序列的长度
+    height = block_size #高度，也就是第一个序列中截取的长度，与分块大小相等
+    width  = block_size * 3 #宽度，也就是第二个序列中截取的长度，为了更长的上下文，还需要考虑上一个分块和下一个分块
+    #创建用来遮挡未来信息的标准掩码
+    #i越大，可见的部分越多，j相反，+block_size是因为上一个分块可见。
+    rela_dist = get_relative_mat(height,width,k=block_size+i-j)
+    #边界超出处理
+    #下超出
+    down_out = max(0,i+height-i_end)
+    #左超出
+    left_out = max(0,block_size-j)
+    #右超出
+    right_out = max(0,j+block_size*2-j_end)
+    #边界内截取
+    rela_dist = rela_dist[:height-down_out,left_out:width-right_out]
+    return rela_dist.astype(np.float32)
+
+#用于添加绝对位置信息的掩码
+def get_absolute_mask(i,j,block_size,i_end,j_end):
+    if block_size == 0:
+        assert i==0 and j==0 ,"i!=0 or j!=0"
+        return np.triu(np.ones((i_end,j_end),dtype='bool'), k=0)
+    #i,j:当前分块的起始位置
+    #block_size:分块大小
+    #i_end,j_end:序列的长度
+    height = block_size #高度，也就是第一个序列中截取的长度，与分块大小相等
+    width  = block_size * 3 #宽度，也就是第二个序列中截取的长度，为了更长的上下文，还需要考虑上一个分块和下一个分块
+    #创建用来遮挡未来信息的标准掩码
+    #i越大，可见的部分越多，j相反，+block_size是因为上一个分块可见。
+    abs_mask = np.triu(np.ones((height,width),dtype='bool'), k=block_size+i-j)
+    #边界超出处理
+    #下超出
+    down_out = max(0,i+height-i_end)
+    #左超出
+    left_out = max(0,block_size-j)
+    #右超出
+    right_out = max(0,j+block_size*2-j_end)
+    #边界内截取
+    abs_mask = abs_mask[:height-down_out,left_out:width-right_out]
+    return abs_mask
+
+#用于遮挡未来信息的标准掩码
+def get_std_mask(i,j,block_size,i_end,j_end):
+    if block_size == 0:
+        assert i==0 and j==0 ,"i!=0 or j!=0"
+        return np.triu(np.ones((i_end,j_end),dtype='bool'), k=1) == False
+    #i,j:当前分块的起始位置
+    #block_size:分块大小
+    #i_end,j_end:序列的长度
+    height = block_size #高度，也就是第一个序列中截取的长度，与分块大小相等
+    width  = block_size * 3 #宽度，也就是第二个序列中截取的长度，为了更长的上下文，还需要考虑上一个分块和下一个分块
+    #创建用来遮挡未来信息的标准掩码
+    #i越大，可见的部分越多，j相反，+block_size是因为上一个分块可见。
+    std_mask = np.triu(np.ones((height,width),dtype='bool'), k=1+block_size+i-j)
+    #边界超出处理
+    #下超出
+    down_out = max(0,i+height-i_end)
+    #左超出
+    left_out = max(0,block_size-j)
+    #右超出
+    right_out = max(0,j+block_size*2-j_end)
+    #边界内截取
+    std_mask = std_mask[:height-down_out,left_out:width-right_out]
+    return std_mask == False
+
+#标记一个需要多次使用的tensor
+def ident(p_list):
+    i,j,block_size,i_end,j_end = p_list[1:]
+    ret = [p_list[0]]
+    if p_list[0]=='r' or p_list[0]=='a':
+        if block_size == 0:
+            ret += [i_end,j_end,0]
+        else:
+            height = block_size
+            width  = block_size * 3
+            ret += [height,width,block_size+i-j]
+            down_out = max(0,i+height-i_end)
+            left_out = max(0,block_size-j)
+            right_out = max(0,j+block_size*2-j_end)
+            ret += [height-down_out,left_out,width-right_out]
+    else:
+        if block_size == 0:
+            ret += [i_end,j_end,1]
+        else:
+            height = block_size
+            width  = block_size * 3
+            ret += [height,width,1+block_size+i-j]
+            down_out = max(0,i+height-i_end)
+            left_out = max(0,block_size-j)
+            right_out = max(0,j+block_size*2-j_end)
+            ret += [height-down_out,left_out,width-right_out]
+    return str(ret)
+
+#缓存字典与定时器
+reg_dict = dict()
+reg_timer = dict()
+
+#查看是否未注册
+def un_reg(p):
+    return not p in reg_dict
+
+#注册需要重复使用的tensor
+def reg(p,v):
+    #找缓冲中用的最少的
+    keys = [k for k in reg_dict]
+    time_min = 0
+    if len(keys) != 0:
+        key_min = keys[0]
+        time_min = reg_timer[key_min]
+        for k in keys:
+            if reg_timer[k]<time_min:
+                key_min = k
+                time_min = reg_timer[key_min]
+    #计数
+    if not p in reg_timer:
+        reg_timer[p] = 1
+    else:
+        reg_timer[p] += 1
+    #缓冲满了就删掉最少用的
+    if len(keys) > 12:
+        del reg_dict[key_min]
+    #比最小的值大就保留
+    if reg_timer[p] > time_min or len(keys) < 12:
+        reg_dict[p] = v
+
+#从缓冲区中获取可重复使用的张量
+def get_reg(p):
+    reg_timer[p] += 1
+    return reg_dict[p]
+
+
+#多头注意力
+class MultiHeadAttention(nn.Module):
+    def __init__(self,embedding_dim,key_dim,head_number,position_information_type,enable_affine,enable_talking_head, \
+                 self_attention_block_size,dropout_rate,enable_el_cache):
+        super(MultiHeadAttention, self).__init__()
+        self.embedding_dim              = embedding_dim
+        self.key_dim                    = key_dim
+        self.head_number                = head_number
+        self.position_information_type  = position_information_type
+        self.enable_talking_head        = enable_talking_head
+        self.self_attention_block_size  = self_attention_block_size
+        self.dropout_layer              = nn.Dropout(p=dropout_rate)
+        self.enable_affine              = enable_affine
+
+        self.query_w = nn.Linear(embedding_dim,key_dim*head_number,bias=False)
+        self.key_w   = nn.Linear(embedding_dim,key_dim*head_number,bias=False)
+        self.value_w = nn.Linear(embedding_dim,key_dim*head_number,bias=False)
+        self.out_w   = nn.Linear(key_dim*head_number,embedding_dim,bias=False)
+
+        self.enable_el_cache = enable_el_cache
+        # 带有自动垃圾回收的字典
+        self.kv_cache = None
+        self.temp     = None
+        self.cnt      = None
+
+        if enable_affine == True:
+            self.query_a = Affine(1.0)
+            self.key_a   = Affine(1.0)
+            self.value_a = Affine(1.0)
+            self.out_a   = Affine(1.0)
+
+        if enable_talking_head == True:
+            self.talking_before_softmax = nn.Linear(head_number,head_number,bias=False)
+            self.talking_after_softmax  = nn.Linear(head_number,head_number,bias=False)
+        else:
+            self.talking_before_softmax = None
+            self.talking_after_softmax  = None
+
+        if position_information_type == "mask":
+            self.absolute_affine = Affine(1.0,grad_factor=1.0)
+            self.relative_affine = Affine(0.1,grad_factor=1.0)
+        else:
+            self.absolute_affine = None
+            self.relative_affine = None
+
+    #注意力运算
+    def attention(self, query, q_mask, key_value, mask_future, session_id):
+        #为了使用EL-Attention需要修改参数传递方式
+        absolute_affine = self.absolute_affine
+        relative_affine = self.relative_affine
+        talking_before_softmax = self.talking_before_softmax
+        talking_after_softmax = self.talking_after_softmax
+        block_size = self.self_attention_block_size
+        #提前调整q_mask的形状，方便广播
+        #query:[batch,head,query_len,emb_dim]
+        #q_mask:[batch,query_len]
+        #q_mask:[batch,query_len]->[batch,1,query_len]
+        #q_mask:[batch,1,query_len]->[batch,head,query_len]
+        q_mask = q_mask.unsqueeze(1).expand(*(query.size()[:-1]))
+        #判断是否需要分块运算
+        if block_size == 0:
+            #不进行分块
+            #计算scores
+            scores = torch.matmul(query,key_value.transpose(-1,-2))
+            if self.enable_affine == True:
+                scores = scores+self.temp[session_id]
+            scores = scores/math.sqrt(self.key_dim)
+            #尝试添加相对位置信息
+            if relative_affine is not None:
+                if self.enable_el_cache and query.size(-2) == 1:
+                    p = ident(['er',0,0,0,query.size(-2),key_value.size(-2)])
+                    if un_reg(p):
+                        rela_dist = np.arange(self.cnt[session_id],-1,-1).reshape(1,-1)
+                        rela_dist = torch.from_numpy(rela_dist).detach().to(query.device)
+                        reg(p,rela_dist)
+                    else:
+                        rela_dist = get_reg(p)
+                else:
+                    p = ident(['r',0,0,0,query.size(-2),key_value.size(-2)])
+                    if un_reg(p):
+                        rela_dist = get_relative_dist(0,0,0,query.size(-2),key_value.size(-2))
+                        #直接广播更高效
+                        rela_dist = torch.from_numpy(rela_dist).detach().to(query.device)
+                        reg(p,rela_dist)
+                    else:
+                        rela_dist = get_reg(p)
+                dist_decay= rela_dist.mul(relative_affine(1.0)).add(1.0).reciprocal()
+                scores    = scores.mul(dist_decay)
+            #尝试添加绝对位置信息
+            if absolute_affine is not None:
+                if self.enable_el_cache and query.size(-2) == 1:
+                    p = ident(['ea',0,0,0,query.size(-2),key_value.size(-2)])
+                    if un_reg(p):
+                        abs_mask = np.array([[False]*(self.cnt[session_id])+[True]])
+                        abs_mask = torch.from_numpy(abs_mask).unsqueeze_(0).unsqueeze_(0).detach().to(query.device)
+                        reg(p,abs_mask)
+                    else:
+                        abs_mask  = get_reg(p)
+                else:
+                    p = ident(['a',0,0,0,query.size(-2),key_value.size(-2)])
+                    if un_reg(p):
+                        abs_mask = get_absolute_mask(0,0,0,query.size(-2),key_value.size(-2))
+                        #mask:[query_len,key_len]->[batch,head,query_len,key_len]
+                        abs_mask = torch.from_numpy(abs_mask).unsqueeze_(0).unsqueeze_(0).detach().to(query.device)
+                        reg(p,abs_mask)
+                    else:
+                        abs_mask  = get_reg(p)
+                abs_mask  = abs_mask.expand(*(scores.size()))
+                value_to_sub = absolute_affine(1.0)
+                scores = torch.where(abs_mask == 0, scores - value_to_sub, scores)
+            #遮挡信息之前先talk，这样数值稳定
+            if talking_before_softmax is not None:
+                scores = talking_before_softmax(scores.transpose(-1,-3)).transpose(-1,-3)
+            #是否需要遮挡未来信息
+            if mask_future == True:
+                p = ident(['f',0,0,0,query.size(-2),key_value.size(-2)])
+                if un_reg(p):
+                    #创建遮挡未来信息的掩码
+                    #mask:[query_len,key_len]->[batch,head,query_len,key_len]
+                    std_mask = get_std_mask(0,0,0,query.size(-2),key_value.size(-2))
+                    std_mask = torch.from_numpy(std_mask).unsqueeze_(0).unsqueeze_(0).detach().to(query.device)
+                    reg(p,std_mask)
+                else:
+                    std_mask = get_reg(p)
+                std_mask = std_mask.expand(*(scores.size()))
+                #q_mask:[batch,head,query_len]->[batch,head,query_len,key_len]
+                std_mask = q_mask.unsqueeze_(-1).expand(*(std_mask.size())) & std_mask
+                scores.masked_fill_(std_mask == 0.0,-1e3)
+            #计算概率权重
+            p_attn = F.softmax(scores, dim = -1)
+            #权重talk
+            if talking_after_softmax is not None:
+                p_attn = talking_after_softmax(p_attn.transpose(-1,-3)).transpose(-1,-3)
+            if self.enable_affine:
+                temp = p_attn.sum(dim=-1,keepdim=True)*self.value_a.bias*self.value_a.grad_factor
+            #计算加权求和的结果
+            ret = torch.matmul(p_attn, key_value)
+        else:
+            #分块时需要一个空间存放最终计算结果
+            ret = torch.zeros_like(query)
+            temp = torch.zeros_like(query[...,:1])
+            #分块操作
+            for i in range(0,query.size(-2),block_size):
+                #进行分块
+                query_block  =  query[...,i:i+block_size,:]
+                q_mask_block = q_mask[...,i:i+block_size]
+                key_value_block  =  key_value[...,max(0,i-block_size):i+block_size*2,:]
+                #计算scores
+                scores = torch.matmul(query_block,key_value_block.transpose(-1,-2))
+                if self.enable_affine == True:
+                    scores = scores+self.temp[session_id][:,:,i:i+block_size]
+                scores = scores/math.sqrt(self.key_dim)
+                #尝试添加相对位置信息
+                if relative_affine is not None: 
+                    p = ident(['r',i,i,block_size,query.size(-2),key_value.size(-2)])
+                    if un_reg(p):
+                        rela_dist = get_relative_dist(i,i,block_size,query.size(-2),key_value.size(-2))
+                        rela_dist = torch.from_numpy(rela_dist).detach().to(query.device)
+                        reg(p,rela_dist)
+                    else:
+                        rela_dist = get_reg(p)
+                    # dist_decay= 1.0 / (1 + rela_dist*relative_affine(1.0))
+                    dist_decay= rela_dist.mul(relative_affine(1.0)).add(1.0).reciprocal()
+                    scores    = scores.mul(dist_decay)
+                    
+                #尝试添加绝对位置信息
+                if absolute_affine is not None:
+                    p = ident(['a',i,i,block_size,query.size(-2),key_value.size(-2)])
+                    if un_reg(p):
+                        abs_mask  = get_absolute_mask(i,i,block_size,query.size(-2),key_value.size(-2))
+                        abs_mask  = torch.from_numpy(abs_mask).unsqueeze_(0).unsqueeze_(0).detach().to(query.device)
+                        reg(p,abs_mask)
+                    else:
+                        abs_mask = get_reg(p)
+                    abs_mask  = abs_mask.expand(*(scores.size()))
+                    value_to_sub = absolute_affine(1.0)
+                    scores = torch.where(abs_mask == 0, scores - value_to_sub, scores)
+        
+                #遮挡信息之前先talk，这样数值稳定
+                if talking_before_softmax is not None:
+                    scores = talking_before_softmax(scores.transpose(-1,-3)).transpose(-1,-3)
+                    
+                #是否需要遮挡未来信息
+                if mask_future == True:
+                    p = ident(['f',i,i,block_size,query.size(-2),key_value.size(-2)])
+                    if un_reg(p):
+                        #创建遮挡未来信息的掩码，因为是批次操作，需要进行升维
+                        std_mask = get_std_mask(i,i,block_size,query.size(-2),key_value.size(-2))
+                        std_mask = torch.from_numpy(std_mask).unsqueeze_(0).unsqueeze_(0).detach().to(query.device)
+                        reg(p,std_mask)
+                    else:
+                        std_mask = get_reg(p)
+                    std_mask = std_mask.expand(*(scores.size()))
+                    std_mask = q_mask_block.unsqueeze(-1).expand(*(std_mask.size())) & std_mask
+                    scores.masked_fill_(std_mask == 0.0,-1e3)
+        
+                #计算概率权重
+                p_attn = F.softmax(scores, dim = -1)
+        
+                #权重talk
+                if talking_after_softmax is not None:
+                    p_attn = talking_after_softmax(p_attn.transpose(-1,-3)).transpose(-1,-3)
+                if self.enable_affine:
+                    temp[...,i:i+block_size,:] = p_attn.sum(dim=-1,keepdim=True)*self.value_a.bias*self.value_a.grad_factor
+                #计算加权求和的结果
+                ret[...,i:i+block_size,:] = torch.matmul(p_attn, key_value_block)
+        if self.enable_affine:
+            ret = ret * self.value_a.value * self.value_a.grad_factor
+        ret = torch.matmul(ret,self.value_w.weight.view(self.head_number,self.key_dim,self.embedding_dim).transpose(1,2)) + temp
+        return ret
+        
+    def forward(self, query, q_mask, key_value, mask_future, session_id):
+        #采用EL-Attention方案
+        if self.enable_el_cache:
+            if self.kv_cache is None:
+                self.kv_cache = ExpiringDict(ttl=600)
+                self.kv_cache.start_auto_cleanup()
+                self.temp     = ExpiringDict(ttl=600)
+                self.temp.start_auto_cleanup()
+                self.cnt      = ExpiringDefaultDict(int, ttl=600)
+                self.cnt.start_auto_cleanup()
+            if query.size(-2) > 1:
+                self.cnt[session_id] = query.size(-2) - 1
+                self.kv_cache[session_id] = key_value
+            else:
+                self.cnt[session_id] += 1
+                self.kv_cache[session_id] = torch.cat((self.kv_cache[session_id],key_value),1)
+                key_value = self.kv_cache[session_id]
+                mask_future = False
+        #经过线性变换得到真正的QKV
+        query = self.query_w(query)
+        batch_size = query.size(0)
+        query = query.view(batch_size, -1, self.head_number, self.key_dim).transpose(1,2)
+        #进行仿射变换，加快训练速度
+        if self.enable_affine == True:
+            query = self.query_a(query)
+            self.temp[session_id] = query.sum(dim=-1,keepdim=True)*self.key_a.bias*self.key_a.grad_factor
+            query = query*self.key_a.value*self.key_a.grad_factor
+        #划分注意力头
+        query = torch.matmul(query,self.key_w.weight.view(self.head_number, self.key_dim, self.embedding_dim))
+        key_value = key_value.view(batch_size,-1,1,self.embedding_dim).transpose(1,2)
+        #query:[batch,head,seq_len,emd_dim]
+        #key_value:[batch,1,seq_len,emd_dim]
+        #计算多头注意力
+        out = self.attention(query, q_mask, key_value, mask_future, session_id)
+        self.temp[session_id] = None
+        #将计算完注意力的结果拼接回去
+        out = out.transpose(1,2).contiguous().view(batch_size, -1, self.head_number * self.key_dim)
+        if self.enable_affine:
+            return self.dropout_layer(self.out_a(self.out_w(out)))
+        else:
+            return self.dropout_layer(self.out_w(out))
             

app.py

@@ -1,221 +1,295 @@
-import time
-import uuid
-import html
-import threading
-import numpy as np
-import gradio as gr
-from queue import Queue
-from tokenizer import tokenizer,vocab_size,token2str
-
-import torch
-import torch.nn as nn
-from make_model import make_model
-from train_and_use import El_text_continue_stream
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-model = make_model(
-    #token是从1开始的，0填充，剩下的用来覆盖全部字节
-    vocab_size = vocab_size+1+255,
-    embedding_dim = 768,
-    key_dim = 128,
-    head_number = 12,
-    position_information_type = "mask",
-    enable_affine = True,
-    enable_talking_head = True,
-    use_diff = False,
-    self_attention_block_size = 0,
-    feed_forward_dim = 1536,
-    enable_layer_norm = True,
-    deep = 12,
-    dropout_rate = 0.1,
-    enable_el_cache = True
-).to(device)
-model.load_state_dict(torch.load('large_model_instruct_27567.weight',map_location=device,weights_only=True))
-model = model.eval()
-
-# 全局字典，存 per-session 的不可 deepcopy 对象 / 状态
-user_queues = {}             # session_id -> Queue()
-user_stop_flags = {}         # session_id -> bool (True 表示停止)
-user_current_sessions = {}   # session_id -> 最后一个 session (list), 可 deepcopy
-
-# token包装器
-def token_wapper(token):
-    return f'<span style="background-color: #FFD580; padding: 2px 4px; border-radius: 6px; margin: 1px; display: inline-block;">{html.escape(token)}</span>'
-    
-def token_split_wapper(token):
-    safe_token = html.escape(token)
-    text = f"({safe_token})[单字多token]"
-    return f'<span style="background-color: #FF0000; padding: 2px 4px; border-radius: 6px; margin: 1px; display: inline-block;">{text}</span>'
-
-# 后台生成函数（只访问全局字典，通过 session_id 定位）
-def generate_text(user_message, session_id, temperature, repeat_penalty, max_length, decay):
-    out = ""
-    q = user_queues.get(session_id)
-    if q is None:
-        return
-    # 通过分词器转化为token
-    user_tokens = tokenizer(user_message,5.0)
-    # 将token还原并进行包装
-    words = []
-    temp = []
-    for token in user_tokens:
-        if token > 0:
-            if len(temp):
-                words += [token_split_wapper(token2str(temp))]
-                temp = []
-            words += [token_wapper(token2str([token]))]
-        else:
-            temp += [token]
-    if len(temp):
-        words += [token_split_wapper(token2str(temp))]
-    user_tokens = ''.join(words)
-    # 准备模型输入
-    if len(tokenizer(user_message,5.0)) < 2:
-        user_message = f' {user_message}'
-    tokens_batch = [tokenizer('<|im_start|>user '+user_message+'<|im_end|><|im_start|>assistant ',5.0)]
-    tokens_batch = np.array(tokens_batch,dtype=np.int64)+255
-    inputs = torch.from_numpy(tokens_batch).to(device).data
-    last_len = -1
-    # 模型输出
-    with torch.no_grad():
-        for o in El_text_continue_stream(
-            model,inputs,out_length=max_length,
-            repeat_penalty_value=repeat_penalty,
-            temperature=temperature,
-            decay=decay
-        ):
-            split = ''
-            if o[0,-1] > 255: #确保是完整的字符才可以输出
-                out += token2str(o[0][last_len:].cpu().numpy()-255,split=split)
-                last_len = -1
-                sess = [
-                    {"role": "user", "content": user_tokens},
-                    {"role": "assistant", "content": out},
-                ]
-                user_current_sessions[session_id] = sess
-                try:
-                    q.put(sess, block=False)
-                except:
-                    # 极少情况：队列放入失败（一般不会发生），忽略
-                    pass
-            else:
-                last_len -= 1
-            if user_stop_flags.get(session_id, True):
-                break
-            if '<|im_end|>' in out:
-                out = out.split('<|im_end|>')[0]
-                sess = [
-                    {"role": "user", "content": user_tokens},
-                    {"role": "assistant", "content": out},
-                ]
-                user_current_sessions[session_id] = sess
-                try:
-                    q.put(sess, block=False)
-                except:
-                    # 极少情况：队列放入失败（一般不会发生），忽略
-                    pass
-                break
-
-# 点击按钮的处理逻辑：start / stop / clear
-def click_process(sess, label, user_message, state, stop_flag_state, session_id, temperature, repeat_penalty, max_length, decay):
-    # 安全检查
-    if session_id is None or session_id not in user_queues:
-        # session 还没初始化好，直接返回不改变 UI
-        return "", "发送消息", state or {"current_session": []}, stop_flag_state or {"stop": True}, session_id
-
-    # 如果现在处于"停止"状态并且有用户输入 -> 启动生成线程
-    if stop_flag_state.get("stop", True) and user_message and sess == []:
-        user_stop_flags[session_id] = False
-        thread = threading.Thread(target=generate_text, args=(user_message, session_id, temperature, repeat_penalty, max_length, decay))
-        thread.daemon = True
-        thread.start()
-        # 更新返回给前端的 state/stop_flag（gradio 会把这些值保存到 session state）
-        return "",  "终止输出", {"current_session": user_current_sessions.get(session_id, [])}, {"stop": False}, session_id
-
-    # 如果正在输出 -> 终止
-    elif not stop_flag_state.get("stop", True) and label != "清空会话":
-        user_stop_flags[session_id] = True
-        return user_message, "清空会话", {"current_session": user_current_sessions.get(session_id, [])}, {"stop": True}, session_id
-
-    # 否则清空会话
-    else:
-        user_stop_flags[session_id] = True
-        user_current_sessions[session_id] = []
-        q = user_queues.get(session_id)
-        if q:
-            while not q.empty():
-                try:
-                    q.get_nowait()
-                except:
-                    break
-        return user_message, "发送消息", {"current_session": []}, {"stop": True}, session_id
-
-# 流式输出 generator（只需触发一次即可一直运行）
-def stream_output(state, stop_flag_state):
-    global user_queues,user_stop_flags,user_current_sessions
-    # 页面加载时初始化 session（返回可 deepcopy 的 state 值和 session_id）
-    session_id = str(uuid.uuid4())
-    user_queues[session_id] = Queue()
-    user_stop_flags[session_id] = True
-    user_current_sessions[session_id] = []  # 初始为空会话
-    # 返回给 gradio 的 state 值（这些都是 deepcopy-friendly）
-    yield gr.update(), gr.update(), {"current_session": []}, {"stop": True}, session_id
-    t0 = time.time()
-    while True:
-        q = user_queues[session_id]
-        stopped = user_stop_flags.get(session_id, True)
-        # 优先处理队列中的消息（FIFO）
-        if (not stopped) and (not q.empty()):
-            t0 = time.time()
-            while q.qsize() > 5:
-                sess = q.get()
-            sess = q.get()
-            # 更新 chatbot（返回的 sess 是 [{"role":...}, ...]）
-            # 同时把 state 返回为 deepcopy-friendly 字典（gr.State 需要可 deepcopied）
-            yield sess, "终止输出", {"current_session": sess}, gr.update(), session_id
-        else:
-            last = user_current_sessions.get(session_id, [])
-            if last == []:
-                yield last, gr.update(), {"current_session": last},gr.update(), session_id
-            else:
-                if time.time() - t0 > 3:
-                    yield last, "清空会话", {"current_session": last},gr.update(), session_id
-        time.sleep(0.1)  # 防止 busy-wait 占满 CPU
-
-# ========== Gradio UI ==========
-with gr.Blocks() as demo:
-    gr.Markdown("# LLM 在线体验（指令微调版）")
-    
-    chatbot = gr.Chatbot(type="messages", label="输入/输出", autoscroll=False, show_copy_button=False)
-    msg = gr.Textbox(placeholder="请输入你的问题。", label="用户问题输入", lines=4)
-    
-    with gr.Row():
-        temperature = gr.Slider(0.0001, 3.0001, value=0.0001, step=0.1, label="Temperature")
-        repeat_penalty = gr.Slider(0.0, 5.0, value=2.5, step=0.1, label="Repeat Penalty")
-        max_length = gr.Slider(64, 8192, value=512, step=64, label="Max Length")
-        decay = gr.Slider(0.90, 1.0, value=0.98, step=0.01, label="Repeat Penalty Decay Rate")
-    
-    btn = gr.Button("发送消息")
-
-    # gr.State 用来在前端保存可 deepcopied 的 session 值
-    state = gr.State()
-    stop_flag_state = gr.State()
-    session_id = gr.State()
-
-    # 点击按钮处理 - 使用 session_id 定位用户资源
-    btn.click(
-        click_process,
-        inputs=[chatbot, btn, msg, state, stop_flag_state, session_id, temperature, repeat_penalty, max_length, decay],
-        outputs=[msg, btn, state, stop_flag_state, session_id],
-    )
-
-    # 页面加载后再触发 stream_output（只要触发一次，generator 会一直运行）
-    demo.load(
-        stream_output,
-        inputs=[state, stop_flag_state],
-        outputs=[chatbot, btn, state, stop_flag_state, session_id],
-    )
-
-if __name__ == "__main__":
-    demo.queue(max_size=128, default_concurrency_limit=128)
-    demo.launch(share=False)
+# 公开库
+import time
+import html
+import uuid
+import torch
+import threading
+import numpy as np
+import gradio as gr
+# 私有库
+from queue import Queue
+from make_model import make_model
+from LazyCache import ExpiringDict
+from train_and_use import El_text_continue_stream
+from tokenizer import tokenizer,vocab_size,token2str
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# 加载模型
+model = make_model(
+    #token是从1开始的，0填充，剩下的用来覆盖全部字节
+    vocab_size = vocab_size+1+255,
+    embedding_dim = 768,
+    key_dim = 128,
+    head_number = 12,
+    position_information_type = "mask",
+    enable_affine = True,
+    enable_talking_head = True,
+    use_diff = False,
+    self_attention_block_size = 0,
+    feed_forward_dim = 1536,
+    enable_layer_norm = True,
+    deep = 12,
+    dropout_rate = 0.1,
+    enable_el_cache = True
+).to(device)
+model.load_state_dict(torch.load('large_model_instruct_09271556.weight',map_location=device,weights_only=True))
+model = model.eval()
+
+# token包装函数 - 使用HTML span标签确保每个token在独立矩形中
+def token_wapper(token):
+    # 对特殊字符进行HTML转义处理
+    escaped_token = html.escape(token)
+    return f'<span class="token-box">{escaped_token}</span>'
+
+# 多token包装函数 - 使用HTML span标签确保每个token在独立矩形中
+def token_split_wapper(token):
+    # 对特殊字符进行HTML转义处理
+    escaped_token = html.escape(token)
+    return f'<span class="multi-token-box">({escaped_token})[多token]</span>'
+
+# 处理用户输入的token，返回安全的显示格式
+def process_user_tokens(user_message):
+    # 通过分词器转化为token
+    user_tokens = tokenizer(user_message, 5.0)
+    
+    # 将token还原并进行安全包装
+    words = [] # token列表
+    temp = [] # token是特殊字节，要合并
+    for token in user_tokens:
+        if token > 0:
+            # 将合并成功的加入列表
+            if len(temp):
+                words.append(token_split_wapper(token2str(temp)))
+                temp = []
+            # 将新的token加入列表
+            words.append(token_wapper(token2str([token])))
+        else:
+            # 将字节送去合并
+            temp.append(token)
+    # 结束的时候要进行收尾
+    if len(temp):
+        words.append(token_split_wapper(token2str(temp)))
+    # 返回包装好的token列表
+    return ''.join(words)
+
+# 全局字典，存 per-session 的不可 deepcopy 对象 / 状态
+user_queues = ExpiringDict(ttl=550)             # session_id -> Queue(list([string,string]))，用于流式输出
+user_queues.start_auto_cleanup()
+user_stop_flags = ExpiringDict(ttl=550)         # session_id -> bool (True 表示停止)
+user_stop_flags.start_auto_cleanup()
+user_history_sessions_show = ExpiringDict(ttl=550)   # session_id -> 用于显示的历史记录，list([string,string])
+user_history_sessions_show.start_auto_cleanup()
+user_history_sessions_text = ExpiringDict(ttl=550)   # session_id -> 纯文本历史记录，string
+user_history_sessions_text.start_auto_cleanup()
+
+# 后台生成函数（只访问全局字典，通过 session_id 定位）
+def generate_text(sess, user_message, session_id, temperature, repeat_penalty, max_length, decay):
+    out = ""
+    q = user_queues.get(session_id)
+    # 立即刷出用户问题
+    q.put(out, block=False)
+    # 构建完整的对话历史输入
+    if len(sess) == 1:
+        user_history_sessions_text[session_id] = f"<|im_start|>user {user_message}<|im_end|><|im_start|>assistant "
+    else:
+        user_history_sessions_text[session_id] += f"<|im_start|>user {user_message}<|im_end|><|im_start|>assistant "
+    # 转换为模型输入格式
+    tokens_batch = [tokenizer(user_history_sessions_text[session_id], 5.0)]
+    tokens_batch = np.array(tokens_batch, dtype=np.int64) + 255
+    inputs = torch.from_numpy(tokens_batch).to(device).data
+    last_len = -1
+    # 模型输出
+    with torch.no_grad():
+        for o in El_text_continue_stream(
+            model, inputs, out_length=max_length,
+            repeat_penalty_value=repeat_penalty,
+            temperature=temperature,decay=decay,session_id=session_id):
+            # 如果当前位置可以完整解码
+            if o[0,-1] > 255:
+                # 将未解码的部分一起解码
+                temp = token2str(o[0][last_len:].cpu().numpy()-255)
+                out += temp
+                user_history_sessions_text[session_id] += temp
+                # 重置为解码光标
+                last_len = -1    
+                q.put(out, block=False)
+            else:
+                # 无法解码，光标固定
+                last_len -= 1
+            # 如果用户主动断开连接，停止生成,去除潜在标记
+            if user_stop_flags.get(session_id, True):
+                if '<' + out.split('<')[-1] in '<|im_end|>':
+                    # 显示的部分去除标记
+                    out = '<'+'<'.join(out.split('<')[:-1])
+                    # 历史的部分保留标记
+                    user_history_sessions_text[session_id] = '<'+'<'.join(user_history_sessions_text[session_id].split('<')[:-1])+'<|im_end|>'
+                break
+            
+            # 如果是输出终止标记
+            if '<|im_end|>' in out:
+                # 显示的部分，去除标记
+                out = out.split('<|im_end|>')[0]
+                q.put(out, block=False)
+                break
+            # 如果用户中断
+            if user_stop_flags[session_id] == True:
+                break
+        # 更新标记为暂停
+        user_stop_flags[session_id] = True
+
+# 按钮处理逻辑：发送消息 / 停止生成 / 清空会话
+def send_message(sess, btn_label, user_message, session_id, temperature, repeat_penalty, max_length, decay):
+    # 发送消息按钮 - 启动生成线程
+    if btn_label == "发送消息" and user_message:
+        # 设置当前用户正在生成的标志
+        user_stop_flags[session_id] = False
+        # 立即在UI中显示用户消息
+        user_tokens_display = process_user_tokens(user_message)
+        # 添加用户消息到当前会话
+        user_history_sessions_show[session_id] = sess
+        user_history_sessions_show[session_id] += [[user_tokens_display, ""]]
+        if session_id not in user_history_sessions_text:
+            return "", "会话过期！"
+        # 在这里开始流式输出
+        thread = threading.Thread(target=generate_text, args=(sess, user_message, session_id, temperature, repeat_penalty, max_length, decay))
+        thread.daemon = True #主进程退出时退出
+        thread.start() #启动
+        user_stop_flags[session_id] = False
+        # 更新返回给前端的 state/stop_flag
+        return "", "停止生成"
+    else:
+        # 停止生成按钮 - 设置标志位
+        user_stop_flags[session_id] = True
+        # 更新返回给前端的 state/stop_flag
+        return user_message, "发送消息"
+
+# 清空会话
+def clear_session():
+    return []
+
+# 流式输出,无限循环刷新页面
+def stream_output(sess):
+    global user_queues, user_stop_flags, user_history_sessions_show, user_history_sessions_text
+    # 页面加载时初始化 session
+    session_id = str(uuid.uuid4())
+    user_queues[session_id] = Queue()
+    user_stop_flags[session_id] = True
+    user_history_sessions_show[session_id] = []  # 初始化历史会话记录，用于显示
+    user_history_sessions_text[session_id] = ""  # 初始化历史会话记录，用于文本存储
+    # 返回初始状态
+    yield [], "发送消息", session_id
+    # 不断刷新
+    while True:
+        time.sleep(0.01)  # 防止 busy-wait 占满 CPU
+        # 等待队列有数据
+        q = user_queues.get(session_id)
+        if q is None:
+            continue
+        # 处理队列中的消息
+        if not q.empty():
+            # 取到最后一个加入的数据
+            while q.qsize() > 1:
+                q.get()
+            out = q.get()
+            sess = user_history_sessions_show[session_id]
+            sess[-1][1] = out
+            # 更新UI状态
+            current_stopped = user_stop_flags.get(session_id, True)
+            button_label = "停止生成" if not current_stopped else "发送消息"
+            yield sess, button_label, session_id
+
+# UI美化
+css = """
+/* 大标题居中 */
+.title {
+    text-align: center;
+}
+/* 高级选项字体居中 */
+#adv-param button {
+    justify-content: center;
+}
+/* 高级选项字体放大 */
+#adv-param > button > span {
+    font-size: 16px !important;
+    font-weight: 600 !important;
+}
+/* 自定义token样式 */
+.token-box {
+    display: inline-block;
+    background-color: #f0f0f0;
+    border: 1px solid #ddd;
+    border-radius: 4px;
+    padding: 2px 4px;
+    margin: 2px;
+    font-family: monospace;
+}
+.multi-token-box {
+    display: inline-block;
+    background-color: #e6f7ff;
+    border: 1px solid #91d5ff;
+    border-radius: 4px;
+    padding: 2px 4px;
+    margin: 2px;
+    font-family: monospace;
+}
+"""
+# ========== Gradio UI ==========
+with gr.Blocks(css=css) as demo:
+    with gr.Column(elem_classes="container"):
+        gr.Markdown("# 0.18B中文大语言模型在线体验", elem_classes="title")
+        # 聊天界面
+        chatbot = gr.Chatbot(
+            label="对话", 
+            autoscroll=False,
+            show_copy_button=True,
+            elem_classes="chatbox",
+            type="tuples",
+            height=400
+        )
+        # 输入区域
+        with gr.Column(elem_classes="input-area"):
+            msg = gr.Textbox(
+                placeholder="请输入你的问题...", 
+                label="", 
+                lines=3,
+                show_label=False
+            )
+            # 按钮区域
+            with gr.Row(elem_classes="button-row"):
+                send_btn = gr.Button("发送消息", elem_classes="send-btn")
+                clear_btn = gr.Button("清空会话", elem_classes="clear-btn")
+        # 参数设置区域（可折叠）
+        with gr.Accordion("高级参数设置", open=False, elem_classes="parameter-row", elem_id="adv-param"):
+            with gr.Row():
+                temperature = gr.Slider(0.0001, 3.0001, value=0.0001, step=0.1, label="Temperature")
+                repeat_penalty = gr.Slider(0.0, 5.0, value=2.5, step=0.1, label="Repeat Penalty")
+            with gr.Row():
+                max_length = gr.Slider(64, 8192, value=512, step=64, label="Max Length")
+                decay = gr.Slider(0.90, 1.0, value=0.98, step=0.01, label="Repeat Penalty Decay Rate")
+    # gr.State 用来在前端保存可 deepcopied 的 session 值
+    session_id = gr.State()
+    # 发送按钮处理
+    send_btn.click(
+        send_message,
+        inputs=[chatbot, send_btn, msg, session_id, temperature, repeat_penalty, max_length, decay],
+        outputs=[msg, send_btn],
+    )
+
+    clear_btn.click(
+        clear_session,
+        inputs=[],
+        outputs=[chatbot],
+    )
+
+    # 无限循环，一直更新聊天界面
+    demo.load(
+        stream_output,
+        inputs=[chatbot],
+        outputs=[chatbot, send_btn, session_id],
+    )
+if __name__ == "__main__":
+    """主函数：启动Gradio界面"""
+    # 设置队列参数以提高并发处理能力
+    demo.queue(max_size=128, default_concurrency_limit=128)
+    # 启动Gradio应用，不公开分享，并应用CSS样式
+    demo.launch(share=False)

train_and_use.py

@@ -1,444 +1,444 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import numpy as np
-import time
-import threading
-import copy
-class Batch:
-    def __init__(self,input_sequences):
-        self.data_type = "generator"
-        self.query = input_sequences[...,:-1]
-        self.label = input_sequences[...,1:]
-        self.q_mask = self.query != 0
-        self.ntokens = float((self.label != 0).sum())
-
-#交叉熵损失，“0”填充特殊处理
-class CrossEntropyLoss(nn.Module):
-    def __init__(self):
-        super(CrossEntropyLoss, self).__init__()
-        # 使用KL散度损失函数（接受对数概率分布x和概率分布y,并不是简单的KL散度计算）
-        self.criterion = nn.KLDivLoss(reduction='sum')
-        
-    def forward(self, model_output_dist, target_sequence):
-        #根据模型输出的分布与标签的分布计算交叉熵损失
-        #为目标分布分配和模型输出形状、类型一样的空间，默认不追踪梯度，写明更清晰
-        true_dist = torch.zeros_like(model_output_dist,requires_grad=False)
-        #使用置信度填充目标词的位置（true_dist是词表那么长的概率分布）
-        #目标序列升维，target_sequence:[batch*len]->[batch*len,1]
-        #true_dist:[batch*len,vocab]
-        #在vocab的维度上用标签值当作索引，找到对应元素，填充1.0
-        true_dist.scatter_(1, target_sequence.data.unsqueeze(1), 1.0)
-        #将填充位置概率设为0
-        true_dist[:,0] = 0.0
-        #计算模型输出分布与平目标序列标签平滑后的分布之间的交叉熵
-        #model_output_dist是对数概率分布,应由F.log_softmax(self.project(x),dim=-1)产生
-        #但实际上为了压缩softmax的值域已达到自动丢弃异常值的效果，在Generator.Projector进行了特殊实现
-        return self.criterion(model_output_dist, true_dist)
-
-class AdamOptimizerWithBase:
-    "带有Base的自适应矩估计优化器"
-    def __init__(self, params, base, half_life, betas, eps):
-        self.beta1 = betas[0]
-        self.beta2 = betas[1]
-        self.beta3 = (1/2)**(1/half_life)
-        self.epsilon = eps
-        self.t = 0
-        self.param_groups = []
-        for p,b in zip(params,base):
-            self.param_groups.append({
-                'params': p, 
-                'lr'    : 0.0, 
-                'm'     : torch.zeros_like(p).detach(), 
-                'v'     : torch.zeros_like(p).detach(), 
-                'b'     : b.clone().detach()
-            })
-
-    def step(self):
-        self.t += 1
-        for group in self.param_groups:
-            # 获取梯度
-            grad = group['params'].grad
-            if grad is None:
-                continue
-            with torch.no_grad():
-                # 历史衰减
-                group['m'].mul_(self.beta1).add_(grad, alpha = 1 - self.beta1)
-                group['v'].mul_(self.beta2).addcmul_(grad, grad, value = 1 - self.beta2)
-                # 偏差纠正
-                m_hat = group['m'] / (1 - self.beta1 ** self.t)
-                v_hat = group['v'] / (1 - self.beta2 ** self.t)
-                # 参数更新
-                group['params'].sub_(group['lr'] / (v_hat.sqrt() + self.epsilon) * m_hat).mul_(self.beta3).add_(group['b'],alpha = 1 - self.beta3)
-
-    def zero_grad(self):
-        for group in self.param_groups:
-            if group['params'].grad is not None:
-                group['params'].grad.detach_()
-                group['params'].grad.zero_()
-
-    def refresh(self):
-        for group in self.param_groups:
-            group['m'] = torch.zeros_like(group['params']).detach()
-            group['v'] = torch.zeros_like(group['params']).detach()
-            group['b'] = group['params'].clone().detach()
-            self.t = 0
-
-class SimpleAdamOptimizer:
-    "简单的自适应矩估计优化器"
-    def __init__(self, params, betas, eps):
-        self.beta1 = betas[0]
-        self.beta2 = betas[1]
-        self.epsilon = eps
-        self.t = 0
-        self.param_groups = []
-        for p in params:
-            self.param_groups.append({
-                'params': p, 
-                'lr'    : 0.0, 
-                'm'     : torch.zeros_like(p).detach(), 
-                'v'     : torch.zeros_like(p).detach()
-            })
-
-    def step(self):
-        self.t += 1
-        for group in self.param_groups:
-            # 获取梯度
-            grad = group['params'].grad
-            if grad is None:
-                continue
-            grad[grad!=grad] = 0.0
-            grad[grad>100] = 100.0
-            grad[grad<-100] = -100.0
-            with torch.no_grad():
-                # 历史衰减
-                group['m'].mul_(self.beta1).add_(grad, alpha = 1 - self.beta1)
-                group['v'].mul_(self.beta2).addcmul_(grad, grad, value = 1 - self.beta2)
-                # 偏差纠正
-                m_hat = group['m'] / (1 - self.beta1 ** self.t)
-                v_hat = group['v'] / (1 - self.beta2 ** self.t)
-                # 参数更新
-                group['params'].sub_(group['lr'] / (v_hat.sqrt() + self.epsilon) * m_hat)
-
-    def zero_grad(self):
-        for group in self.param_groups:
-            if group['params'].grad is not None:
-                group['params'].grad.detach_()
-                group['params'].grad.zero_()
-                
-def get_lrate(start_step,total_step,lr_from,lr_to,transition,enable_wave):
-    assert transition > 0 and transition % 2 == 0, "Need transition lt 0 and transition mod 2 eq 0."
-    mid_transition = transition // 2
-    half_lr_gap = (lr_to - lr_from)/2
-    if total_step >= start_step + transition:
-        ret = lr_to
-    elif total_step < start_step + mid_transition:
-        ret = lr_from + half_lr_gap * (total_step - start_step)**2 / mid_transition**2
-    else:
-        ret =  lr_to - half_lr_gap * (start_step + transition - total_step)**2 / mid_transition**2
-    #最后的时候震荡，否则有危害
-    if ret != lr_to or enable_wave == False or lr_to > 2e-4:
-        return ret
-    else:
-        return ret + np.sin((total_step - start_step) * np.pi / mid_transition) * lr_to * 0.9
-
-record = {
-    "loss_line" : [],
-    "lr_line" : []
-}
-
-class OptimizerWrapper:
-    def __init__(self, optimizer, warm_up, lr, enable_wave = False):
-        self.lr_from   = 0                 #初始学习率
-        self.lr_to     = lr                #目标学习率
-        self.warm_up   = warm_up           #预热步数
-        self.start_step= 0                 #起始步数
-        self.total_step= 0                 #总步数
-        self.optimizer = optimizer         #优化器，用于执行梯度下降
-        self.enable_wave = enable_wave     #学习率波动
-
-    def update(self):
-        global record
-        #设置优化器中每个参数组的学习率并执行梯度下降
-        lrate = self.lrate()
-        record["lr_line"] += [lrate]
-        for parameters in self.optimizer.param_groups:
-            parameters['lr'] = lrate
-        self.optimizer.step()
-        self.optimizer.zero_grad()
-
-    def lrate(self):
-        self.total_step += 1
-        return get_lrate(
-            self.start_step,
-            self.total_step,
-            self.lr_from,
-            self.lr_to,
-            self.warm_up,
-            self.enable_wave)
-        
-    def set_lrate(self,lrate,transition):
-        self.lr_from = self.lr_to
-        self.lr_to = lrate
-        self.warm_up = transition
-        self.start_step = self.total_step
-
-stop = False
-pause = False
-
-def run_epoch(model,data_iter,caculate_size,loss_f,optimizer,epoch,use_amp):
-    global stop
-    global pause
-    global record
-    for step, batch in enumerate(data_iter):
-        if stop:
-            break
-        while pause:
-            time.sleep(0.5)
-        total_loss = 0
-        t_start = time.time()
-        for i in range(0,batch.query.size(0),caculate_size):
-            if use_amp:
-                with torch.amp.autocast("cuda"):
-                    model_output = model(batch.query[i:i+caculate_size], batch.q_mask[i:i+caculate_size])
-                    loss = loss_f(torch.log(F.softmax(model_output,dim=-1).mul(0.99).add(5e-3)).view(-1,model_output.size(-1)),
-                                batch.label[i:i+caculate_size].reshape(-1))/ batch.ntokens
-                    loss.backward()
-                    total_loss += float(loss) * batch.ntokens
-            else:
-                model_output = model(batch.query[i:i+caculate_size], batch.q_mask[i:i+caculate_size])
-                loss = loss_f(torch.log(F.softmax(model_output,dim=-1).mul(0.99).add(5e-3)).view(-1,model_output.size(-1)),
-                            batch.label[i:i+caculate_size].reshape(-1))/ batch.ntokens
-                loss.backward()
-                total_loss += float(loss) * batch.ntokens
-        optimizer.update()
-        mean_loss = total_loss/batch.ntokens
-        record["loss_line"] += [mean_loss]
-        t_end = time.time()
-        print('\repoch:',epoch,'\tstep:',step,'\tloss:',str(mean_loss)[:5],'\tspeed:',str(batch.ntokens/(t_end - t_start))[:7],'tokens/s',end = ' '*20)
-
-            
-#训练函数以服务模式运行，可以随时手动调整
-def train(model,data_generator,batch_size,caculate_size,loss_f,optimizer,use_amp):
-    global stop
-    epoch = 0
-    while(True):
-        if stop:
-            break
-        run_epoch(model,data_generator(batch_size),caculate_size,loss_f,optimizer,epoch,use_amp)
-        epoch += 1
-
-#启动训练服务
-def train_server_start(model,generator_batch_pair,split_n,loss_f,optimizer,use_amp = False):
-    assert generator_batch_pair[1] % split_n == 0, "Need batch_size mod split_n eq 0."
-    data_generator,batch_size = generator_batch_pair
-    thread = threading.Thread(target=train,args=(model,data_generator,batch_size,batch_size//split_n,loss_f,optimizer,use_amp))
-    thread.start()
-
-def TOGGLE():
-    global pause
-    pause = not pause
-    print("pause:",pause)
-    
-def STOP():
-    global stop
-    stop = True
-
-#贪婪解码
-def greedy_decode(model,inputs,out_length):
-    if model.model_type == "generator":
-        for _ in range(out_length):
-            query = model.embedding(inputs)
-            prob_dist = model.projector(model.encoder(query,inputs==inputs)[:,-1:,:])
-            next_token = torch.max(prob_dist, dim = -1)[1]
-            inputs = torch.cat([inputs,next_token.to(inputs.device)], dim=-1)
-        return inputs
-
-def El_greedy_decode(model,inputs,out_length):
-    if model.model_type == "generator":
-        assert len(inputs[0]) > 1, "初始序列长度必须大于1，与增量续写进行区分"
-        query = model.embedding(inputs)
-        prob_dist = model.projector(model.encoder(query,inputs==inputs)[:,-1:,:])
-        next_token = torch.max(prob_dist, dim = -1)[1]
-        inputs = torch.cat([inputs,next_token.to(inputs.device)], dim=-1)
-        for _ in range(0,out_length-1,1):
-            query = model.embedding(inputs[:,[-1]])
-            prob_dist = model.projector(model.encoder(query,(inputs==inputs)[:,[-1]])[:,-1:,:])
-            next_token = torch.max(prob_dist, dim = -1)[1]
-            inputs = torch.cat([inputs,next_token.to(inputs.device)], dim=-1)
-        return inputs
-
-#概率解码
-def sampling_decode(model,inputs,out_length):
-    if model.model_type == "generator":
-        for _ in range(out_length):
-            query = model.embedding(inputs)
-            prob_dist = model.projector(model.encoder(query,inputs==inputs)[:,-1,:])
-            next_token = torch.multinomial(F.softmax(prob_dist, dim = -1), num_samples = 1)
-            inputs = torch.cat([inputs,next_token.to(inputs.device)], dim=-1)
-        return inputs
-
-def El_sampling_decode(model,inputs,out_length):
-    if model.model_type == "generator":
-        assert len(inputs[0]) > 1, "初始序列长度必须大于1，与增量续写进行区分"
-        query = model.embedding(inputs)
-        prob_dist = model.projector(model.encoder(query,inputs==inputs)[:,-1,:])
-        next_token = torch.multinomial(F.softmax(prob_dist, dim = -1), num_samples = 1)
-        inputs = torch.cat([inputs,next_token.to(inputs.device)], dim=-1)
-        for _ in range(0,out_length-1,1):
-            query = model.embedding(inputs[:,[-1]])
-            prob_dist = model.projector(model.encoder(query,(inputs==inputs)[:,[-1]])[:,-1,:])
-            next_token = torch.multinomial(F.softmax(prob_dist, dim = -1), num_samples = 1)
-            inputs = torch.cat([inputs,next_token.to(inputs.device)], dim=-1)
-        return inputs
-
-#更可控的文本续写工具
-def text_continue(model,inputs,out_length,repeat_penalty_value,temperature,decay=0.98):
-    if model.model_type == "generator":
-        repeat_penalty = None
-        for _ in range(out_length):
-            query = model.embedding(inputs)
-            prob_dist = model.projector(model.encoder(query,inputs==inputs)[:,-1,:])
-            if repeat_penalty is None:
-                repeat_penalty = torch.zeros_like(prob_dist, device=inputs.device)
-                for index in range(inputs.size(1)):
-                    for line in range(inputs.size(0)):
-                        repeat_penalty[line][inputs[line][index]] -= repeat_penalty_value
-                    repeat_penalty *= decay
-            else:
-                repeat_penalty *= decay
-            prob_dist += repeat_penalty
-            next_token = torch.multinomial(F.softmax(prob_dist/temperature, dim = -1), num_samples = 1)
-            inputs = torch.cat([inputs,next_token.to(inputs.device)], dim=-1)
-            for i in range(next_token.size(0)):
-                repeat_penalty[i][next_token[i]] -= repeat_penalty_value
-        return inputs
-
-def El_text_continue(model,inputs,out_length,repeat_penalty_value,temperature,decay=0.98):
-    if model.model_type == "generator":
-        assert len(inputs[0]) > 1, "初始序列长度必须大于1，与增量续写进行区分"
-        query = model.embedding(inputs)
-        prob_dist = model.projector(model.encoder(query,inputs==inputs)[:,-1,:])
-        repeat_penalty = torch.zeros_like(prob_dist, device=inputs.device)
-        for index in range(inputs.size(1)):
-            for line in range(inputs.size(0)):
-                repeat_penalty[line][inputs[line][index]] -= repeat_penalty_value
-            repeat_penalty *= decay
-        prob_dist += repeat_penalty
-        next_token = torch.multinomial(F.softmax(prob_dist/temperature, dim = -1), num_samples = 1)
-        inputs = torch.cat([inputs,next_token.to(inputs.device)], dim=-1)
-        for i in range(next_token.size(0)):
-            repeat_penalty[i][next_token[i]] -= repeat_penalty_value
-        for _ in range(0,out_length-1,1):
-            query = model.embedding(inputs[:,[-1]])
-            prob_dist = model.projector(model.encoder(query,(inputs==inputs)[:,[-1]])[:,-1,:])
-            repeat_penalty *= decay
-            prob_dist += repeat_penalty
-            next_token = torch.multinomial(F.softmax(prob_dist/temperature, dim = -1), num_samples = 1)
-            inputs = torch.cat([inputs,next_token.to(inputs.device)], dim=-1)
-            for i in range(next_token.size(0)):
-                repeat_penalty[i][next_token[i]] -= repeat_penalty_value
-        return inputs
-
-def El_text_continue_stream(model,inputs,out_length,repeat_penalty_value,temperature,decay=0.98):
-    if model.model_type == "generator":
-        assert len(inputs[0]) > 1, "初始序列长度必须大于1，与增量续写进行区分"
-        query = model.embedding(inputs)
-        prob_dist = model.projector(model.encoder(query,inputs==inputs)[:,-1,:])
-        repeat_penalty = torch.zeros_like(prob_dist, device=inputs.device)
-        for index in range(inputs.size(1)):
-            for line in range(inputs.size(0)):
-                repeat_penalty[line][inputs[line][index]] -= repeat_penalty_value
-            repeat_penalty *= decay
-        prob_dist += repeat_penalty
-        next_token = torch.multinomial(F.softmax(prob_dist/temperature, dim = -1), num_samples = 1)
-        inputs = torch.cat([inputs,next_token.to(inputs.device)], dim=-1)[:,-4:]
-        yield inputs
-        for i in range(next_token.size(0)):
-            repeat_penalty[i][next_token[i]] -= repeat_penalty_value
-        for _ in range(0,out_length-1,1):
-            query = model.embedding(inputs[:,[-1]])
-            prob_dist = model.projector(model.encoder(query,(inputs==inputs)[:,[-1]])[:,-1,:])
-            repeat_penalty *= decay
-            prob_dist += repeat_penalty
-            next_token = torch.multinomial(F.softmax(prob_dist/temperature, dim = -1), num_samples = 1)
-            inputs = torch.cat([inputs,next_token.to(inputs.device)], dim=-1)[:,-4:] #留下最后4个字就足够了(utf-8最长是4字节)
-            for i in range(next_token.size(0)):
-                repeat_penalty[i][next_token[i]] -= repeat_penalty_value
-            yield inputs
-
-#值函数，给基于蒙特卡洛树的续写用
-def text_continue_value(model,inputs,out_length,repeat_penalty,repeat_penalty_value,temperature,decay):
-    if model.model_type == "generator":
-        ret = 0
-        assert len(inputs[0]) > 1,"初始序列长度必须大于1，与增量续写进行区分"
-        query = model.embedding(inputs)
-        prob_dist = model.projector(model.encoder(query,inputs==inputs)[:,-1,:])
-        prob_dist += repeat_penalty
-        repeat_penalty *= decay
-        prob_dist = F.softmax(prob_dist/temperature, dim = -1)
-        next_token = torch.multinomial(prob_dist, num_samples = 1)
-        inputs = torch.cat([inputs,next_token.to(inputs.device)], dim=-1)
-        for i in range(next_token.size(0)):
-            repeat_penalty[i][next_token[i]] -= repeat_penalty_value
-            ret += prob_dist[i,next_token[i]]
-        for _ in range(0,out_length-1,1):
-            query = model.embedding(inputs[:,[-1]])
-            prob_dist = model.projector(model.encoder(query,(inputs==inputs)[:,[-1]])[:,-1,:])
-            prob_dist += repeat_penalty
-            repeat_penalty *= decay
-            prob_dist = F.softmax(prob_dist/temperature, dim = -1)
-            next_token = torch.multinomial(prob_dist, num_samples = 1)
-            inputs = torch.cat([inputs,next_token.to(inputs.device)], dim=-1)
-            for i in range(next_token.size(0)):
-                repeat_penalty[i][next_token[i]] -= repeat_penalty_value
-                ret += prob_dist[i,next_token[i]]
-        return ret
-        
-#基于蒙特卡洛树的续写
-def MC_continue(model,inputs,out_length,repeat_penalty_value,temperature,try_n,acc_n,deep_n,decay=0.98):
-    if model.model_type == "generator":
-        repeat_penalty = None
-        assert inputs.dim() == 1, "不支持并行续写！Need inputs.dim eq 1"
-        #复制多份进行树搜索
-        values = [0] * try_n
-        inputs = inputs.repeat(try_n,1)
-        query = model.embedding(inputs)
-        prob_dist = model.projector(model.encoder(query,inputs==inputs)[:,-1,:])
-        repeat_penalty = torch.zeros_like(prob_dist, device=inputs.device)
-        for index in range(inputs.size(1)):
-            for line in range(inputs.size(0)):
-                repeat_penalty[line][inputs[line][index]] -= repeat_penalty_value
-            repeat_penalty *= decay
-        prob_dist += repeat_penalty
-        prob_dist = F.softmax(prob_dist/temperature, dim = -1)
-        next_token = torch.multinomial(prob_dist, num_samples = 1)
-        inputs = torch.cat([inputs,next_token.to(inputs.device)], dim=-1)
-        for i in range(try_n):
-            repeat_penalty[i][next_token[i]] -= repeat_penalty_value
-            values[i] += prob_dist[i,next_token[i]]
-        for cur in range(0,out_length-1,1):
-            query = model.embedding(inputs[:,[-1]])
-            prob_dist = model.projector(model.encoder(query,(inputs==inputs)[:,[-1]])[:,-1,:])
-            repeat_penalty *= decay
-            prob_dist += repeat_penalty
-            prob_dist = F.softmax(prob_dist/temperature, dim = -1)
-            next_token = torch.multinomial(prob_dist, num_samples = 1)
-            inputs = torch.cat([inputs,next_token.to(inputs.device)], dim=-1)
-            for i in range(try_n):
-                repeat_penalty[i][next_token[i]] -= repeat_penalty_value
-                values[i] += prob_dist[i,next_token[i]]
-        max_v = 0.0
-        max_i = 0
-        cnt = 0
-        for test_input,test_repeat_penalty,value in zip(inputs,repeat_penalty,values):
-            test_input = test_input.repeat(acc_n,1)
-            test_repeat_penalty = test_repeat_penalty.repeat(acc_n,1)
-            value += float(text_continue_value(
-                model,test_input,deep_n,test_repeat_penalty,repeat_penalty_value,temperature,decay
-            ))/(acc_n*deep_n)
-            if value > max_v:
-                max_v = value
-                max_i = cnt
-            cnt += 1
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+import time
+import threading
+import copy
+class Batch:
+    def __init__(self,input_sequences):
+        self.data_type = "generator"
+        self.query = input_sequences[...,:-1]
+        self.label = input_sequences[...,1:]
+        self.q_mask = self.query != 0
+        self.ntokens = float((self.label != 0).sum())
+
+#交叉熵损失，“0”填充特殊处理
+class CrossEntropyLoss(nn.Module):
+    def __init__(self):
+        super(CrossEntropyLoss, self).__init__()
+        # 使用KL散度损失函数（接受对数概率分布x和概率分布y,并不是简单的KL散度计算）
+        self.criterion = nn.KLDivLoss(reduction='sum')
+        
+    def forward(self, model_output_dist, target_sequence):
+        #根据模型输出的分布与标签的分布计算交叉熵损失
+        #为目标分布分配和模型输出形状、类型一样的空间，默认不追踪梯度，写明更清晰
+        true_dist = torch.zeros_like(model_output_dist,requires_grad=False)
+        #使用置信度填充目标词的位置（true_dist是词表那么长的概率分布）
+        #目标序列升维，target_sequence:[batch*len]->[batch*len,1]
+        #true_dist:[batch*len,vocab]
+        #在vocab的维度上用标签值当作索引，找到对应元素，填充1.0
+        true_dist.scatter_(1, target_sequence.data.unsqueeze(1), 1.0)
+        #将填充位置概率设为0
+        true_dist[:,0] = 0.0
+        #计算模型输出分布与平目标序列标签平滑后的分布之间的交叉熵
+        #model_output_dist是对数概率分布,应由F.log_softmax(self.project(x),dim=-1)产生
+        #但实际上为了压缩softmax的值域已达到自动丢弃异常值的效果，在Generator.Projector进行了特殊实现
+        return self.criterion(model_output_dist, true_dist)
+
+class AdamOptimizerWithBase:
+    "带有Base的自适应矩估计优化器"
+    def __init__(self, params, base, half_life, betas, eps):
+        self.beta1 = betas[0]
+        self.beta2 = betas[1]
+        self.beta3 = (1/2)**(1/half_life)
+        self.epsilon = eps
+        self.t = 0
+        self.param_groups = []
+        for p,b in zip(params,base):
+            self.param_groups.append({
+                'params': p, 
+                'lr'    : 0.0, 
+                'm'     : torch.zeros_like(p).detach(), 
+                'v'     : torch.zeros_like(p).detach(), 
+                'b'     : b.clone().detach()
+            })
+
+    def step(self):
+        self.t += 1
+        for group in self.param_groups:
+            # 获取梯度
+            grad = group['params'].grad
+            if grad is None:
+                continue
+            with torch.no_grad():
+                # 历史衰减
+                group['m'].mul_(self.beta1).add_(grad, alpha = 1 - self.beta1)
+                group['v'].mul_(self.beta2).addcmul_(grad, grad, value = 1 - self.beta2)
+                # 偏差纠正
+                m_hat = group['m'] / (1 - self.beta1 ** self.t)
+                v_hat = group['v'] / (1 - self.beta2 ** self.t)
+                # 参数更新
+                group['params'].sub_(group['lr'] / (v_hat.sqrt() + self.epsilon) * m_hat).mul_(self.beta3).add_(group['b'],alpha = 1 - self.beta3)
+
+    def zero_grad(self):
+        for group in self.param_groups:
+            if group['params'].grad is not None:
+                group['params'].grad.detach_()
+                group['params'].grad.zero_()
+
+    def refresh(self):
+        for group in self.param_groups:
+            group['m'] = torch.zeros_like(group['params']).detach()
+            group['v'] = torch.zeros_like(group['params']).detach()
+            group['b'] = group['params'].clone().detach()
+            self.t = 0
+
+class SimpleAdamOptimizer:
+    "简单的自适应矩估计优化器"
+    def __init__(self, params, betas, eps):
+        self.beta1 = betas[0]
+        self.beta2 = betas[1]
+        self.epsilon = eps
+        self.t = 0
+        self.param_groups = []
+        for p in params:
+            self.param_groups.append({
+                'params': p, 
+                'lr'    : 0.0, 
+                'm'     : torch.zeros_like(p).detach(), 
+                'v'     : torch.zeros_like(p).detach()
+            })
+
+    def step(self):
+        self.t += 1
+        for group in self.param_groups:
+            # 获取梯度
+            grad = group['params'].grad
+            if grad is None:
+                continue
+            grad[grad!=grad] = 0.0
+            grad[grad>100] = 100.0
+            grad[grad<-100] = -100.0
+            with torch.no_grad():
+                # 历史衰减
+                group['m'].mul_(self.beta1).add_(grad, alpha = 1 - self.beta1)
+                group['v'].mul_(self.beta2).addcmul_(grad, grad, value = 1 - self.beta2)
+                # 偏差纠正
+                m_hat = group['m'] / (1 - self.beta1 ** self.t)
+                v_hat = group['v'] / (1 - self.beta2 ** self.t)
+                # 参数更新
+                group['params'].sub_(group['lr'] / (v_hat.sqrt() + self.epsilon) * m_hat)
+
+    def zero_grad(self):
+        for group in self.param_groups:
+            if group['params'].grad is not None:
+                group['params'].grad.detach_()
+                group['params'].grad.zero_()
+                
+def get_lrate(start_step,total_step,lr_from,lr_to,transition,enable_wave):
+    assert transition > 0 and transition % 2 == 0, "Need transition lt 0 and transition mod 2 eq 0."
+    mid_transition = transition // 2
+    half_lr_gap = (lr_to - lr_from)/2
+    if total_step >= start_step + transition:
+        ret = lr_to
+    elif total_step < start_step + mid_transition:
+        ret = lr_from + half_lr_gap * (total_step - start_step)**2 / mid_transition**2
+    else:
+        ret =  lr_to - half_lr_gap * (start_step + transition - total_step)**2 / mid_transition**2
+    #最后的时候震荡，否则有危害
+    if ret != lr_to or enable_wave == False or lr_to > 2e-4:
+        return ret
+    else:
+        return ret + np.sin((total_step - start_step) * np.pi / mid_transition) * lr_to * 0.9
+
+record = {
+    "loss_line" : [],
+    "lr_line" : []
+}
+
+class OptimizerWrapper:
+    def __init__(self, optimizer, warm_up, lr, enable_wave = False):
+        self.lr_from   = 0                 #初始学习率
+        self.lr_to     = lr                #目标学习率
+        self.warm_up   = warm_up           #预热步数
+        self.start_step= 0                 #起始步数
+        self.total_step= 0                 #总步数
+        self.optimizer = optimizer         #优化器，用于执行梯度下降
+        self.enable_wave = enable_wave     #学习率波动
+
+    def update(self):
+        global record
+        #设置优化器中每个参数组的学习率并执行梯度下降
+        lrate = self.lrate()
+        record["lr_line"] += [lrate]
+        for parameters in self.optimizer.param_groups:
+            parameters['lr'] = lrate
+        self.optimizer.step()
+        self.optimizer.zero_grad()
+
+    def lrate(self):
+        self.total_step += 1
+        return get_lrate(
+            self.start_step,
+            self.total_step,
+            self.lr_from,
+            self.lr_to,
+            self.warm_up,
+            self.enable_wave)
+        
+    def set_lrate(self,lrate,transition):
+        self.lr_from = self.lr_to
+        self.lr_to = lrate
+        self.warm_up = transition
+        self.start_step = self.total_step
+
+stop = False
+pause = False
+
+def run_epoch(model,data_iter,caculate_size,loss_f,optimizer,epoch,use_amp):
+    global stop
+    global pause
+    global record
+    for step, batch in enumerate(data_iter):
+        if stop:
+            break
+        while pause:
+            time.sleep(0.5)
+        total_loss = 0
+        t_start = time.time()
+        for i in range(0,batch.query.size(0),caculate_size):
+            if use_amp:
+                with torch.amp.autocast("cuda"):
+                    model_output = model(batch.query[i:i+caculate_size], batch.q_mask[i:i+caculate_size])
+                    loss = loss_f(torch.log(F.softmax(model_output,dim=-1).mul(0.99).add(5e-3)).view(-1,model_output.size(-1)),
+                                batch.label[i:i+caculate_size].reshape(-1))/ batch.ntokens
+                    loss.backward()
+                    total_loss += float(loss) * batch.ntokens
+            else:
+                model_output = model(batch.query[i:i+caculate_size], batch.q_mask[i:i+caculate_size])
+                loss = loss_f(torch.log(F.softmax(model_output,dim=-1).mul(0.99).add(5e-3)).view(-1,model_output.size(-1)),
+                            batch.label[i:i+caculate_size].reshape(-1))/ batch.ntokens
+                loss.backward()
+                total_loss += float(loss) * batch.ntokens
+        optimizer.update()
+        mean_loss = total_loss/batch.ntokens
+        record["loss_line"] += [mean_loss]
+        t_end = time.time()
+        print('\repoch:',epoch,'\tstep:',step,'\tloss:',str(mean_loss)[:5],'\tspeed:',str(batch.ntokens/(t_end - t_start))[:7],'tokens/s',end = ' '*20)
+
+            
+#训练函数以服务模式运行，可以随时手动调整
+def train(model,data_generator,batch_size,caculate_size,loss_f,optimizer,use_amp):
+    global stop
+    epoch = 0
+    while(True):
+        if stop:
+            break
+        run_epoch(model,data_generator(batch_size),caculate_size,loss_f,optimizer,epoch,use_amp)
+        epoch += 1
+
+#启动训练服务
+def train_server_start(model,generator_batch_pair,split_n,loss_f,optimizer,use_amp = False):
+    assert generator_batch_pair[1] % split_n == 0, "Need batch_size mod split_n eq 0."
+    data_generator,batch_size = generator_batch_pair
+    thread = threading.Thread(target=train,args=(model,data_generator,batch_size,batch_size//split_n,loss_f,optimizer,use_amp))
+    thread.start()
+
+def TOGGLE():
+    global pause
+    pause = not pause
+    print("pause:",pause)
+    
+def STOP():
+    global stop
+    stop = True
+
+#贪婪解码
+def greedy_decode(model,inputs,out_length):
+    if model.model_type == "generator":
+        for _ in range(out_length):
+            query = model.embedding(inputs)
+            prob_dist = model.projector(model.encoder(query,inputs==inputs)[:,-1:,:])
+            next_token = torch.max(prob_dist, dim = -1)[1]
+            inputs = torch.cat([inputs,next_token.to(inputs.device)], dim=-1)
+        return inputs
+
+def El_greedy_decode(model,inputs,out_length):
+    if model.model_type == "generator":
+        assert len(inputs[0]) > 1, "初始序列长度必须大于1，与增量续写进行区分"
+        query = model.embedding(inputs)
+        prob_dist = model.projector(model.encoder(query,inputs==inputs)[:,-1:,:])
+        next_token = torch.max(prob_dist, dim = -1)[1]
+        inputs = torch.cat([inputs,next_token.to(inputs.device)], dim=-1)
+        for _ in range(0,out_length-1,1):
+            query = model.embedding(inputs[:,[-1]])
+            prob_dist = model.projector(model.encoder(query,(inputs==inputs)[:,[-1]])[:,-1:,:])
+            next_token = torch.max(prob_dist, dim = -1)[1]
+            inputs = torch.cat([inputs,next_token.to(inputs.device)], dim=-1)
+        return inputs
+
+#概率解码
+def sampling_decode(model,inputs,out_length):
+    if model.model_type == "generator":
+        for _ in range(out_length):
+            query = model.embedding(inputs)
+            prob_dist = model.projector(model.encoder(query,inputs==inputs)[:,-1,:])
+            next_token = torch.multinomial(F.softmax(prob_dist, dim = -1), num_samples = 1)
+            inputs = torch.cat([inputs,next_token.to(inputs.device)], dim=-1)
+        return inputs
+
+def El_sampling_decode(model,inputs,out_length):
+    if model.model_type == "generator":
+        assert len(inputs[0]) > 1, "初始序列长度必须大于1，与增量续写进行区分"
+        query = model.embedding(inputs)
+        prob_dist = model.projector(model.encoder(query,inputs==inputs)[:,-1,:])
+        next_token = torch.multinomial(F.softmax(prob_dist, dim = -1), num_samples = 1)
+        inputs = torch.cat([inputs,next_token.to(inputs.device)], dim=-1)
+        for _ in range(0,out_length-1,1):
+            query = model.embedding(inputs[:,[-1]])
+            prob_dist = model.projector(model.encoder(query,(inputs==inputs)[:,[-1]])[:,-1,:])
+            next_token = torch.multinomial(F.softmax(prob_dist, dim = -1), num_samples = 1)
+            inputs = torch.cat([inputs,next_token.to(inputs.device)], dim=-1)
+        return inputs
+
+#更可控的文本续写工具
+def text_continue(model,inputs,out_length,repeat_penalty_value,temperature,decay=0.98):
+    if model.model_type == "generator":
+        repeat_penalty = None
+        for _ in range(out_length):
+            query = model.embedding(inputs)
+            prob_dist = model.projector(model.encoder(query,inputs==inputs)[:,-1,:])
+            if repeat_penalty is None:
+                repeat_penalty = torch.zeros_like(prob_dist, device=inputs.device)
+                for index in range(inputs.size(1)):
+                    for line in range(inputs.size(0)):
+                        repeat_penalty[line][inputs[line][index]] -= repeat_penalty_value
+                    repeat_penalty *= decay
+            else:
+                repeat_penalty *= decay
+            prob_dist += repeat_penalty
+            next_token = torch.multinomial(F.softmax(prob_dist/temperature, dim = -1), num_samples = 1)
+            inputs = torch.cat([inputs,next_token.to(inputs.device)], dim=-1)
+            for i in range(next_token.size(0)):
+                repeat_penalty[i][next_token[i]] -= repeat_penalty_value
+        return inputs
+
+def El_text_continue(model,inputs,out_length,repeat_penalty_value,temperature,decay=0.98):
+    if model.model_type == "generator":
+        assert len(inputs[0]) > 1, "初始序列长度必须大于1，与增量续写进行区分"
+        query = model.embedding(inputs)
+        prob_dist = model.projector(model.encoder(query,inputs==inputs)[:,-1,:])
+        repeat_penalty = torch.zeros_like(prob_dist, device=inputs.device)
+        for index in range(inputs.size(1)):
+            for line in range(inputs.size(0)):
+                repeat_penalty[line][inputs[line][index]] -= repeat_penalty_value
+            repeat_penalty *= decay
+        prob_dist += repeat_penalty
+        next_token = torch.multinomial(F.softmax(prob_dist/temperature, dim = -1), num_samples = 1)
+        inputs = torch.cat([inputs,next_token.to(inputs.device)], dim=-1)
+        for i in range(next_token.size(0)):
+            repeat_penalty[i][next_token[i]] -= repeat_penalty_value
+        for _ in range(0,out_length-1,1):
+            query = model.embedding(inputs[:,[-1]])
+            prob_dist = model.projector(model.encoder(query,(inputs==inputs)[:,[-1]])[:,-1,:])
+            repeat_penalty *= decay
+            prob_dist += repeat_penalty
+            next_token = torch.multinomial(F.softmax(prob_dist/temperature, dim = -1), num_samples = 1)
+            inputs = torch.cat([inputs,next_token.to(inputs.device)], dim=-1)
+            for i in range(next_token.size(0)):
+                repeat_penalty[i][next_token[i]] -= repeat_penalty_value
+        return inputs
+
+def El_text_continue_stream(model,inputs,out_length,repeat_penalty_value,temperature,decay=0.98,session_id='0'):
+    if model.model_type == "generator":
+        assert len(inputs[0]) > 1, "初始序列长度必须大于1，与增量续写进行区分"
+        query = model.embedding(inputs)
+        prob_dist = model.projector(model.encoder(query,inputs==inputs,session_id)[:,-1,:])
+        repeat_penalty = torch.zeros_like(prob_dist, device=inputs.device)
+        for index in range(inputs.size(1)):
+            for line in range(inputs.size(0)):
+                repeat_penalty[line][inputs[line][index]] -= repeat_penalty_value
+            repeat_penalty *= decay
+        prob_dist += repeat_penalty
+        next_token = torch.multinomial(F.softmax(prob_dist/temperature, dim = -1), num_samples = 1)
+        inputs = torch.cat([inputs,next_token.to(inputs.device)], dim=-1)[:,-4:]
+        yield inputs
+        for i in range(next_token.size(0)):
+            repeat_penalty[i][next_token[i]] -= repeat_penalty_value
+        for _ in range(0,out_length-1,1):
+            query = model.embedding(inputs[:,[-1]])
+            prob_dist = model.projector(model.encoder(query,(inputs==inputs)[:,[-1]],session_id)[:,-1,:])
+            repeat_penalty *= decay
+            prob_dist += repeat_penalty
+            next_token = torch.multinomial(F.softmax(prob_dist/temperature, dim = -1), num_samples = 1)
+            inputs = torch.cat([inputs,next_token.to(inputs.device)], dim=-1)[:,-4:] #留下最后4个字就足够了(utf-8最长是4字节)
+            for i in range(next_token.size(0)):
+                repeat_penalty[i][next_token[i]] -= repeat_penalty_value
+            yield inputs
+
+#值函数，给基于蒙特卡洛树的续写用
+def text_continue_value(model,inputs,out_length,repeat_penalty,repeat_penalty_value,temperature,decay):
+    if model.model_type == "generator":
+        ret = 0
+        assert len(inputs[0]) > 1,"初始序列长度必须大于1，与增量续写进行区分"
+        query = model.embedding(inputs)
+        prob_dist = model.projector(model.encoder(query,inputs==inputs)[:,-1,:])
+        prob_dist += repeat_penalty
+        repeat_penalty *= decay
+        prob_dist = F.softmax(prob_dist/temperature, dim = -1)
+        next_token = torch.multinomial(prob_dist, num_samples = 1)
+        inputs = torch.cat([inputs,next_token.to(inputs.device)], dim=-1)
+        for i in range(next_token.size(0)):
+            repeat_penalty[i][next_token[i]] -= repeat_penalty_value
+            ret += prob_dist[i,next_token[i]]
+        for _ in range(0,out_length-1,1):
+            query = model.embedding(inputs[:,[-1]])
+            prob_dist = model.projector(model.encoder(query,(inputs==inputs)[:,[-1]])[:,-1,:])
+            prob_dist += repeat_penalty
+            repeat_penalty *= decay
+            prob_dist = F.softmax(prob_dist/temperature, dim = -1)
+            next_token = torch.multinomial(prob_dist, num_samples = 1)
+            inputs = torch.cat([inputs,next_token.to(inputs.device)], dim=-1)
+            for i in range(next_token.size(0)):
+                repeat_penalty[i][next_token[i]] -= repeat_penalty_value
+                ret += prob_dist[i,next_token[i]]
+        return ret
+        
+#基于蒙特卡洛树的续写
+def MC_continue(model,inputs,out_length,repeat_penalty_value,temperature,try_n,acc_n,deep_n,decay=0.98):
+    if model.model_type == "generator":
+        repeat_penalty = None
+        assert inputs.dim() == 1, "不支持并行续写！Need inputs.dim eq 1"
+        #复制多份进行树搜索
+        values = [0] * try_n
+        inputs = inputs.repeat(try_n,1)
+        query = model.embedding(inputs)
+        prob_dist = model.projector(model.encoder(query,inputs==inputs)[:,-1,:])
+        repeat_penalty = torch.zeros_like(prob_dist, device=inputs.device)
+        for index in range(inputs.size(1)):
+            for line in range(inputs.size(0)):
+                repeat_penalty[line][inputs[line][index]] -= repeat_penalty_value
+            repeat_penalty *= decay
+        prob_dist += repeat_penalty
+        prob_dist = F.softmax(prob_dist/temperature, dim = -1)
+        next_token = torch.multinomial(prob_dist, num_samples = 1)
+        inputs = torch.cat([inputs,next_token.to(inputs.device)], dim=-1)
+        for i in range(try_n):
+            repeat_penalty[i][next_token[i]] -= repeat_penalty_value
+            values[i] += prob_dist[i,next_token[i]]
+        for cur in range(0,out_length-1,1):
+            query = model.embedding(inputs[:,[-1]])
+            prob_dist = model.projector(model.encoder(query,(inputs==inputs)[:,[-1]])[:,-1,:])
+            repeat_penalty *= decay
+            prob_dist += repeat_penalty
+            prob_dist = F.softmax(prob_dist/temperature, dim = -1)
+            next_token = torch.multinomial(prob_dist, num_samples = 1)
+            inputs = torch.cat([inputs,next_token.to(inputs.device)], dim=-1)
+            for i in range(try_n):
+                repeat_penalty[i][next_token[i]] -= repeat_penalty_value
+                values[i] += prob_dist[i,next_token[i]]
+        max_v = 0.0
+        max_i = 0
+        cnt = 0
+        for test_input,test_repeat_penalty,value in zip(inputs,repeat_penalty,values):
+            test_input = test_input.repeat(acc_n,1)
+            test_repeat_penalty = test_repeat_penalty.repeat(acc_n,1)
+            value += float(text_continue_value(
+                model,test_input,deep_n,test_repeat_penalty,repeat_penalty_value,temperature,decay
+            ))/(acc_n*deep_n)
+            if value > max_v:
+                max_v = value
+                max_i = cnt
+            cnt += 1
         return inputs[max_i]