git_issue/word2vec-test.py

75 lines
2.1 KiB
Python

__author__ = 'qiangge'
from sklearn.feature_extraction.text import CountVectorizer, TfidfVectorizer
from nltk.tokenize import TweetTokenizer
import cPickle as pickle
import helper
import pymysql
import dao
import logging
from gensim.models import word2vec
logging.basicConfig(format='%(asctime)s : %(levelname)s : %(message)s', level=logging.INFO)
# path_oringan = 'result2'
# preprocess of data
def data_preprocess(project_name, methold_name = ''):
# project_name = "owncloud"
path = 'result_wordembedding/'+project_name+'/'
helper.mkdir(path)
# if len(methold_name):
# data_path = path + methold_name + '/data/'
# else:
# data_path = path + 'data/'
# helper.mkdir(data_path)
print('=' * 80)
print("get data: ")
cur = dao.get_data(project_name)
fetchall = cur.fetchall()
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: word embedding")
# 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]))
# TF-IDF
# vectorizer = TfidfVectorizer(sublinear_tf=True, max_df=0.5,
# stop_words='english', tokenizer=helper.tokenize_help)
# X = vectorizer.fit_transform(train_data)
sentences = []
for text in train_data:
sentences.append(helper.tokenize_help(text))
model = word2vec.Word2Vec(sentences,size=200,sg=1)
model.save("ttt")
# model = word2vec.Word2Vec.load('ttt')
# for sentence in sentences:
# for word in sentence:
# print(word)
# if model[word].any():
# print('get')
print("done")
model = word2vec.Word2Vec.load('ttt')
# data_preprocess('24444')
print('finish')