]> code.communitydata.science - cdsc_reddit.git/blob - similarities/tfidf.py
5b1c0c94d450da2d7259bc828404c779c734d62c
[cdsc_reddit.git] / 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/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/comment_terms.parquet',
59                 topN=25000):
60
61     return tfidf_weekly("/gscratch/comdata/output/reddit_ngrams/comment_terms.parquet",
62                  outpath,
63                  topN,
64                  'term',
65                  []
66                  )
67
68
69 if __name__ == "__main__":
70     fire.Fire({'authors':tfidf_authors,
71                'terms':tfidf_terms,
72                'authors_weekly':tfidf_authors_weekly,
73                'terms_weekly':tfidf_terms_weekly})

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