generate_noisy_pairs.py

import os
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from tqdm import tqdm
import argparse
import matplotlib.pyplot as plt
from scipy.stats import multivariate_normal
import random

class VAE(nn.Module):
    def __init__(self, input_dim, latent_dim=32, hidden_dims=[256, 128]):
        super(VAE, self).__init__()
        
        # Encoder
        encoder_layers = []
        prev_dim = input_dim
        for hidden_dim in hidden_dims:
            encoder_layers.extend([
                nn.Linear(prev_dim, hidden_dim),
                #nn.LayerNorm(hidden_dim),
                nn.ReLU()
            ])
            prev_dim = hidden_dim
        
        self.encoder = nn.Sequential(*encoder_layers)
        
        # Mean and log-variance heads
        self.fc_mu = nn.Linear(hidden_dims[-1], latent_dim)
        self.fc_var = nn.Linear(hidden_dims[-1], latent_dim)
        
        # Decoder
        decoder_layers = []
        prev_dim = latent_dim
        for hidden_dim in reversed(hidden_dims):
            decoder_layers.extend([
                nn.Linear(prev_dim, hidden_dim),
                #nn.LayerNorm(hidden_dim),
                nn.ReLU()
            ])
            prev_dim = hidden_dim
        
        decoder_layers.append(nn.Linear(hidden_dims[0], input_dim))
        self.decoder = nn.Sequential(*decoder_layers)
        
    def encode(self, x):
        x = self.encoder(x)
        mu = self.fc_mu(x)
        log_var = self.fc_var(x)
        return mu, log_var
    
    def reparameterize(self, mu, log_var):
        std = torch.exp(0.5 * log_var)
        eps = torch.randn_like(std)
        return mu + eps * std
    
    def decode(self, z):
        return self.decoder(z)
    
    def forward(self, x):
        mu, log_var = self.encode(x)
        z = self.reparameterize(mu, log_var)
        return self.decode(z), mu, log_var

class SharedVAE(nn.Module):
    def __init__(self, input_dim, latent_dim=32, hidden_dims=[256, 128], temperature=0.5):
        super(SharedVAE, self).__init__()
        
        # Shared encoder/decoder
        encoder_layers = []
        prev_dim = input_dim
        for hidden_dim in hidden_dims:
            encoder_layers.extend([
                nn.Linear(prev_dim, hidden_dim),
                #nn.LayerNorm(hidden_dim),
                nn.ReLU()
            ])
            prev_dim = hidden_dim
        
        self.encoder = nn.Sequential(*encoder_layers)
        
        # Mean/logvar for positive samples
        self.fc_mu_pos = nn.Linear(hidden_dims[-1], latent_dim)
        self.fc_var_pos = nn.Linear(hidden_dims[-1], latent_dim)
        
        # Mean/logvar for negative samples
        self.fc_mu_neg = nn.Linear(hidden_dims[-1], latent_dim)
        self.fc_var_neg = nn.Linear(hidden_dims[-1], latent_dim)
        
        # Decoder
        decoder_layers = []
        prev_dim = latent_dim
        for hidden_dim in reversed(hidden_dims):
            decoder_layers.extend([
                nn.Linear(prev_dim, hidden_dim),
                #nn.LayerNorm(hidden_dim),
                nn.ReLU()
            ])
            prev_dim = hidden_dim
        
        decoder_layers.append(nn.Linear(hidden_dims[0], input_dim))
        self.decoder = nn.Sequential(*decoder_layers)
        
        self.temperature = temperature
        
    def encode_pos(self, x):
        x = self.encoder(x)
        mu = self.fc_mu_pos(x)
        log_var = self.fc_var_pos(x)
        return mu, log_var
    
    def encode_neg(self, x):
        x = self.encoder(x)
        mu = self.fc_mu_neg(x)
        log_var = self.fc_var_neg(x)
        return mu, log_var
    
    def reparameterize(self, mu, log_var):
        std = torch.exp(0.5 * log_var)
        eps = torch.randn_like(std)
        return mu + eps * std
    
    def decode(self, z):
        return self.decoder(z)
    
    def forward(self, x_pos, x_neg):
        # 正样本编码
        mu_pos, log_var_pos = self.encode_pos(x_pos)
        z_pos = self.reparameterize(mu_pos, log_var_pos)
        recon_pos = self.decode(z_pos)
        
        # 负样本编码
        mu_neg, log_var_neg = self.encode_neg(x_neg)
        z_neg = self.reparameterize(mu_neg, log_var_neg)
        recon_neg = self.decode(z_neg)
        
        return (recon_pos, mu_pos, log_var_pos), (recon_neg, mu_neg, log_var_neg)
    
    def get_latent(self, x_pos, x_neg):
        mu_pos, _ = self.encode_pos(x_pos)
        mu_neg, _ = self.encode_neg(x_neg)
        return mu_pos, mu_neg

class DualVAE(nn.Module):
    def __init__(self, input_dim, latent_dim=32, hidden_dims=[256, 128], temperature=0.5):
        super(DualVAE, self).__init__()
        self.positive_vae = VAE(input_dim, latent_dim, hidden_dims)
        self.negative_vae = VAE(input_dim, latent_dim, hidden_dims)
        self.temperature = temperature
        
    def forward(self, x_pos, x_neg):
        recon_pos, mu_pos, log_var_pos = self.positive_vae(x_pos)
        recon_neg, mu_neg, log_var_neg = self.negative_vae(x_neg)
        return (recon_pos, mu_pos, log_var_pos), (recon_neg, mu_neg, log_var_neg)
    
    def get_latent(self, x_pos, x_neg):
        mu_pos, _ = self.positive_vae.encode(x_pos)
        mu_neg, _ = self.negative_vae.encode(x_neg)
        return mu_pos, mu_neg

class EmbeddingDataset(Dataset):
    def __init__(self, positive_features, negative_features):
        self.positive_features = torch.FloatTensor(positive_features)
        self.negative_features = torch.FloatTensor(negative_features)
        
    def __len__(self):
        return len(self.positive_features)
    
    def __getitem__(self, idx):
        return self.positive_features[idx], self.negative_features[idx]

def generate_noisy_pairs(model, positive_features, negative_features, noise_std=0.5, batch_size=128, num_variants=1, n_noise=10, use_shared_vae=False, include_original_samples=True):
    """
    Generate new pairs by sampling noisy variants in the latent space

    Args:
        model: Trained VAE (DualVAE or SharedVAE)
        positive_features: Positive features
        negative_features: Negative features
        noise_std: Std of Gaussian noise
        batch_size: Batch size
        num_variants: Number of variants per sample
        n_noise: Candidates per sample, select top-k by VAE likelihood
        use_shared_vae: Use SharedVAE
        include_original_samples: Include original samples in outputs
    """
    device = next(model.parameters()).device
    model.eval()
    
    # Dataset and loader
    dataset = EmbeddingDataset(positive_features, negative_features)
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
    
    generated_positive = []
    generated_negative = []
    
    with torch.no_grad():
        for pos_batch, neg_batch in tqdm(dataloader, desc="Generating noisy pairs"):
            pos_batch = pos_batch.to(device)
            neg_batch = neg_batch.to(device)
            
            # Latent representation and VAE distribution parameters
            if use_shared_vae:
                mu_pos, log_var_pos = model.encode_pos(pos_batch)
                mu_neg, log_var_neg = model.encode_neg(neg_batch)
                z_pos = model.reparameterize(mu_pos, log_var_pos)
                z_neg = model.reparameterize(mu_neg, log_var_neg)
            else:
                mu_pos, log_var_pos = model.positive_vae.encode(pos_batch)
                mu_neg, log_var_neg = model.negative_vae.encode(neg_batch)
                z_pos = model.positive_vae.reparameterize(mu_pos, log_var_pos)
                z_neg = model.negative_vae.reparameterize(mu_neg, log_var_neg)
            
            # Generate variants
            pos_variants = []
            neg_variants = []
            
            # Optionally include original samples
            if include_original_samples:
                pos_variants.append(pos_batch)
                neg_variants.append(neg_batch)
            
            # Positive: vectorized noisy variants
            # Expand z_pos to [n_noise, batch, latent_dim]
            z_pos_expanded = z_pos.unsqueeze(0).expand(n_noise, -1, -1)
            # Noise [n_noise, batch, latent_dim]
            noise_pos = torch.randn_like(z_pos_expanded) * noise_std - 0.5*noise_std
            z_pos_noisy = z_pos_expanded + noise_pos
            
            # VAE likelihood (vectorized)
            std_pos = torch.exp(0.5 * log_var_pos)
            # Expand mu/std to [n_noise, batch, latent_dim]
            mu_pos_expanded = mu_pos.unsqueeze(0).expand(n_noise, -1, -1)
            std_pos_expanded = std_pos.unsqueeze(0).expand(n_noise, -1, -1)
            
            # Log-likelihood [n_noise, batch]
            pos_log_likelihood = -0.5 * (
                torch.sum(((z_pos_noisy - mu_pos_expanded) / std_pos_expanded) ** 2, dim=2) + 
                torch.sum(torch.log(2 * np.pi * std_pos_expanded ** 2), dim=2)
            )
            
            # Negative: vectorized noisy variants
            z_neg_expanded = z_neg.unsqueeze(0).expand(n_noise, -1, -1)
            noise_neg = torch.randn_like(z_neg_expanded) * noise_std - 0.5*noise_std
            z_neg_noisy = z_neg_expanded + noise_neg
            
            # Negative log-likelihood
            std_neg = torch.exp(0.5 * log_var_neg)
            mu_neg_expanded = mu_neg.unsqueeze(0).expand(n_noise, -1, -1)
            std_neg_expanded = std_neg.unsqueeze(0).expand(n_noise, -1, -1)
            
            neg_log_likelihood = -0.5 * (
                torch.sum(((z_neg_noisy - mu_neg_expanded) / std_neg_expanded) ** 2, dim=2) + 
                torch.sum(torch.log(2 * np.pi * std_neg_expanded ** 2), dim=2)
            )
            
            # Select top-k by VAE likelihood
            for i in range(min(num_variants, n_noise)):
                # Top-k for positives
                top_pos_idx = torch.argmax(pos_log_likelihood, dim=0)  # [batch_size]
                z_pos_noisy_selected = z_pos_noisy[top_pos_idx, torch.arange(len(top_pos_idx))]
                pos_log_likelihood[top_pos_idx, torch.arange(len(top_pos_idx))] = float('-inf')
                
                # Top-k for negatives
                top_neg_idx = torch.argmax(neg_log_likelihood, dim=0)  # [batch_size]
                z_neg_noisy_selected = z_neg_noisy[top_neg_idx, torch.arange(len(top_neg_idx))]
                neg_log_likelihood[top_neg_idx, torch.arange(len(top_neg_idx))] = float('-inf')
                
                # Decode back to feature space
                if use_shared_vae:
                    recon_pos = model.decode(z_pos_noisy_selected)
                    recon_neg = model.decode(z_neg_noisy_selected)
                else:
                    recon_pos = model.positive_vae.decode(z_pos_noisy_selected)
                    recon_neg = model.negative_vae.decode(z_neg_noisy_selected)
                
                pos_variants.append(recon_pos)
                neg_variants.append(recon_neg)
            
            # Produce all pairwise combinations
            for i in range(len(pos_variants)):
                for j in range(len(neg_variants)):
                    generated_positive.append(pos_variants[i].cpu().numpy())
                    generated_negative.append(neg_variants[j].cpu().numpy())
    
    # Concatenate all batches
    generated_positive = np.concatenate(generated_positive, axis=0)
    generated_negative = np.concatenate(generated_negative, axis=0)
    
    return generated_positive, generated_negative



def evaluate_vae_distribution(model, positive_features, negative_features, batch_size=128, use_shared_vae=False):
    """
    Evaluate whether VAE means/variances describe the data distribution well

    Args:
        model: Trained VAE (DualVAE or SharedVAE)
        positive_features: Positive features
        negative_features: Negative features
        batch_size: Batch size
        use_shared_vae: Whether to use SharedVAE

    Returns:
        dict with evaluation metrics
    """
    device = next(model.parameters()).device
    model.eval()
    
    # Dataset and loader
    dataset = EmbeddingDataset(positive_features, negative_features)
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
    
    # Collect latent samples and VAE parameters
    all_pos_latent = []
    all_neg_latent = []
    all_pos_mu = []
    all_pos_log_var = []
    all_neg_mu = []
    all_neg_log_var = []
    
    with torch.no_grad():
        for pos_batch, neg_batch in tqdm(dataloader, desc="Collecting latents and VAE params"):
            pos_batch = pos_batch.to(device)
            neg_batch = neg_batch.to(device)
            
            if use_shared_vae:
                # SharedVAE means/logvars
                mu_pos, log_var_pos = model.encode_pos(pos_batch)
                mu_neg, log_var_neg = model.encode_neg(neg_batch)
                z_pos = model.reparameterize(mu_pos, log_var_pos)
                z_neg = model.reparameterize(mu_neg, log_var_neg)
            else:
                # DualVAE means/logvars
                mu_pos, log_var_pos = model.positive_vae.encode(pos_batch)
                mu_neg, log_var_neg = model.negative_vae.encode(neg_batch)
                z_pos = model.positive_vae.reparameterize(mu_pos, log_var_pos)
                z_neg = model.negative_vae.reparameterize(mu_neg, log_var_neg)
            
            all_pos_latent.append(z_pos.cpu().numpy())
            all_neg_latent.append(z_neg.cpu().numpy())
            all_pos_mu.append(mu_pos.cpu().numpy())
            all_pos_log_var.append(log_var_pos.cpu().numpy())
            all_neg_mu.append(mu_neg.cpu().numpy())
            all_neg_log_var.append(log_var_neg.cpu().numpy())
    
    # Concatenate all batches
    all_pos_latent = np.concatenate(all_pos_latent, axis=0)
    all_neg_latent = np.concatenate(all_neg_latent, axis=0)
    all_pos_mu = np.concatenate(all_pos_mu, axis=0)
    all_pos_log_var = np.concatenate(all_pos_log_var, axis=0)
    all_neg_mu = np.concatenate(all_neg_mu, axis=0)
    all_neg_log_var = np.concatenate(all_neg_log_var, axis=0)
    
    # Compute means/covariances
    pos_mean = np.mean(all_pos_latent, axis=0)
    pos_cov = np.cov(all_pos_latent, rowvar=False)
    
    neg_mean = np.mean(all_neg_latent, axis=0)
    neg_cov = np.cov(all_neg_latent, rowvar=False)
    
    # Average learned means/variances
    pos_mu_mean = np.mean(all_pos_mu, axis=0)
    pos_var_mean = np.mean(np.exp(all_pos_log_var), axis=0)
    
    neg_mu_mean = np.mean(all_neg_mu, axis=0)
    neg_var_mean = np.mean(np.exp(all_neg_log_var), axis=0)
    
    # Differences vs empirical distribution
    pos_mu_diff = np.mean(np.abs(pos_mu_mean - pos_mean))
    pos_var_diff = np.mean(np.abs(pos_var_mean - np.diag(pos_cov)))
    
    neg_mu_diff = np.mean(np.abs(neg_mu_mean - neg_mean))
    neg_var_diff = np.mean(np.abs(neg_var_mean - np.diag(neg_cov)))
    
    # Likelihood under learned distributions (vectorized)
    # Add small epsilon for stability
    eps = 1e-4
    
    # Log-likelihoods
    pos_var = np.exp(all_pos_log_var) + eps
    pos_log_likelihood = -0.5 * (
        np.sum(((all_pos_latent - all_pos_mu) ** 2) / pos_var, axis=1) +
        np.sum(np.log(2 * np.pi * pos_var), axis=1)
    )
    
    # ... for negatives
    neg_var = np.exp(all_neg_log_var) + eps
    neg_log_likelihood = -0.5 * (
        np.sum(((all_neg_latent - all_neg_mu) ** 2) / neg_var, axis=1) +
        np.sum(np.log(2 * np.pi * neg_var), axis=1)
    )
    
    # Positives under negative distribution
    pos_in_neg_log_likelihood = -0.5 * (
        np.sum(((all_pos_latent - all_neg_mu) ** 2) / neg_var, axis=1) +
        np.sum(np.log(2 * np.pi * neg_var), axis=1)
    )
    
    # Negatives under positive distribution
    neg_in_pos_log_likelihood = -0.5 * (
        np.sum(((all_neg_latent - all_pos_mu) ** 2) / pos_var, axis=1) +
        np.sum(np.log(2 * np.pi * pos_var), axis=1)
    )
    
    # Separation via Mahalanobis distance
    diff = pos_mean - neg_mean
    inv_cov = np.linalg.inv((pos_cov + neg_cov) / 2)
    mahalanobis_dist = np.sqrt(np.dot(np.dot(diff, inv_cov), diff))
    
    # Optional visualization when latent dim is 2
    if all_pos_latent.shape[1] == 2:
        plt.figure(figsize=(12, 6))
        
        # Scatter positive/negative latents
        plt.subplot(1, 2, 1)
        plt.scatter(all_pos_latent[:, 0], all_pos_latent[:, 1], alpha=0.5, label='Positive')
        plt.scatter(all_neg_latent[:, 0], all_neg_latent[:, 1], alpha=0.5, label='Negative')
        
        # Contours for learned distributions
        x, y = np.mgrid[all_pos_latent[:, 0].min():all_pos_latent[:, 0].max():100j, 
                        all_pos_latent[:, 1].min():all_pos_latent[:, 1].max():100j]
        pos = np.dstack((x, y))
        
        # Using learned mean/var
        rv_pos = multivariate_normal(pos_mu_mean, np.diag(pos_var_mean))
        rv_neg = multivariate_normal(neg_mu_mean, np.diag(neg_var_mean))
        
        plt.contour(x, y, rv_pos.pdf(pos), levels=5, colors='blue', alpha=0.5)
        plt.contour(x, y, rv_neg.pdf(pos), levels=5, colors='red', alpha=0.5)
        
        plt.title('Learned VAE distributions')
        plt.legend()
        
        # Log-likelihood histograms
        plt.subplot(1, 2, 2)
        plt.hist(pos_log_likelihood, bins=30, alpha=0.5, label='Pos in pos-dist loglik')
        plt.hist(neg_log_likelihood, bins=30, alpha=0.5, label='Neg in neg-dist loglik')
        plt.hist(pos_in_neg_log_likelihood, bins=30, alpha=0.5, label='Pos in neg-dist loglik')
        plt.hist(neg_in_pos_log_likelihood, bins=30, alpha=0.5, label='Neg in pos-dist loglik')
        plt.title('VAE log-likelihoods')
        plt.legend()
        
        plt.tight_layout()
        plt.savefig('vae_distribution_evaluation.png')
        plt.close()
    
    # 返回评估结果
    results = {
        'pos_mean': pos_mean,
        'pos_cov': pos_cov,
        'neg_mean': neg_mean,
        'neg_cov': neg_cov,
        'pos_mu_mean': pos_mu_mean,
        'pos_var_mean': pos_var_mean,
        'neg_mu_mean': neg_mu_mean,
        'neg_var_mean': neg_var_mean,
        'pos_mu_diff': pos_mu_diff,
        'pos_var_diff': pos_var_diff,
        'neg_mu_diff': neg_mu_diff,
        'neg_var_diff': neg_var_diff,
        'pos_log_likelihood_mean': np.mean(pos_log_likelihood),
        'neg_log_likelihood_mean': np.mean(neg_log_likelihood),
        'pos_in_neg_log_likelihood_mean': np.mean(pos_in_neg_log_likelihood),
        'neg_in_pos_log_likelihood_mean': np.mean(neg_in_pos_log_likelihood),
        'mahalanobis_distance': mahalanobis_dist
    }
    
    return results

def main():
    # 设置随机种子
    seed = 42
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    random.seed(seed)
    
    parser = argparse.ArgumentParser(description='Generate noisy pairs using trained VAE')
    parser.add_argument('--model_path', type=str, default='', help='Path to the trained VAE model')
    parser.add_argument('--input_features', type=str, default='/nobackup2/taoleitian/rm/embeddings/hh_rlhf/llama_sft_100k/train_10k.npy', help='Path to the input features file')
    parser.add_argument('--output_dir', type=str, default='/nobackup2/taoleitian/rm/vae_results/hh_rlhf/llama_sft_10k_simple/generated_pairs', help='Output directory for generated pairs')
    parser.add_argument('--noise_std', type=float, default=0.01, help='Standard deviation of Gaussian noise')
    parser.add_argument('--batch_size', type=int, default=128, help='Batch size for generation')
    parser.add_argument('--latent_dim', type=int, default=32, help='Latent dimension of VAE')
    parser.add_argument('--hidden_dims', type=int, nargs='+', default=[128], help='Hidden layer dimensions')
    parser.add_argument('--num_variants', type=int, default=2, help='Number of variants per sample')
    parser.add_argument('--n_noise', type=int, default=20, help='Number of noisy candidates per sample; pick top-k by VAE likelihood')
    parser.add_argument('--use_shared_vae', action='store_true', help='Use SharedVAE instead of DualVAE')
    parser.add_argument('--evaluate_vae', action='store_true', help='Evaluate learned VAE distributions')
    parser.add_argument('--n_components', type=int, default=1, help='Number of GMM components')
    parser.add_argument('--seed', type=int, default=42, help='Random seed')
    parser.add_argument('--include_original_samples', type=str, default='false', choices=['true', 'false'], help='Include original samples in outputs')
    parser.add_argument('--output_suffix', type=str, default='', help='Suffix for output filename (e.g., to denote noise level)')
    args = parser.parse_args()
    
    # 将字符串转换为布尔值
    args.include_original_samples = args.include_original_samples.lower() == 'true'
    
    # 使用命令行参数中的种子
    seed = args.seed
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    random.seed(seed)
    
    # 创建输出目录
    os.makedirs(args.output_dir, exist_ok=True)
    
    # 加载特征文件
    print("Loading feature file...")
    features = np.load(args.input_features)
    positive_features = features[:, 0, -1, :].astype(np.float32)
    negative_features = features[:, 1, -1, :].astype(np.float32)
    
    print(f"Input positive shape: {positive_features.shape}")
    print(f"Input negative shape: {negative_features.shape}")
    
    # 加载模型
    print("\nLoading VAE model...")
    input_dim = positive_features.shape[1]
    
    if args.use_shared_vae:
        print("Using SharedVAE")
        model = SharedVAE(input_dim, args.latent_dim, args.hidden_dims)
    else:
        print("Using DualVAE")
        model = DualVAE(input_dim, args.latent_dim, args.hidden_dims)
    
    model.load_state_dict(torch.load(args.model_path))
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = model.to(device)
    

    
    # 评估VAE学到的分布
    if args.evaluate_vae:
        print("\nEvaluating learned VAE distributions...")
        vae_results = evaluate_vae_distribution(
            model,
            positive_features,
            negative_features,
            batch_size=args.batch_size,
            use_shared_vae=args.use_shared_vae
        )
        
        print("\nVAE distribution evaluation:")
        print("\nPositive distribution:")
        print("Mean:")
        print(vae_results['pos_mu_mean'])
        print("\nVariance:")
        print(vae_results['pos_var_mean'])
        
        print("\nNegative distribution:")
        print("Mean:")
        print(vae_results['neg_mu_mean'])
        print("\nVariance:")
        print(vae_results['neg_var_mean'])
        
        print(f"\nPos mean diff: {vae_results['pos_mu_diff']:.4f}")
        print(f"Pos var diff: {vae_results['pos_var_diff']:.4f}")
        print(f"Neg mean diff: {vae_results['neg_mu_diff']:.4f}")
        print(f"Neg var diff: {vae_results['neg_var_diff']:.4f}")
        print(f"Pos in pos-dist loglik (mean): {vae_results['pos_log_likelihood_mean']:.4f}")
        print(f"Neg in neg-dist loglik (mean): {vae_results['neg_log_likelihood_mean']:.4f}")
        print(f"Pos in neg-dist loglik (mean): {vae_results['pos_in_neg_log_likelihood_mean']:.4f}")
        print(f"Neg in pos-dist loglik (mean): {vae_results['neg_in_pos_log_likelihood_mean']:.4f}")
        print(f"Mahalanobis distance: {vae_results['mahalanobis_distance']:.4f}")
        
        # 保存评估结果
        np.save(os.path.join(args.output_dir, "vae_distribution_evaluation_results.npy"), vae_results)
        print(f"Saved VAE eval results to {os.path.join(args.output_dir, 'vae_distribution_evaluation_results.npy')}")
    
    # 生成带噪声的样本对
    print("\nGenerating noisy pairs...")
    generated_positive, generated_negative = generate_noisy_pairs(
        model,
        positive_features,
        negative_features,
        noise_std=args.noise_std,
        batch_size=args.batch_size,
        num_variants=args.num_variants,
        n_noise=args.n_noise,
        use_shared_vae=args.use_shared_vae,
        include_original_samples=args.include_original_samples
    )
    
    # Compute number of combinations per original sample
    num_combinations = (args.num_variants + (1 if args.include_original_samples else 0)) * (args.num_variants + (1 if args.include_original_samples else 0))
    
    # Create array with original feature layout
    generated_features = np.zeros((len(generated_positive), 2, 1, positive_features.shape[1]))
    generated_features[:, 0, -1, :] = generated_positive
    generated_features[:, 1, -1, :] = generated_negative
    
    # Save generated pairs
    output_filename = "generated_noisy_pairs"
    if args.output_suffix:
        output_filename += f"_{args.output_suffix}"
    output_path = os.path.join(args.output_dir, f"{output_filename}.npy")
    np.save(output_path, generated_features)
    
    print(f"\nGenerated pairs saved to {output_path}")
    print(f"Generated features shape: {generated_features.shape}")
    print(f"Combinations per original sample: {num_combinations}")
    print(f"Include original samples: {args.include_original_samples}")
    print(f"Noise std: {args.noise_std}")

if __name__ == "__main__":
    main()