.gitignore

@@ -2,10 +2,4 @@
 /dataset
 /out
 wandb/
-**/*.log
-models/sentence_transformers/
-models/sentence_transformers_cache/
-**/*.pyc
-qwen2-1.7B/
-images/
-cache/
+**/*.log

.vscode/launch.json

@@ -1,102 +0,0 @@
-{
-    "version": "0.2.0",
-    "configurations": [
-        {
-            "name": "Debug Train Pretrain Accelerate",
-            "type": "python",
-            "request": "launch",
-            "program": "${workspaceFolder}/train_pretrain_accelerate.py",
-            "console": "integratedTerminal",
-            "python": "/opt/conda/envs/mini/bin/python",
-            "cwd": "${workspaceFolder}",
-            "env": {
-                "PYTHONPATH": "${workspaceFolder}",
-                "CUDA_VISIBLE_DEVICES": "0"
-            },
-            "justMyCode": false,
-            "stopOnEntry": false,
-            "redirectOutput": true
-        },
-        {
-            "name": "Debug Train Pretrain Accelerate (Multi-GPU)",
-            "type": "python",
-            "request": "launch",
-            "program": "${workspaceFolder}/train_pretrain_accelerate.py",
-            "console": "integratedTerminal",
-            "python": "/opt/conda/envs/mini/bin/python",
-            "args": [
-                "--hidden_size", "512",
-                "--max_seq_len", "512", 
-                "--n_layers", "8",
-                "--batch_size", "8",
-                "--epochs", "1"
-            ],
-            "cwd": "${workspaceFolder}",
-            "env": {
-                "PYTHONPATH": "${workspaceFolder}",
-                "CUDA_VISIBLE_DEVICES": "0,1"
-            },
-            "justMyCode": false,
-            "stopOnEntry": false,
-            "redirectOutput": true
-        },
-        {
-            "name": "Debug Train Pretrain Accelerate (Small Test)",
-            "type": "python",
-            "request": "launch",
-            "program": "${workspaceFolder}/train_pretrain_accelerate.py",
-            "console": "integratedTerminal",
-            "python": "/opt/conda/envs/mini/bin/python",
-            "args": [
-                "--hidden_size", "512",
-                "--max_seq_len", "512", 
-                "--n_layers", "8",
-                "--batch_size", "2",
-                "--epochs", "1",
-                "--log_interval", "10",
-                "--save_interval", "100",
-                "--accumulation_steps", "4"
-            ],
-            "cwd": "${workspaceFolder}",
-            "env": {
-                "PYTHONPATH": "${workspaceFolder}",
-                "CUDA_VISIBLE_DEVICES": "0",
-                "WANDB_MODE": "offline"
-            },
-            "justMyCode": false,
-            "stopOnEntry": false,
-            "redirectOutput": true
-        },
-        {
-            "name": "Debug ExtractDB Comparison",
-            "type": "python",
-            "request": "launch",
-            "program": "${workspaceFolder}/train_pretrain_accelerate.py",
-            "console": "integratedTerminal",
-            "python": "/opt/conda/envs/mini/bin/python",
-            "args": [
-                "--hidden_size", "512",
-                "--max_seq_len", "256", 
-                "--n_layers", "4",
-                "--batch_size", "2",
-                "--epochs", "1",
-                "--log_interval", "10",
-                "--save_interval", "200",
-                "--accumulation_steps", "2",
-                "--comparison_mode",
-                "--knowledge_num", "256",
-                "--knowledge_length", "64",
-                "--comparison_mode"
-            ],
-            "cwd": "${workspaceFolder}",
-            "env": {
-                "PYTHONPATH": "${workspaceFolder}",
-                "CUDA_VISIBLE_DEVICES": "0",
-                "WANDB_MODE": "offline"
-            },
-            "justMyCode": false,
-            "stopOnEntry": false,
-            "redirectOutput": true
-        }
-    ]
-} 

.vscode/settings.json

@@ -1,18 +0,0 @@
-{
-    "python.pythonPath": "/home/iomgaa/miniconda3/envs/accelerate/bin/python",
-    "python.defaultInterpreterPath": "/home/iomgaa/miniconda3/envs/accelerate/bin/python",
-    "python.terminal.activateEnvironment": true,
-    "python.terminal.activateEnvInCurrentTerminal": true,
-    "python.linting.enabled": true,
-    "python.linting.pylintEnabled": false,
-    "python.linting.flake8Enabled": true,
-    "python.formatting.provider": "black",
-    "python.analysis.autoImportCompletions": true,
-    "python.analysis.typeCheckingMode": "off",
-    "files.exclude": {
-        "**/__pycache__": true,
-        "**/*.pyc": true,
-        "**/.git": false,
-        "**/wandb": false
-    }
-} 

model/LMConfig.py

@@ -19,7 +19,6 @@ class LMConfig(PretrainedConfig):
             rope_theta: int = 1e6,
             dropout: float = 0.0,
             flash_attn: bool = True,
-            embeddings_epoch: int = 2,
             ####################################################
             # DB related configurations
             ####################################################
@@ -40,7 +39,6 @@ class LMConfig(PretrainedConfig):
             ####################################################
             knowledge_num: int = 64*64,
             knowledge_length: int = 8,
-            knowledge_dim: int = 128,
             **kwargs,
     ):
         self.dim = dim
@@ -55,7 +53,6 @@ class LMConfig(PretrainedConfig):
         self.rope_theta = rope_theta
         self.dropout = dropout
         self.flash_attn = flash_attn
-        self.embeddings_epoch = embeddings_epoch
         ####################################################
         # DB related configurations
         ####################################################
@@ -75,5 +72,4 @@ class LMConfig(PretrainedConfig):
         ####################################################
         self.knowledge_num = knowledge_num
         self.knowledge_length = knowledge_length
-        self.knowledge_dim = knowledge_dim
         super().__init__(**kwargs)

model/model.py

@@ -11,9 +11,14 @@ import torch.nn.functional as F
 from torch import nn
 from transformers import PreTrainedModel
 from transformers.modeling_outputs import CausalLMOutputWithPast
+from torch import nn, einsum
+from einops import rearrange, repeat
+
+def exists(val):
+    return val is not None
 
 
-
+# RMSNorm 类定义了一个用于归一化输入张量的模块。
 class RMSNorm(torch.nn.Module):
     def __init__(self, dim: int, eps: float = 1e-6):
         super().__init__()
@@ -26,7 +31,7 @@ class RMSNorm(torch.nn.Module):
     def forward(self, x):
         return self.weight * self._norm(x.float()).type_as(x)
 
-
+# precompute_pos_cis 函数用于预计算位置编码（复数版本）。
 def precompute_pos_cis(dim: int, end: int = int(32 * 1024), theta: float = 1e6):
     freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
     t = torch.arange(end, device=freqs.device)  # type: ignore
@@ -34,7 +39,7 @@ def precompute_pos_cis(dim: int, end: int = int(32 * 1024), theta: float = 1e6):
     pos_cis = torch.polar(torch.ones_like(freqs), freqs)  # complex64
     return pos_cis
 
-
+# apply_rotary_emb 函数用于应用旋转位置编码（复数版本）。
 def apply_rotary_emb(xq, xk, pos_cis):
     def unite_shape(pos_cis, x):
         ndim = x.ndim
@@ -50,194 +55,200 @@ def apply_rotary_emb(xq, xk, pos_cis):
     xk_out = torch.view_as_real(xk_ * pos_cis).flatten(3)
     return xq_out.type_as(xq), xk_out.type_as(xk)
 
-class KnowledgeDataset(nn.Module):
-    def __init__(self, params, tok_embeddings, is_train=True):
+# precompute_pos_cis_real 函数用于预计算位置编码（实数版本）。
+def precompute_pos_cis_real(dim: int, end: int = int(32 * 1024), theta: float = 1e6):
+    """使用实数张量实现位置编码，避免使用复数张量
+
+    这个函数与precompute_pos_cis完全等价，但使用实数张量而非复数张量。
+    原始函数生成形状为[seq_len, dim//2]的复数张量，其中实部全为1，虚部为旋转角度。
+    这个函数生成形状为[seq_len, dim]的实数张量，其中偶数索引是cos(角度)，奇数索引是sin(角度)。
+    """
+    # 确保dim是偶数
+    if dim % 2 != 0:
+        raise ValueError(f"维度必须是偶数，但得到了 {dim}")
+
+    # 复制原始函数的频率计算逻辑
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
+    t = torch.arange(end, device=freqs.device)
+    freqs = torch.outer(t, freqs).float()
+
+    # 计算cos和sin值
+    # 在复数版本中，pos_cis = torch.polar(torch.ones_like(freqs), freqs)
+    # 等价于 cos(freqs) + i*sin(freqs)
+    cos = torch.cos(freqs)
+    sin = torch.sin(freqs)
+
+    # 创建实数张量，交错排列cos和sin
+    pos_emb = torch.zeros((end, dim), device=freqs.device)
+    pos_emb[:, 0::2] = cos  # 偶数索引放cos
+    pos_emb[:, 1::2] = sin  # 奇数索引放sin
+
+    return pos_emb
+
+# apply_rotary_emb_real 函数用于应用旋转位置编码（实数版本）。
+def apply_rotary_emb_real(xq, xk, pos_emb):
+    """使用实数张量实现旋转位置编码，避免使用复数张量
+
+    这个函数与apply_rotary_emb完全等价，但使用实数张量而非复数张量。
+    原始函数将输入张量转换为复数形式，与位置编码相乘，然后再转回实数形式。
+    这个函数直接使用实数运算实现相同的旋转操作。
+    """
+    # 获取形状信息
+    bsz, seq_len, n_heads, head_dim = xq.shape
+
+    # 确保pos_emb形状正确
+    assert pos_emb.shape[0] >= seq_len, f"位置编码长度 {pos_emb.shape[0]} 小于序列长度 {seq_len}"
+    assert pos_emb.shape[1] == head_dim, f"位置编码维度 {pos_emb.shape[1]} 与头维度 {head_dim} 不匹配"
+
+    # 截取需要的位置编码长度
+    pos_emb = pos_emb[:seq_len]
+
+    # 将pos_emb调整为广播形状 [1, seq_len, 1, head_dim]
+    pos_emb = pos_emb.unsqueeze(0).unsqueeze(2)
+
+    # 将head_dim分成两半
+    half_head_dim = head_dim // 2
+
+    # 提取cos和sin值（偶数索引是cos，奇数索引是sin）
+    cos = pos_emb[..., 0::2]
+    sin = pos_emb[..., 1::2]
+
+    # 将xq和xk重新排列，以便进行旋转操作
+    # 原始复数版本中，xq和xk被重塑为复数张量，其中实部和虚部交错排列
+    # 在实数版本中，我们需要将偶数索引和奇数索引分开处理
+
+    # 分离偶数和奇数索引
+    xq_even = xq[..., 0::2]  # 偶数索引，对应复数的实部
+    xq_odd = xq[..., 1::2]   # 奇数索引，对应复数的虚部
+    xk_even = xk[..., 0::2]
+    xk_odd = xk[..., 1::2]
+
+    # 应用旋转（等价于复数乘法）
+    # (a + bi)(cos + sin*i) = (a*cos - b*sin) + (a*sin + b*cos)i
+    # 其中a是偶数索引，b是奇数索引
+    xq_out_even = xq_even * cos - xq_odd * sin  # 新的偶数索引（实部）
+    xq_out_odd = xq_even * sin + xq_odd * cos   # 新的奇数索引（虚部）
+    xk_out_even = xk_even * cos - xk_odd * sin
+    xk_out_odd = xk_even * sin + xk_odd * cos
+
+    # 重新组合偶数和奇数索引
+    xq_out = torch.zeros_like(xq)
+    xk_out = torch.zeros_like(xk)
+    xq_out[..., 0::2] = xq_out_even
+    xq_out[..., 1::2] = xq_out_odd
+    xk_out[..., 0::2] = xk_out_even
+    xk_out[..., 1::2] = xk_out_odd
+
+    return xq_out.type_as(xq), xk_out.type_as(xk)
+
+# repeat_kv 函数用于重复键值对。
+def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """torch.repeat_interleave(x, dim=2, repeats=n_rep)"""
+    bs, slen, n_kv_heads, head_dim = x.shape
+    if n_rep == 1:
+        return x
+    return (
+        x[:, :, :, None, :]
+        .expand(bs, slen, n_kv_heads, n_rep, head_dim)
+        .reshape(bs, slen, n_kv_heads * n_rep, head_dim)
+    )
+
+
+class Attention(nn.Module):
+    def __init__(self, args: LMConfig):
         super().__init__()
-        self.is_train = is_train
-        self.params = params
-        self.tok_embeddings = tok_embeddings
+        self.n_kv_heads = args.n_heads if args.n_kv_heads is None else args.n_kv_heads
+        assert args.n_heads % self.n_kv_heads == 0
+        self.n_local_heads = args.n_heads
+        self.n_local_kv_heads = self.n_kv_heads
+        self.n_rep = self.n_local_heads // self.n_local_kv_heads
+        self.head_dim = args.dim // args.n_heads
+        self.wq = nn.Linear(args.dim, args.n_heads * self.head_dim, bias=False)
+        self.wk = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False)
+        self.wv = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False)
+        self.wo = nn.Linear(args.n_heads * self.head_dim, args.dim, bias=False)
+        self.attn_dropout = nn.Dropout(args.dropout)
+        self.resid_dropout = nn.Dropout(args.dropout)
+        self.dropout = args.dropout
+        self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention') and args.flash_attn
+        # print("WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0")
+        mask = torch.full((1, 1, args.max_seq_len, args.max_seq_len), float("-inf"))
+        mask = torch.triu(mask, diagonal=1)
+        self.register_buffer("mask", mask, persistent=False)
 
-        # 嵌入参数
-        self.knowledge_dim = params.knowledge_dim
-        self.key_dim = self.knowledge_dim // 2
-        self.to_queries = nn.Sequential(
-                nn.Linear(params.dim, self.knowledge_dim, bias=False),
+    def forward(self,
+                x: torch.Tensor,
+                pos_cis: torch.Tensor,
+                db_value=None):
+        bsz, seq_len, _ = x.shape #bsz: 批量大小, seq_len: 序列长度, _: 隐藏维度
+        xq, xk, xv = self.wq(x), self.wk(x), self.wv(x) #将输入张量x分别通过线性层wq, wk, wv进行变换，得到查询、键和值。
+        xq = xq.view(bsz, seq_len, self.n_local_heads, self.head_dim) #将变换后的张量xq重塑为形状为(bsz, seq_len, n_local_heads, head_dim)的形状。
+        xk = xk.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim) #将变换后的张量xk重塑为形状为(bsz, seq_len, n_local_kv_heads, head_dim)的形状。
+        xv = xv.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim) #将变换后的张量xv重塑为形状为(bsz, seq_len, n_local_kv_heads, head_dim)的形状。
+
+        # 应用旋转位置编码（使用实数版本）
+        xq, xk = apply_rotary_emb_real(xq, xk, pos_cis)
+        # kv_cache实现 REMOVED
+        # if past_key_value is not None:
+        #     xk = torch.cat([past_key_value[0], xk], dim=1)
+        #     xv = torch.cat([past_key_value[1], xv], dim=1)
+        # past_kv = (xk, xv) if use_cache else None
+
+        # 重复键值对
+        xq, xk, xv = (
+            xq.transpose(1, 2),
+            repeat_kv(xk, self.n_rep).transpose(1, 2),
+            repeat_kv(xv, self.n_rep).transpose(1, 2)
         )
 
-        ## 数据库参数
-        self.knowledge_num = params.knowledge_num
-        self.knowledge_length = params.knowledge_length
-        self.keys = nn.Parameter(torch.randn(self.knowledge_num, self.knowledge_dim) * 0.02, requires_grad=True)
-        self.product_key_topk = min(16, self.knowledge_num)
-        
-        # 使用频率统计 - 使用register_buffer以便在GPU/CPU间正确移动
-        self.register_buffer('has_update_keys', torch.zeros(self.knowledge_num))
+        # 如果提供了db_value，根据头的数量调整它的形状并与xv合并
+        if db_value is not None:
+            # 确保db_value的形状与xv兼容，假设db_value形状为[B, N, H, D]
+            if db_value.ndim == 4:  # [B, N, H, D]
+                db_value = db_value.transpose(1, 2)  # -> [B, H, N, D]
 
-        # 知识库存储 - 使用register_buffer因为这是整数索引，不需要梯度
-        self.register_buffer('knowledge_dataset', 
-            torch.randint(low=0, high=params.vocab_size, size=(self.knowledge_num, self.knowledge_length), dtype=torch.long)
-        )
+                # 检查是否需要调整D维度
+                if db_value.shape[-1] != xv.shape[-1]:
+                    # 如果db_value的维度与xv不同，可以添加一个投影层
+                    # 或者在这里使用简单的调整方法
+                    # 这里我们简单地通过均值池化或重复来调整维度
+                    if db_value.shape[-1] > xv.shape[-1]:
+                        # 降维
+                        factor = db_value.shape[-1] // xv.shape[-1]
+                        db_value = db_value.view(bsz, self.n_local_heads, seq_len, factor, xv.shape[-1])
+                        db_value = db_value.mean(dim=3)
+                    else:
+                        # 升维
+                        factor = xv.shape[-1] // db_value.shape[-1]
+                        db_value = db_value.unsqueeze(-1).repeat(1, 1, 1, 1, factor)
+                        db_value = db_value.view(bsz, self.n_local_heads, seq_len, xv.shape[-1])
 
-        # 计算step数目，用于动态调整权重
-        self.step_counter = 0
+                # 将db_value与xv相加或融合
+                # 这里我们简单地将它们相加，但你也可以使用其他融合方法
+                xv = xv + db_value
 
-        self.freeze_embedding = False
+        # 使用Flash Attention
+        if self.flash and seq_len != 1:
+            dropout_p = self.dropout if self.training else 0.0
+            output = F.scaled_dot_product_attention(
+                xq, xk, xv,
+                attn_mask=None,
+                dropout_p=dropout_p,
+                is_causal=True
+            )
+        else:
+            scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(self.head_dim)
+            scores += self.mask[:, :, :seq_len, :seq_len]
+            scores = F.softmax(scores.float(), dim=-1).type_as(xq)
+            scores = self.attn_dropout(scores)
+            output = scores @ xv
+
+        output = output.transpose(1, 2).reshape(bsz, seq_len, -1)
+        output = self.resid_dropout(self.wo(output))
+        return output
 
-        
-
-    def intelligent_selection(self, query, all_scores, all_indices):
-        """智能分层选择策略"""
-        if self.is_train == False:
-            return all_scores, all_indices
-        
-        batch_size = all_scores.size(0)
-        device = all_scores.device
-        dtype = all_scores.dtype
-
-        # 对每个batch进行分层选择
-        enhanced_scores = all_scores.clone()
-        query_features = query.mean(dim=1)  # [batch_size, dim]
-
-        # 预先计算所有候选条目的嵌入（批量优化）
-        all_candidate_indices = torch.cat([all_indices[i] for i in range(batch_size)], dim=0)
-        unique_indices, inverse_indices = torch.unique(all_candidate_indices, return_inverse=True)
-
-        # 批量计算唯一候选条目的嵌入
-        candidate_tokens = self.knowledge_dataset[unique_indices]
-        flat_tokens = candidate_tokens.view(-1)
-        flat_embeddings = self.tok_embeddings(flat_tokens)
-        #获取flat_tokens对应的index
-        pre_update_indices = unique_indices.view(-1)
-        pre_update_embeddings = flat_embeddings.view(
-            len(unique_indices), self.knowledge_length, -1
-        )
-
-        unique_candidate_features = flat_embeddings.view(
-            len(unique_indices), self.knowledge_length, -1
-        ).mean(dim=1)  # [num_unique_candidates, dim]
-            
-        # 归一化候选特征（优化相似度计算）
-        normalized_candidates = F.normalize(unique_candidate_features, dim=-1)
-        normalized_queries = F.normalize(query_features, dim=-1)
-
-        # 收集所有batch的best_tokens
-        batch_best_tokens = []
-        batch_best_tokens_embeddings = []
-
-        for batch_idx in range(batch_size):
-            indices = all_indices[batch_idx]
-                
-            # 获取当前batch候选条目对应的特征索引
-            start_idx = batch_idx * len(indices)
-            end_idx = start_idx + len(indices)
-            batch_inverse_indices = inverse_indices[start_idx:end_idx]
-                
-            # 使用预计算的归一化特征进行优化相似度计算
-            batch_candidate_features = normalized_candidates[batch_inverse_indices]
-            query_feature = normalized_queries[batch_idx]
-                
-            # 使用矩阵乘法计算余弦相似度
-            similarity_scores = torch.mv(batch_candidate_features, query_feature)
-            
-            # 找到最大相似度分数的索引
-            max_similarity_idx = torch.argmax(similarity_scores)
-            
-            # 获取最大相似度对应的候选条目索引
-            best_candidate_idx = indices[max_similarity_idx]
-            
-            # 获取对应的tokens
-            best_tokens = self.knowledge_dataset[best_candidate_idx]
-            best_tokens_embeddings = self.tok_embeddings(best_tokens)
-            
-            # 将当前batch的best_tokens添加到列表中
-            batch_best_tokens.append(best_tokens)
-            batch_best_tokens_embeddings.append(best_tokens_embeddings)
-
-        # 将所有batch的best_tokens堆叠成一个张量
-        # [batch_size, knowledge_length]
-        all_best_tokens = torch.stack(batch_best_tokens, dim=0)
-        all_best_tokens_embeddings = torch.stack(batch_best_tokens_embeddings, dim=0)
-
-        # 获取
-        
-        # 使用重新计算的embeddings更新self.keys
-        if self.is_train:
-            self._update_keys_with_embeddings(pre_update_indices, pre_update_embeddings)
-        
-        # 更新被修改过的key
-        with torch.no_grad():
-            self.has_update_keys[pre_update_indices] = 1
-        
-        return all_best_tokens, all_best_tokens_embeddings
-    
-    def _update_keys_with_embeddings(self, pre_update_indices, pre_update_embeddings):
-        if self.freeze_embedding:
-            return
-        # 使用pre_update_embeddings更新self.keys
-        with torch.no_grad():
-            pre_update_embeddings = pre_update_embeddings.mean(dim=1)  # [337, 512]
-            pre_update_embeddings = self.to_queries(pre_update_embeddings)
-            self.keys[pre_update_indices] = pre_update_embeddings
-    
-    def search_index(self,x):
-        batch_size, seq_len, dim = x.shape
-
-        # collapse sequence dimension by averaging
-        x_flat = x.mean(dim=1)  # [batch_size, dim]
-
-        queries = self.to_queries(x_flat) # [batch_size, 2*dim_key]
-        # queries = queries.reshape(batch_size, 2, self.key_dim)
-        # queries = queries.permute(1, 0, 2)
-
-        # 2. 计算queries与keys的相似度
-        sim = torch.einsum('b d, k d -> b k', queries, self.keys)
-
-        # 3. 在两个子空间分别做top-k
-        scores_and_indices = sim.topk(self.product_key_topk, dim=-1)
-        scores, indices = scores_and_indices[0], scores_and_indices[1]
-
-        # 5. 应用智能分层选择策略
-        best_tokens, best_tokens_embeddings = self.intelligent_selection(x, scores, indices)
-
-        # 6. 更新1%的keys
-        if self.is_train:
-            # 获取未更新过的keys的索引
-            not_updated_indices = torch.where(self.has_update_keys == 0)[0]
-                
-            # 如果有未更新的keys，随机选择num_update_keys个进行更新
-            if len(not_updated_indices) > 0:
-                num_update_keys = int(self.knowledge_num * 0.01)
-                perm = torch.randperm(len(not_updated_indices))[:num_update_keys]
-                perm_num = perm.shape[0]
-                pre_update_indices = not_updated_indices[perm]
-                pre_update_tokens = self.knowledge_dataset[pre_update_indices]
-                pre_update_embeddings = self.tok_embeddings(pre_update_tokens.view(-1))
-                pre_update_embeddings = pre_update_embeddings.view(perm_num, self.knowledge_length, -1)
-                self._update_keys_with_embeddings(pre_update_indices, pre_update_embeddings)
-                # 更新被修改过的key
-                with torch.no_grad():
-                    self.has_update_keys[pre_update_indices] = 1
-            else:
-                print("all keys are updated")
-                # 重置所有keys的更新状态
-                self.has_update_keys.zero_()
-                # 重新获取所有可更新的索引
-                not_updated_indices = torch.arange(len(self.has_update_keys), device=self.has_update_keys.device)
-                num_update_keys = int(self.knowledge_num * 0.01)
-                perm = torch.randperm(len(not_updated_indices))[:num_update_keys]
-                pre_update_indices = not_updated_indices[perm]
-                pre_update_tokens = self.knowledge_dataset[pre_update_indices]
-                pre_update_embeddings = self.tok_embeddings(pre_update_tokens.view(-1))
-                pre_update_embeddings = pre_update_embeddings.view(num_update_keys, self.knowledge_length, -1)
-                self._update_keys_with_embeddings(pre_update_indices, pre_update_embeddings)
-                # 更新被修改过的key
-                with torch.no_grad():
-                    self.has_update_keys[pre_update_indices] = 1
-            
 
 
-        
-        return best_tokens, best_tokens_embeddings
 
 class CrossAttention(nn.Module):
     def __init__(
@@ -284,58 +295,6 @@ class CrossAttention(nn.Module):
 
         return context
 
-class Attention(nn.Module):
-    def __init__(self, args: LMConfig):
-        super().__init__()
-        self.n_kv_heads = args.n_heads if args.n_kv_heads is None else args.n_kv_heads
-        assert args.n_heads % self.n_kv_heads == 0
-        self.n_local_heads = args.n_heads
-        self.n_local_kv_heads = self.n_kv_heads
-        self.n_rep = self.n_local_heads // self.n_local_kv_heads
-        self.head_dim = args.dim // args.n_heads
-        self.wq = nn.Linear(args.dim, args.n_heads * self.head_dim, bias=False)
-        self.wk = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False)
-        self.wv = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False)
-        self.wo = nn.Linear(args.n_heads * self.head_dim, args.dim, bias=False)
-        self.attn_dropout = nn.Dropout(args.dropout)
-        self.resid_dropout = nn.Dropout(args.dropout)
-        self.dropout = args.dropout
-        self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention') and args.flash_attn
-        # print("WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0")
-        mask = torch.full((1, 1, args.max_seq_len, args.max_seq_len), float("-inf"))
-        mask = torch.triu(mask, diagonal=1)
-        self.register_buffer("mask", mask, persistent=False)
-
-    def forward(self,
-                x: torch.Tensor,
-                pos_cis: torch.Tensor):
-        bsz, seq_len, _ = x.shape
-        xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
-        xq = xq.view(bsz, seq_len, self.n_local_heads, self.head_dim)
-        xk = xk.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)
-        xv = xv.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)
-
-        xq, xk = apply_rotary_emb(xq, xk, pos_cis)
-        if self.flash and seq_len != 1:
-            dropout_p = self.dropout if self.training else 0.0
-            output = F.scaled_dot_product_attention(
-                xq, xk, xv,
-                attn_mask=None,
-                dropout_p=dropout_p,
-                is_causal=True
-            )
-        else:
-            scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(self.head_dim)
-            scores += self.mask[:, :, :seq_len, :seq_len]
-            scores = F.softmax(scores.float(), dim=-1).type_as(xq)
-            scores = self.attn_dropout(scores)
-            output = scores @ xv
-
-        output = output.transpose(1, 2).reshape(bsz, seq_len, -1)
-        output = self.resid_dropout(self.wo(output))
-        return output
-
-
 class FeedForward(nn.Module):
     def __init__(self, config: LMConfig):
         super().__init__()
@@ -468,30 +427,169 @@ class MOEFeedForward(nn.Module):
 
 
 class MiniMindBlock(nn.Module):
-    def __init__(self, layer_id: int, config: LMConfig, knowledge_dataset: KnowledgeDataset):
+    def __init__(self, layer_id: int, config: LMConfig):
         super().__init__()
         self.n_heads = config.n_heads
         self.dim = config.dim
         self.head_dim = config.dim // config.n_heads
-        self.self_attention = Attention(config)
-        self.cross_attention = CrossAttention(config)
-        self.knowledge_dataset = knowledge_dataset
+        self.attention = Attention(config)
+        self.cross_att = CrossAttention(config)
 
         self.layer_id = layer_id
         self.attention_norm = RMSNorm(config.dim, eps=config.norm_eps)
         self.ffn_norm = RMSNorm(config.dim, eps=config.norm_eps)
         self.feed_forward = FeedForward(config) if not config.use_moe else MOEFeedForward(config)
 
-    def forward(self, x, pos_cis):
-        h_attn = self.self_attention(
+        # 假设num_experts是已定义的总专家数量的平方根
+
+
+        # 查询生成的参数
+
+
+        # 创建查询生成模块
+        # if weight_down_embed is not None:
+        #     self.to_queries = nn.Sequential(
+        #         nn.Linear(config.dim, self.dim_key * 2, bias=False),
+        #         # nn.Unflatten(2, (2, self.n_heads, self.dim_key))  # 替代Rearrange
+        #     )
+
+        #     # 超参数
+        #     self.product_key_topk = min(16, self.num_keys)  # 确保不超过num_keys
+        #     self.num_experts_per_head_topk = 1  # 最终每个头选取的专家数
+
+    def forward(self, x, db_value, pos_cis):
+        # import pdb;pdb.set_trace()
+        # db_value = None
+
+        # # 如果有weight_down_embed，使用Product Key机制
+        # if self.weight_down_embed is not None:
+        #     # 1. 生成queries
+        #     batch_size, seq_len, dim = x.shape
+
+        #     # collapse sequence dimension by averaging
+        #     x_flat = x.mean(dim=1)  # [batch_size, dim]
+        #     queries = self.to_queries(x_flat)  # [batch_size, 2*dim_key]
+        #     queries = queries.reshape(batch_size, 2, self.dim_key)  # [batch_size, 2, dim_key]
+        #     queries = queries.permute(1, 0, 2)  # [2, batch_size, dim_key]
+
+        #     # 2. 计算queries与keys的相似度
+        #     sim = torch.einsum('p b d, k p d -> p b k', queries, self.keys)
+
+        #     # 3. 在两个子空间分别做top-k
+        #     scores_and_indices = [sim[p].topk(self.product_key_topk, dim=-1) for p in range(2)]
+        #     scores_x, scores_y = scores_and_indices[0][0], scores_and_indices[1][0]
+        #     indices_x, indices_y = scores_and_indices[0][1], scores_and_indices[1][1]
+
+        #     # 4. 组合两个子空间的分数和索引
+        #     all_scores = scores_x.unsqueeze(-1) + scores_y.unsqueeze(-2)
+        #     all_scores = all_scores.view(*all_scores.shape[:-2], -1)
+
+        #     all_indices = (indices_x.unsqueeze(-1) * self.num_keys) + indices_y.unsqueeze(-2)
+        #     all_indices = all_indices.view(*all_indices.shape[:-2], -1)
+
+        #     # 5. 最终top-k选择
+        #     scores, pk_indices = all_scores.topk(self.num_experts_per_head_topk, dim=-1)
+        #     indices = all_indices.gather(-1, pk_indices)
+
+        #     # 6. 从embedding中获取专家值
+
+        #     # 从embedding中获取值
+        #     flat_indices = indices.view(-1)  # 将索引展平为一维张量
+        #     db_values = self.weight_down_embed(flat_indices)
+
+        #     # 重塑回原始形状
+        #     db_value = db_values.view(batch_size, -1, dim)
+
+
+        # 注意力计算
+        h_attn = self.attention(
             self.attention_norm(x),
-            pos_cis
+            pos_cis,
+            db_value=db_value
         )
-        db, db_embeddings = self.knowledge_dataset.search_index(h_attn)
-        h_attn = self.cross_attention(h_attn, db_embeddings)
+
+        h_attn = self.cross_att(h_attn, db_value)
+
+        # 残差连接
         h = x + h_attn
+
+        # 前馈神经网络
         out = h + self.feed_forward(self.ffn_norm(h))
-        return out
+        return out 
+
+class ExtractDB(nn.Module):
+    def __init__(self,params):
+         # 修改专家数量和知识维度，确保能开方
+        super().__init__()
+        self.batch_size = None
+        self.dim = params.dim
+        self.dim_key = self.dim // 2
+        self.knowledge_num = params.knowledge_num  # 100专家，确保是完全平方数
+        # 将knowledge_dim设置为与head_dim相同，以便在attention中直接使用
+        self.head_dim = params.dim // params.n_heads
+        self.knowledge_length = params.knowledge_length
+
+        # 使用register_buffer代替nn.Parameter，避免梯度问题
+        # self.register_buffer('weight_down_embed', torch.randn(self.knowledge_num, self.knowledge_length) * 0.02)
+        self.register_buffer('weight_down_embed',torch.randint(low=0,high=6400, size=(self.knowledge_num, self.knowledge_length),dtype=torch.long))
+    
+    
+
+
+        self.num_keys = int(math.sqrt(self.knowledge_num)) if self.knowledge_num > 0 else 0
+        self.product_key_topk = min(16, self.num_keys)
+        self.keys = nn.Parameter(torch.randn(self.num_keys, 2, self.dim_key) * 0.02)
+        self.num_experts_per_head_topk = 1
+        self.to_queries = nn.Sequential(
+                nn.Linear(params.dim, self.dim_key * 2, bias=False),
+        )
+ 
+    def q_to_k(self,x):
+        # 1. 生成queries
+            self.batch_size, seq_len, dim = x.shape
+
+            # collapse sequence dimension by averaging
+            x_flat = x.mean(dim=1)  # [batch_size, dim]
+
+            queries = self.to_queries(x_flat)  # [batch_size, 2*dim_key]
+            queries = queries.reshape(self.batch_size, 2, self.dim_key)  # [batch_size, 2, dim_key]
+            queries = queries.permute(1, 0, 2)  # [2, batch_size, dim_key]
+
+            # 2. 计算queries与keys的相似度
+            sim = torch.einsum('p b d, k p d -> p b k', queries, self.keys)
+
+            # 3. 在两个子空间分别做top-k
+            scores_and_indices = [sim[p].topk(self.product_key_topk, dim=-1) for p in range(2)]
+            scores_x, scores_y = scores_and_indices[0][0], scores_and_indices[1][0]
+            indices_x, indices_y = scores_and_indices[0][1], scores_and_indices[1][1]
+
+            # 4. 组合两个子空间的分数和索引
+            all_scores = scores_x.unsqueeze(-1) + scores_y.unsqueeze(-2)
+            all_scores = all_scores.view(*all_scores.shape[:-2], -1)
+
+            all_indices = (indices_x.unsqueeze(-1) * self.num_keys) + indices_y.unsqueeze(-2)
+            all_indices = all_indices.view(*all_indices.shape[:-2], -1)
+
+            # 5. 最终top-k选择
+            scores, pk_indices = all_scores.topk(self.num_experts_per_head_topk, dim=-1)
+            indices = all_indices.gather(-1, pk_indices)
+            flat_indices = indices.view(-1)
+            return flat_indices
+
+    def get_data(self, index):
+        # 直接从GPU获取embedding
+        db_values = self.weight_down_embed[index]#变成token了所以是1,后续再过emb
+        # db_value = db_values.view(self.batch_size,-1) 
+        return db_values
+
+    @torch.no_grad()
+    def updata_value(self, k, v):#要加一个从向量返回index的过程
+        # 直接更新buffer上的值 (不需要梯度)
+        v_reshaped = v.view(v.size(0), -1)
+        # 确保数据类型匹配
+        v_reshaped = v_reshaped.to(dtype=self.weight_down_embed.dtype)
+        self.weight_down_embed[k] = v_reshaped
+
 
 
 class MiniMindLM(PreTrainedModel):
@@ -503,36 +601,110 @@ class MiniMindLM(PreTrainedModel):
         self.vocab_size, self.n_layers = params.vocab_size, params.n_layers
         self.tok_embeddings = nn.Embedding(params.vocab_size, params.dim)
         self.dropout = nn.Dropout(params.dropout)
-        self.knowledge_dataset = KnowledgeDataset(params, self.tok_embeddings)
-        self.layers = nn.ModuleList([MiniMindBlock(l, params, self.knowledge_dataset) for l in range(self.n_layers)])
+        # 移除旧的weight_down_embed声明
+        self.extract_db = ExtractDB(self.params)
+
+        # 将self.weight_down_embed传递给每个MiniMindBlock
+        self.layers = nn.ModuleList([MiniMindBlock(l, params) for l in range(self.n_layers)])
         self.norm = RMSNorm(params.dim, eps=params.norm_eps)
         self.output = nn.Linear(params.dim, params.vocab_size, bias=False)
+        self.database_output = nn.Linear(params.dim, params.knowledge_length, bias=False)
         self.tok_embeddings.weight = self.output.weight
-        self.register_buffer("pos_cis",
-                             precompute_pos_cis(dim=params.dim // params.n_heads, theta=params.rope_theta),
+        self.database_output.weight = self.output.weight
+
+        # Calculate input dimension
+        input_dim = (self.params.max_seq_len-1)*self.params.n_layers
+        # Use a bottleneck architecture to reduce parameters
+        bottleneck_dim = 256  # Significantly smaller bottleneck dimension
+
+        # Factorized shared downsampling using two smaller convolutions
+        self.shared_downsample = nn.Sequential(
+            # First reduce input dimension to bottleneck
+            nn.Conv1d(input_dim, bottleneck_dim, kernel_size=1, padding='same'),
+            nn.ReLU(),  # Non-linearity to improve representation capacity
+            # Then expand to target dimension
+            nn.Conv1d(bottleneck_dim, 128*8, kernel_size=1, padding='same')
+        )
+
+        # Specific layers for v path
+        self.downsample_v_specific = nn.Sequential(
+            nn.Conv1d(128*8, 128, kernel_size=1, padding='same'),
+            nn.Conv1d(128, self.params.knowledge_length, kernel_size=1, padding='same')
+        )
+        
+        # Specific layers for q path
+        self.downsample_q_specific = nn.Sequential(
+            nn.Conv1d(128*8, 512, kernel_size=1, padding='same')
+        )
+        # 使用实数版本的位置编码，避免复数张量可能导致的段错误
+        self.register_buffer("pos_cis_real",
+                             precompute_pos_cis_real(dim=params.dim // params.n_heads, theta=params.rope_theta),
                              persistent=False)
-        self.OUT = CausalLMOutputWithPast()
-        self.freeze_embedding = False
+        self.params = params
 
     def forward(self,
                 input_ids: Optional[torch.Tensor] = None,
                 logits_to_keep: Union[int, torch.Tensor] = 0,
-                step: int = 0,
                 **args):
         start_pos = args.get('start_pos', 0)
-        if self.freeze_embedding and step == 0:
-            self.tok_embeddings.weight.requires_grad = False
-            # 同时冻结KnowledgeDataset的嵌入更新
-            self.knowledge_dataset.freeze_embedding = True
-            print("tok_embeddings.weight.requires_grad: ", self.tok_embeddings.weight.requires_grad)
-            print("knowledge_dataset.freeze_embedding: ", self.knowledge_dataset.freeze_embedding)
         h = self.dropout(self.tok_embeddings(input_ids))
-        pos_cis = self.pos_cis[start_pos:start_pos + input_ids.size(1)]
+        pos_cis_real = self.pos_cis_real[start_pos:start_pos + input_ids.size(1)]
+        h_list = []
+
         for l, layer in enumerate(self.layers):
+            # 禁用数据库模式，使用固定值替代数据库查询
+            if self.params.disable_db:
+                # 创建一个形状为[batch_size, n_layers, dim]的tensor，所有元素值为1e-4
+                batch_size = h.size(0)
+                db_value = torch.full((batch_size, self.n_layers, self.params.dim), 1e-4,
+                                      dtype=h.dtype, device=h.device)
+            else:
+                # 正常模式，使用数据库查询
+                # import pdb;pdb.set_trace()
+                index = self.extract_db.q_to_k(h)
+            
+                token_idx = self.extract_db.get_data(index) #这里是index
+
+                db_value =self.tok_embeddings(token_idx)
+                
             h = layer(
-                h, pos_cis
+                h, db_value, pos_cis_real
             )
 
+            h_list.append(h.unsqueeze(0))
+
+        h_tensor = torch.cat(h_list, dim=0).permute(1, 0, 2, 3)
+
+        # 只在非禁用数据库模式下执行数据库更新逻辑
+        if not self.params.disable_db:
+            # 使用detach()分离计算图，避免多次反向传播
+            h_tensor_detached = h_tensor.detach()
+            h_tensor_detached = h_tensor_detached.reshape(h_tensor_detached.shape[0], -1, self.params.dim)
+
+            # 数据库更新逻辑与主计算图分离
+            with torch.no_grad():
+                
+                # Compute shared downsampling layer once
+                shared_features = self.shared_downsample(h_tensor_detached)
+                
+                # Get features from v path - now we output embedding-dimension vectors
+                z_v_features = self.downsample_v_specific(shared_features)
+                batch_z, seq_len, dim_z = z_v_features.shape
+                
+                # Reshape to batch_size * knowledge_length, dim
+                z_v_flat = z_v_features.reshape(-1, dim_z)
+                
+                # Direct token prediction - like the main language model head
+                token_logits = self.database_output(z_v_flat)  # [batch_z * seq_len, vocab_size]
+                # Get token indices directly from logits
+                token_indices_flat = torch.argmax(token_logits, dim=-1)
+                token_indices = token_indices_flat.reshape(batch_z, -1)
+                
+                # Process query path as before
+                z_q = self.downsample_q_specific(shared_features)
+                z_k = self.extract_db.q_to_k(z_q)
+                self.extract_db.updata_value(z_k, token_indices)
+
         slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
         logits = self.output(self.norm(h)[:, slice_indices, :])
         aux_loss = sum(l.feed_forward.aux_loss for l in self.layers if isinstance(l.feed_forward, MOEFeedForward))
@@ -545,6 +717,12 @@ class MiniMindLM(PreTrainedModel):
 
         output.aux_loss = aux_loss
 
+        # 尝试添加其他属性（如果支持的话）
+        # try:
+        #     output.hidden_states = h
+        # except:
+        #     pass
+
         return output
 
     @torch.inference_mode()
@@ -577,14 +755,13 @@ class MiniMindLM(PreTrainedModel):
         return res
 
     def _stream(self, input_ids, eos_token_id, max_new_tokens, temperature, top_p, rp, **args):
-        start, first_seq, past_kvs = input_ids.shape[1], True, None
+        start, first_seq = input_ids.shape[1], True
         while input_ids.shape[1] < max_new_tokens - 1:
             if first_seq:
                 out, first_seq = self(input_ids, **args), False
             else:
-                out = self(input_ids[:, -1:],
-                           start_pos=input_ids.shape[1] - 1, **args)
-            logits, past_kvs = out.logits[:, -1, :], out.past_key_values
+                out = self(input_ids[:, -1:], start_pos=input_ids.shape[1] - 1, **args)
+            logits = out.logits[:, -1, :]
             logits[:, list(set(input_ids.tolist()[0]))] /= rp
             logits /= (temperature + 1e-9)
             if top_p is not None and top_p < 1.0:
@@ -600,4 +777,4 @@ class MiniMindLM(PreTrainedModel):
             input_ids = torch.cat((input_ids, input_ids_next), dim=1)
             yield input_ids[:, start:]
             if input_ids_next.item() == eos_token_id:
-                break
+                break

preprocessing/README_trex_processor.md

@@ -1,154 +0,0 @@
-# TREx 数据集处理工具使用说明
-
-这个工具支持两步骤处理 TREx 数据集：
-1. **句子提取**：从 TREx 数据集提取三元组并转换为自然语言句子
-2. **LLM 处理**：使用 ollama qwen3:4b 模型进行句子修正和重要性评分
-
-## 🆕 防卡死机制
-
-为了解决LLM处理时可能出现的卡死问题，新增了以下功能：
-
-### 超时和重试机制
-- **超时时间**：每个LLM请求60秒超时
-- **重试机制**：失败后最多重试2次，采用指数退避策略
-- **并发控制**：降低并发数至4个，减少服务器压力
-
-### 心跳监控系统
-- **实时监控**：每30秒检查一次LLM响应状态
-- **异常警告**：超过30秒无成功响应时发出警告
-- **服务检测**：自动检查ollama服务状态
-- **详细统计**：实时显示成功率、超时率等统计信息
-
-### 日志系统
-- **详细日志**：所有操作都记录在 `logs/` 目录下
-- **双重输出**：同时输出到日志文件和控制台
-- **时间戳标记**：日志文件包含启动时间戳
-
-### 改进的错误处理
-- **异常恢复**：LLM处理失败时使用原句子和默认评分
-- **状态监控**：处理前检查ollama服务状态
-- **批次间休息**：批次之间休息5秒，避免过度压力
-
-## 安装依赖
-
-```bash
-pip install agno asyncio pydantic requests
-```
-
-确保已安装并启动 ollama，并下载 qwen3:4b 模型：
-```bash
-ollama pull qwen3:4b
-```
-
-## 使用方法
-
-### 1. 完整流程（两步骤连续执行）
-
-```bash
-python trex_to_sentences_simple.py --step all --input_dir dataset/TREx --max_files 2
-```
-
-### 2. 分步骤执行
-
-#### 步骤1：仅提取句子
-```bash
-python trex_to_sentences_simple.py --step extract --input_dir dataset/TREx --sentences_json my_sentences.json --max_files 2
-```
-
-#### 步骤2：仅LLM处理
-```bash
-python trex_to_sentences_simple.py --step llm --sentences_json my_sentences.json --output_file final_output.txt
-```
-
-## 主要参数说明
-
-- `--step`: 运行步骤
-  - `extract`: 仅提取句子
-  - `llm`: 仅LLM处理  
-  - `all`: 完整流程（默认）
-
-- `--input_dir`: TREx数据集目录（默认：`dataset/TREx`）
-- `--sentences_json`: 提取的句子JSON文件（默认：`extracted_sentences.json`）
-- `--output_file`: 最终输出文件（默认：`trex_sentences_enhanced.txt`）
-- `--max_files`: 最大处理文件数（用于测试）
-- `--no_llm`: 禁用LLM处理
-
-## 输出文件
-
-**注意：所有输出文件都会自动保存在相应目录中**
-
-### 句子提取输出
-- `output/extracted_sentences.json`: 提取的原始句子，包含元数据
-
-### LLM处理输出  
-- `output/{output_file}.txt`: 修正后的句子文本文件
-- `output/{output_file}.json`: 完整的处理结果（包含原句、修正句、评分）
-- `output/{output_file}_sorted_by_importance.txt`: 按重要性评分排序的句子
-
-### 检查点文件
-- `output/{output_file}_checkpoint_{数量}.json`: 每1000条句子自动保存的检查点
-
-### 日志文件
-- `logs/trex_processor_{时间戳}.log`: 详细的处理日志
-
-## 🆕 故障诊断
-
-### 如果遇到卡死问题：
-
-1. **检查日志文件**：查看 `logs/` 目录下的最新日志
-2. **观察心跳监控**：注意控制台的心跳警告信息
-3. **检查ollama服务**：
-   ```bash
-   ps aux | grep ollama
-   curl http://localhost:11434/api/tags
-   ```
-4. **重启ollama服务**（如果需要）：
-   ```bash
-   pkill ollama
-   ollama serve &
-   ```
-
-### 常见警告信息：
-
-- `⚠️ 心跳检测`: 30秒无成功响应（正常情况下会自动恢复）
-- `❌ 严重警告`: 90秒无成功响应（可能需要检查服务）
-- `💀 Ollama服务异常`: ollama服务可能已停止
-- `💀 致命错误`: 连续多次警告（建议重启程序）
-
-## 检查点恢复机制
-
-- 步骤2会自动检测已有的检查点文件（在 `output/` 目录中）
-- 只处理尚未处理的句子，避免重复工作
-- 如果所有句子都已处理，会直接生成最终输出文件
-- 中断后重新运行会自动从最新检查点继续
-
-## 示例工作流
-
-```bash
-# 1. 先提取句子（可以快速完成）
-python trex_to_sentences_simple.py --step extract --max_files 5
-
-# 2. 后续进行LLM处理（耗时较长，支持断点续传）
-python trex_to_sentences_simple.py --step llm
-
-# 如果中途中断，再次运行步骤2会自动从检查点恢复
-python trex_to_sentences_simple.py --step llm
-```
-
-## 性能特点
-
-- **保守的并发**: 最大4个并发LLM请求（降低卡死风险）
-- **检查点保存**: 每1000条句子自动保存，支持断点续传
-- **智能监控**: 详细的处理进度和时间预估
-- **健壮的错误处理**: LLM请求失败时使用原句子和默认评分
-- **服务监控**: 自动检测ollama服务状态
-
-## 注意事项
-
-1. 首次运行步骤2前，必须先完成步骤1
-2. 检查点文件会占用额外磁盘空间（每个都包含所有已处理数据）  
-3. LLM处理速度取决于模型性能和网络状况
-4. 建议先用`--max_files`参数测试小批量数据
-5. **新增**：如果遇到卡死，查看日志文件和心跳监控信息
-6. **新增**：程序会自动检测并报告ollama服务状态
-7. **新增**：所有处理过程都有详细日志记录，便于问题诊断 

preprocessing/merge_output_json.py

@@ -1,225 +0,0 @@
-#!/usr/bin/env python3
-"""
-JSON文件合并脚本
-读取多个JSON文件并合并为一个JSON文件
-"""
-
-import json
-import os
-from typing import Dict, List, Any, Union
-
-# 需要合并的JSON文件列表
-JSON_FILES_TO_MERGE = [
-    "output/trex_sentences_enhanced_checkpoint_360000.json"
-]
-for i in range(1, 1010):
-    JSON_FILES_TO_MERGE.append(f"output/trex_sentences_enhanced_batch_{i}.json")
-
-def load_json_file(file_path: str) -> Union[Dict, List, None]:
-    """加载JSON文件"""
-    if not os.path.exists(file_path):
-        print(f"警告: 文件 {file_path} 不存在")
-        return None
-    
-    try:
-        with open(file_path, 'r', encoding='utf-8') as f:
-            data = json.load(f)
-        print(f"成功加载: {file_path}")
-        return data
-    except json.JSONDecodeError as e:
-        print(f"错误: 无法解析JSON文件 {file_path} - {e}")
-        return None
-    except Exception as e:
-        print(f"错误: 读取文件 {file_path} 失败 - {e}")
-        return None
-
-def merge_json_data(data1: Union[Dict, List], data2: Union[Dict, List]) -> Union[Dict, List]:
-    """合并两个JSON数据结构"""
-    
-    # 如果两个都是列表，直接合并
-    if isinstance(data1, list) and isinstance(data2, list):
-        print(f"合并两个列表: {len(data1)} + {len(data2)} = {len(data1) + len(data2)} 项")
-        return data1 + data2
-    
-    # 如果两个都是字典
-    elif isinstance(data1, dict) and isinstance(data2, dict):
-        print("合并两个字典结构")
-        merged = data1.copy()
-        
-        # 特殊处理：如果都有'sentences'字段且为列表，合并sentences
-        if 'sentences' in data1 and 'sentences' in data2:
-            if isinstance(data1['sentences'], list) and isinstance(data2['sentences'], list):
-                print(f"合并sentences字段: {len(data1['sentences'])} + {len(data2['sentences'])} = {len(data1['sentences']) + len(data2['sentences'])} 项")
-                merged['sentences'] = data1['sentences'] + data2['sentences']
-                
-                # 更新metadata if exists
-                if 'metadata' in merged:
-                    if isinstance(merged['metadata'], dict):
-                        merged['metadata']['total_sentences'] = len(merged['sentences'])
-                        merged['metadata']['merged_from'] = [os.path.basename(f) for f in JSON_FILES_TO_MERGE if os.path.exists(f)]
-                
-                # 合并其他字段
-                for key, value in data2.items():
-                    if key != 'sentences' and key not in merged:
-                        merged[key] = value
-                        
-                return merged
-        
-        # 普通字典合并
-        for key, value in data2.items():
-            if key in merged:
-                # 如果key重复且都是列表，合并列表
-                if isinstance(merged[key], list) and isinstance(value, list):
-                    merged[key] = merged[key] + value
-                # 如果key重复且都是字典，递归合并
-                elif isinstance(merged[key], dict) and isinstance(value, dict):
-                    merged[key] = merge_json_data(merged[key], value)
-                else:
-                    # 其他情况保留第二个文件的值
-                    merged[key] = value
-                    print(f"字段 '{key}' 被覆盖")
-            else:
-                merged[key] = value
-        
-        return merged
-    
-    # 类型不匹配的情况，创建一个包含两者的新结构
-    else:
-        print("数据类型不匹配，创建包含两者的新结构")
-        return {
-            "data_from_save.json": data1,
-            "data_from_save2.json": data2,
-            "merged_at": "test.py"
-        }
-
-def save_merged_json(data: Union[Dict, List], output_path: str):
-    """保存合并后的JSON数据"""
-    try:
-        # 确保输出目录存在
-        os.makedirs(os.path.dirname(output_path), exist_ok=True)
-        
-        with open(output_path, 'w', encoding='utf-8') as f:
-            json.dump(data, f, ensure_ascii=False, indent=2)
-        
-        print(f"合并结果已保存到: {output_path}")
-        
-        # 显示统计信息
-        if isinstance(data, dict):
-            if 'sentences' in data and isinstance(data['sentences'], list):
-                print(f"总计句子数: {len(data['sentences'])}")
-            print(f"总计字段数: {len(data)}")
-        elif isinstance(data, list):
-            print(f"总计列表项数: {len(data)}")
-            
-    except Exception as e:
-        print(f"错误: 保存文件失败 - {e}")
-
-def remove_duplicates_from_sentences(data: Union[Dict, List]) -> Union[Dict, List]:
-    """从合并结果中移除重复的句子（基于句子内容）"""
-    if isinstance(data, dict) and 'sentences' in data:
-        if isinstance(data['sentences'], list):
-            original_count = len(data['sentences'])
-            seen_sentences = set()
-            unique_sentences = []
-            
-            for item in data['sentences']:
-                if isinstance(item, dict):
-                    # 如果是字典，使用sentence字段或corrected_sentence字段作为唯一标识
-                    sentence_key = item.get('sentence') or item.get('corrected_sentence') or item.get('original_sentence')
-                elif isinstance(item, str):
-                    sentence_key = item
-                else:
-                    sentence_key = str(item)
-                
-                if sentence_key and sentence_key not in seen_sentences:
-                    seen_sentences.add(sentence_key)
-                    unique_sentences.append(item)
-            
-            data['sentences'] = unique_sentences
-            
-            # 更新metadata
-            if 'metadata' in data and isinstance(data['metadata'], dict):
-                data['metadata']['total_sentences'] = len(unique_sentences)
-                data['metadata']['duplicates_removed'] = original_count - len(unique_sentences)
-            
-            print(f"去重完成: {original_count} -> {len(unique_sentences)} (移除了 {original_count - len(unique_sentences)} 个重复项)")
-    
-    return data
-
-def merge_multiple_json_data(data_list: List[Union[Dict, List]]) -> Union[Dict, List]:
-    """合并多个JSON数据结构"""
-    if not data_list:
-        return {}
-    
-    if len(data_list) == 1:
-        return data_list[0]
-    
-    print(f"准备合并 {len(data_list)} 个JSON数据结构")
-    
-    # 从第一个数据开始，逐步合并其他数据
-    merged_data = data_list[0]
-    
-    for i, data in enumerate(data_list[1:], 1):
-        print(f"正在合并第 {i+1} 个数据结构...")
-        merged_data = merge_json_data(merged_data, data)
-    
-    return merged_data
-
-def main():
-    """主函数"""
-    print("=== JSON文件合并脚本 ===")
-    
-    # 输出路径
-    output_path = "output/merged.json"
-    
-    print(f"准备合并以下文件:")
-    for i, file_path in enumerate(JSON_FILES_TO_MERGE, 1):
-        print(f"  {i}. {file_path}")
-    print(f"输出文件: {output_path}")
-    print()
-    
-    # 加载所有文件
-    loaded_data = []
-    successfully_loaded = []
-    
-    for file_path in JSON_FILES_TO_MERGE:
-        data = load_json_file(file_path)
-        if data is not None:
-            loaded_data.append(data)
-            successfully_loaded.append(file_path)
-    
-    # 检查是否至少有一个文件加载成功
-    if not loaded_data:
-        print("错误: 没有文件能够成功加载，退出")
-        return
-    
-    print(f"成功加载了 {len(loaded_data)} 个文件:")
-    for file_path in successfully_loaded:
-        print(f"  ✓ {file_path}")
-    
-    if len(loaded_data) < len(JSON_FILES_TO_MERGE):
-        failed_count = len(JSON_FILES_TO_MERGE) - len(loaded_data)
-        print(f"警告: {failed_count} 个文件加载失败")
-    print()
-    
-    # 合并所有数据
-    if len(loaded_data) == 1:
-        print("只有一个文件可用，直接使用...")
-        merged_data = loaded_data[0]
-    else:
-        print("开始合并所有文件...")
-        merged_data = merge_multiple_json_data(loaded_data)
-    
-    # 去重处理
-    print("\n检查并去除重复项...")
-    merged_data = remove_duplicates_from_sentences(merged_data)
-    
-    # 保存合并结果
-    print("\n保存合并结果...")
-    save_merged_json(merged_data, output_path)
-    
-    print("\n=== 合并完成 ===")
-    print(f"合并了 {len(successfully_loaded)} 个文件的数据")
-
-if __name__ == "__main__":
-    main()

run_file/DynamicKV-LLM_Mini_Minimind.sh

@@ -1,8 +1,8 @@
 #!/bin/bash
 
 # 激活conda环境
-source $(conda info --base)/etc/profile.d/conda.sh
-conda activate mini
+# source $(conda info --base)/etc/profile.d/conda.sh
+# conda activate ycz_accelerate
 
 # 设置环境变量以帮助调试
 export NCCL_DEBUG=INFO
@@ -26,9 +26,24 @@ export PYTHONFAULTHANDLER=1
 #     --profile_interval 10
 
 # 方法2: 使用命令行参数直接配置accelerate
-CUDA_VISIBLE_DEVICES=0 /opt/conda/envs/mini/bin/python -m accelerate.commands.launch \
+CUDA_VISIBLE_DEVICES=0 accelerate launch \
     --num_processes=1 \
     --mixed_precision=bf16 \
     --main_process_port=29500 \
     train_pretrain_accelerate.py \
-
+    --epochs 3 \
+    --batch_size 24 \
+    --learning_rate 2e-4 \
+    --dtype bfloat16 \
+    --accumulation_steps 32 \
+    --grad_clip 1.0 \
+    --log_interval 100 \
+    --save_interval 10000 \
+    --dim 512 \
+    --n_layers 12 \
+    --max_seq_len 512 \
+    --use_flash_attn \
+    --profile \
+    --profile_interval 10\
+    --knowledge_num 4096 \
+    --knowledge_length 8

train_pretrain_accelerate.py

@@ -3,7 +3,6 @@ import os
 os.environ["WANDB_MODE"] = "offline"  # 或者使用 "dryrun"
 import platform
 import argparse
-from tqdm import tqdm
 import time
 import math
 import warnings
@@ -19,10 +18,8 @@ from accelerate.utils import set_seed
 from accelerate.utils import DeepSpeedPlugin
 from accelerate.utils import DistributedDataParallelKwargs
 from transformers import AutoTokenizer, get_cosine_schedule_with_warmup
-import numpy as np
-from sklearn.metrics.pairwise import cosine_similarity
 
-from model.model import MiniMindLM, RMSNorm
+from model.model import MiniMindLM
 from model.LMConfig import LMConfig
 from model.dataset import PretrainDataset
 
@@ -44,40 +41,9 @@ def get_lr(it, num_iters, learning_rate):
     return learning_rate * 0.5 * (1.0 + math.cos(math.pi * it / num_iters))
 
 # 初始化模型函数
-def init_model(lm_config, pretrained_embedding_path=None, database_init_path=None, args=None):
+def init_model(lm_config, pretrained_embedding_path=None):
     tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer')
     model = MiniMindLM(lm_config)
-    
-    # 默认模型初始化
-    Logger("Performing default model initialization...")
-    
-    # 初始化嵌入层权重
-    nn.init.normal_(model.tok_embeddings.weight, mean=0.0, std=0.02)
-    
-    # 初始化输出层权重（如果不共享权重的话）
-    if not hasattr(model.tok_embeddings, 'weight') or model.output.weight is not model.tok_embeddings.weight:
-        nn.init.normal_(model.output.weight, mean=0.0, std=0.02)
-    
-    # 初始化所有线性层
-    for name, module in model.named_modules():
-        if isinstance(module, nn.Linear):
-            # 使用Xavier/Glorot初始化
-            nn.init.xavier_uniform_(module.weight)
-            if module.bias is not None:
-                nn.init.zeros_(module.bias)
-        elif isinstance(module, nn.Embedding):
-            # 嵌入层使用正态分布初始化
-            nn.init.normal_(module.weight, mean=0.0, std=0.02)
-        elif isinstance(module, RMSNorm):
-            # RMSNorm的权重初始化为1
-            if hasattr(module, 'weight'):
-                nn.init.ones_(module.weight)
-    
-    # 初始化位置编码相关参数
-    if hasattr(model.knowledge_dataset, 'keys'):
-        nn.init.normal_(model.knowledge_dataset.keys, mean=0.0, std=0.02)
-    
-    Logger("Default model initialization completed")
 
     # 如果提供了预训练的嵌入权重，加载它们
     if pretrained_embedding_path:
@@ -85,136 +51,7 @@ def init_model(lm_config, pretrained_embedding_path=None, database_init_path=Non
         pretrained_embeddings = torch.load(pretrained_embedding_path)
         model.tok_embeddings.weight.data.copy_(pretrained_embeddings)
         model.output.weight.data.copy_(pretrained_embeddings)  # 共享权重
-    
-    if database_init_path:
-        import json
-        import os
-        
-        # 数据库参数
-        knowledge_num = args.knowledge_num
-        knowledge_length = args.knowledge_length
-        
-        # 检查是否使用缓存
-        cache_dir = os.path.dirname(args.cluster_cache_path)
-        if cache_dir:
-            os.makedirs(cache_dir, exist_ok=True)
-        
-        processed_tensor = None
-        
-        # 尝试加载缓存的处理结果
-        if not args.recompute_clusters and os.path.exists(args.cluster_cache_path):
-            try:
-                Logger(f"Loading cached processed results from {args.cluster_cache_path}")
-                processed_tensor = torch.load(args.cluster_cache_path)
-                
-                # 验证缓存文件的形状是否可用
-                cached_knowledge_num, cached_knowledge_length = processed_tensor.shape
-                
-                if cached_knowledge_length == knowledge_length:
-                    if cached_knowledge_num >= knowledge_num:
-                        # 缓存足够大，可以截取使用
-                        processed_tensor = processed_tensor[:knowledge_num, :]
-                        Logger(f"Successfully loaded cached data with shape {processed_tensor.shape}")
-                        Logger(f"Truncated from cached shape ({cached_knowledge_num}, {cached_knowledge_length}) to required shape ({knowledge_num}, {knowledge_length})")
-                        Logger("Skipping database initialization - using cached results")
-                    else:
-                        # 缓存太小，需要重新计算
-                        Logger(f"Cached knowledge_num ({cached_knowledge_num}) < required knowledge_num ({knowledge_num}), recomputing...")
-                        processed_tensor = None
-                else:
-                    # knowledge_length不匹配，需要重新计算
-                    Logger(f"Cached knowledge_length ({cached_knowledge_length}) != required knowledge_length ({knowledge_length}), recomputing...")
-                    processed_tensor = None
-            except Exception as e:
-                Logger(f"Failed to load cached data: {e}, recomputing...")
-                processed_tensor = None
-        
-        # 只有在没有有效缓存时才进行数据库初始化和处理
-        if processed_tensor is None:
-            Logger(f"Loading database initialization data from {database_init_path}")
-            
-            # 1. 加载JSON文件
-            with open(database_init_path, 'r', encoding='utf-8') as f:
-                database_data = json.load(f)
-            
-            # 提取sentences列表
-            sentences_data = database_data.get('sentences', [])
-            Logger(f"Loaded {len(sentences_data)} sentences from database")
-            
-            # 2. 按照importance_score进行排序（从高到低）
-            sorted_sentences = sorted(sentences_data, key=lambda x: x.get('importance_score', 0.0), reverse=True)
-            Logger(f"Sorted sentences by importance score (highest: {sorted_sentences[0].get('importance_score', 0.0)}, lowest: {sorted_sentences[-1].get('importance_score', 0.0)})")
-            
-            # 3. 处理每条数据，不进行聚类
-            Logger("Processing individual sentences...")
-            processed_rows = []
-            
-            # 获取空token的id（用于填充）
-            pad_token_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0
-            
-            # 处理所需数量的句子
-            num_to_process = min(knowledge_num, len(sorted_sentences))
-            
-            for i in range(num_to_process):
-                sentence_data = sorted_sentences[i]
-                sentence = sentence_data.get('corrected_sentence', '')
-                
-                # 将句子转换为tokens
-                sentence_tokens = tokenizer.encode(sentence, add_special_tokens=False)
-                
-                # 截断或填充到knowledge_length
-                if len(sentence_tokens) > knowledge_length:
-                    # 如果超过长度，截断
-                    sentence_tokens = sentence_tokens[:knowledge_length]
-                    Logger(f"Sentence {i+1} truncated from {len(tokenizer.encode(sentence, add_special_tokens=False))} to {knowledge_length} tokens")
-                else:
-                    # 如果不足长度，用空token填充
-                    original_length = len(sentence_tokens)
-                    sentence_tokens.extend([pad_token_id] * (knowledge_length - len(sentence_tokens)))
-                    if original_length < knowledge_length:
-                        Logger(f"Sentence {i+1} padded from {original_length} to {knowledge_length} tokens")
-                
-                processed_rows.append(sentence_tokens)
-                
-                if (i + 1) % 1000 == 0:
-                    Logger(f"Processed {i + 1}/{num_to_process} sentences")
-            
-            # 如果句子数量不足，用空token填充剩余位置
-            while len(processed_rows) < knowledge_num:
-                empty_tokens = [pad_token_id] * knowledge_length
-                processed_rows.append(empty_tokens)
-                if len(processed_rows) % 1000 == 0:
-                    Logger(f"Added empty entry {len(processed_rows)}/{knowledge_num}")
-            
-            Logger(f"Finished adding empty entries. Total: {len(processed_rows)}/{knowledge_num}")
-            
-            # 转换为tensor
-            processed_tensor = torch.tensor(processed_rows, dtype=torch.long)
-            
-            Logger(f"Data processing completed:")
-            Logger(f"  - Processed {num_to_process} sentences")
-            Logger(f"  - Added {knowledge_num - num_to_process} empty entries")
-            Logger(f"  - Final shape: {processed_tensor.shape}")
-            Logger(f"  - Expected shape: ({knowledge_num}, {knowledge_length})")
-            
-            # 保存处理结果到缓存文件
-            try:
-                torch.save(processed_tensor, args.cluster_cache_path)
-                Logger(f"Processed results saved to {args.cluster_cache_path}")
-            except Exception as e:
-                Logger(f"Failed to save processed results: {e}")
-        
-        # 4. 初始化模型的knowledge_dataset
-        if hasattr(model, 'knowledge_dataset') and hasattr(model.knowledge_dataset, 'knowledge_dataset'):
-            model.knowledge_dataset.knowledge_dataset.data.copy_(processed_tensor)
-            Logger("Successfully initialized model.knowledge_dataset.knowledge_dataset with processed data")
-        else:
-            Logger("Warning: Could not find model.knowledge_dataset.knowledge_dataset to initialize")
-            # 存储为全局变量作为备选
-            globals()['processed_database'] = processed_tensor
 
-    Logger(f"Database embeddings and sentences stored in model")
-    
     Logger(f'LLM总参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万')
     return model, tokenizer
 
@@ -224,7 +61,6 @@ def train_epoch(epoch, accelerator, model, train_loader, optimizer, scheduler, a
     total_steps_in_epoch = len(train_loader)
     total_training_steps = args.epochs * total_steps_in_epoch
     moe_path = '_moe' if args.use_moe else ''
-    best_loss = float('10000')
 
     # 添加CUDA事件来分析性能 (只在主进程进行)
     if args.profile and accelerator.is_main_process:
@@ -288,12 +124,7 @@ def train_epoch(epoch, accelerator, model, train_loader, optimizer, scheduler, a
 
             # 前向传播
             with ctx:
-                if step == 0 and args.embedding_epoch == epoch:
-                    # 需要设置原始模型的freeze_embedding属性，而不是包装后的模型
-                    unwrapped_model = accelerator.unwrap_model(model)
-                    unwrapped_model.freeze_embedding = True
-                    Logger(f"Set freeze_embedding=True for epoch {epoch}, step {step}", accelerator)
-                res = model(X, step=step)
+                res = model(X)
                 loss = loss_fct(
                     res.logits.view(-1, res.logits.size(-1)),
                     Y.view(-1)
@@ -417,9 +248,7 @@ def train_epoch(epoch, accelerator, model, train_loader, optimizer, scheduler, a
                     wandb.log(log_dict)
 
             # 保存模型 (只在主进程进行)
-            loss_total = loss.item() * args.accumulation_steps
-            if best_loss > loss_total and accelerator.is_main_process:
-                best_loss = loss_total
+            if (step + 1) % args.save_interval == 0 and accelerator.is_main_process:
                 # 使用函数开始处定义的moe_path变量
                 ckp = f'{args.save_dir}/pretrain_{args.dim}{moe_path}.pth'
 
@@ -438,22 +267,21 @@ def train_epoch(epoch, accelerator, model, train_loader, optimizer, scheduler, a
 def main():
     parser = argparse.ArgumentParser(description="MiniMind Pretraining with Accelerate")
     parser.add_argument("--out_dir", type=str, default="out")
-    parser.add_argument("--epochs", type=int, default=4)
-    parser.add_argument("--embedding_epoch", type=int, default=2, help="embedding训练的epoch数")
-    parser.add_argument("--batch_size", type=int, default=64)
+    parser.add_argument("--epochs", type=int, default=3)
+    parser.add_argument("--batch_size", type=int, default=24)
     parser.add_argument("--learning_rate", type=float, default=2e-4)
     parser.add_argument("--dtype", type=str, default="bfloat16")
     parser.add_argument("--use_wandb", default=True, action="store_true")
     parser.add_argument("--wandb_project", type=str, default="MiniMind-Pretrain")
-    parser.add_argument("--num_workers", type=int, default=8)
+    parser.add_argument("--num_workers", type=int, default=48)
     parser.add_argument("--accumulation_steps", type=int, default=32)
     parser.add_argument("--grad_clip", type=float, default=1.0)
     parser.add_argument("--warmup_iters", type=int, default=0)
     parser.add_argument("--log_interval", type=int, default=100)
     parser.add_argument("--save_interval", type=int, default=10000)
-    parser.add_argument('--dim', default=512, type=int)
-    parser.add_argument('--n_layers', default=8, type=int)
-    parser.add_argument('--max_seq_len', default=512, type=int)
+    parser.add_argument('--dim', default=1024, type=int)
+    parser.add_argument('--n_layers', default=32, type=int)
+    parser.add_argument('--max_seq_len', default=1024, type=int)
     parser.add_argument('--use_moe', default=False, type=bool)
     parser.add_argument('--disable_db', action='store_true', help="禁用数据库功能，使用固定值1e-4替代")
     parser.add_argument("--data_path", type=str, default="./dataset/pretrain_hq.jsonl")
@@ -461,12 +289,8 @@ def main():
     parser.add_argument("--profile", action="store_true", default=True, help="启用性能分析")
     parser.add_argument("--profile_interval", type=int, default=10, help="性能分析打印间隔（步数）")
     parser.add_argument("--use_flash_attn", action="store_true", default=True, help="启用FlashAttention")
-    parser.add_argument("--knowledge_num", type=int, default=8192,help="知识库的数据数目")
-    parser.add_argument("--knowledge_length", type=int, default=32,help="知识库的句子长度")
-    parser.add_argument("--database_init_path", type=str, default="./dataset/database_init.json", help="数据库初始化路径")
-    parser.add_argument("--fast_clustering", action="store_true", default=True, help="使用快速近似聚类算法（适用于大数据集）")
-    parser.add_argument("--cluster_cache_path", type=str, default="./cache/cluster_tokens_single.pt", help="聚类结果缓存文件路径")
-    parser.add_argument("--recompute_clusters", action="store_true", default=False, help="强制重新计算聚类，忽略缓存文件")
+    parser.add_argument("--knowledge_num", type=int, default=64*64,help="知识库的数据数目")
+    parser.add_argument("--knowledge_length", type=int, default=8,help="知识库的句子长度")
     args = parser.parse_args()
 
     #########################################################
@@ -504,8 +328,7 @@ def main():
         disable_db=args.disable_db,
         flash_attn=args.use_flash_attn,
         knowledge_num=args.knowledge_num,
-        knowledge_length=args.knowledge_length,
-        embeddings_epoch=args.embedding_epoch
+        knowledge_length=args.knowledge_length
     )
 
     #########################################################
@@ -556,7 +379,7 @@ def main():
     #########################################################
     # 初始化模型和tokenizer
     #########################################################
-    model, tokenizer = init_model(lm_config, args.pretrained_embedding_path, args.database_init_path, args)
+    model, tokenizer = init_model(lm_config, args.pretrained_embedding_path)
     # 将accelerator传递给init_model函数中的Logger调用
     Logger(f'模型初始化完成', accelerator)