优化Pandas生成NLP语料库关键词索引的高效方法

针对你的NLP语料库关键词生成代码，核心瓶颈在于低效的循环遍历、内存频繁扩容和线性时间复杂度的成员检查，以下是具体优化措施：

1. 提前预处理静态资源 #

把停用词转成集合，避免每次循环重复调用stopwords.words('english')，集合的in操作是O(1)，远快于列表的O(n)：

# 在类初始化或方法开头提前加载
self.stopwords = set(stopwords.words('english'))

2. 替换iterrows()为apply() #

iterrows()是Pandas中最慢的遍历方式之一，改用apply()结合自定义函数处理每一行，效率提升显著。同时用集合替代列表做去重，i not in s的检查从O(n)降到O(1)。

3. 避免逐行扩容DataFrame #

原代码中self.keyword_index.loc[len(...)]每次添加行都会重新分配内存，改用列表收集所有(paper_id, keyword)对，最后一次性生成DataFrame，这是Pandas批量处理的标准高效做法。

4. 用集合维护全局关键词映射 #

self.keyword_map改用集合存储，避免重复添加，最后再转成列表生成CSV，成员检查从O(n)降到O(1)。

优化后的完整代码 #

import pandas as pd
from tqdm import tqdm
import nltk
from nltk.corpus import stopwords
# 首次运行需下载停用词
# nltk.download('stopwords')

def isHyperlink(text):
    # 保留你的原有实现
    return 'http' in text or 'www' in text

class Doc_Finder():
    def __init__(self, corpus_file):
        self.corpus = pd.read_csv(corpus_file, dtype={'paper_id': str})
        self.corpus_dir = corpus_file
        # 提前加载停用词为集合
        self.stopwords = set(stopwords.words('english'))

    def make_keywords(self, save=False, load=False, **kwargs):
        if save and load:
            raise Exception("Invalid Parameters")
        if save:
            total = kwargs['size']
            # 批量收集关键词索引数据
            keyword_index_data = []
            # 用集合维护全局关键词，避免重复
            keyword_map_set = set()

            def process_row(row):
                text = str(row['full_text_cleaned'])
                doc_keywords = set()
                for token in text.split(' '):
                    if (token not in doc_keywords
                        and not token.isnumeric()
                        and token not in self.stopwords
                        and token.isalnum()
                        and not isHyperlink(token)):
                        doc_keywords.add(token)
                        keyword_map_set.add(token)
                # 生成当前文档的关键词配对
                for kw in doc_keywords:
                    keyword_index_data.append([row['paper_id'], kw])

            # 遍历并显示进度
            _ = tqdm(self.corpus.head(total).apply(process_row, axis=1), total=total)

            # 一次性生成DataFrame
            self.keyword_index = pd.DataFrame(keyword_index_data, columns=['id', 'keyword'])
            self.keyword_map = list(keyword_map_set)

            # 保存文件
            idx_dir = './keyword_index_' + str(total) + '.csv'
            self.keyword_index.to_csv(idx_dir, index=False)
            tempdf = pd.DataFrame(data=self.keyword_map, columns=['keyword'], dtype=str)
            map_dir = './keyword_map_' + str(total) + '.csv'
            tempdf.to_csv(map_dir, index=False)
        elif load:
            self.keyword_index = pd.read_csv(kwargs['index_dir'])
            self.keyword_map = pd.read_csv(kwargs['map_dir'])['keyword'].tolist()

额外加速建议 #

1. 多进程并行处理：如果机器有多个CPU核心，可使用swifter库让apply()自动并行化，进一步提升速度：

   # 安装swifter：pip install swifter
   import swifter
   
   # 替换原apply为swifter.apply
   _ = tqdm(self.corpus.head(total).swifter.apply(process_row, axis=1), total=total)
   

2. 优化文本分割：如果full_text_cleaned存在连续空格，用正则分割替代split(' ')，避免空字符串干扰：

   import re
   for token in re.split(r'\s+', text):
   

3. 减少内存占用：读取CSV时明确指定列类型，避免Pandas自动推断占用额外内存，比如self.corpus = pd.read_csv(corpus_file, dtype={'paper_id': str})。

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