从零实现BERTopic:OpenAI Embedding + UMAP + HDBSCAN + c-TF-IDF法律文书主题建模

核心技术: 不依赖bertopic库,从零实现四步管线
论文灵感: BERTopic (Grootendorst, 2022) + UMAP (McInnes, 2018) + HDBSCAN (Campello, 2013)
特色: 纯API + NumPy实现,无GPU,优雅降级


一、前言:为什么法律文书需要主题建模?

当律所/法院的文书库积累到数百甚至上千份时,简单的"民事/刑事/行政"分类已经远远不够:

  • 民事案件里可能包含"合同纠纷"、“房产纠纷”、“婚姻家事”、"知识产权"等几十个子类
  • 相似案由的文书在不同时期可能有完全不同的裁判趋势
  • 新兴领域(如数据隐私、AI责任)的案件无法用预定义分类覆盖

BERTopic 的优势

传统的 LDA 主题建模基于词袋模型,主题连贯性差。BERTopic 利用预训练语言模型的语义嵌入,能发现更有意义的主题:

对比维度 LDA BERTopic
输入表示 BoW (词袋) 语义向量 (BERT/OpenAI)
主题连贯性
可解释性 一般 c-TF-IDF词清晰
自动确定主题数 否 (需指定K) 是 (HDBSCAN自适应)
处理长文本

二、不依赖bertopic库的原因

官方 bertopic 库虽然好用,但依赖链很深:

bertopic → sentence-transformers → torch → transformers → ...

仅 PyTorch 就需要 2GB+ 磁盘空间和 GPU 加速。对于一个 Web 应用,这太重了。

我们的方案:用 OpenAI API 替代本地 sentence-transformers,用纯 NumPy 实现 UMAP/HDBSCAN 的降级方案,从零复现 BERTopic 的核心管线。


三、四步管线详解

3.1 整体架构

Step 1: OpenAI Embedding              Step 2: UMAP 降维
┌─────────────────────┐    1536维    ┌──────────────────┐
│ text-embedding-3-   │ ─────────→  │ UMAP(n=5)        │    5维
│ small               │             │ or PCA fallback   │ ────┐
└─────────────────────┘             └──────────────────┘     │
                                                              ▼
Step 4: c-TF-IDF 主题词            Step 3: HDBSCAN 聚类
┌─────────────────────┐            ┌──────────────────┐
│ jieba分词            │ ←────────  │ HDBSCAN           │
│ TF * IDF per topic  │   labels   │ or cosine cluster │
└─────────────────────┘            └──────────────────┘

3.2 Step 1: OpenAI Embedding

def _get_embeddings_batch(texts: list[str], client: OpenAI) -> np.ndarray:
    """批量获取 embedding"""
    embeddings = []
    for start in range(0, len(texts), 100):
        batch = texts[start:start+100]
        try:
            resp = client.embeddings.create(
                input=batch,
                model="text-embedding-3-small"
            )
            for item in resp.data:
                embeddings.append(item.embedding)
        except Exception:
            # fallback: 随机向量(测试/离线模式)
            for _ in batch:
                embeddings.append(np.random.randn(1536).tolist())
    return np.array(embeddings)

设计考量

  • text-embedding-3-small: 1536维,$0.02/1M tokens,性价比最高
  • 批处理100条/批,避免API限流
  • 随机向量fallback保证系统可用性

3.3 Step 2: UMAP 降维(优雅降级至PCA)

UMAP (Uniform Manifold Approximation and Projection) 是BERTopic的标配降维方法,但安装 umap-learn 可能有依赖问题。我们实现了优雅降级:

def _simple_umap(embeddings: np.ndarray, n_components: int = 2, 
                  n_neighbors: int = 5) -> np.ndarray:
    n_samples = embeddings.shape[0]
    if n_samples <= n_components:
        return np.zeros((n_samples, n_components))

    # 优先使用 umap-learn
    try:
        import umap
        n_neighbors_actual = min(n_neighbors, n_samples - 1)
        reducer = umap.UMAP(
            n_components=n_components,
            n_neighbors=max(2, n_neighbors_actual),
            metric='cosine',
            random_state=42,
        )
        return reducer.fit_transform(embeddings)
    except ImportError:
        pass

    # Fallback: PCA (纯NumPy实现)
    centered = embeddings - embeddings.mean(axis=0)
    if n_samples < embeddings.shape[1]:
        # 小样本: 用 n×n 协方差矩阵
        cov = centered @ centered.T
        eigvals, eigvecs = np.linalg.eigh(cov)
        idx = np.argsort(eigvals)[::-1][:n_components]
        components = centered.T @ eigvecs[:, idx]
        norms = np.linalg.norm(components, axis=0, keepdims=True)
        norms[norms == 0] = 1
        components = components / norms
        reduced = centered @ components
    else:
        # 大样本: 用 d×d 协方差矩阵
        cov = centered.T @ centered / n_samples
        eigvals, eigvecs = np.linalg.eigh(cov)
        idx = np.argsort(eigvals)[::-1][:n_components]
        reduced = centered @ eigvecs[:, idx]
    return reduced

PCA vs UMAP 的差异

  • PCA 保留全局线性结构,UMAP 保留局部拓扑结构
  • 对于文档数 < 50 的场景,PCA 和 UMAP 的聚类效果差异不大
  • PCA 是纯数学运算,无额外依赖

3.4 Step 3: HDBSCAN 密度聚类(优雅降级至余弦聚类)

HDBSCAN 是一种层次密度聚类算法,核心优势:自动确定簇数量,能识别噪声点

def _simple_hdbscan(embeddings: np.ndarray, min_cluster_size: int = 2) -> np.ndarray:
    n = embeddings.shape[0]
    if n < 2:
        return np.array([-1] * n)

    # 优先使用 hdbscan 库
    try:
        import hdbscan
        clusterer = hdbscan.HDBSCAN(
            min_cluster_size=min_cluster_size,
            min_samples=1,
            metric='cosine',
        )
        return clusterer.fit_predict(embeddings)
    except ImportError:
        pass

    # Fallback: 贪心余弦相似度聚类
    norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
    norms[norms == 0] = 1
    normalized = embeddings / norms
    sim_matrix = normalized @ normalized.T  # n×n 相似度矩阵

    labels = np.full(n, -1)
    cluster_id = 0
    for i in range(n):
        if labels[i] >= 0:
            continue
        similar = np.where(sim_matrix[i] > 0.7)[0]  # 阈值0.7
        if len(similar) >= min_cluster_size:
            for j in similar:
                if labels[j] < 0:
                    labels[j] = cluster_id
            cluster_id += 1
    return labels

聚类阈值 0.7 的选择

  • 法律文书同主题的余弦相似度通常在 0.7-0.9
  • 不同主题在 0.3-0.6
  • 0.7 是一个较保守的值,宁可漏聚不可错聚

3.5 Step 4: c-TF-IDF 主题词提取

c-TF-IDF (class-based TF-IDF) 是 BERTopic 的核心创新——它不是对每个文档计算TF-IDF,而是对**每个主题(簇)**计算:

c-TFIDF(w,c)=TF(w,c)×log⁡(1+Aˉfw)c\text{-}TFIDF(w, c) = TF(w, c) \times \log\left(1 + \frac{\bar{A}}{f_w}\right)c-TFIDF(w,c)=TF(w,c)×log(1+fwAˉ)

其中:

  • TF(w,c)TF(w, c)TF(w,c): 词 www 在主题 ccc 的所有文档中的总频率
  • Aˉ\bar{A}Aˉ: 所有主题的平均词数
  • fwf_wfw: 词 www 在所有文档中的全局频率
def _extract_topic_words(texts: list[str], labels: np.ndarray, top_n: int = 8):
    try:
        import jieba
        def tokenize(text):
            return [w.strip() for w in jieba.cut(text) if len(w.strip()) >= 2]
    except ImportError:
        def tokenize(text):
            return re.findall(r'[\u4e00-\u9fff]{2,}', text)

    stop_words = {'的', '了', '在', '是', '我', '有', '和', '就', '不', '人', ...}

    # 按主题分组文档
    topic_docs = {}
    for i, label in enumerate(labels):
        if label < 0: continue
        topic_docs.setdefault(int(label), []).append(texts[i])

    # 全局词频
    global_counter = Counter()
    for text in texts:
        words = [w for w in tokenize(text) if w not in stop_words]
        global_counter.update(words)

    total_words = sum(global_counter.values())
    avg_per_topic = total_words / max(len(set(labels) - {-1}), 1)

    topic_words = {}
    for label, docs in topic_docs.items():
        topic_counter = Counter()
        for doc in docs:
            words = [w for w in tokenize(doc) if w not in stop_words]
            topic_counter.update(words)

        scored = []
        for word, tf in topic_counter.items():
            idf = np.log(1 + avg_per_topic / max(global_counter[word], 1))
            score = tf * idf
            scored.append({"word": word, "score": round(float(score), 4), "count": tf})

        scored.sort(key=lambda x: x["score"], reverse=True)
        topic_words[label] = scored[:top_n]

    return topic_words

为什么 c-TF-IDF 比普通 TF-IDF 好?

  • 普通TF-IDF在文档级计算,短文档的词权重被放大
  • c-TF-IDF把同一主题的所有文档合并,消除了文档长度差异
  • 提取出的词更能代表整个主题,而非某一篇文档

四、主函数与输出格式

def discover_topics() -> dict:
    doc_infos = _get_doc_texts()
    if len(doc_infos) < 2:
        return {"topics": [], "message": "至少需要2篇文档"}

    texts = [d["text"] for d in doc_infos]
    client = _get_client()

    # Step 1-4
    embeddings = _get_embeddings_batch(texts, client)
    reduced_cluster = _simple_umap(embeddings, n_components=min(5, len(texts)-1))
    reduced_viz = _simple_umap(embeddings, n_components=2)  # 2D用于可视化
    labels = _simple_hdbscan(reduced_cluster, min_cluster_size=2)
    topic_words = _extract_topic_words(texts, labels)

    return {
        "topics": [{
            "id": label,
            "label": " / ".join(w["word"] for w in words[:3]),  # 自动标签
            "words": words,
            "doc_count": len([...]),
            "doc_ids": [...],
        } for label, words in sorted(topic_words.items())],
        "scatter_data": [{
            "doc_id": "...", "x": float, "y": float,
            "topic_id": int, "filename": "..."
        } for each doc],
        "doc_topic_map": {"doc_id": topic_id},
        "noise_count": int,
    }

五、前端可视化

5.1 ECharts 散点图

const scatterOption = {
  title: { text: '文档主题分布(UMAP降维)' },
  series: topicIds.map((tid, idx) => ({
    name: tid === -1 ? '噪声' : `主题${tid}`,
    type: 'scatter',
    data: scatterData.filter(d => d.topic === tid).map(d => [d.x, d.y]),
    symbolSize: tid === -1 ? 6 : 10,  // 噪声点更小
    itemStyle: { 
      color: tid === -1 ? '#555' : palette[idx % palette.length],
      opacity: tid === -1 ? 0.4 : 0.85  // 噪声点更透明
    },
  })),
}

5.2 主题卡片

每个主题显示:

  • 主题编号 + 自动标签(前3个c-TF-IDF词)
  • 文档数量
  • 关键词列表(opacity反映权重)
<div v-for="t in topics" :key="t.id" class="topic-card">
  <span class="topic-id">主题 {{ t.id }}</span>
  <span class="topic-count">{{ t.doc_count }} 篇文档</span>
  <div class="topic-words">
    <span v-for="w in t.words" :key="w.word" 
          :style="{ opacity: 0.5 + 0.5 * w.weight }">
      {{ w.word }}
    </span>
  </div>
</div>

六、依赖管理策略

依赖 用途 必需? 降级方案
numpy 数值计算
openai Embedding API 随机向量
jieba 中文分词 正则 r'[\u4e00-\u9fff]{2,}'
umap-learn UMAP降维 PCA (numpy)
hdbscan 密度聚类 余弦阈值聚类 (numpy)

最小安装: pip install numpy openai 即可运行。
推荐安装: pip install numpy openai jieba umap-learn hdbscan 获得最佳效果。


七、常见问题

Q1: 为什么UMAP比t-SNE更适合BERTopic?
A: UMAP保留了全局结构(不同簇之间的距离有意义),而t-SNE倾向于均匀分散所有簇。对聚类来说,保留全局结构更重要。

Q2: 文档数量 < 10 时效果如何?
A: PCA降级方案仍然可用,但主题数会很少(1-2个)。建议至少20篇文档才能发现有意义的主题分布。

Q3: 如何评估主题质量?
A: 可用 Topic Coherence (C_V) 指标。直觉方法:看主题词是否能被人类理解为一个有意义的类别。例如 “合同/违约/赔偿” 是好主题,“一个/进行/情况” 是差主题。

Q4: 上传新文档后需要重新建模吗?
A: 目前是全量重算。优化方向:用 Online UMAP + 增量 HDBSCAN 实现在线更新。


八、总结

本文贡献

  1. 从零实现BERTopic四步管线,不依赖bertopic/torch/transformers
  2. 三层优雅降级:UMAP→PCA, HDBSCAN→cosine clustering, jieba→regex
  3. 完整的前端可视化(散点图 + 主题卡片)

与官方BERTopic的差异

对比项 官方BERTopic 本实现
Embedding sentence-transformers (本地) OpenAI API (云端)
降维 UMAP (必需) UMAP / PCA (可选)
聚类 HDBSCAN (必需) HDBSCAN / cosine (可选)
主题表示 c-TF-IDF + MaximalMarginalRelevance c-TF-IDF
依赖大小 ~2GB (含PyTorch) ~10MB (NumPy+OpenAI)
GPU需求 推荐 不需要

优化方向

  • 增量更新: Online UMAP + 增量HDBSCAN
  • 主题标签优化: 用LLM给每个主题生成自然语言标签
  • 时间维度: 追踪主题随时间的演化趋势
  • 交互式探索: 点击散点图的点 → 查看文档详情

Logo

有“AI”的1024 = 2048,欢迎大家加入2048 AI社区

更多推荐