X-Git-Url: https://code.communitydata.science/cdsc_reddit.git/blobdiff_plain/4dc949de5fb8d3eac04bae125c819100002c9522..65deba5e4e4ad9e3f23e82573491f7d6b190e644:/similarities/tfidf.py?ds=inline diff --git a/similarities/tfidf.py b/similarities/tfidf.py index 885dae2..3356299 100644 --- a/similarities/tfidf.py +++ b/similarities/tfidf.py @@ -1,69 +1,110 @@ 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 +from functools import partial - -def _tfidf_wrapper(func, inpath, outpath, topN, term_colname, exclude): +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) df = df.filter(~ f.col(term_colname).isin(exclude)) - include_subs = select_topN_subreddits(topN) + if included_subreddits is not None: + include_subs = set(map(str.strip,open(included_subreddits))) + else: + include_subs = select_topN_subreddits(topN) + + include_subs = spark.sparkContext.broadcast(include_subs) + + # term_id = term_colname + "_id" - df = func(df, include_subs, term_colname) + 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') - df.write.parquet(outpath,mode='overwrite',compression='snappy') + 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): - return _tfidf_wrapper(build_tfidf_dataset, inpath, outpath, topN, term_colname, exclude) +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, topN, term_colname, exclude): - return _tfidf_wrapper(build_weekly_tfidf_dataset, inpath, outpath, topN, term_colname, exclude) +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_authors(outpath='/gscratch/comdata/output/reddit_similarity/tfidf/comment_authors.parquet', - topN=25000): - return tfidf("/gscratch/comdata/output/reddit_ngrams/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, + min_df=None, + max_df=None): + + return tfidf(inpath, outpath, topN, 'author', - ['[deleted]','AutoModerator'] + ['[deleted]','AutoModerator'], + 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', - 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, + 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, + min_df=min_df, + max_df=max_df ) -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", + 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", - outpath, - topN, - 'author', - ['[deleted]','AutoModerator'] - ) + 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', - topN=25000): +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', - [] + [], + included_subreddits=included_subreddits )