run_wiki_flow.py

import torch
import torch.optim as optim
from transformers import AutoTokenizer
from tqdm import tqdm
import torch.nn.functional as F
import os
import evaluate 
import sacrebleu

from src.config import ModelConfig, TrainConfig
from src.models.autoencoder import ReshapedAutoencoder
from src.models.dit import PatchedFlowDiT
from src.trainer import Trainer
from src.utils.data_utils import prepare_data

### 加上判断eos的函数
def _pick_stop_id(tokenizer):
    # BERT/Jina 系通常 eos_token_id=None，用 sep_token_id 作为终止符
    return tokenizer.eos_token_id if tokenizer.eos_token_id is not None else tokenizer.sep_token_id

def _first_pos(x_1d, token_id, default):
    # x_1d: [L]
    idx = (x_1d == token_id).nonzero(as_tuple=True)[0]
    return idx[0].item() if idx.numel() > 0 else default


def calculate_metrics(sources, predictions, references):
    """
    计算 SARI, BLEU, 和 压缩比
    """
    ## 这里尝试去huggingface hub 去下载 BLEU的评估脚本，但是因为网络问题没有找到
    # sari_metric = evaluate.load("sari")
    # bleu_metric = evaluate.load("bleu")
    # SARI 需要 sources
    # sari_score = sari_metric.compute(sources=sources, predictions=predictions, references=[[r] for r in references])
    # # BLEU
    # bleu_score = bleu_metric.compute(predictions=predictions, references=[[r] for r in references])
    
    
    # 1. BLEU
    # sacrebleu 期望 references 是 List[List[str]] (多个参考)
    # 这里的 references 是 List[str] (单个参考)
    # 所以需要 transpose 一下： [[ref1, ref2, ...]]
    bleu = sacrebleu.corpus_bleu(predictions, [references])
    
    # 2. SARI
    try:
        # corpus_sari 返回值就是一个 SARI 对象，它的 score 属性是 float
        sari = sacrebleu.corpus_sari(sources, predictions, [references])
        sari_score = sari.score
    except Exception as e:
        print(f"SARI calculation failed: {e}")
        sari_score = 0.0
    
    # 3. Compression Ratio
    ratios = [len(p) / len(s) if len(s) > 0 else 0 for p, s in zip(predictions, sources)]
    avg_ratio = sum(ratios) / len(ratios)
    
    return {
        "SARI": sari_score, # 直接使用 float
        "BLEU": bleu.score,
        "Compression Ratio": avg_ratio
    }

@torch.no_grad()
def inference_batch(ae, flow, loader, tokenizer, device, steps=10, save_path="results.txt",use_oneshot=True):
    ae.eval()
    flow.eval()

    stop_id = _pick_stop_id(tokenizer)
    pad_id = tokenizer.pad_token_id
    
    print(f"\n>>> Running Inference on {len(loader.dataset)} examples...")
    
    all_sources = []
    all_targets = []
    all_generated = []

    scale = getattr(ae, "latent_scale", 10.0)
    
    with open(save_path, "w", encoding="utf-8") as f:
        f.write("Source\tTarget\tGenerated\n")
        
        for batch in tqdm(loader, desc="Inferencing"):
            src_ids = batch['src_ids'].to(device)
            src_mask = batch['src_mask'].to(device)
            tgt_ids = batch['tgt_ids'].to(device)
            B, L = src_ids.shape

            z_curr = ae.encode(src_ids, src_mask)
            z_cond = z_curr.clone()

            ## 这里分别采用 one-shot 和多布采样
            if use_oneshot:
                # x-pred 最稳：直接 t=0 one-shot
                t0 = torch.zeros(B, device=device)
                z_curr = flow(z_curr, t0, condition=z_cond).float()
            else:
                dt = 1.0 / steps
                for i in range(steps):
                    t_val = i / steps
                    # 避免 t=1 时的除零错误 (虽不常见但要防范)
                    if t_val >= 0.999: break
                    t = torch.ones(z_curr.shape[0], device=device) * t_val
                    
                    ## from v to z
                    # v = flow(z_curr, t, condition=z_cond).float()
                    # z_curr = z_curr + v * dt
                    ## from z to v to zcur
                
                    pred_z1 = flow(z_curr, t, condition=z_cond).float()
                    ## maybe optimize: 1 - t_val -> 1
                    v = (pred_z1 - z_curr) / (1.0 - t_val + + 1e-4) # add epilson
                    z_curr = z_curr + v * dt
                    z_curr = F.normalize(z_curr, p=2, dim=-1) * scale
                z_curr = pred_z1 # 最后一次终点预测直接使用
            
            z_curr = torch.nn.functional.normalize(z_curr, p=2, dim=-1) * scale ## scaling 对齐
            
            # ---- 3) two-pass decode to determine length by EOS ----
            full_mask = torch.ones(B, L, device=device)  # 允许增长：全长都“可生成”
            
            # Pass-1: decode with full mask
            logits1 = ae.decode(z_curr, attention_mask=full_mask)
            ids1 = logits1.argmax(dim=-1)  # [B, L]

            # find stop positions and build gen_mask
            stop_pos = []
            for i in range(B):
                # 如果没预测到 stop，就用 L-1 当作“最大长度”
                pos = _first_pos(ids1[i], stop_id, default=L - 1)
                stop_pos.append(pos)
            stop_pos = torch.tensor(stop_pos, device=device)

            gen_mask = torch.zeros(B, L, device=device)
            for i in range(B):
                gen_mask[i, : stop_pos[i].item() + 1] = 1.0

                # Pass-2: decode again with gen_mask, reducing tail interference
            logits2 = ae.decode(z_curr, attention_mask=gen_mask)
            ids2 = logits2.argmax(dim=-1)

            # enforce pad after stop for clean decoding
            ids2 = ids2.masked_fill(gen_mask == 0, pad_id)

            # ---- 4) decode to text with truncation ----
            src_texts = tokenizer.batch_decode(src_ids, skip_special_tokens=True)
            tgt_texts = tokenizer.batch_decode(tgt_ids, skip_special_tokens=True)

            gen_texts = []
            for i in range(B):
                end = stop_pos[i].item() + 1
                ids_cut = ids2[i, :end]
                gen_texts.append(tokenizer.decode(ids_cut, skip_special_tokens=True))
    
            # Save & Collect
            for s, t, g in zip(src_texts, tgt_texts, gen_texts):
                # 简单的后处理：去掉换行符以便存成 TSV
                s_clean = s.replace("\n", " ")
                t_clean = t.replace("\n", " ")
                g_clean = g.replace("\n", " ")
                
                f.write(f"{s_clean}\t{t_clean}\t{g_clean}\n")
                
                all_sources.append(s_clean)
                all_targets.append(t_clean)
                all_generated.append(g_clean)
                
    return all_sources, all_targets, all_generated

### add saving ckpts
def main():

    ckpt_dir = "checkpoints"
    os.makedirs(ckpt_dir, exist_ok=True)
    print(f"Checkpoints will be saved to: {ckpt_dir}")

    # Config
    m_cfg = ModelConfig(
        encoder_name='../jina-embeddings-v2-base-code',
        latent_dim=512, 
        max_seq_len=128 # Wiki 任务文本短，用 128 足够且快
    )
    
    t_cfg = TrainConfig(
        batch_size=16, # 推理时可以大一点
        num_epochs_ae=20,   # 增加一点 AE 训练
        num_epochs_flow=35, # 增加 Flow 训练
        grad_accum_steps=4,
        use_amp=False
    )
    
    tokenizer = AutoTokenizer.from_pretrained(m_cfg.encoder_name, trust_remote_code=True)
    
    # 1. Load Data (Train & Test)
    train_loader = prepare_data("wiki", tokenizer, m_cfg.max_seq_len, t_cfg.batch_size, split="train")
    test_loader = prepare_data("wiki", tokenizer, m_cfg.max_seq_len, t_cfg.batch_size, split="test")
    
    # Init
    ae = ReshapedAutoencoder(m_cfg).to(t_cfg.device).float()
    flow = PatchedFlowDiT(m_cfg).to(t_cfg.device).float()
    
    if ae.encoder.config.pad_token_id is None:
        ae.encoder.config.pad_token_id = tokenizer.pad_token_id
        
    # trainer = Trainer(ae, flow, t_cfg, train_loader)
    ## 加上pad_id 和 stop_id
    trainer = Trainer(ae, flow, t_cfg, train_loader, pad_id=tokenizer.pad_token_id, stop_id=_pick_stop_id(tokenizer))

    
    # 2. Train AE
    opt_ae = optim.AdamW(filter(lambda p: p.requires_grad, ae.parameters()), lr=t_cfg.lr_ae)
    best_ae_loss = float('inf')
    print("\n>>> Start Training Autoencoder...")
    for epoch in range(t_cfg.num_epochs_ae):
        loss = trainer.train_ae(opt_ae)
        print(f"AE Epoch {epoch}: Loss {loss:.4f}")
        
        # 保存 Best
        if loss < best_ae_loss:
            best_ae_loss = loss
            torch.save(ae.state_dict(), os.path.join(ckpt_dir, "ae_best.pt"))
            # print(f"  Saved Best AE (Loss {loss:.4f})")
        
        # 保存 Last (每个 epoch 覆盖，用于断点续训或检查)
        torch.save(ae.state_dict(), os.path.join(ckpt_dir, "ae_last.pt"))
    
    print(f"AE Training Done. Best Loss: {best_ae_loss:.4f}")

    # 3. Train Flow
    opt_flow = optim.AdamW(flow.parameters(), lr=t_cfg.lr_flow)
    best_flow_loss = float('inf')
    print("\n>>> Start Training Flow DiT...")
    for epoch in range(t_cfg.num_epochs_flow):
        loss = trainer.train_flow(opt_flow)
        print(f"Flow Epoch {epoch}: Loss {loss:.4f}")

        # 保存 Best
        if loss < best_flow_loss:
            best_flow_loss = loss
            torch.save(flow.state_dict(), os.path.join(ckpt_dir, "flow_best.pt"))
            # print(f"  Saved Best Flow (Loss {loss:.4f})")
            
        # 保存 Last
        torch.save(flow.state_dict(), os.path.join(ckpt_dir, "flow_last.pt"))
    
    print(f"Flow Training Done. Best Loss: {best_flow_loss:.4f}")
    
    # 4. Evaluation
    # 加载最佳权重
    ae_path = os.path.join(ckpt_dir, "ae_best.pt")
    flow_path = os.path.join(ckpt_dir, "flow_best.pt")
    
    if os.path.exists(ae_path):
        ae.load_state_dict(torch.load(ae_path, map_location=t_cfg.device))
        print("Loaded AE Best.")
    else:
        print("Warning: AE Best ckpt not found, using last state.")

    if os.path.exists(flow_path):
        flow.load_state_dict(torch.load(flow_path, map_location=t_cfg.device))
        print("Loaded Flow Best.")
    else:
        print("Warning: Flow Best ckpt not found, using last state.")

    print("\n--- Starting Inference ---")
    sources, targets, gens = inference_batch(
        ae, flow, test_loader, tokenizer, t_cfg.device, 
        steps=10, 
        save_path="wiki_results.tsv"
    )
    
    # Calculate Metrics
    metrics = calculate_metrics(sources, gens, targets)
    print("\n=== Metrics ===")
    for k, v in metrics.items():
        print(f"{k}: {v:.4f}")
    
    print(f"\nResults saved to wiki_results.tsv")

if __name__ == "__main__":
    main()