X-Git-Url: https://code.communitydata.science/cdsc_reddit.git/blobdiff_plain/f0176d9f0de93f0e4f3ab1d676c852c2e5fad3b3..541e125b28dbca5c06d2160a5cd59ce112657b2a:/similarities/tfidf.py

diff --git a/similarities/tfidf.py b/similarities/tfidf.py
index 30033a8..01b0b20 100644
--- a/similarities/tfidf.py
+++ b/similarities/tfidf.py
@@ -1,75 +1,82 @@
 import fire
 from pyspark.sql import SparkSession
 from pyspark.sql import functions as f
-from similarities_helper import build_tfidf_dataset, build_weekly_tfidf_dataset, select_topN_subreddits
+from similarities_helper import tfidf_dataset, build_weekly_tfidf_dataset, select_topN_subreddits
 
 def _tfidf_wrapper(func, inpath, outpath, topN, term_colname, exclude, included_subreddits):
-    spark = SparkSession.builder.getOrCreate()
+    spark = SparkSession.builder.getOrCreate()y
 
     df = spark.read.parquet(inpath)
 
     df = df.filter(~ f.col(term_colname).isin(exclude))
 
     if included_subreddits is not None:
-        include_subs = list(open(included_subreddits))
+        include_subs = set(map(str.strip,open(included_subreddits)))
     else:
         include_subs = select_topN_subreddits(topN)
 
-    df = func(df, include_subs, term_colname)
-
-    df.write.parquet(outpath,mode='overwrite',compression='snappy')
+    dfwriter = func(df, include_subs, term_colname)
 
+    dfwriter.parquet(outpath,mode='overwrite',compression='snappy')
     spark.stop()
 
 def tfidf(inpath, outpath, topN, term_colname, exclude, included_subreddits):
-    return _tfidf_wrapper(build_tfidf_dataset, inpath, outpath, topN, term_colname, exclude, included_subreddits)
+    return _tfidf_wrapper(tfidf_dataset, inpath, outpath, topN, term_colname, exclude, included_subreddits)
 
 def tfidf_weekly(inpath, outpath, topN, term_colname, exclude, included_subreddits):
     return _tfidf_wrapper(build_weekly_tfidf_dataset, inpath, outpath, topN, term_colname, exclude, included_subreddits)
 
-def tfidf_authors(outpath='/gscratch/comdata/output/reddit_similarity/tfidf/comment_authors.parquet',
-                  topN=25000):
+def tfidf_authors(inpath="/gscratch/comdata/output/reddit_ngrams/comment_authors.parquet",
+                  outpath='/gscratch/comdata/output/reddit_similarity/tfidf/comment_authors.parquet',
+                  topN=None,
+                  included_subreddits=None):
 
-    return tfidf("/gscratch/comdata/output/reddit_ngrams/comment_authors.parquet",
+    return tfidf(inpath,
                  outpath,
                  topN,
                  'author',
                  ['[deleted]','AutoModerator'],
-                 included_subreddits=None
+                 included_subreddits=included_subreddits
                  )
 
-def tfidf_terms(outpath='/gscratch/comdata/output/reddit_similarity/tfidf/comment_terms.parquet',
-                topN=25000):
+def tfidf_terms(inpath="/gscratch/comdata/output/reddit_ngrams/comment_terms.parquet",
+                outpath='/gscratch/comdata/output/reddit_similarity/tfidf/comment_terms.parquet',
+                topN=None,
+                included_subreddits=None):
 
-    return tfidf("/gscratch/comdata/output/reddit_ngrams/comment_terms.parquet",
+    return tfidf(inpath,
                  outpath,
                  topN,
                  'term',
                  [],
-                 included_subreddits=None
+                 included_subreddits=included_subreddits
                  )
 
-def tfidf_authors_weekly(outpath='/gscratch/comdata/output/reddit_similarity/tfidf_weekly/comment_authors.parquet',
-                         topN=25000):
+def tfidf_authors_weekly(inpath="/gscratch/comdata/output/reddit_ngrams/comment_authors.parquet",
+                         outpath='/gscratch/comdata/output/reddit_similarity/tfidf_weekly/comment_authors.parquet',
+                         topN=None,
+                         included_subreddits=None):
 
-    return tfidf_weekly("/gscratch/comdata/output/reddit_ngrams/comment_authors.parquet",
+    return tfidf_weekly(inpath,
                         outpath,
                         topN,
                         'author',
                         ['[deleted]','AutoModerator'],
-                        included_subreddits=None
+                        included_subreddits=included_subreddits
                         )
 
-def tfidf_terms_weekly(outpath='/gscratch/comdata/output/reddit_similarity/tfidf_weekly/comment_terms.parquet',
-                       topN=25000):
+def tfidf_terms_weekly(inpath="/gscratch/comdata/output/reddit_ngrams/comment_terms.parquet",
+                       outpath='/gscratch/comdata/output/reddit_similarity/tfidf_weekly/comment_terms.parquet',
+                       topN=None,
+                       included_subreddits=None):
 
 
-    return tfidf_weekly("/gscratch/comdata/output/reddit_ngrams/comment_terms.parquet",
+    return tfidf_weekly(inpath,
                         outpath,
                         topN,
                         'term',
                         [],
-                        included_subreddits=None
+                        included_subreddits=included_subreddits
                         )