Update model.py

model.py

@@ -13,6 +13,7 @@ from encoders import (
 )
 
 class MultiModalDenseTransformer(nn.Module):
+
     def __init__(
         self,
         model_dim: int = 2048,
@@ -25,37 +26,31 @@ class MultiModalDenseTransformer(nn.Module):
         dropout: float = 0.0,
         attn_dropout: float = 0.0,
         
-        # MoE配置
         use_moe: bool = False,
         num_experts: int = 8,
         moe_top_k: int = 2,
         moe_layers: Optional[List[int]] = None,
         
-        # PEFT配置
         use_adapter: bool = False,
         adapter_dim: int = 64,
         use_lora: bool = False,
         lora_rank: int = 8,
         
-        # 训练配置
         use_gradient_checkpointing: bool = False,
         use_parallel_residual: bool = False,
         
-        # 位置编码
         rope_scaling_factor: float = 1.0,
         rope_scaling_type: str = "yarn",
         sliding_window: Optional[int] = None,
         
-        # 规范化
         norm_eps: float = 1e-6,
         initializer_range: float = 0.02,
         ffn_dim_multiplier: Optional[float] = None,
         tie_word_embeddings: bool = True,
         
-        # 多模态配置
-        use_multimodal_fusion: bool = True,
+        use_multimodal_fusion: bool = False,
         fusion_layers: int = 4,
-        use_contrastive: bool = True,
+        use_contrastive: bool = False,
         vision_depth: int = 24,
         audio_depth: int = 12,
         video_spatial_depth: int = 12,
@@ -71,7 +66,6 @@ class MultiModalDenseTransformer(nn.Module):
         self.tie_word_embeddings = tie_word_embeddings
         self.use_multimodal_fusion = use_multimodal_fusion
         
-        # Token embedding
         self.token_embedding = nn.Embedding(vocab_size, model_dim)
         self.modality_embedding = nn.Embedding(4, model_dim)
         self.embed_dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity()
@@ -104,7 +98,6 @@ class MultiModalDenseTransformer(nn.Module):
             adapter_dim=adapter_dim
         )
         
-        # 多模态融合模块
         if use_multimodal_fusion:
             self.fusion_module = MultiModalFusionModule(
                 dim=model_dim,
@@ -157,24 +150,8 @@ class MultiModalDenseTransformer(nn.Module):
         
         self.n_params = sum(p.numel() for p in self.parameters())
         trainable_params = sum(p.numel() for p in self.parameters() if p.requires_grad)
-        
-        print(f"\n{'='*80}")
-        print(f"Improved Model Configuration:")
-        print(f"  Model Dimension: {model_dim}")
-        print(f"  Vocab Size: {vocab_size}")
-        print(f"  Layers: {n_layers}")
-        print(f"  Attention Heads: {n_heads}")
-        print(f"  KV Heads: {n_kv_heads if n_kv_heads else n_heads}")
-        print(f"  Max Sequence Length: {max_seq_len}")
-        print(f"  Multimodal Fusion: {use_multimodal_fusion}")
-        print(f"  Contrastive Learning: {use_contrastive}")
-        print(f"  MoE: {use_moe} (Experts: {num_experts}, Top-K: {moe_top_k})")
-        print(f"  Total Parameters: {self.n_params / 1e9:.2f}B")
-        print(f"  Trainable Parameters: {trainable_params / 1e9:.2f}B")
-        print(f"{'='*80}\n")
 
     def _init_weights(self, module):
-        """权重初始化"""
         if isinstance(module, nn.Linear):
             torch.nn.init.normal_(module.weight, mean=0.0, std=self.initializer_range)
             if module.bias is not None:
@@ -185,12 +162,10 @@ class MultiModalDenseTransformer(nn.Module):
                 module.weight.data[module.padding_idx].zero_()
 
     def _init_lm_head(self):
-        """初始化LM head"""
         std = self.initializer_range / math.sqrt(2 * self.n_layers)
         torch.nn.init.normal_(self.lm_head.weight, mean=0.0, std=std)
 
     def _encode_modality(self, segment: Dict) -> torch.Tensor:
-        """编码单个模态"""
         seg_type = segment['type']
         seg_data = segment['data']
         
@@ -217,15 +192,11 @@ class MultiModalDenseTransformer(nn.Module):
         output_hidden_states: bool = False,
         compute_contrastive: bool = False
     ) -> Dict:
-        """前向传播"""
         device = self.token_embedding.weight.device
         
-        # 编码每个模态
         encoded_segments = []
         for segment in input_data.get('segments', []):
             encoded = self._encode_modality(segment)
-            
-            # 添加模态嵌入
             modality_id = segment.get('modality_id', 0)
             modality_embeds = self.modality_embedding(
                 torch.tensor([modality_id], device=device)
@@ -237,7 +208,6 @@ class MultiModalDenseTransformer(nn.Module):
                 'modality_id': modality_id
             })
         
-        # 多模态融合
         contrastive_losses = {}
         if self.use_multimodal_fusion and len(encoded_segments) > 1:
             fusion_output = self.fusion_module(
@@ -247,7 +217,6 @@ class MultiModalDenseTransformer(nn.Module):
             x = fusion_output['fused_features']
             contrastive_losses = fusion_output.get('contrastive_losses', {})
         else:
-            # 简单拼接
             all_embeddings = [seg['data'] for seg in encoded_segments]
             x = torch.cat(all_embeddings, dim=1) if all_embeddings else torch.zeros(
                 1, 1, self.model_dim, device=device
@@ -256,21 +225,16 @@ class MultiModalDenseTransformer(nn.Module):
         x = self.embed_dropout(x)
         if position_ids is None:
             if past_key_values is not None:
-                # 缓存的长度 (KV cache 的 shape 是 [B, H, SeqLen, D])
                 past_length = past_key_values[0][0].size(2)
-                # 当前输入的长度
                 seq_length = x.shape[1]
-                # 生成正确的位置索引: [past_length, past_length + 1, ...]
                 position_ids = torch.arange(
                     past_length, past_length + seq_length, dtype=torch.long, device=device
                 ).unsqueeze(0).expand(x.shape[0], -1)
             else:
-                # 如果没有缓存，从 0 开始
                 seq_length = x.shape[1]
                 position_ids = torch.arange(
                     0, seq_length, dtype=torch.long, device=device
                 ).unsqueeze(0).expand(x.shape[0], -1)
-        # Transformer层
         present_key_values = [] if use_cache else None
         all_hidden_states = [] if output_hidden_states else None
         all_attentions = [] if output_attentions else None
@@ -332,7 +296,6 @@ class MultiModalDenseTransformer(nn.Module):
         if output_hidden_states:
             all_hidden_states.append(hidden_states)
         
-        # 组装输出
         outputs = {
             'logits': logits,
             'moe_aux_loss': moe_aux_loss,
@@ -370,7 +333,6 @@ class MultiModalDenseTransformer(nn.Module):
         do_sample: bool = True,
         num_beams: int = 1
     ) -> torch.Tensor:
-        """改进的生成方法"""
         self.eval()
         device = next(self.parameters()).device
         
@@ -406,7 +368,6 @@ class MultiModalDenseTransformer(nn.Module):
             }
             
             if step > 0 and use_cache:
-                # 添加当前 token 的 mask (1)
                 new_mask = torch.ones(batch_size,1,dtype=torch.long,device=device)
                 attention_mask = torch.cat([attention_mask, new_mask], dim=1)
                 current_positions = (attention_mask.sum(dim=1 , keepdim=True) -1).clamp(min=0)
@@ -415,7 +376,7 @@ class MultiModalDenseTransformer(nn.Module):
                 current_positions_ids=position_ids
             outputs = self.forward(
                 current_input_data,
-                attention_mask=attention_mask, # <--- 传入 Mask
+                attention_mask=attention_mask, 
                 position_ids=current_positions_ids,
                 use_cache=use_cache,
                 past_key_values=past_key_values
@@ -427,7 +388,6 @@ class MultiModalDenseTransformer(nn.Module):
             
             next_token_logits = logits[:, -1, :] / max(temperature, 1e-5)
             
-            # Repetition penalty
             if repetition_penalty != 1.0 and len(generated_tokens) > 0:
                 prev_generated = torch.cat(generated_tokens, dim=1)
                 score = torch.gather(next_token_logits, 1, prev_generated)
@@ -438,11 +398,9 @@ class MultiModalDenseTransformer(nn.Module):
                 )
                 next_token_logits.scatter_(1, prev_generated, score)
             
-            # Min length constraint
             if step < min_length:
                 next_token_logits[:, eos_token_id] = float('-inf')
             
-            # Sampling
             if do_sample:
                 if top_k > 0:
                     top_k_vals, _ = torch.topk(next_token_logits, top_k)
@@ -464,7 +422,6 @@ class MultiModalDenseTransformer(nn.Module):
             else:
                 next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)
             
-            # Apply unfinished mask
             next_token = next_token * unfinished_sequences[:, None] + pad_token_id * (1 - unfinished_sequences[:, None])
             
             generated_tokens.append(next_token)
@@ -475,7 +432,6 @@ class MultiModalDenseTransformer(nn.Module):
             else:
                 current_tokens = next_token
             
-            # Update unfinished sequences
             unfinished_sequences = unfinished_sequences.mul(
                 (next_token.squeeze(-1) != eos_token_id).long()
             )