import torch import torch.optim as optim from transformers import AutoTokenizer from tqdm import tqdm import torch.nn.functional as F import os import argparse 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 # --- Helper Functions for Inference (复制过来以便独立运行) --- def _pick_stop_id(tokenizer): 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): 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): bleu = sacrebleu.corpus_bleu(predictions, [references]) try: sari = sacrebleu.corpus_sari(sources, predictions, [references]) sari_score = sari.score except Exception: sari_score = 0.0 ratios = [len(p) / len(s) if len(s) > 0 else 0 for p, s in zip(predictions, sources)] avg_ratio = sum(ratios) / len(ratios) if ratios else 0 return {"SARI": sari_score, "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 # Encode z_curr = ae.encode(src_ids, src_mask) z_cond = z_curr.clone() # Flow Sampling if use_oneshot: 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 if t_val >= 0.999: break t = torch.ones(B, device=device) * t_val pred_z1 = flow(z_curr, t, condition=z_cond).float() v = (pred_z1 - z_curr) / (1.0 - t_val + 1e-4) z_curr = z_curr + v * dt z_curr = pred_z1 # Decode (Pass 1: Detect Length) full_mask = torch.ones(B, L, device=device) logits1 = ae.decode(z_curr, attention_mask=full_mask) ids1 = logits1.argmax(dim=-1) stop_pos = [] for i in range(B): pos = _first_pos(ids1[i], stop_id, default=L - 1) stop_pos.append(pos) # Decode (Pass 2: Clean Decode) gen_mask = torch.zeros(B, L, device=device) for i in range(B): gen_mask[i, : stop_pos[i] + 1] = 1.0 logits2 = ae.decode(z_curr, attention_mask=gen_mask) ids2 = logits2.argmax(dim=-1) ids2 = ids2.masked_fill(gen_mask == 0, pad_id) # Convert to Text 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] + 1 ids_cut = ids2[i, :end] gen_texts.append(tokenizer.decode(ids_cut, skip_special_tokens=True)) for s, t, g in zip(src_texts, tgt_texts, gen_texts): s_c = s.replace("\n", " ") t_c = t.replace("\n", " ") g_c = g.replace("\n", " ") f.write(f"{s_c}\t{t_c}\t{g_c}\n") all_sources.append(s_c) all_targets.append(t_c) all_generated.append(g_c) return all_sources, all_targets, all_generated def main(): parser = argparse.ArgumentParser() parser.add_argument("--ae_ckpt", type=str, default="/mnt/hdfs/user/lixinyu.222/CodeFlow/residual_robust_checkpoints/ae_best.pt", help="Path to pre-trained AE checkpoint") parser.add_argument("--save_dir", type=str, default="residual_robust_checkpoints", help="Directory to save flow checkpoints") parser.add_argument("--use_oneshot", action="store_true", default=True, help="Use one-shot sampling for inference") args = parser.parse_args() os.makedirs(args.save_dir, exist_ok=True) # --- Config --- m_cfg = ModelConfig( encoder_name='../jina-embeddings-v2-base-code', latent_dim=512, max_seq_len=128 ) t_cfg = TrainConfig( batch_size=16, num_epochs_flow=35, # 只关注 Flow 的 epoch grad_accum_steps=4, use_amp=False, lr_flow=2e-4 ) # --- Tokenizer & Data --- tokenizer = AutoTokenizer.from_pretrained(m_cfg.encoder_name,local_files_only=True, trust_remote_code=False) 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") # --- Load AE (Pre-trained) --- print(f"\n>>> Loading Pre-trained Autoencoder from {args.ae_ckpt} ...") ae = ReshapedAutoencoder(m_cfg).to(t_cfg.device).float() if not os.path.exists(args.ae_ckpt): raise FileNotFoundError(f"AE checkpoint not found at {args.ae_ckpt}. Please run train_ae.py first.") ae.load_state_dict(torch.load(args.ae_ckpt, map_location=t_cfg.device)) # 冻结 AE 的所有参数，Flow 训练时不更新 AE ae.eval() for param in ae.parameters(): param.requires_grad = False print(">>> Autoencoder loaded and frozen.") if ae.encoder.config.pad_token_id is None: ae.encoder.config.pad_token_id = tokenizer.pad_token_id # --- Initialize Flow --- flow = PatchedFlowDiT(m_cfg).to(t_cfg.device).float() # --- Trainer --- trainer = Trainer( ae=ae, flow=flow, cfg=t_cfg, loader=train_loader, pad_id=tokenizer.pad_token_id, stop_id=_pick_stop_id(tokenizer) ) # --- Optimizer --- opt_flow = optim.AdamW(flow.parameters(), lr=t_cfg.lr_flow) # --- Training Loop --- best_flow_loss = float('inf') print("\n>>> Start Training Flow DiT...") for epoch in range(t_cfg.num_epochs_flow): # 传入 opt_flow 训练 Flow loss = trainer.train_flow(opt_flow) print(f"Flow Epoch {epoch}: Loss {loss:.4f}") # Save Best if loss < best_flow_loss: best_flow_loss = loss save_path = os.path.join(args.save_dir, "flow_best.pt") torch.save(flow.state_dict(), save_path) # print(f" Saved Best Flow to {save_path}") # Save Last torch.save(flow.state_dict(), os.path.join(args.save_dir, "flow_last.pt")) print(f"Flow Training Done. Best Loss: {best_flow_loss:.4f}") # --- Inference / Evaluation --- print("\n>>> Loading Best Flow Checkpoint for Evaluation...") best_flow_path = os.path.join(args.save_dir, "flow_best.pt") if os.path.exists(best_flow_path): flow.load_state_dict(torch.load(best_flow_path, map_location=t_cfg.device)) else: print("Warning: Best checkpoint not found, utilizing last epoch weights.") print("\n--- Starting Inference ---") sources, targets, gens = inference_batch( ae, flow, test_loader, tokenizer, t_cfg.device, steps=10, save_path="wiki_results.tsv", use_oneshot=args.use_oneshot ) # 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()