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 build_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 = list(open(included_subreddits))
  15     else:
  16         include_subs = select_topN_subreddits(topN)
  17
  18     df = func(df, include_subs, term_colname)
  19
  20     df.write.parquet(outpath,mode='overwrite',compression='snappy')
  21
  22     spark.stop()
  23
  24 def tfidf(inpath, outpath, topN, term_colname, exclude, included_subreddits):
  25     return _tfidf_wrapper(build_tfidf_dataset, inpath, outpath, topN, term_colname, exclude, included_subreddits)
  26
  27 def tfidf_weekly(inpath, outpath, topN, term_colname, exclude, included_subreddits):
  28     return _tfidf_wrapper(build_weekly_tfidf_dataset, inpath, outpath, topN, term_colname, exclude, included_subreddits)
  29
  30 def tfidf_authors(outpath='/gscratch/comdata/output/reddit_similarity/tfidf/comment_authors.parquet',
  31                   topN=25000):
  32
  33     return tfidf("/gscratch/comdata/output/reddit_ngrams/comment_authors.parquet",
  34                  outpath,
  35                  topN,
  36                  'author',
  37                  ['[deleted]','AutoModerator'],
  38                  included_subreddits=None
  39                  )
  40
  41 def tfidf_terms(outpath='/gscratch/comdata/output/reddit_similarity/tfidf/comment_terms.parquet',
  42                 topN=25000):
  43
  44     return tfidf("/gscratch/comdata/output/reddit_ngrams/comment_terms.parquet",
  45                  outpath,
  46                  topN,
  47                  'term',
  48                  [],
  49                  included_subreddits=None
  50                  )
  51
  52 def tfidf_authors_weekly(outpath='/gscratch/comdata/output/reddit_similarity/tfidf_weekly/comment_authors.parquet',
  53                          topN=25000):
  54
  55     return tfidf_weekly("/gscratch/comdata/output/reddit_ngrams/comment_authors.parquet",
  56                         outpath,
  57                         topN,
  58                         'author',
  59                         ['[deleted]','AutoModerator'],
  60                         included_subreddits=None
  61                         )
  62
  63 def tfidf_terms_weekly(outpath='/gscratch/comdata/output/reddit_similarity/tfidf_weekly/comment_terms.parquet',
  64                        topN=25000):
  65
  66
  67     return tfidf_weekly("/gscratch/comdata/output/reddit_ngrams/comment_terms.parquet",
  68                         outpath,
  69                         topN,
  70                         'term',
  71                         [],
  72                         included_subreddits=None
  73                         )
  74
  75
  76 if __name__ == "__main__":
  77     fire.Fire({'authors':tfidf_authors,
  78                'terms':tfidf_terms,
  79                'authors_weekly':tfidf_authors_weekly,
  80                'terms_weekly':tfidf_terms_weekly})