]> code.communitydata.science - cdsc_reddit.git/blob - similarities/tfidf.py
update pushshift dumps.
[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 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()y
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})

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