]> code.communitydata.science - cdsc_reddit.git/blob - tf_reddit_comments.py
3596062f318a1acc458a50d49a7e6b25cf55400a
[cdsc_reddit.git] / tf_reddit_comments.py
1 import pyarrow as pa
2 import pyarrow.dataset as ds
3 import pyarrow.parquet as pq
4 from itertools import groupby, islice, chain
5 import fire
6 from collections import Counter
7 import pandas as pd
8 import os
9 import datetime
10 import re
11 from nltk import wordpunct_tokenize, MWETokenizer, sent_tokenize
12 from nltk.corpus import stopwords
13 from nltk.util import ngrams
14 import string
15 from random import random
16
17 # remove urls
18 # taken from https://stackoverflow.com/questions/3809401/what-is-a-good-regular-expression-to-match-a-url
19 urlregex = re.compile(r"[-a-zA-Z0-9@:%._\+~#=]{1,256}\.[a-zA-Z0-9()]{1,6}\b([-a-zA-Z0-9()@:%_\+.~#?&//=]*)")
20
21 # compute term frequencies for comments in each subreddit by week
22 def weekly_tf(partition, mwe_pass = 'first'):
23     dataset = ds.dataset(f'/gscratch/comdata/output/reddit_comments_by_subreddit.parquet/{partition}', format='parquet')
24
25     if not os.path.exists("/gscratch/comdata/users/nathante/reddit_comment_ngrams_10p_sample/"):
26         os.mkdir("/gscratch/comdata/users/nathante/reddit_comment_ngrams_10p_sample/")
27
28     if not os.path.exists("/gscratch/comdata/users/nathante/reddit_tfidf_test_authors.parquet_temp/"):
29         os.mkdir("/gscratch/comdata/users/nathante/reddit_tfidf_test_authors.parquet_temp/")
30
31     ngram_output = partition.replace("parquet","txt")
32
33     if os.path.exists(f"/gscratch/comdata/users/nathante/reddit_comment_ngrams_10p_sample/{ngram_output}"):
34         os.remove(f"/gscratch/comdata/users/nathante/reddit_comment_ngrams_10p_sample/{ngram_output}")
35     
36     batches = dataset.to_batches(columns=['CreatedAt','subreddit','body','author'])
37
38     schema = pa.schema([pa.field('subreddit', pa.string(), nullable=False),
39                         pa.field('term', pa.string(), nullable=False),
40                         pa.field('week', pa.date32(), nullable=False),
41                         pa.field('tf', pa.int64(), nullable=False)]
42     )
43
44     author_schema = pa.schema([pa.field('subreddit', pa.string(), nullable=False),
45                                pa.field('author', pa.string(), nullable=False),
46                                pa.field('week', pa.date32(), nullable=False),
47                                pa.field('tf', pa.int64(), nullable=False)]
48     )
49
50     dfs = (b.to_pandas() for b in batches)
51
52     def add_week(df):
53         df['week'] = (df.CreatedAt - pd.to_timedelta(df.CreatedAt.dt.dayofweek, unit='d')).dt.date
54         return(df)
55
56     dfs = (add_week(df) for df in dfs)
57
58     def iterate_rows(dfs):
59         for df in dfs:
60             for row in df.itertuples():
61                 yield row
62
63     rows = iterate_rows(dfs)
64
65     subreddit_weeks = groupby(rows, lambda r: (r.subreddit, r.week))
66
67     mwe_tokenize = MWETokenizer().tokenize
68
69     def remove_punct(sentence):
70         new_sentence = []
71         for token in sentence:
72             new_token = ''
73             for c in token:
74                 if c not in string.punctuation:
75                     new_token += c
76             if len(new_token) > 0:
77                 new_sentence.append(new_token)
78         return new_sentence
79
80
81     stopWords = set(stopwords.words('english'))
82
83     # we follow the approach described in datta, phelan, adar 2017
84     def my_tokenizer(text):
85         # remove stopwords, punctuation, urls, lower case
86         # lowercase        
87         text = text.lower()
88
89         # remove urls
90         text = urlregex.sub("", text)
91
92         # sentence tokenize
93         sentences = sent_tokenize(text)
94
95         # wordpunct_tokenize
96         sentences = map(wordpunct_tokenize, sentences)
97
98         # remove punctuation
99                         
100         sentences = map(remove_punct, sentences)
101
102         # remove sentences with less than 2 words
103         sentences = filter(lambda sentence: len(sentence) > 2, sentences)
104
105         # datta et al. select relatively common phrases from the reddit corpus, but they don't really explain how. We'll try that in a second phase.
106         # they say that the extract 1-4 grams from 10% of the sentences and then find phrases that appear often relative to the original terms
107         # here we take a 10 percent sample of sentences 
108         if mwe_pass == 'first':
109             sentences = list(sentences)
110             for sentence in sentences:
111                 if random() <= 0.1:
112                     grams = list(chain(*map(lambda i : ngrams(sentence,i),range(4))))
113                     with open(f'/gscratch/comdata/users/nathante/reddit_comment_ngrams_10p_sample/{ngram_output}','a') as gram_file:
114                         for ng in grams:
115                             gram_file.write(' '.join(ng) + '\n')
116                 for token in sentence:
117                     if token not in stopWords:
118                         yield token
119
120         else:
121             # remove stopWords
122             sentences = map(lambda s: filter(lambda token: token not in stopWords, s), sentences)
123             return chain(* sentences)
124
125     def tf_comments(subreddit_weeks):
126         for key, posts in subreddit_weeks:
127             subreddit, week = key
128             tfs = Counter([])
129             authors = Counter([])
130             for post in posts:
131                 tokens = my_tokenizer(post.body)
132                 tfs.update(tokens)
133                 authors.update([post.author])
134
135             for term, tf in tfs.items():
136                 yield [True, subreddit, term, week, tf]
137
138             for author, tf in authors.items():
139                 yield [False, subreddit, author, week, tf]
140
141     outrows = tf_comments(subreddit_weeks)
142
143     outchunksize = 10000
144
145     with pq.ParquetWriter("/gscratch/comdata/users/nathante/reddit_tfidf_test.parquet_temp/{partition}",schema=schema,compression='snappy',flavor='spark') as writer, pq.ParquetWriter("/gscratch/comdata/users/nathante/reddit_tfidf_test_authors.parquet_temp/{partition}",schema=author_schema,compression='snappy',flavor='spark') as author_writer:
146         while True:
147             chunk = islice(outrows,outchunksize)
148             pddf = pd.DataFrame(chunk, columns=["is_token"] + schema.names)
149             print(pddf)
150             author_pddf = pddf.loc[pddf.is_token == False]
151             author_pddf = author_pddf.rename({'term':'author'}, axis='columns')
152             author_pddf = author_pddf.loc[:,author_schema.names]
153             
154             pddf = pddf.loc[pddf.is_token == True, schema.names]
155
156             print(pddf)
157             print(author_pddf)
158             table = pa.Table.from_pandas(pddf,schema=schema)
159             author_table = pa.Table.from_pandas(author_pddf,schema=author_schema)
160             if table.shape[0] == 0:
161                 break
162             writer.write_table(table)
163             author_writer.write_table(author_table)
164             
165         writer.close()
166         author_writer.close()
167
168
169 def gen_task_list():
170     files = os.listdir("/gscratch/comdata/output/reddit_comments_by_subreddit.parquet/")
171     with open("tf_task_list",'w') as outfile:
172         for f in files:
173             if f.endswith(".parquet"):
174                 outfile.write(f"python3 tf_reddit_comments.py weekly_tf {f}\n")
175
176 if __name__ == "__main__":
177     fire.Fire({"gen_task_list":gen_task_list,
178                "weekly_tf":weekly_tf})

Community Data Science Collective || Want to submit a patch?