src/trainer.py

import torch
import torch.nn.functional as F
from tqdm import tqdm

class Trainer:
    def __init__(self, ae, flow, cfg, loader, pad_id, stop_id):
        self.ae = ae.to(cfg.device)
        self.flow = flow.to(cfg.device) if flow else None
        self.cfg = cfg
        self.loader = loader
        self.device = cfg.device
        self.pad_id = pad_id
        self.stop_id = stop_id

    def train_ae(self, optimizer):
        self.ae.train()
        total_loss = 0
        pbar = tqdm(self.loader, desc="Train AE")
        optimizer.zero_grad()
        
        for step, batch in enumerate(pbar):
            tgt = batch['tgt_ids'].to(self.device)
            mask = batch['tgt_mask'].to(self.device)
            
            # logits, z = self.ae(tgt, mask)

            # ## 不太明白这里的mask 机制
            # labels = tgt.masked_fill(mask == 0, -100)
            # loss = F.cross_entropy(
            #     logits.view(-1, logits.size(-1)),
            #     labels.view(-1),
            #     ignore_index=-100
            # )
            # Reconstruction Loss
            # loss = F.cross_entropy(logits.view(-1, logits.size(-1)), tgt.view(-1), ignore_index=1)
            logits, z = self.ae(tgt, mask)   # decoder_mask 默认 = mask

            V = logits.size(-1)
            B, L = tgt.shape3

            # 1) token loss：只看 mask==1
            labels_tok = tgt.masked_fill(mask == 0, -100)
            loss_tok = F.cross_entropy(
                logits.view(-1, V),
                labels_tok.view(-1),
                ignore_index=-100,
                reduction="mean"
            )

            # 2) pad loss：mask==0 的位置强制预测 PAD（轻权重）
            pad_pos = (mask == 0)
            if pad_pos.any():
                # 每个位置的 CE
                ce_all = F.cross_entropy(
                    logits.view(-1, V),
                    tgt.new_full((B * L,), self.pad_id),
                    reduction="none"
                ).view(B, L)
                loss_pad = (ce_all * pad_pos).sum() / (pad_pos.sum() + 1e-6)
            else:
                loss_pad = logits.new_tensor(0.0)

            # 3) 可选：stop 位置加权（让 SEP 更稳）
            stop_pos = ((tgt == self.stop_id) & (mask == 1))
            if stop_pos.any():
                ce_tok = F.cross_entropy(
                    logits.view(-1, V),
                    tgt.view(-1),
                    reduction="none"
                ).view(B, L)
                loss_stop = (ce_tok * stop_pos).sum() / (stop_pos.sum() + 1e-6)
            else:
                loss_stop = logits.new_tensor(0.0)

            # 合成：pad/stop 的权重别太大
            lambda_pad = 0.1
            lambda_stop = 0.2
            loss = loss_tok + lambda_pad * loss_pad + lambda_stop * loss_stop

            loss = loss / self.cfg.grad_accum_steps
            loss.backward()
            
            if (step + 1) % self.cfg.grad_accum_steps == 0:
                optimizer.step()
                optimizer.zero_grad()
                
            total_loss += loss.item() * self.cfg.grad_accum_steps
            pbar.set_postfix(loss=loss.item() * self.cfg.grad_accum_steps)
            
        return total_loss / len(self.loader)

    def train_robust_ae(self, optimizer):
        
        self.ae.train()
        total_loss = 0
        noise_std = 0.05 

        for batch in self.loader:
            tgt_ids = batch['tgt_ids'].to(self.device)
            tgt_mask = batch['tgt_mask'].to(self.device)
            
            # 1. get normal z
            with torch.no_grad():
                z_clean = self.ae.encode(tgt_ids, tgt_mask)
            
            # 2. add noise (Denoising Training)
            #  Decoder -> like z 
            noise = torch.randn_like(z_clean) * noise_std
            z_noisy = z_clean + noise
            
            # 3. Decode
            logits = self.ae.decode(z_noisy, attention_mask=tgt_mask)

            # 4. Loss
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), 
                                   tgt_ids.view(-1), 
                                   reduction='none')
            loss = (loss * tgt_mask.view(-1)).sum() / tgt_mask.sum()

            # Backward
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            
            total_loss += loss.item()
            
        return total_loss / len(self.loader)


    def train_ae_combined(self, optimizer, epoch, max_epochs):
        """
        结合了 基础重建 + 课程去噪 + Pad/Stop 优化
        """
        self.ae.train()
        total_loss = 0
        
        # --- 课程噪声调度 (Curriculum Noise) ---
        # 前 20% 的 Epoch 不加噪声，先学好重建
        # 后面线性增加到 0.1
        if epoch < max_epochs * 0.2:
            current_noise = 0.0
        else:
            progress = (epoch - max_epochs * 0.2) / (max_epochs * 0.8)
            current_noise = 0.1 * progress # 最大噪声 0.1

        pbar = tqdm(self.loader, desc=f"Train AE (Noise={current_noise:.4f})")
        
        for step, batch in enumerate(pbar):
            tgt = batch['tgt_ids'].to(self.device)
            mask = batch['tgt_mask'].to(self.device)
            
            # 1. Encode Clean
            with torch.no_grad():
                z_clean = self.ae.encode(tgt, mask)
            
            # 2. Add Noise (如果 noise > 0)
            if current_noise > 0:
                noise = torch.randn_like(z_clean) * current_noise
                z_input = z_clean + noise
            else:
                z_input = z_clean
            
            # 3. Decode
            logits = self.ae.decode(z_input, attention_mask=mask)
            
            # 4. Calculate Advanced Loss (Copy from your original code)
            V = logits.size(-1)
            B, L = tgt.shape

            # Token Loss (只看 mask==1)
            labels_tok = tgt.masked_fill(mask == 0, -100)
            loss_tok = F.cross_entropy(logits.view(-1, V), labels_tok.view(-1), ignore_index=-100)

            # Pad Loss (mask==0)
            pad_pos = (mask == 0)
            if pad_pos.any():
                ce_pad = F.cross_entropy(logits.view(-1, V), tgt.new_full((B*L,), self.pad_id), reduction='none').view(B,L)
                loss_pad = (ce_pad * pad_pos).sum() / (pad_pos.sum() + 1e-6)
            else:
                loss_pad = torch.tensor(0.0, device=self.device)
                
            # Stop Loss
            stop_pos = ((tgt == self.stop_id) & (mask == 1))
            if stop_pos.any():
                ce_stop = F.cross_entropy(logits.view(-1, V), tgt.view(-1), reduction='none').view(B,L)
                loss_stop = (ce_stop * stop_pos).sum() / (stop_pos.sum() + 1e-6)
            else:
                loss_stop = torch.tensor(0.0, device=self.device)

            # 合并 Loss
            loss = loss_tok + 0.1 * loss_pad + 0.5 * loss_stop # 提高一点 stop 的权重
            
            # Backward
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            
            total_loss += loss.item()
            pbar.set_postfix(loss=loss.item())
            
        return total_loss / len(self.loader)

    def train_flow(self, optimizer):
        self.flow.train()
        self.ae.eval()
        total_loss = 0
        pbar = tqdm(self.loader, desc="Train Flow")
        optimizer.zero_grad()

        scale = getattr(self.ae, "latent_scale", 10.0)
        
        for step, batch in enumerate(pbar):
            src = batch['src_ids'].to(self.device)
            src_mask = batch['src_mask'].to(self.device)
            tgt = batch['tgt_ids'].to(self.device)
            tgt_mask = batch['tgt_mask'].to(self.device)
            
            with torch.no_grad():
                z_bad = self.ae.encode(src, src_mask) # norm ~ scale
                z_good = self.ae.encode(tgt, tgt_mask) # norm ~ scale
                
            # Rectified Flow
            bs = z_bad.shape[0]
            t = torch.rand(bs, device=self.device).view(bs, 1, 1)
            
            
            # Interpolation: Bad -> Good, modify-> push back to sphere
            z_t_linear = (1 - t) * z_bad + t * z_good
            ## test before or after
            # z_t = F.normalize(z_t_linear, p=2, dim=-1) * scale
            z_t = z_t_linear

            # Modify: pred_v to pred_x
            # target_v = z_good - z_bad
            # pred_v = self.flow(z_t, t.squeeze(), condition=z_bad)
            # loss = F.mse_loss(pred_v, target_v)
            
            # to predict z_good (Target)
            pred_z1 = self.flow(z_t, t, condition=z_bad)
            # 3) (强烈建议) 把输出也投影回同一球面，避免 off-manifold -> 都不要normalize
            pred_z1 = pred_z1
            # pred_z1 = F.normalize(pred_z1, p=2, dim=-1) * scale
            # Loss 直接算与 z_good 的距离
            ## 修改：loss必须按照mask 算有效token
            mse = (pred_z1 - z_good).pow(2).mean(dim=-1)  # [B,L]
            w = tgt_mask.float()

            # stop 位置加权
            stop_pos = ((tgt == self.stop_id) & (tgt_mask == 1))
            w = w + stop_pos.float() * 2.0   # 让 SEP 位置权重更大（比如 +2）

            loss = (mse * w).sum() / (w.sum() + 1e-6)

            # loss = (mse * tgt_mask).sum() / (tgt_mask.sum() + 1e-6)
            # loss = F.mse_loss(pred_z1, z_good)

            loss = loss / self.cfg.grad_accum_steps
            loss.backward()
            
            if (step + 1) % self.cfg.grad_accum_steps == 0:
                optimizer.step()
                optimizer.zero_grad()
                
            total_loss += loss.item() * self.cfg.grad_accum_steps
            pbar.set_postfix(loss=loss.item() * self.cfg.grad_accum_steps)
            
        return total_loss / len(self.loader)