Top2Vec模型在Colab运行abcnews数据集时停滞的解决方法咨询

问题背景 #

使用Top2Vec处理小型JSON数据集（vol7.json）时代码运行正常，但切换到abcnews-date-text.csv数据集时，Colab陷入无限运行状态，核心代码如下：

# Extract the text data from the dataset
documents = data['headline_text'].tolist()

# Initialize Top2Vec model
top2vec_model = Top2Vec(documents, embedding_model="distiluse-base-multilingual-cased")

# Get the number of topics
num_topics = 5  # You can adjust this number according to your preference

# Get the top topics
top_topics = top2vec_model.get_topics(num_topics)

# Print the top topics
for i, topic in enumerate(top_topics):
    print(f"Topic {i+1}: {', '.join(topic)}")

核心原因与解决思路 #

abcnews-date-text.csv包含超百万条新闻标题，远大于原JSON数据集的规模，Top2Vec全量处理会耗尽Colab资源导致卡顿或无限运行，可从以下方向优化：

• 缩小数据集规模先做测试
  先抽取数据集的子集（比如前10000条）验证代码逻辑，确认模型能正常运行后再逐步扩大规模：

  # 抽取前10000条数据测试
  documents = data['headline_text'].head(10000).tolist()
  

• 切换轻量嵌入模型
  distiluse-base-multilingual-cased是大模型，处理海量数据时计算量极大。换成更轻量的英文专用模型（如all-MiniLM-L6-v2），能大幅降低计算负载：

  top2vec_model = Top2Vec(documents, embedding_model="all-MiniLM-L6-v2")
  

• 启用Colab更高配置
  在Colab中切换到GPU或TPU运行时：

  1. 点击菜单栏「Runtime」→「Change runtime type」
  2. 选择「Hardware accelerator」为GPU/TPU，保存后重新运行代码

• 分步拆解模型训练流程
  把Top2Vec的训练拆分为嵌入计算、降维、聚类三个步骤，单独监控每个步骤的运行状态，定位卡顿环节：

  # 先单独计算文本嵌入
  from sentence_transformers import SentenceTransformer
  model = SentenceTransformer("all-MiniLM-L6-v2")
  embeddings = model.encode(documents, show_progress_bar=True)
  
  # 再初始化Top2Vec并传入预计算的嵌入
  top2vec_model = Top2Vec(documents, embeddings=embeddings)
  

• 设置聚类参数限制计算量
  Top2Vec默认的聚类算法会自适应处理数据，可手动设置speed参数为"fast-learn"，强制使用更快的聚类逻辑：

  top2vec_model = Top2Vec(documents, embedding_model="all-MiniLM-L6-v2", speed="fast-learn")
  

内容的提问来源于stack exchange，提问作者PS Nayak