如何序列化含自定义分词器的scikit-learn TfidfVectorizer？

你遇到的问题本质是pickle对函数的序列化机制导致的：pickle不会保存函数的代码本身，只保存函数的名称和它所在的模块路径。当你加载时，Python需要在当前运行环境中找到完全同名、同模块的函数定义，否则就会抛出AttributeError。

下面给你几个可行的解决方案，按推荐程度排序：

1. 使用Cloudpickle替代Pickle（最省心） #

Cloudpickle是pickle的增强版，能直接序列化自定义函数、lambda表达式这类pickle搞不定的对象，不需要额外调整代码结构。

步骤： #

• 先安装cloudpickle：

pip install cloudpickle

• 保存TfidfVectorizer时替换pickle为cloudpickle：

import cloudpickle
import os

dest = os.path.join('classifier','pkl_object')
# 保存整个tfidf对象（包含tokenizer）
with open(os.path.join(dest,'vect.pkl'),'wb') as f:
    cloudpickle.dump(tfidf, f)

• 加载时同样用cloudpickle：

import cloudpickle
import os

with open(os.path.join('classifier','pkl_object','vect.pkl'),'rb') as file:
    vect = cloudpickle.load(file)

这样不需要担心函数定义的位置问题，cloudpickle会把函数的代码一起序列化。

2. 将自定义Tokenizer放到独立模块中（最规范） #

把custom_tokenizer函数放到一个单独的Python模块里，训练和加载时都从这个模块导入，确保pickle保存的函数引用在加载环境中存在。

步骤： #

• 创建一个名为tokenizers.py的文件，把tokenizer和依赖的对象都放进去：

# tokenizers.py
# 确保这里导入vect_tokenizer和lemmatizer，比如：
# from nltk.tokenize import word_tokenize as vect_tokenizer
# from nltk.stem import WordNetLemmatizer
# lemmatizer = WordNetLemmatizer()

def custom_tokenizer(doc):
    tokens = vect_tokenizer(doc)
    return [lemmatizer.lemmatize(token) for token in tokens]

• 训练时从模块导入函数：

from tokenizers import custom_tokenizer
tfidf = TfidfVectorizer(tokenizer=custom_tokenizer, stop_words="english")

# 训练完成后保存tfidf
import pickle
import os
dest = os.path.join('classifier','pkl_object')
with open(os.path.join(dest,'vect.pkl'),'wb') as f:
    pickle.dump(tfidf, f, protocol=4)

• 加载时先导入模块中的函数，再加载pickle文件：

from tokenizers import custom_tokenizer  # 必须先导入，让Python能找到这个函数
import pickle
import os

with open(os.path.join('classifier','pkl_object','vect.pkl'),'rb') as file:
    vect = pickle.load(file)

3. 用类实现Tokenizer（更灵活） #

把tokenizer的逻辑包装成一个类，实现__call__方法，这样pickle可以序列化类实例（会保存实例的状态），不需要依赖外部函数定义。

步骤： #

• 定义Tokenizer类：

class CustomTokenizer:
    def __init__(self, vect_tokenizer, lemmatizer):
        self.vect_tokenizer = vect_tokenizer
        self.lemmatizer = lemmatizer
    
    def __call__(self, doc):
        tokens = self.vect_tokenizer(doc)
        return [self.lemmatizer.lemmatize(token) for token in tokens]

• 训练时创建类实例并传入TfidfVectorizer：

# 假设你已经有vect_tokenizer和lemmatizer的实例
tokenizer = CustomTokenizer(vect_tokenizer, lemmatizer)
tfidf = TfidfVectorizer(tokenizer=tokenizer, stop_words="english")

# 训练后保存tfidf
import pickle
import os
dest = os.path.join('classifier','pkl_object')
with open(os.path.join(dest,'vect.pkl'),'wb') as f:
    pickle.dump(tfidf, f, protocol=4)

• 加载时直接用pickle加载即可，不需要额外导入：

import pickle
import os

with open(os.path.join('classifier','pkl_object','vect.pkl'),'rb') as file:
    vect = pickle.load(file)

补充：你之前代码的小问题 #

你保存Vectorizer时，先dump了custom_tokenizer又dump了best_vector，但加载时只做了一次pickle.load()，应该对应两次load才能拿到两个对象：

# 错误的加载方式
with open(..., 'rb') as file:
    vect = pickle.load(file)  # 只拿到了custom_tokenizer，不是tfidf对象

# 正确的加载方式（如果要保留你原来的保存逻辑）
with open(..., 'rb') as file:
    loaded_tokenizer = pickle.load(file)
    loaded_vect = pickle.load(file)

不过更推荐直接保存整个tfidf对象，而不是分开保存tokenizer和vectorizer。

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