TabGAN

高质量合成表格数据生成

概述

TabGAN 提供了一个统一的 Python 接口，用于使用多种最先进的生成方法生成合成表格数据：

所有生成器都共享一个通用流程：生成 → 后处理 → 对抗性过滤，确保合成数据与真实数据分布保持一致。

基于论文：不均衡分布下的表格 GAN (arXiv:2010.00638)

核心特性

• 统一 API — 通过更改单个参数即可在 GAN、扩散模型和 LLM 之间切换
• 对抗性过滤 — 内置基于 LightGBM 的验证机制，使合成样本分布一致
• 混合数据类型 — 原生支持连续、分类和自由文本列
• 条件生成 — 通过 LLM 提示生成基于分类属性的文本
• LLM API 支持 — 可与 LM Studio、OpenAI、Ollama 或任何兼容 OpenAI 的端点集成
• 质量验证 — 通过一次函数调用即可比较原始和合成分布
• AutoSynth — 自动运行所有生成器，比较质量和隐私，选择最佳方案
• HuggingFace 集成 — 一键合成任意 HF 数据集，并将结果推回 Hub
• 在线演示 — 在 HuggingFace Spaces 中浏览器体验

安装

pip install tabgan

快速入门

import pandas as pd
import numpy as np
from tabgan.sampler import GANGenerator

train = pd.DataFrame(np.random.randint(-10, 150, size=(150, 4)), columns=list("ABCD"))
target = pd.DataFrame(np.random.randint(0, 2, size=(150, 1)), columns=list("Y"))
test = pd.DataFrame(np.random.randint(0, 100, size=(100, 4)), columns=list("ABCD"))

new_train, new_target = GANGenerator().generate_data_pipe(train, target, test)

可用生成器

API 参考

常见参数

所有生成器均接受以下参数：

生成器特定参数 (gen_params)

GANGenerator:

{"batch_size": 500, "patience": 25, "epochs": 500}

LLMGenerator:

{"batch_size": 32, "epochs": 4, "llm": "distilgpt2", "max_length": 500}

generate_data_pipe 方法

new_train, new_target = generator.generate_data_pipe(
    train_df,           # pd.DataFrame - 训练特征
    target,             # pd.DataFrame - 目标变量（或 None）
    test_df,            # pd.DataFrame - 测试特征，用于分布对齐
    deep_copy=True,     # bool - 复制输入 DataFrame
    only_adversarial=False,  # bool - 跳过生成，仅进行过滤
    use_adversarial=True,    # bool - 启用对抗性过滤
)

返回值: Tuple[pd.DataFrame, pd.DataFrame] — (new_train, new_target)

数据格式

TabGAN 接受 pandas.DataFrame 输入，包含：

• 连续列 — 任何实数值的数值数据
• 分类列 — 具有有限取值集合的离散列

注意: TabGAN 内部以浮点数处理数值。若需整数输出，请在生成后进行四舍五入。

示例

所有生成器的基本用法

from tabgan.sampler import (
    OriginalGenerator, GANGenerator, ForestDiffusionGenerator,
    BayesianGenerator, LLMGenerator,
)
import pandas as pd
import numpy as np

train = pd.DataFrame(np.random.randint(-10, 150, size=(150, 4)), columns=list("ABCD"))
target = pd.DataFrame(np.random.randint(0, 2, size=(150, 1)), columns=list("Y"))
test = pd.DataFrame(np.random.randint(0, 100, size=(100, 4)), columns=list("ABCD"))

new_train1, new_target1 = OriginalGenerator().generate_data_pipe(train, target, test)
new_train2, new_target2 = GANGenerator(
    gen_params={"batch_size": 500, "epochs": 10, "patience": 5}
).generate_data_pipe(train, target, test)
new_train3, new_target3 = ForestDiffusionGenerator().generate_data_pipe(train, target, test)
new_train4, new_target4 = BayesianGenerator().generate_data_pipe(train, target, test)
new_train5, new_target5 = LLMGenerator(
    gen_params={"batch_size": 32, "epochs": 4, "llm": "distilgpt2", "max_length": 500}
).generate_data_pipe(train, target, test)

完整参数示例

new_train, new_target = GANGenerator(
    gen_x_times=1.1,
    cat_cols=None,
    bot_filter_quantile=0.001,
    top_filter_quantile=0.999,
    is_post_process=True,
    adversarial_model_params={
        "metrics": "AUC", "max_depth": 2, "max_bin": 100,
        "learning_rate": 0.02, "random_state": 42, "n_estimators": 100,
    },
    pregeneration_frac=2,
    only_generated_data=False,
    gen_params={"batch_size": 500, "patience": 25, "epochs": 500},
).generate_data_pipe(
    train, target, test,
    deep_copy=True,
    only_adversarial=False,
    use_adversarial=True,
)

LLM 条件文本生成

根据分类属性生成包含新颖文本值的合成行：

import pandas as pd
from tabgan.sampler import LLMGenerator

train = pd.DataFrame({
    "Name": ["Anna", "Maria", "Ivan", "Sergey", "Olga", "Boris"],
    "Gender": ["F", "F", "M", "M", "F", "M"],
    "Age": [25, 30, 35, 40, 28, 32],
    "Occupation": ["Engineer", "Doctor", "Artist", "Teacher", "Manager", "Pilot"],
})

new_train, _ = LLMGenerator(
    gen_x_times=1.5,
    text_generating_columns=["Name"],      # 需要生成新颖文本的列
    conditional_columns=["Gender"],         # 用于条件化文本生成的列
    gen_params={"batch_size": 32, "epochs": 4, "llm": "distilgpt2", "max_length": 500},
    is_post_process=False,
).generate_data_pipe(train, target=None, test_df=None, only_generated_data=True)

工作原理：

1. 从条件列的经验分布中采样值。
2. 使用拟合的 GReaT 模型插补其余非文本列。
3. 通过基于提示的生成方式生成新颖的文本。
4. 确保生成的文本值与原始数据不同。

基于 LLM API 的文本生成

使用外部 LLM API（LM Studio、OpenAI、Ollama）代替本地模型：

import pandas as pd
from tabgan.sampler import LLMGenerator
from tabgan.llm_config import LLMAPIConfig

train = pd.DataFrame({
    "Name": ["Anna", "Maria", "Ivan", "Sergey", "Olga", "Boris"],
    "Gender": ["F", "F", "M", "M", "F", "M"],
    "Age": [25, 30, 35, 40, 28, 32],
    "Occupation": ["Engineer", "Doctor", "Artist", "Teacher", "Manager", "Pilot"],
})

# LM Studio
api_config = LLMAPIConfig.from_lm_studio(
    base_url="http://localhost:1234",
    model="google/gemma-3-12b",
    timeout=90,
)

# 或 OpenAI：LLMAPIConfig.from_openai(api_key="...", model="gpt-4")
# 或 Ollama：LLMAPIConfig.from_ollama(model="llama3")

new_train, _ = LLMGenerator(
    gen_x_times=1.5,
    text_generating_columns=["Name"],
    conditional_columns=["Gender"],
    gen_params={"batch_size": 32, "epochs": 4, "llm": "distilgpt2", "max_length": 500},
    llm_api_config=api_config,
    is_post_process=False,
).generate_data_pipe(train, target=None, test_df=None, only_generated_data=True)

from tabgan.llm_config import LLMAPIConfig
from tabgan.llm_api_client import LLMAPIClient

config = LLMAPIConfig.from_lm_studio()
with LLMAPIClient(config) as client:
    print(f"API 可用：{client.check_connection()}")
    print(f"生成结果：{client.generate('生成一个女性名字：')}")

提升模型性能

import sklearn
import pandas as pd
from tabgan.sampler import GANGenerator

def evaluate(clf, X_train, y_train, X_test, y_test):
    clf.fit(X_train, y_train)
    return sklearn.metrics.roc_auc_score(y_test, clf.predict_proba(X_test)[:, 1])

dataset = sklearn.datasets.load_breast_cancer()
clf = sklearn.ensemble.RandomForestClassifier(n_estimators=25, max_depth=6)
X_train, X_test, y_train, y_test = sklearn.model_selection.train_test_split(
    pd.DataFrame(dataset.data),
    pd.DataFrame(dataset.target, columns=["target"]),
    test_size=0.33, random_state=42,
)

print("基线：", evaluate(clf, X_train, y_train, X_test, y_test))

new_train, new_target = GANGenerator().generate_data_pipe(X_train, y_train, X_test)
print("使用 GAN 后：", evaluate(clf, new_train, new_target, X_test, y_test))

时间序列数据生成

import pandas as pd
import numpy as np
from tabgan.utils import get_year_mnth_dt_from_date, collect_dates
from tabgan.sampler import GANGenerator

train = pd.DataFrame(np.random.randint(-10, 150, size=(100, 4)), columns=list("ABCD"))
min_date, max_date = pd.to_datetime("2019-01-01"), pd.to_datetime("2021-12-31")
d = (max_date - min_date).days + 1
train["Date"] = min_date + pd.to_timedelta(np.random.randint(d, size=100), unit="d")
train = get_year_mnth_dt_from_date(train, "Date")

new_train, _ = GANGenerator(
    gen_x_times=1.1, cat_cols=["year"],
    bot_filter_quantile=0.001, top_filter_quantile=0.999,
    is_post_process=True, pregeneration_frac=2,
).generate_data_pipe(train.drop("Date", axis=1), None, train.drop("Date", axis=1))

new_train = collect_dates(new_train)

质量报告

生成一份自包含的 HTML 报告，从多个质量维度比较原始数据和合成数据：列统计信息、PSI、相关性热图、分布图以及机器学习效用（TSTR 与 TRTR）。

from tabgan import QualityReport

report = QualityReport(
    original_df, synthetic_df,
    cat_cols=["gender"],
    target_col="target",      # 启用机器学习效用评估
).compute()

# 导出为单个 HTML 文件（图表以内嵌 base64 格式）
report.to_html("quality_report.html")

# 或者以编程方式访问指标
summary = report.summary()
print(f"总体评分: {summary['overall_score']}")
print(f"PSI 均值: {summary['psi']['mean']}")
print(f"ML 效用比: {summary['ml_utility']['utility_ratio']}")

若只需快速比较而无需完整报告，可以使用以下代码：

from tabgan.utils import compare_dataframes

score = compare_dataframes(original_df, generated_df)  # 0.0（差）至 1.0（优）

约束条件

对生成的数据强制执行业务规则。约束条件在数据生成后应用——无效行会被修复或过滤掉。

from tabgan import GANGenerator, RangeConstraint, UniqueConstraint, FormulaConstraint, RegexConstraint

new_train, new_target = GANGenerator(gen_x_times=1.5).generate_data_pipe(
    train, target, test,
    constraints=[
        RangeConstraint("age", min_val=0, max_val=120),
        UniqueConstraint("email"),
        FormulaConstraint("end_date > start_date"),
        RegexConstraint("zip_code", r"\d{5}"),
    ],
)

可用约束条件：

ConstraintEngine 支持两种策略：“fix”（先修复再过滤）和“filter”（仅丢弃违规行）：

from tabgan import ConstraintEngine, RangeConstraint

engine = ConstraintEngine(
    constraints=[RangeConstraint("price", min_val=0)],
    strategy="fix",  # 或 "filter"
)
cleaned_df = engine.apply(generated_df)

隐私度量

在共享合成数据之前，评估其重新识别风险。包括最近邻距离（DCR）、最近邻距离比（NNDR）以及成员推理风险。

from tabgan import PrivacyMetrics

pm = PrivacyMetrics(original_df, synthetic_df, cat_cols=["gender"])
summary = pm.summary()

print(f"总体隐私评分: {summary['overall_privacy_score']}")  # 0（有风险）至 1（隐私性强）
print(f"DCR 均值: {summary['dcr']['mean']}")
print(f"NNDR 均值: {summary['nndr']['mean']}")
print(f"成员推理 AUC: {summary['membership_inference']['auc']}")  # 接近 0.5 表示更好

度量解释：

sklearn 管道集成

使用 TabGANTransformer 将合成数据增强插入到 sklearn Pipeline 中：

from sklearn.pipeline import Pipeline
from sklearn.ensemble import RandomForestClassifier
from tabgan import TabGANTransformer

pipe = Pipeline([
    ("augment", TabGANTransformer(gen_x_times=1.5, cat_cols=["gender"])),
    ("model", RandomForestClassifier()),
])

# fit() 会生成合成数据，并在增强后的数据上训练模型
pipe.fit(X_train, y_train)

适用于任何生成器，并支持约束条件：

from tabgan import TabGANTransformer、GANGenerator 和 RangeConstraint

transformer = TabGANTransformer(
    generator_class=GANGenerator,
    gen_x_times=2.0,
    gen_params={"batch_size": 500, "epochs": 10, "patience": 5},
    constraints=[RangeConstraint("age", min_val=0, max_val=120)],
)

X_augmented = transformer.fit_transform(X_train, y_train)
y_augmented = transformer.get_augmented_target()

AutoSynth

不知道哪种生成器最适合您的数据？AutoSynth 会运行所有生成器，并根据质量和隐私评分选出最佳方案：

from tabgan import AutoSynth

result = AutoSynth(df, target_col="label").run()

print(result.report)
#   生成器          状态  评分  质量  隐私  行数  时间 (s)
# 0 GAN (CTGAN)        OK      0.847  0.891    0.743    165   12.3
# 1 Forest Diffusion   OK      0.812  0.834    0.761    165   45.1
# 2 Random Baseline    OK      0.654  0.621    0.732    165   0.1

best_synthetic = result.best_data
print(f"获胜者: {result.best_name}")

自定义评分权重：

result = AutoSynth(
    df,
    target_col="label",
    quality_weight=0.5,   # 平等权重
    privacy_weight=0.5,
).run()

HuggingFace Hub 集成

只需一次调用，即可从 HuggingFace Hub 合成任意表格数据集：

from tabgan import synthesize_hf_dataset

# 自动加载 → 生成 → 评估
result = synthesize_hf_dataset("scikit-learn/iris", target_col="target")
print(result.synthetic_df.head())
print(f"质量: {result.quality_summary['overall_score']}")

# 将合成数据推回 Hub
result = synthesize_hf_dataset(
    "scikit-learn/iris",
    target_col="target",
    push_to_hub=True,
    hub_repo_id="your-username/iris-synthetic",
)

命令行界面

tabgan-generate \
    --input-csv train.csv \
    --target-col target \
    --generator gan \
    --gen-x-times 1.5 \
    --cat-cols year,gender \
    --output-csv synthetic_train.csv

流程架构

输入（train_df、target、test_df）
  |
  v
[预处理] --> 验证数据框，准备列
  |
  v
[生成]  --> CTGAN / ForestDiffusion / GReaT LLM / 随机采样
  |
  v
[后处理] --> 基于分位数的过滤，使其符合测试分布
  |
  v
[对抗性过滤] --> LightGBM 分类器移除不相似样本
  |
  v
输出（synthetic_df、synthetic_target）

基准测试结果

归一化 ROC AUC 分数（越高越好）：

## 引用

```bibtex
@misc{ashrapov2020tabular,
    title={用于不均衡分布的表格 GAN},
    author={Insaf Ashrapov},
    year={2020},
    eprint={2010.00638},
    archivePrefix={arXiv},
    primaryClass={cs.LG}
}

参考文献

1. 徐亮，维拉马切内尼，K.（2018）。使用生成对抗网络合成表格数据。arXiv:1811.11264。
2. 若利库尔-马蒂诺，A.，法特拉斯，K.，卡赫曼，T.（2023）。通过扩散模型和基于流的梯度提升树生成与填补表格数据。SamsungSAILMontreal/ForestDiffusion。
3. 徐亮，斯库拉里杜，M.，库埃斯塔-因凡特，A.，维拉马切内尼，K.（2019）。使用条件生成对抗网络建模表格数据。NeurIPS。
4. 鲍里索夫，V.，塞斯勒，K.，莱曼，T.，帕韦尔奇克，M.，卡斯内奇，G.（2023）。语言模型是真实的表格数据生成器。ICLR。

许可证

Apache许可证2.0 — 详情请参阅LICENSE文件。

TabGAN 快速上手指南

TabGAN 是一个用于生成高质量合成表格数据的 Python 库，支持 GAN、扩散模型和大语言模型（LLM）等多种生成方法。它提供统一的 API，能够处理混合数据类型（连续值、类别值、自由文本），并内置对抗性过滤机制以确保生成数据分布与真实数据一致。

环境准备

• 操作系统：Linux, macOS, Windows
• Python 版本：3.8 - 3.12
• 核心依赖：
  • pandas
  • numpy
  • scikit-learn
  • torch (PyTorch，用于 GAN 和扩散模型)
  • transformers (可选，用于 LLM 生成)

提示：建议在使用前确保已安装 PyTorch。若需加速下载，可使用国内镜像源配置 pip。

安装步骤

使用 pip 直接安装：

pip install tabgan

推荐使用国内镜像源加速安装：

pip install tabgan -i https://pypi.tuna.tsinghua.edu.cn/simple

若需使用基于大语言模型的生成器（LLMGenerator），请额外安装相关依赖：

pip install transformers accelerate -i https://pypi.tuna.tsinghua.edu.cn/simple

基本使用

以下示例展示如何使用默认的 GANGenerator 快速生成合成数据。

1. 导入库并准备数据

import pandas as pd
import numpy as np
from tabgan.sampler import GANGenerator

# 构造示例训练数据 (特征)
train = pd.DataFrame(np.random.randint(-10, 150, size=(150, 4)), columns=list("ABCD"))

# 构造示例目标变量
target = pd.DataFrame(np.random.randint(0, 2, size=(150, 1)), columns=list("Y"))

# 构造示例测试数据 (用于分布对齐)
test = pd.DataFrame(np.random.randint(0, 100, size=(100, 4)), columns=list("ABCD"))

2. 生成合成数据

调用 generate_data_pipe 方法即可完成“生成 -> 后处理 -> 对抗过滤”的全流程：

# 初始化生成器并生成数据
new_train, new_target = GANGenerator().generate_data_pipe(train, target, test)

# 查看生成结果
print(f"原始数据形状：{train.shape}")
print(f"合成数据形状：{new_train.shape}")
print(new_train.head())

3. 切换其他生成器（可选）

TabGAN 支持通过更换类名轻松切换底层算法，API 保持一致：

from tabgan.sampler import ForestDiffusionGenerator, BayesianGenerator, LLMGenerator

# 使用基于树的扩散模型
new_train_diff, _ = ForestDiffusionGenerator().generate_data_pipe(train, target, test)

# 使用贝叶斯方法 (速度快，保留相关性)
new_train_bayes, _ = BayesianGenerator().generate_data_pipe(train, target, test)

# 使用大语言模型 (适合包含文本列的数据)
# new_train_llm, _ = LLMGenerator().generate_data_pipe(train, target, test)

注意：生成的数据内部以浮点数处理。如果原始数据包含整数，建议在生成后对结果进行取整操作。