similarities/tfidf.py

   1 import fire
   2 from pyspark.sql import SparkSession
   3 from pyspark.sql import functions as f
   4 from similarities_helper import tfidf_dataset, build_weekly_tfidf_dataset, select_topN_subreddits
   5
   6 def _tfidf_wrapper(func, inpath, outpath, topN, term_colname, exclude, included_subreddits):
   7     spark = SparkSession.builder.getOrCreate()
   8
   9     df = spark.read.parquet(inpath)
  10
  11     df = df.filter(~ f.col(term_colname).isin(exclude))
  12
  13     if included_subreddits is not None:
  14         include_subs = set(map(str.strip,open(included_subreddits)))
  15     else:
  16         include_subs = select_topN_subreddits(topN)
  17
  18     dfwriter = func(df, include_subs, term_colname)
  19
  20     dfwriter.parquet(outpath,mode='overwrite',compression='snappy')
  21     spark.stop()
  22
  23 def tfidf(inpath, outpath, topN, term_colname, exclude, included_subreddits):
  24     return _tfidf_wrapper(tfidf_dataset, inpath, outpath, topN, term_colname, exclude, included_subreddits)
  25
  26 def tfidf_weekly(inpath, outpath, topN, term_colname, exclude, included_subreddits):
  27     return _tfidf_wrapper(build_weekly_tfidf_dataset, inpath, outpath, topN, term_colname, exclude, included_subreddits)
  28
  29 def tfidf_authors(inpath="/gscratch/comdata/output/reddit_ngrams/comment_authors.parquet",
  30                   outpath='/gscratch/comdata/output/reddit_similarity/tfidf/comment_authors.parquet',
  31                   topN=None,
  32                   included_subreddits=None):
  33
  34     return tfidf(inpath,
  35                  outpath,
  36                  topN,
  37                  'author',
  38                  ['[deleted]','AutoModerator'],
  39                  included_subreddits=included_subreddits
  40                  )
  41
  42 def tfidf_terms(inpath="/gscratch/comdata/output/reddit_ngrams/comment_terms.parquet",
  43                 outpath='/gscratch/comdata/output/reddit_similarity/tfidf/comment_terms.parquet',
  44                 topN=None,
  45                 included_subreddits=None):
  46
  47     return tfidf(inpath,
  48                  outpath,
  49                  topN,
  50                  'term',
  51                  [],
  52                  included_subreddits=included_subreddits
  53                  )
  54
  55 def tfidf_authors_weekly(inpath="/gscratch/comdata/output/reddit_ngrams/comment_authors.parquet",
  56                          outpath='/gscratch/comdata/output/reddit_similarity/tfidf_weekly/comment_authors.parquet',
  57                          topN=None,
  58                          included_subreddits=None):
  59
  60     return tfidf_weekly(inpath,
  61                         outpath,
  62                         topN,
  63                         'author',
  64                         ['[deleted]','AutoModerator'],
  65                         included_subreddits=included_subreddits
  66                         )
  67
  68 def tfidf_terms_weekly(inpath="/gscratch/comdata/output/reddit_ngrams/comment_terms.parquet",
  69                        outpath='/gscratch/comdata/output/reddit_similarity/tfidf_weekly/comment_terms.parquet',
  70                        topN=None,
  71                        included_subreddits=None):
  72
  73
  74     return tfidf_weekly(inpath,
  75                         outpath,
  76                         topN,
  77                         'term',
  78                         [],
  79                         included_subreddits=included_subreddits
  80                         )
  81
  82
  83 if __name__ == "__main__":
  84     fire.Fire({'authors':tfidf_authors,
  85                'terms':tfidf_terms,
  86                'authors_weekly':tfidf_authors_weekly,
  87                'terms_weekly':tfidf_terms_weekly})