'''
將原始資料的word特徵數字化為countvector特徵，並將結果儲存到本地

article特徵可做類似處理

'''
import pandas as pd
from sklearn.feature_extraction.text import countvectororizer
import pickle
import time

t_start = time.time()

"""=====================================================================================================================
1 資料預處理
"""
# 讀取原始資料train和test檔案
df_train=pd.read_csv('train_set.csv')
df_test=pd.read_csv('test_set.csv')

# 刪除特徵article，只保留特徵word
df_train.drop(columns='article', inplace=True)
df_test.drop(columns='article', inplace=True)

# 按行拼接df_train和df_test
df_all = pd.concat(objs=[df_train, df_test], axis=0, sort=True)

# 獲取train檔案中的特徵class
y_train = (df_train['class'] - 1).values

"""=====================================================================================================================
2 特徵工程
"""
print('2 特徵工程')
# 將原始資料數字化為countvector特徵

vectorizer = countvectororizer(ngram_range=(1, 2), min_df=100, max_df=0.8)
vectorizer.fit(df_all['word_seg'])
x_train = vectorizer.transform(df_train['word_seg'])
x_test = vectorizer.transform(df_test['word_seg'])

"""=====================================================================================================================
3 儲存至本地
"""
print('3 儲存特徵')
data = (x_train, y_train, x_test)
with open('countvector_word.pkl', 'wb') as f:
	pickle.dump(data,f)

t_end = time.time()
print("共耗時：{}min".format((t_end-t_start)/60))