topic-based-retrieval/function.py at master · ntson2002/topic-based-retrieval · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
import glob
import pickle
from sklearn import manifold
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.feature_extraction.text import CountVectorizer
from nltk.corpus import stopwords
from sklearn.metrics.pairwise import cosine_similarity
import json


def getStopWords():
    list = stopwords.words('english')
    list.append("iv")
    list.append("v")
    list.append("iii")
    list.append("rrb")
    list.append("lrb")
    list.append("shall")
    list.append("may")
    list.append("article")
    list.append("paragraph")
    return list


def convert_jsoncorpus_to_tfidf_matrix(path):
    train_set = []
    datas = {}

    # load json file which contains all documents
    with open(path) as data_file:
        jsonData = json.load(data_file)

    files = []  # list of file id or file name
    for i in range(len(jsonData)):
        files.append(jsonData[i]['article_id'])
        sentences = [item['text'] for item in jsonData[i]['sentences']]
        sentences.append(' ' + jsonData[i]['description'])
        content = ' '.join(sentences)
        train_set.append(content)
        datas[jsonData[i]['article_id']] = content

    # convert into TF-IDF vectors
    stopWords = getStopWords()
    vectorizer = TfidfVectorizer(stop_words=stopWords, use_idf=True, sublinear_tf=False, norm='l2', smooth_idf=True)

    # vectorizer = CountVectorizer(stop_words = stopWords, min_df=1)
    trainVectorizerArray = vectorizer.fit_transform(train_set)

    return (files, vectorizer, trainVectorizerArray.toarray(), datas)


def convert_jsoncorpus_to_count_matrix(path):
    train_set = []
    datas = {}

    with open(path) as data_file:
        jsonData = json.load(data_file)

    files = []
    for i in range(len(jsonData)):
        files.append(jsonData[i]['article_id'])
        sentences = [item['text'] for item in jsonData[i]['sentences']]
        sentences.append(' ' + jsonData[i]['description'])
        content = ' '.join(sentences)
        train_set.append(content)
        datas[jsonData[i]['article_id']] = content

    stopWords = getStopWords()
    vectorizer = CountVectorizer(stop_words=stopWords, min_df=1)
    trainVectorizerArray = vectorizer.fit_transform(train_set)
    return (files, vectorizer, trainVectorizerArray.toarray(), datas)


def convert_to_tfidf_matrix(path):
    train_set = []
    datas = {}
    files = glob.glob(path)
    print 'Total files:', len(files)

    for file in files:
        f = open(file, 'r')
        data = f.readlines();
        content = ' '.join(data)  # join all line in a doc
        train_set.append(content)
        datas[file] = content
        f.close()

    stopWords = getStopWords()

    # vectorizer = TfidfVectorizer(stop_words = stopWords, use_idf=True, sublinear_tf=False , norm='l2', smooth_idf=True)
    vectorizer = CountVectorizer(stop_words=stopWords, min_df=1)
    trainVectorizerArray = vectorizer.fit_transform(train_set)
    return (files, vectorizer, trainVectorizerArray.toarray(), datas)


def query(test_query, files, vectorizer, matrix):
    testVectorizerArray = vectorizer.transform(test_query).toarray()
    scores = cosine_similarity(testVectorizerArray, matrix)
    sorted_by_second = sorted(zip(files, scores[0, :]), key=lambda tup: tup[1], reverse=True)
    return sorted_by_second


def transform_query_to_vector(vectorizer, test_query):
    return vectorizer.transform(test_query).toarray()


def convert_to_mds(matrix, dimensions=50):
    mds = manifold.MDS(n_components=dimensions, dissimilarity='precomputed')
    mdsMatrix = mds.fit_transform(matrix)
    return mdsMatrix


def save_object(obj, filename):
    with open(filename, 'wb') as output:
        pickle.dump(obj, output, pickle.HIGHEST_PROTOCOL)