kse-01/search-data.py

import re
import argparse
import os
import pandas as pd
import nltk
import numpy as np
from nltk.corpus import stopwords
from gensim.similarities import SparseMatrixSimilarity, MatrixSimilarity
from gensim.models import TfidfModel, LsiModel, LdaModel
from gensim.corpora import Dictionary
from collections import defaultdict

nltk.download('stopwords')

SCRIPT_DIR = os.path.abspath(os.path.dirname(__file__))
IN_DATASET = os.path.join(SCRIPT_DIR, "data.csv")

# using ntlk stop words and example words for now
STOP_WORDS = set(stopwords.words('english')) \
    .union(['test', 'tests', 'main', 'this'])


def find_all(regex, word):
    matches = re.finditer(regex, word)
    return [m.group(0).lower() for m in matches]


# https://stackoverflow.com/a/29920015
def camel_case_split(word):
    return find_all('.+?(?:(?<=[a-z])(?=[A-Z])|(?<=[A-Z])(?=[A-Z][a-z])|$)', word)


def identifier_split(identifier):
    return [y for x in identifier.split("_") for y in camel_case_split(x)]


def comment_split(comment):
    return find_all('[A-Za-z0-9]+', comment)


def remove_stopwords(input_bow_list):
    return [word for word in input_bow_list if word not in STOP_WORDS]


def get_bow(data, split_f):
    if data is None or (type(data) == float and np.isnan(data)):
        return []
    return remove_stopwords(split_f(data))


def print_sims(corpus, query, df, dictionary):
    index = SparseMatrixSimilarity(corpus, num_features=len(dictionary))
    sims = index[query]
    pick_top = 5
    
    for idx, score in sorted(enumerate(sims), key=lambda x: x[1], reverse=True)[:pick_top]:
        row = df.loc[idx]
        print("Similarity: {s:2.02f}%".format(s=score*100))
        print("Python {feat}: {name}\nFile: {file}\nLine: {line}\n" \
              .format(feat=row["type"], name=row["name"], file=row["file"], line=row["line"]))


def search(query, method):
    df = pd.read_csv(IN_DATASET)
    df["name_bow"] = df["name"].apply(lambda n: get_bow(n, identifier_split))
    df["comment_bow"] = df["comment"].apply(lambda c: get_bow(c, comment_split))

    corpus_list = []
    for idx, row in df.iterrows():
        document_words = row["name_bow"] + row["comment_bow"]
        corpus_list.append(document_words)

    dictionary = Dictionary(corpus_list)
    corpus_bow = [dictionary.doc2bow(text) for text in corpus_list]
    query_bow = dictionary.doc2bow(get_bow(query, comment_split))
   
    if method == "tfidf":
        tfidf = TfidfModel(corpus_bow)
        print_sims(tfidf[corpus_bow], tfidf[query_bow], df, dictionary)
    elif method == "freq":
        print_sims(corpus_bow, query_bow, df, dictionary)
    elif method == "lsi":
        lsi = LsiModel(corpus_bow)
        print_sims(lsi[corpus_bow], lsi[query_bow], df, dictionary)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("method", help="the method to compare similarities with", type=str)
    parser.add_argument("query", help="the query to search the corpus with", type=str)
    args = parser.parse_args()
    search(args.query, args.method)


if __name__ == "__main__":
    main()