Merge remote-tracking branch 'refs/remotes/origin/excise_reindex' into excise_reindex

[cdsc_reddit.git] / similarities / tfidf.py

diff --git a/similarities/tfidf.py b/similarities/tfidf.py
index 19d30138457843df62ef440d3a75acc45b41df87..c44fd0ddbf14d49f7c96e9f4be92c03bcd5b4c96 100644 (file)
--- a/similarities/tfidf.py
+++ b/similarities/tfidf.py
@@ -2,8 +2,11 @@ import fire
 from pyspark.sql import SparkSession
 from pyspark.sql import functions as f
 from similarities_helper import tfidf_dataset, build_weekly_tfidf_dataset, select_topN_subreddits
+from functools import partial
 
-def _tfidf_wrapper(func, inpath, outpath, topN, term_colname, exclude, included_subreddits):
+inpath = '/gscratch/comdata/users/nathante/competitive_exclusion_reddit/data/tfidf/comment_authors_compex.parquet'
+# include_terms is a path to a parquet file that contains a column of term_colname + '_id' to include.
+def _tfidf_wrapper(func, inpath, outpath, topN, term_colname, exclude, included_subreddits, included_terms=None, min_df=None, max_df=None):
     spark = SparkSession.builder.getOrCreate()
 
     df = spark.read.parquet(inpath)
@@ -15,60 +18,87 @@ def _tfidf_wrapper(func, inpath, outpath, topN, term_colname, exclude, included_
     else:
         include_subs = select_topN_subreddits(topN)
 
-    dfwriter = func(df, include_subs, term_colname)
+    include_subs = spark.sparkContext.broadcast(include_subs)
+
+    #    term_id = term_colname + "_id"
+
+    if included_terms is not None:
+        terms_df = spark.read.parquet(included_terms)
+        terms_df = terms_df.select(term_colname).distinct()
+        df = df.join(terms_df, on=term_colname, how='left_semi')
+
+    dfwriter = func(df, include_subs.value, term_colname)
 
     dfwriter.parquet(outpath,mode='overwrite',compression='snappy')
     spark.stop()
 
-def tfidf(inpath, outpath, topN, term_colname, exclude, included_subreddits):
-    return _tfidf_wrapper(tfidf_dataset, inpath, outpath, topN, term_colname, exclude, included_subreddits)
+def tfidf(inpath, outpath, topN, term_colname, exclude, included_subreddits, min_df, max_df):
+    tfidf_func = partial(tfidf_dataset, max_df=max_df, min_df=min_df)
+    return _tfidf_wrapper(tfidf_func, inpath, outpath, topN, term_colname, exclude, included_subreddits)
+
+def tfidf_weekly(inpath, outpath, static_tfidf_path, topN, term_colname, exclude, included_subreddits):
+    return _tfidf_wrapper(build_weekly_tfidf_dataset, inpath, outpath, topN, term_colname, exclude, included_subreddits, included_terms=static_tfidf_path)
 
-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',
+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):
+                  included_subreddits=None,
+                  min_df=None,
+                  max_df=None):
 
-    return tfidf("/gscratch/comdata/output/reddit_ngrams/comment_authors.parquet",
+    return tfidf(inpath,
                  outpath,
                  topN,
                  'author',
                  ['[deleted]','AutoModerator'],
-                 included_subreddits=included_subreddits
+                 included_subreddits=included_subreddits,
+                 min_df=min_df,
+                 max_df=max_df
                  )
 
-def tfidf_terms(outpath='/gscratch/comdata/output/reddit_similarity/tfidf/comment_terms.parquet',
+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):
+                included_subreddits=None,
+                min_df=None,
+                max_df=None):
 
-    return tfidf("/gscratch/comdata/output/reddit_ngrams/comment_terms.parquet",
+    return tfidf(inpath,
                  outpath,
                  topN,
                  'term',
                  [],
-                 included_subreddits=included_subreddits
+                 included_subreddits=included_subreddits,
+                 min_df=min_df,
+                 max_df=max_df
                  )
 
-def tfidf_authors_weekly(outpath='/gscratch/comdata/output/reddit_similarity/tfidf_weekly/comment_authors.parquet',
+def tfidf_authors_weekly(inpath="/gscratch/comdata/output/reddit_ngrams/comment_authors.parquet",
+                         static_tfidf_path="/gscratch/comdata/output/reddit_similarity/tfidf/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,
+                        static_tfidf_path,
                         topN,
                         'author',
                         ['[deleted]','AutoModerator'],
                         included_subreddits=included_subreddits
                         )
 
-def tfidf_terms_weekly(outpath='/gscratch/comdata/output/reddit_similarity/tfidf_weekly/comment_terms.parquet',
+def tfidf_terms_weekly(inpath="/gscratch/comdata/output/reddit_ngrams/comment_terms.parquet",
+                       static_tfidf_path="/gscratch/comdata/output/reddit_similarity/tfidf/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,
+                        static_tfidf_path,
                         topN,
                         'term',
                         [],