python sklearn 主题分类

文件下载地址：
链接: https://pan.baidu.com/s/1Rlokc8Vj3GQnBM-6Dt2h7w 提取码: 7x6n

"""
    自然语言处理 新闻文件主题识别
"""
import sklearn.datasets as sd
import sklearn.feature_extraction.text as ft
import sklearn.naive_bayes as nb
import numpy as np

# 加载新闻数据
train = sd.load_files("./20news", encoding="latin1", shuffle=True,
                      random_state=7)

# 数据类型转换
train_data = np.array(train.data)   # 文本
train_y = np.array(train.target)    # 类型编码值
categrories = np.array(train.target_names)  # 类型名称

print(train_data.shape, train_y.shape, categrories.shape)
# print(train_data[0], train_y[0], categrories)

# 构建tfdif模型 词袋模型 统计单词出现的次数
cv = ft.CountVectorizer()
# 统计词频
bow = cv.fit_transform(train_data)

tt = ft.TfidfTransformer()
# 逆文档频率
tfdif = tt.fit_transform(bow)

# 模型建立 训练基于多项分布的朴素贝叶斯
model = nb.MultinomialNB()
model.fit(tfdif, train_y)

# 测试数据
test_data = [
    'The curveballs of right handed pitchers tend to curve to the left',
    'Caesar cipher is an ancient form of encryption',
    'This two-wheeler is really good on slippery roads']
test_bow = cv.transform(test_data)
test_tfdif = tt.transform(test_bow)

# 预测数据
prd_test_y = model.predict(test_tfdif)

print(categrories[prd_test_y])