AI ninja project [day 4] AI RPA系统--名片model建立篇

好的，我们继续利用上一篇文章建立的csv资料及来进行训练，
首先我们建立一个功能来帮助从字串提取特徵:
strcount.py

# -*- coding:utf8 -*-
import string
import numpy as np

def str_count(s,mode):
    count_en = count_dg = count_sp = count_zh = count_pu = count_low=0
    s_len = len(s)
    for c in s:
        # 英文
        if c in string.ascii_letters:
            count_en += 1
            #小写
            if c.islower():
                count_low+=1

        # 数字
        elif c.isdigit():
            count_dg += 1
        # 空格
        elif c.isspace():
            count_sp += 1
        # 非英文其他字元
        elif c.isalpha():
            count_zh += 1
        # 特殊字符
        else:
            count_pu += 1

   
    if mode == 'dict':    
        rtn_dict ={'total':s_len, 'zh':count_zh ,'en':count_en ,'digit':count_dg,'punc':count_pu}
        return rtn_dict
    elif mode == 'numpy' :
        rtn_np =  np.array([[count_en,count_dg,count_pu,count_zh]], dtype=object) 
        return rtn_np
    else:
        return NONE

因次这里我以一个字串有多少英文字(程序中能区分大小写)、特殊字元、
中文字(目前程序是计算排除了英文的其他字元)、
数字、空格(上一篇前处理已经排除，所以回传没有列入)。

create_model.py

import pandas as pd
from strcount import str_count
from sklearn import preprocessing
from sklearn import tree
from joblib import dump, load

df = pd.read_csv('train.csv')

le = preprocessing.LabelEncoder()

df['label'] = le.fit_transform(df['label'])

# print(df.head())

enList=[]
digList=[]
puncList=[]
zhList=[]

for cont_str in df["content"]:
    count_dict= str_count(cont_str,'dict')
    # print(count_dict)
    enList.append(count_dict['en'])
    digList.append(count_dict['digit'])
    puncList.append(count_dict['punc'])
    zhList.append(count_dict['zh'])
    

df['en'] = enList
df['digit'] = digList
df['punc'] = puncList
df['zh'] = zhList

print(df.head())


X = df.iloc[:, 2:]
y = df.iloc[:, 1]

decision_tree = tree.DecisionTreeClassifier(criterion='gini')

decision_tree.fit(X, y)

再来引入套件(pandas、sklearn、上面写的strcount、joblib)
我们把每个类别使用LabelEncoder这个功能 mapping 到某个数值。

将特徵数值加入到新的栏位中。
使用pandas分隔特徵(X)与标签(Y)，
使用决策树演算法来拟合训练。

测试模型是否能正常预测。

print(le.inverse_transform(decision_tree.predict([(10,0,5,3)])))

顺边查看类别对应的数值。

print(le.inverse_transform([0]))
print(le.inverse_transform([1]))
print(le.inverse_transform([2]))
print(le.inverse_transform([3]))
print(le.inverse_transform([4]))
print(le.inverse_transform([5]))
print(le.inverse_transform([6]))
print(le.inverse_transform([7]))
print(le.inverse_transform([8]))

我们用joblib套件将模型储存成joblib物件。

dump(decision_tree, 'tree.joblib')

未来使用:
load_model.py

import os
import pymongo
# import openpyxl
from strcount import str_count
import pandas as pd
from joblib import dump, load
import io
from sklearn import preprocessing


label_convert_dict = {0:'address',1:'company',2:'contact_info',3:'email',4:'line',5:'name',6:'position',7:'tax',8:'unknown'}

clf = load('tree.joblib') 

count=str_count('苏韦文',mode='numpy')
print(count)


my_prediction = clf.predict(count)
print(my_prediction[0])
predict_answer =  my_prediction[0]
print(label_convert_dict[predict_answer])

可以使用一个字典来对应预测数值与答案，
载入模型进行预测答案。
我建立的模型是有可能无法区别职务以及姓名的，
所以可以多加一些字串规则或是使用正规表达式来进一步筛选与分类，
帮助未来预测上的表现。

至於用哪一种演算法，
也可以把决策树更换为KNN或是SVM测试看看哪一种效果比较好，
差异并没有很大。