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