git_issue/kmeans.py

88 lines
2.7 KiB
Python

from sklearn.cluster import KMeans
from sklearn.feature_extraction.text import TfidfVectorizer
import helper
import cPickle as pickle
import dao
def data_preprocess(project_name, methold_name = ''):
# project_name = "owncloud"
path = 'result/'+project_name+'/'
helper.mkdir(path)
# if len(methold_name):
# data_path = path + methold_name + '/data/'
# else:
# data_path = path + 'data/'
# helper.mkdir(data_path)
f_train = path +'train.pkl'
f_target = path + 'target.pkl'
f_id = path + 'id.pkl'
f_vect = path + 'vect.pkl'
f_train_data = path + 'train_data.pkl'
f_cur_all = path + 'cur_all.pkl'
print('=' * 80)
print("get data: ")
cur = dao.get_data(project_name)
fetchall = cur.fetchall()
with open(f_cur_all, 'w') as f: # open file with write-mode
pickle.dump(fetchall, f)
train_data = []
train_target = []
x_id = []
for r in fetchall:
str = helper.filter_str(r[1]+".\n"+r[2])
# str = helper.filter_code(str)
train_data.append(str)
train_target.append(r[0])
x_id.append(r[3])
print("data length is : ", len(train_target))
cur.close()
print('_' * 80)
print("processing data: TF-IDF")
# categories = ['bug','enhancement','feature','documentation','question','others']
categories = ['bug','enhancement']
y = [None]*len(train_target)
for i in range(len(train_target)):
y[i] = (categories.index(train_target[i]))
with open(f_train_data, 'w') as f: # open file with write-mode
pickle.dump(train_data, f)
# TF-IDF
vectorizer = TfidfVectorizer(sublinear_tf=True, max_df=0.5,
stop_words='english', tokenizer=helper.tokenize_help)
X = vectorizer.fit_transform(train_data)
# store data preprocess result in file
with open(f_train, 'w') as f: # open file with write-mode
pickle.dump(X, f) # serialize and save object
with open(f_target, 'w') as f: # open file with write-mode
pickle.dump(y, f) # serialize and save object
with open(f_id, 'w') as f: # open file with write-mode
pickle.dump(x_id, f) # serialize and save object
with open(f_vect, 'w') as f: # open file with write-mode
pickle.dump(vectorizer, f) # serialize and save object
print("done")
return X,y,x_id,vectorizer
project_name = "11450"
X,y,x_id,vect = helper.get_tfidf_data(project_name)
# kmeans
y_pred = KMeans(n_clusters= 10).fit_predict(X)
# dao save in db
dao.save_kmeans_result(x_id,y_pred,project_name)