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

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