tf_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     if mwe_pass != 'first':
  68         mwe_dataset = ds.dataset(f'/gscratch/comdata/users/nathante/reddit_comment_ngrams_pwmi.parquet',format='parquet')
  69         mwe_dataset = mwe_dataset.to_pandas(columns=['phrase','phraseCount','phrasePWMI'])
  70         mwe_dataset = mwe_dataset.sort_values(['phrasePWMI'],ascending=False)
  71         mwe_phrases = list(mwe_dataset.phrase[0:1000])
  72
  73
  74         mwe_tokenize = MWETokenizer(mwe_phrases).tokenize
  75
  76
  77     def remove_punct(sentence):
  78         new_sentence = []
  79         for token in sentence:
  80             new_token = ''
  81             for c in token:
  82                 if c not in string.punctuation:
  83                     new_token += c
  84             if len(new_token) > 0:
  85                 new_sentence.append(new_token)
  86         return new_sentence
  87
  88
  89     stopWords = set(stopwords.words('english'))
  90
  91     # we follow the approach described in datta, phelan, adar 2017
  92     def my_tokenizer(text):
  93         # remove stopwords, punctuation, urls, lower case
  94         # lowercase
  95         text = text.lower()
  96
  97         # remove urls
  98         text = urlregex.sub("", text)
  99
 100         # sentence tokenize
 101         sentences = sent_tokenize(text)
 102
 103         # wordpunct_tokenize
 104         sentences = map(wordpunct_tokenize, sentences)
 105
 106         # remove punctuation
 107
 108         sentences = map(remove_punct, sentences)
 109
 110         # remove sentences with less than 2 words
 111         sentences = filter(lambda sentence: len(sentence) > 2, sentences)
 112
 113         # 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.
 114         # 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
 115         # here we take a 10 percent sample of sentences
 116         if mwe_pass == 'first':
 117             sentences = list(sentences)
 118             for sentence in sentences:
 119                 if random() <= 0.1:
 120                     grams = list(chain(*map(lambda i : ngrams(sentence,i),range(4))))
 121                     with open(f'/gscratch/comdata/users/nathante/reddit_comment_ngrams_10p_sample/{ngram_output}','a') as gram_file:
 122                         for ng in grams:
 123                             gram_file.write(' '.join(ng) + '\n')
 124                 for token in sentence:
 125                     if token not in stopWords:
 126                         yield token
 127
 128         else:
 129             # remove stopWords
 130             sentences = map(mwe_tokenize, sentences)
 131             sentences = map(lambda s: filter(lambda token: token not in stopWords, s), sentences)
 132             return chain(* sentences)
 133
 134     def tf_comments(subreddit_weeks):
 135         for key, posts in subreddit_weeks:
 136             subreddit, week = key
 137             tfs = Counter([])
 138             authors = Counter([])
 139             for post in posts:
 140                 tokens = my_tokenizer(post.body)
 141                 tfs.update(tokens)
 142                 authors.update([post.author])
 143
 144             for term, tf in tfs.items():
 145                 yield [True, subreddit, term, week, tf]
 146
 147             for author, tf in authors.items():
 148                 yield [False, subreddit, author, week, tf]
 149
 150     outrows = tf_comments(subreddit_weeks)
 151
 152     outchunksize = 10000
 153
 154     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:
 155         while True:
 156             chunk = islice(outrows,outchunksize)
 157             pddf = pd.DataFrame(chunk, columns=["is_token"] + schema.names)
 158
 159             author_pddf = pddf.loc[pddf.is_token == False, schema.names]
 160             pddf = pddf.loc[pddf.is_token == True, schema.names]
 161
 162             author_pddf = author_pddf.rename({'term':'author'}, axis='columns')
 163             author_pddf = author_pddf.loc[:,author_schema.names]
 164
 165             table = pa.Table.from_pandas(pddf,schema=schema)
 166             author_table = pa.Table.from_pandas(author_pddf,schema=author_schema)
 167             if table.shape[0] == 0:
 168                 break
 169             writer.write_table(table)
 170             author_writer.write_table(author_table)
 171
 172         writer.close()
 173         author_writer.close()
 174
 175
 176 def gen_task_list():
 177     files = os.listdir("/gscratch/comdata/output/reddit_comments_by_subreddit.parquet/")
 178     with open("tf_task_list",'w') as outfile:
 179         for f in files:
 180             if f.endswith(".parquet"):
 181                 outfile.write(f"python3 tf_comments.py weekly_tf {f}\n")
 182
 183 if __name__ == "__main__":
 184     fire.Fire({"gen_task_list":gen_task_list,
 185                "weekly_tf":weekly_tf})