]> code.communitydata.science - cdsc_reddit.git/blob - similarities/tfidf.py
bug fix in affinity clustering
[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 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     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(build_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(outpath='/gscratch/comdata/output/reddit_similarity/tfidf/comment_authors.parquet',
30                   topN=25000):
31
32     return tfidf("/gscratch/comdata/output/reddit_ngrams/comment_authors.parquet",
33                  outpath,
34                  topN,
35                  'author',
36                  ['[deleted]','AutoModerator'],
37                  included_subreddits=None
38                  )
39
40 def tfidf_terms(outpath='/gscratch/comdata/output/reddit_similarity/tfidf/comment_terms.parquet',
41                 topN=25000):
42
43     return tfidf("/gscratch/comdata/output/reddit_ngrams/comment_terms.parquet",
44                  outpath,
45                  topN,
46                  'term',
47                  [],
48                  included_subreddits=None
49                  )
50
51 def tfidf_authors_weekly(outpath='/gscratch/comdata/output/reddit_similarity/tfidf_weekly/comment_authors.parquet',
52                          topN=25000):
53
54     return tfidf_weekly("/gscratch/comdata/output/reddit_ngrams/comment_authors.parquet",
55                         outpath,
56                         topN,
57                         'author',
58                         ['[deleted]','AutoModerator'],
59                         included_subreddits=None
60                         )
61
62 def tfidf_terms_weekly(outpath='/gscratch/comdata/output/reddit_similarity/tfidf_weekly/comment_terms.parquet',
63                        topN=25000):
64
65
66     return tfidf_weekly("/gscratch/comdata/output/reddit_ngrams/comment_terms.parquet",
67                         outpath,
68                         topN,
69                         'term',
70                         [],
71                         included_subreddits=None
72                         )
73
74
75 if __name__ == "__main__":
76     fire.Fire({'authors':tfidf_authors,
77                'terms':tfidf_terms,
78                'authors_weekly':tfidf_authors_weekly,
79                'terms_weekly':tfidf_terms_weekly})

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