MyCourseIndex/generate_simstring_db.py at master · shljessie/MyCourseIndex · 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
import json
import pickle

import boto3
import nltk
from tqdm import tqdm
import html2text

from app.utils.simstring_doc import (
    CharNgramFeatureExtractor,
    RamDatabase,
    CosineSimilarity,
    Searcher
)
from app.utils.toke import sp


s3 = boto3.client('s3')
s3.download_file('cs4300-data-models', 'P03Data_mod.json', 'P03Data.json')

with open("P03Data.json") as fp:
    P03Data = json.load(fp)

db = RamDatabase(CharNgramFeatureExtractor(3))

h = html2text.HTML2Text()
h.ignore_links = True

# map them with triplet keys (course, Piazza/Resource, Key)
# This is encoded as COURSE|RESOURCE|KEY
bigram_2_str = lambda w: w[0] + " " + w[1]
trigram_2_str = lambda w: w[0] + " " + w[1] + " " + w[2]

# Parse Textbook data first
for key, item in tqdm(P03Data["CS 4300"]["Resource"].items(), desc="Resource", leave=True):
    raw = item["raw"]
    tokes = sp(raw.lower())
    tokenized = []
    idxes = []
    punct_count = 0
    for w in tokes:
        if w.is_punct:
            punct_count += 1
        else:
            tokenized.append(w.text)
            idxes.append(w.idx - punct_count)
    # idxes = [w.idx for w in tokenized]
    # tokenized = [w.text for w in tokenized if not w.is_punct]
    bigrams = nltk.bigrams(tokenized)
    trigrams = nltk.trigrams(tokenized)
    location = ("CS 4300", "Piazza", key)
    for idx, w in tqdm(zip(idxes,tokenized), leave=False, desc="Single words"):
        db.add(w, location + (idx,))
    for idx, w in tqdm(zip(idxes[:-1], bigrams), leave=False, desc="Bigrams"):
        db.add(bigram_2_str(w), location + (idx,))
    for idx, w in tqdm(zip(idxes[:-2], trigrams), leave=False, desc="Trigrams"):
        db.add(trigram_2_str(w), location + (idx,))
    # tokenized = [w.text for w in tokenized if not w.is_punct]
    # bigrams = nltk.bigrams(tokenized)
    # trigrams = nltk.trigrams(tokenized)
    # location = "CS 4300|Resource|" + key
    # for w in tqdm(tokenized, leave=False, desc="Single words"):
    #     db.add(w, location)
    # for w in tqdm(bigrams, leave=False, desc="Bigrams"):
    #     db.add(bigram_2_str(w), location)
    # for w in tqdm(trigrams, leave=False, desc="Trigrams"):
    #     db.add(trigram_2_str(w), location)

for post_id, post in tqdm(P03Data["CS 4300"]["Piazza"].items(), leave=True, desc="Piazza"):
    post = post["raw"]
    all_text = ""
    subject = h.handle(post['history'][0]['subject']).replace("\n", " ")
    question = h.handle(post['history'][0]['content']).replace("\n", " ")
    other_text = ""
    for answer in post['children']:
        if answer['type'] == "i_answer":
            other_text += " " + h.handle(answer['history'][0]['content']).replace("\n", " ")
        elif answer['type'] == "s_answer":
            other_text += " " + h.handle(answer['history'][0]['content']).replace("\n", " ")
        elif answer['type'] == "followup":
            other_text += " " + h.handle(answer['subject']).replace("\n", " ")
            for fb in answer['children']:
                other_text += " " + h.handle(fb['subject']).replace("\n", " ")

    all_text = subject + " " + question + " " + other_text
    tokes = sp(all_text.lower())
    tokenized = []
    idxes = []
    punct_count = 0
    for w in tokes:
        if w.is_punct:
            punct_count += 1
        else:
            tokenized.append(w.text)
            idxes.append(w.idx - punct_count)
    # idxes = [w.idx for w in tokenized]
    # tokenized = [w.text for w in tokenized if not w.is_punct]
    bigrams = nltk.bigrams(tokenized)
    trigrams = nltk.trigrams(tokenized)
    location = ("CS 4300", "Piazza", post_id)
    for idx, w in tqdm(zip(idxes,tokenized), leave=False, desc="Single words"):
        db.add(w, location + (idx,))
    for idx, w in tqdm(zip(idxes[:-1], bigrams), leave=False, desc="Bigrams"):
        db.add(bigram_2_str(w), location + (idx,))
    for idx, w in tqdm(zip(idxes[:-2], trigrams), leave=False, desc="Trigrams"):
        db.add(trigram_2_str(w), location + (idx,))

indexer = Searcher(db, CosineSimilarity())

with open("ramDB.pkl", "wb") as fp:
    pickle.dump(db, fp)