├── .gitignore ├── README.md ├── analysis ├── dataset_statistics.py ├── gather_regression.py ├── regression-base.sh ├── regression-large.sh ├── regression_results_base.csv ├── regression_results_large.csv ├── run_glue.py └── time_statistics.py ├── data ├── dev.tsv.gz ├── fullrjokes.json.gz ├── preprocessed.csv.gz ├── submissions.json.gz ├── test.tsv.gz └── train.tsv.gz ├── download_nltk_packages.sh ├── plots ├── all │ ├── count.pdf │ ├── count.png │ ├── logdist.pdf │ ├── logdist.png │ ├── max.pdf │ ├── max.png │ ├── percentiles.csv │ ├── score_distributions.pdf │ ├── score_distributions.png │ ├── sentiment.pdf │ ├── sentiment.png │ ├── statistics.txt │ ├── top_wordcounts.pdf │ ├── top_wordcounts.png │ ├── total_score_distribution.pdf │ ├── total_score_distribution.png │ ├── wordCloud.pdf │ ├── wordCloud.png │ └── wordCloudSmall.pdf └── yearbyyear │ ├── 2008_top_wordcounts.pdf │ ├── 2008_top_wordcounts.png │ ├── 2009_top_wordcounts.pdf │ ├── 2009_top_wordcounts.png │ ├── 2010_top_wordcounts.pdf │ ├── 2010_top_wordcounts.png │ ├── 2011_top_wordcounts.pdf │ ├── 2011_top_wordcounts.png │ ├── 2012_top_wordcounts.pdf │ ├── 2012_top_wordcounts.png │ ├── 2013_top_wordcounts.pdf │ ├── 2013_top_wordcounts.png │ ├── 2014_top_wordcounts.pdf │ ├── 2014_top_wordcounts.png │ ├── 2015_top_wordcounts.pdf │ ├── 2015_top_wordcounts.png │ ├── 2016_top_wordcounts.pdf │ ├── 2016_top_wordcounts.png │ ├── 2017_top_wordcounts.pdf │ ├── 2017_top_wordcounts.png │ ├── 2018_top_wordcounts.pdf │ ├── 2018_top_wordcounts.png │ ├── 2019_top_wordcounts.pdf │ └── 2019_top_wordcounts.png ├── prepare_data ├── gather_reddit_pushshift.py └── preprocess.py └── requirements.txt /.gitignore: -------------------------------------------------------------------------------- 1 | env/ 2 | .vscode 3 | *pycache* 4 | *.bin 5 | *.pt 6 | *cached* 7 | *checkpoint* 8 | analysis/output/* -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # The r/Jokes Dataset: a Large Scale Humor Collection 2 | Code and Datasets from the paper, ["The r/Jokes Dataset: a Large Scale Humor Collection"](https://www.aclweb.org/anthology/2020.lrec-1.753/) by Orion Weller and Kevin Seppi 3 | 4 | Dataset files are located in `data/{train/dev/test}.tsv` for the regression task, while the full unsplit data can be found in `data/preprocessed.tsv`. These files will need to be unzipped after cloning the repo. 5 | 6 | For related projects, see our work on [Humor Detection (separating the humorous jokes from the non-humorous)](https://github.com/orionw/RedditHumorDetection) or [generating humor automatically](https://github.com/orionw/humorTranslate). 7 | 8 | ** **We do not endorse these jokes. Please view at your own risk** ** 9 | 10 | ## License 11 | The data is under the [Reddit License and Terms of Service](https://www.reddit.com/wiki/api-terms) and users must follow the Reddit User Agreement and Privacy Policy, as well as remove any posts if asked to by the original user. For more details on this, please see the link above. 12 | 13 | # Usage 14 | ## Load the Required Packages 15 | 0. Run `pip3 install -r requirements.txt` 16 | 1. Gather the NLTK packages by running `bash download_nltk_packages.sh`. This downloads the packages `averaged_perceptron_tagger`, `words`, `stopwords`, `maxent_ne_chunker`, used for analysis/preprocessing. 17 | 18 | ## Reproduce the current dataset (updated to Jan 1st 2020) 19 | ### We chunk this process into three parts to avoid networking errors 20 | 0. Run `python3 gather_reddit_pushshift.py` after `cd prepare_data` to gather the Reddit post ids. 21 | 1. Run `python3 preprocess.py --update` to update the Reddit post IDs with the full post. 22 | 2. Run `python3 preprocess.py --preprocess` to preprocess the Reddit posts into final datasets 23 | 24 | ## Reproduce plots and analysis from the paper 25 | 0. Run `cd analysis` 26 | 1. Run `python3 time_statistics.py` to gather the statistics that display over time 27 | 2. Run `python3 dataset_statistics.py` to gather the overall dataset statistics 28 | 3. See plots in the `./plots` folder 29 | 30 | ## Re-gather All Jokes and Extend With Newer Jokes 31 | 0. Run the first two commands in the `Reproduce` section above 32 | 1. Update the code in the `preprocess` function of the `preprocess.py` file to NOT remove all jokes after 2020 (line 89). Then run `python3 preprocess.py --preprocess` 33 | 34 | # Reference: 35 | If you found this repository helpful, please cite the following paper: 36 | ``` 37 | @ARTICLE{rjokesData2020, 38 | title={The r/Jokes Dataset: a Large Scale Humor Collection}, 39 | author={Weller, Orion and Seppi, Kevin}, 40 | journal={"Proceedings of the 2020 Conference of Language Resources and Evaluation"}, 41 | month=May, 42 | year = "2020", 43 | } 44 | -------------------------------------------------------------------------------- /analysis/dataset_statistics.py: -------------------------------------------------------------------------------- 1 | import os 2 | import json 3 | import string 4 | import random 5 | 6 | import pandas as pd 7 | import numpy as np 8 | from wordcloud import WordCloud 9 | from sklearn.feature_extraction.text import CountVectorizer 10 | import matplotlib.pyplot as plt 11 | import seaborn as sns 12 | import nltk 13 | from nltk.corpus import stopwords 14 | from nltk.tokenize import word_tokenize 15 | from nltk.stem import PorterStemmer 16 | from nltk.stem import LancasterStemmer 17 | from nltk import ne_chunk, pos_tag, word_tokenize 18 | from nltk.tree import Tree 19 | 20 | from sklearn.decomposition import LatentDirichletAllocation as LDA 21 | 22 | random.seed(42) 23 | np.random.seed(42) 24 | 25 | EXT_TYPES = ["pdf", "png"] 26 | REMOVE_TOKENS = set(stopwords.words('english') + ["n't", "ve"]).union(set(string.punctuation)) 27 | sns.set_style('whitegrid') 28 | 29 | def get_continuous_chunks(text): 30 | chunked = ne_chunk(pos_tag(word_tokenize(text))) 31 | continuous_chunk = [] 32 | current_chunk = [] 33 | for i in chunked: 34 | if type(i) == Tree: 35 | current_chunk.append(" ".join([token for token, pos in i.leaves()])) 36 | elif current_chunk: 37 | named_entity = " ".join(current_chunk) 38 | if named_entity not in continuous_chunk: 39 | continuous_chunk.append(named_entity) 40 | current_chunk = [] 41 | else: 42 | continue 43 | return continuous_chunk 44 | 45 | 46 | def plot_n_most_common_words(text: pd.Series, data_name: str, n: int = 10, specific_year: str = None): 47 | entities = [] 48 | for sentence in text: 49 | entities.extend(get_continuous_chunks(sentence)) 50 | stemmer = PorterStemmer() 51 | filtered_words = [stemmer.stem(w) for w in entities if not w in REMOVE_TOKENS] 52 | # Initialise the count vectorizer with the English stop words 53 | count_vectorizer = CountVectorizer(stop_words='english') 54 | # Fit and transform the processed titles 55 | count_data = count_vectorizer.fit_transform(filtered_words) 56 | # Visualise the 10 most common words 57 | words = count_vectorizer.get_feature_names() 58 | total_counts = np.zeros(len(words)) 59 | for t in count_data: 60 | total_counts+=t.toarray()[0] 61 | 62 | get_lda(count_data, count_vectorizer) 63 | 64 | count_dict = (zip(words, total_counts)) 65 | count_dict = sorted(count_dict, key=lambda x:x[1], reverse=True)[0:n] 66 | words = [w[0] for w in count_dict] 67 | counts = [w[1] for w in count_dict] 68 | x_pos = np.arange(len(words)) 69 | 70 | for ext_type in EXT_TYPES: 71 | plt.figure(2, figsize=(15, 15/1.6180)) 72 | plt.subplot(title='{} most common words'.format(n)) 73 | sns.set_context("notebook", font_scale=1.25, rc={"lines.linewidth": 2.5}) 74 | sns.barplot(x_pos, counts, palette='husl') 75 | plt.xticks(x_pos, words, rotation=90) 76 | plt.xlabel('words') 77 | plt.ylabel('counts') 78 | plt.savefig(os.path.join(os.path.realpath('..'), "plots", data_name, 79 | "{}top_wordcounts.{}".format("" if specific_year is None else str(specific_year) + "_", ext_type))) 80 | plt.close() 81 | 82 | 83 | def run_wordcloud(text: list, data_name: str): 84 | word_tokens = [] 85 | try: 86 | for index, sentence in enumerate(text): 87 | if pd.isnull(sentence): 88 | print("index is null", index) 89 | continue 90 | word_tokens.extend(word_tokenize(sentence)) 91 | except Exception as e: 92 | import pdb; pdb.set_trace() 93 | print(e, sentence) 94 | filtered_words = [w for w in word_tokens if not w in REMOVE_TOKENS] 95 | # Create a WordCloud object 96 | wordcloud = WordCloud(background_color="white", width=800, height=800, max_words=150, contour_width=3, contour_color='steelblue') 97 | # Generate a word cloud 98 | wordcloud.generate(" ".join(filtered_words)) 99 | # Visualize the word cloud 100 | for ext_type in EXT_TYPES: 101 | wordcloud.to_file(os.path.join(os.path.realpath('..'), "plots", data_name, 'wordCloud.{}'.format(ext_type))) 102 | plt.close() 103 | 104 | 105 | def print_topics(model, count_vectorizer, n_top_words): 106 | words = count_vectorizer.get_feature_names() 107 | for topic_idx, topic in enumerate(model.components_): 108 | print("\nTopic #%d:" % topic_idx) 109 | print(" ".join([words[i] 110 | for i in topic.argsort()[:-n_top_words - 1:-1]])) 111 | 112 | def get_lda(count_data, count_vectorizer): 113 | # Tweak the two parameters below 114 | number_topics = 5 115 | number_words = 10 116 | # Create and fit the LDA model 117 | lda = LDA(n_components=number_topics, n_jobs=-1) 118 | lda.fit(count_data) 119 | print("Topics found via LDA:") 120 | print_topics(lda, count_vectorizer, number_words) 121 | 122 | def get_statistics(df: pd.DataFrame, data_name: str): 123 | distrib = df.copy(deep=True) 124 | distrib = distrib[distrib["score"] != 0] 125 | for ext_type in EXT_TYPES: 126 | sns.distplot(distrib["score"]) 127 | plt.savefig(os.path.join(os.path.realpath('..'), "plots", data_name, 'total_score_distribution.{}'.format(ext_type))) 128 | plt.close() 129 | 130 | punch = np.array([len(sentence) for sentence in df["punchline"].values]) 131 | body = np.array([len(sentence) for sentence in df["body"].values]) 132 | joke = np.array([len(sentence) for sentence in df["joke"].values]) 133 | 134 | average_punch = np.mean(punch) 135 | average_body = np.mean(body) 136 | average_joke = np.mean(joke) - 1.0 # Joke Token, `AND` to join body and punchline together 137 | 138 | std_punch = np.nanstd(punch) 139 | std_body = np.nanstd(body) 140 | std_joke = np.nanstd(joke) 141 | 142 | ave_tokens_punch = np.nanmean(np.array([len(nltk.word_tokenize(sentence)) for sentence in df["punchline"].dropna().values])) 143 | ave_tokens_body = np.nanmean(np.array([len(nltk.word_tokenize(sentence)) for sentence in df["body"].dropna().values])) 144 | ave_tokens_joke = np.nanmean(np.array([len(nltk.word_tokenize(sentence)) for sentence in df["joke"].dropna().values])) 145 | 146 | tokens = [] 147 | [tokens.extend(nltk.word_tokenize(joke)) for joke in df["joke"].dropna().values] 148 | total_tokens = len(set(tokens)) 149 | 150 | stat_df = pd.DataFrame([{"ave_punchline_len": average_punch, "ave_body_len": average_body, "ave_joke_len": average_joke, "std_punch": std_punch, 151 | "std_body": std_body, "std_joke": std_joke, "total_tokens": total_tokens}]) 152 | stat_df.to_csv(os.path.join(os.path.realpath('..'), "plots", data_name, 'statistics.txt')) 153 | 154 | 155 | def plot_sentiment(df: pd.DataFrame, data_name: str): 156 | for ext_type in EXT_TYPES: 157 | ax = sns.lineplot(x="date", y="prop", hue="sentiment", data=df, ci=False) 158 | # Find the x,y coordinates for each point 159 | x_coords = [] 160 | y_coords = [] 161 | for point_pair in ax.collections: 162 | for x, y in point_pair.get_offsets(): 163 | x_coords.append(x) 164 | y_coords.append(y) 165 | # create the custom error bars 166 | colors = ['steelblue']*2 + ['coral']*2 167 | ax.errorbar(x_coords, y_coords, yerr=df["std"], 168 | ecolor=colors, fmt=' ', zorder=-1) 169 | ax.savefig(os.path.join(os.path.realpath('..'), "data", data_name, "sentiment_plot.{}".format(ext_type))) 170 | 171 | 172 | 173 | def gather_data(df: pd.DataFrame, data_name: str): 174 | if not os.path.isdir(os.path.join(os.path.realpath('..'), "plots", data_name)): 175 | os.mkdir(os.path.join(os.path.realpath('..'), "plots", data_name)) 176 | run_wordcloud(df["joke"].tolist(), data_name) 177 | plot_n_most_common_words(df["joke"].tolist(), data_name) 178 | get_statistics(df, data_name) 179 | 180 | 181 | def percentiles_upvotes(df: pd.DataFrame, data_name: str) -> pd.DataFrame: 182 | list_of_percentiles = [] 183 | for percentile in [0, 10, 25, 50, 75, 90, 100]: 184 | cur_per = np.percentile(df["score"], percentile) 185 | list_of_percentiles.append({"percentile": percentile, "value": cur_per}) 186 | percent_df = pd.DataFrame(list_of_percentiles) 187 | percent_df.to_csv(os.path.join(os.path.realpath('..'), "plots", data_name, "percentiles.csv")) 188 | return percent_df 189 | 190 | 191 | if __name__ == "__main__": 192 | # NOTE: log-distribution plots are found in the preprocess.py script 193 | df = pd.read_csv(os.path.join(os.path.realpath('..'), "data", "preprocessed.csv"), index_col=None, encoding="UTF-8", keep_default_na=False) 194 | df["date"] = pd.to_numeric(df["date"]) 195 | df["score"] = pd.to_numeric(df["score"]) 196 | df = df[df["date"].isna() == False] 197 | assert df.shape == df.dropna().shape, "was nans that are unaccounted for" 198 | percentiles_upvotes(df, "all") 199 | gather_data(df, "all") 200 | 201 | 202 | 203 | 204 | 205 | 206 | 207 | 208 | 209 | 210 | -------------------------------------------------------------------------------- /analysis/gather_regression.py: -------------------------------------------------------------------------------- 1 | import os 2 | import pandas as pd 3 | import matplotlib.pyplot as plt 4 | import seaborn as sns 5 | import glob 6 | 7 | def gather_data_from_huggingface(root_dir: str, output_dir: str): 8 | all_models = {} 9 | for result_dir in os.listdir(root_dir): 10 | model_name = result_dir.split("/")[-1] 11 | model_results = [] 12 | checkpoint_folder = None 13 | for checkpoint_folder in glob.glob(os.path.join(root_dir, result_dir, "*checkpoint-*/")): 14 | if not os.path.isfile(os.path.join(checkpoint_folder, "eval_results.txt")): 15 | print("No results file in path:", checkpoint_folder) 16 | else: 17 | results_on_check = {} 18 | with open(os.path.join(checkpoint_folder, "eval_results.txt"), "r") as f: 19 | for index, line in enumerate(f): 20 | name, value = line.split("=") 21 | name, value = name.strip(), float(value.strip()) 22 | results_on_check[name] = value 23 | results_on_check["model"] = model_name 24 | results_on_check["checkpoint"] = checkpoint_folder.split("/")[-2] 25 | model_results.append(results_on_check) 26 | results_df = pd.DataFrame(model_results) 27 | if checkpoint_folder is not None: 28 | results_df.to_csv(os.path.join(checkpoint_folder, "model_results.csv")) 29 | ax = sns.lineplot(x="index", y="rmse", data=results_df.reset_index()) 30 | plt.savefig(os.path.join(root_dir, result_dir, "rmse_plot.png")) 31 | plt.close() 32 | all_models[model_name] = results_df 33 | print("Wrote model file to ", os.path.join(result_dir, "rmse_plot.png")) 34 | 35 | if all_models: 36 | # now that all the results have been gathered, let's combine them 37 | full_df = pd.concat(list(all_models.values())).reset_index() 38 | full_df.to_csv(output_dir) 39 | # focus on lowest RMSE (could focus on Pearson, Spearmanr etc.) 40 | min_rmse = full_df.groupby("model")["rmse"].idxmin() 41 | best_results = full_df.iloc[min_rmse, :] 42 | 43 | # plot line plot here 44 | ax = sns.barplot(x="model", y="rmse", data=best_results) 45 | plt.savefig(os.path.join(root_dir, "rmse_plot.png")) 46 | print("Wrote full file to ", os.path.join(root_dir, "rmse_plot.png")) 47 | plt.close() 48 | print(best_results) 49 | 50 | 51 | if __name__ == "__main__": 52 | gather_data_from_huggingface("output/large", "regression_results_large.csv") 53 | gather_data_from_huggingface("output/base/", "regression_results_base.csv") -------------------------------------------------------------------------------- /analysis/regression-base.sh: -------------------------------------------------------------------------------- 1 | export DATA_DIR=../data 2 | # regression task set up 3 | export TASK_NAME=STS-B 4 | declare -a arr=("roberta" "bert" "xlnet") 5 | declare -a arrtwo=("roberta-base" "bert-base-uncased" "xlnet-base-cased") 6 | 7 | 8 | for ((i=0; i<3; i++)); 9 | do 10 | python run_glue.py \ 11 | --model_type "${arr[i]}" \ 12 | --model_name_or_path "${arrtwo[i]}" \ 13 | --task_name $TASK_NAME \ 14 | --do_train \ 15 | --do_eval \ 16 | --do_lower_case \ 17 | --data_dir $DATA_DIR \ 18 | --eval_all_checkpoints \ 19 | --max_seq_length 128 \ 20 | --per_gpu_train_batch_size 96 \ 21 | --learning_rate 2e-5 \ 22 | --num_train_epochs 5.0 \ 23 | --overwrite_output_dir \ 24 | --output_dir output/base/"${arr[i]}" 25 | done 26 | -------------------------------------------------------------------------------- /analysis/regression-large.sh: -------------------------------------------------------------------------------- 1 | export DATA_DIR=../data 2 | # regression task set up 3 | export TASK_NAME=STS-B 4 | declare -a arr=("roberta" "bert" "xlnet") 5 | # I used batch size of 96 with xlnet-large and 128 with others 6 | declare -a arrtwo=("roberta-large" "bert-large-uncased" "xlnet-large-cased") 7 | 8 | 9 | for ((i=0; i<3; i++)); 10 | do 11 | python run_glue.py \ 12 | --model_type "${arr[i]}" \ 13 | --model_name_or_path "${arrtwo[i]}" \ 14 | --task_name $TASK_NAME \ 15 | --do_train \ 16 | --do_eval \ 17 | --do_lower_case \ 18 | --data_dir $DATA_DIR \ 19 | --eval_all_checkpoints \ 20 | --max_seq_length 128 \ 21 | --per_gpu_train_batch_size 128 \ 22 | --learning_rate 2e-5 \ 23 | --num_train_epochs 5.0 \ 24 | --overwrite_output_dir \ 25 | --output_dir output/large/"${arr[i]}" 26 | done 27 | -------------------------------------------------------------------------------- /analysis/regression_results_base.csv: -------------------------------------------------------------------------------- 1 | ,index,corr,pearson,rmse,spearmanr,model,checkpoint 2 | 0,0,0.4443397796905355,0.4651837567754322,1.7541763,0.42349580260563885,roberta,checkpoint-12000 3 | 1,1,0.4313989023673853,0.45019291276869466,1.6386057,0.41260489196607597,roberta,checkpoint-6000 4 | 2,2,0.4452852768941931,0.4658201273950414,1.8800247,0.4247504263933448,roberta,checkpoint-18000 5 | 3,0,0.4318083997703859,0.4531463200754483,1.8907644,0.4104704794653235,xlnet,checkpoint-12000 6 | 4,1,0.42182588036243795,0.4421904039165935,1.6474583,0.4014613568082824,xlnet,checkpoint-6000 7 | 5,2,0.43808309570939125,0.4580278612177484,1.9638016,0.4181383302010341,xlnet,checkpoint-18000 8 | 6,0,0.44518519943639,0.46596080774994675,1.7353058,0.4244095911228332,bert,checkpoint-12000 9 | 7,1,0.43999650071057195,0.4611161397762899,1.6263107,0.41887686164485405,bert,checkpoint-6000 10 | 8,2,0.44383916832300385,0.4650446676294345,1.8012,0.42263366901657323,bert,checkpoint-18000 11 | -------------------------------------------------------------------------------- /analysis/regression_results_large.csv: -------------------------------------------------------------------------------- 1 | ,index,corr,pearson,rmse,spearmanr,model,checkpoint 2 | 0,0,0.4667788714460796,0.4883324584728209,1.7011992,0.4452252844193383,roberta,checkpoint-12000 3 | 1,1,0.4544410080159367,0.47388892823620876,1.6138555,0.43499308779566465,roberta,checkpoint-6000 4 | 2,2,0.4660406387792235,0.48559753527177185,1.8018064,0.4464837422866752,roberta,checkpoint-18000 5 | 3,0,0.43412252294017617,0.4567995734825274,1.7385621,0.411445472397825,xlnet,checkpoint-12000 6 | 4,1,0.37713699244381943,0.38511733976723006,1.9101037,0.3691566451204088,xlnet,checkpoint-6000 7 | 5,2,0.44718551599712086,0.4682513100855499,1.9024026,0.42611972190869185,xlnet,checkpoint-24000 8 | 6,3,0.44616641205790697,0.4687108119213232,1.772816,0.4236220121944907,xlnet,checkpoint-18000 9 | 7,0,0.44932231309459303,0.47116557636851286,1.7352395,0.4274790498206732,bert,checkpoint-12000 10 | 8,1,0.4506944418874647,0.4709748411712956,1.6188005,0.4304140426036338,bert,checkpoint-6000 11 | 9,2,0.4523183259704836,0.47327446616808244,1.8045843,0.43136218577288477,bert,checkpoint-18000 12 | -------------------------------------------------------------------------------- /analysis/run_glue.py: -------------------------------------------------------------------------------- 1 | # coding=utf-8 2 | # Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team. 3 | # Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved. 4 | # 5 | # Licensed under the Apache License, Version 2.0 (the "License"); 6 | # you may not use this file except in compliance with the License. 7 | # You may obtain a copy of the License at 8 | # 9 | # http://www.apache.org/licenses/LICENSE-2.0 10 | # 11 | # Unless required by applicable law or agreed to in writing, software 12 | # distributed under the License is distributed on an "AS IS" BASIS, 13 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | # See the License for the specific language governing permissions and 15 | # limitations under the License. 16 | """ Finetuning the library models for sequence classification on GLUE (Bert, XLM, XLNet, RoBERTa).""" 17 | 18 | from __future__ import absolute_import, division, print_function 19 | 20 | import argparse 21 | import glob 22 | import logging 23 | import os 24 | import random 25 | 26 | import numpy as np 27 | from scipy.stats import pearsonr, spearmanr 28 | from sklearn.metrics import matthews_corrcoef, f1_score, mean_squared_error 29 | import torch 30 | from torch.utils.data import (DataLoader, RandomSampler, SequentialSampler, 31 | TensorDataset) 32 | from torch.utils.data.distributed import DistributedSampler 33 | 34 | from tqdm import tqdm, trange 35 | 36 | from transformers import (WEIGHTS_NAME, BertConfig, 37 | BertForSequenceClassification, BertTokenizer, 38 | RobertaConfig, 39 | RobertaForSequenceClassification, 40 | RobertaTokenizer, 41 | XLMConfig, XLMForSequenceClassification, 42 | XLMTokenizer, XLNetConfig, 43 | XLNetForSequenceClassification, 44 | XLNetTokenizer, 45 | DistilBertConfig, 46 | DistilBertForSequenceClassification, 47 | DistilBertTokenizer) 48 | 49 | from transformers import AdamW, get_linear_schedule_with_warmup 50 | # from transformers import WarmupLinearSchedule as get_linear_schedule_with_warmup 51 | 52 | from transformers import glue_output_modes as output_modes 53 | from transformers import glue_convert_examples_to_features as convert_examples_to_features 54 | from transformers import DataProcessor, InputExample 55 | 56 | logger = logging.getLogger(__name__) 57 | 58 | ALL_MODELS = sum((tuple(conf.pretrained_config_archive_map.keys()) for conf in (BertConfig, XLNetConfig, XLMConfig, 59 | RobertaConfig, DistilBertConfig)), ()) 60 | 61 | MODEL_CLASSES = { 62 | 'bert': (BertConfig, BertForSequenceClassification, BertTokenizer), 63 | 'xlnet': (XLNetConfig, XLNetForSequenceClassification, XLNetTokenizer), 64 | 'xlm': (XLMConfig, XLMForSequenceClassification, XLMTokenizer), 65 | 'roberta': (RobertaConfig, RobertaForSequenceClassification, RobertaTokenizer), 66 | 'distilbert': (DistilBertConfig, DistilBertForSequenceClassification, DistilBertTokenizer) 67 | } 68 | 69 | 70 | def simple_accuracy(preds, labels): 71 | return (preds == labels).mean() 72 | 73 | 74 | def acc_and_f1(preds, labels): 75 | acc = simple_accuracy(preds, labels) 76 | f1 = f1_score(y_true=labels, y_pred=preds) 77 | return { 78 | "acc": acc, 79 | "f1": f1, 80 | "acc_and_f1": (acc + f1) / 2, 81 | } 82 | 83 | 84 | def pearson_and_spearman(preds, labels): 85 | rmse = np.sqrt(mean_squared_error(preds, labels)) 86 | pearson_corr = pearsonr(preds, labels)[0] 87 | spearman_corr = spearmanr(preds, labels)[0] 88 | return { 89 | "pearson": pearson_corr, 90 | "spearmanr": spearman_corr, 91 | "corr": (pearson_corr + spearman_corr) / 2, 92 | "rmse": rmse 93 | } 94 | 95 | 96 | def compute_metrics(task_name, preds, labels): 97 | assert len(preds) == len(labels) 98 | if task_name == "cola": 99 | acc = simple_accuracy(preds, labels) 100 | f1 = f1_score(y_true=labels, y_pred=preds) 101 | mcc = matthews_corrcoef(labels, preds) 102 | return { 103 | "acc": acc, 104 | "f1": f1, 105 | "acc_and_f1": (acc + f1) / 2, 106 | "mcc": mcc 107 | } 108 | elif task_name == "sst-2": 109 | return {"acc": simple_accuracy(preds, labels)} 110 | elif task_name == "mrpc": 111 | return acc_and_f1(preds, labels) 112 | elif task_name == "sts-b": 113 | return pearson_and_spearman(preds, labels) 114 | elif task_name == "qqp": 115 | return acc_and_f1(preds, labels) 116 | elif task_name == "mnli": 117 | return {"acc": simple_accuracy(preds, labels)} 118 | elif task_name == "mnli-mm": 119 | return {"acc": simple_accuracy(preds, labels)} 120 | elif task_name == "qnli": 121 | return {"acc": simple_accuracy(preds, labels)} 122 | elif task_name == "rte": 123 | return {"acc": simple_accuracy(preds, labels)} 124 | elif task_name == "wnli": 125 | return {"acc": simple_accuracy(preds, labels)} 126 | else: 127 | raise KeyError(task_name) 128 | 129 | 130 | class StsbProcessor(DataProcessor): 131 | """OVERRIDE Processor for the STS-B data set (GLUE version).""" 132 | 133 | def get_example_from_tensor_dict(self, tensor_dict): 134 | """See base class.""" 135 | return InputExample(tensor_dict['idx'].numpy(), 136 | tensor_dict['sentence'].numpy().decode('utf-8'), 137 | None, 138 | str(tensor_dict['label'].numpy())) 139 | 140 | def get_train_examples(self, data_dir): 141 | """See base class.""" 142 | return self._create_examples( 143 | self._read_tsv(os.path.join(data_dir, "train.tsv")), "train") 144 | 145 | def get_dev_examples(self, data_dir): 146 | """See base class.""" 147 | return self._create_examples( 148 | self._read_tsv(os.path.join(data_dir, "dev.tsv")), "dev") 149 | 150 | def get_labels(self): 151 | """See base class.""" 152 | return [None] 153 | 154 | def _create_examples(self, lines, set_type): 155 | """Creates examples for the training and dev sets.""" 156 | examples = [] 157 | for (i, line) in enumerate(lines): 158 | try: 159 | if i == 0: 160 | continue 161 | guid = "%s-%s" % (set_type, line[0]) 162 | text_a = line[1] 163 | text_b = None 164 | label = line[0] 165 | examples.append( 166 | InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label)) 167 | except Exception as e: 168 | import pdb; pdb.set_trace() 169 | print(e) 170 | return examples 171 | 172 | # define my own 173 | # from transformers import glue_processors as processors 174 | processors = { 175 | "sts-b": StsbProcessor, 176 | } 177 | 178 | def set_seed(args): 179 | random.seed(args.seed) 180 | np.random.seed(args.seed) 181 | torch.manual_seed(args.seed) 182 | if args.n_gpu > 0: 183 | torch.cuda.manual_seed_all(args.seed) 184 | 185 | 186 | def train(args, train_dataset, model, tokenizer): 187 | """ Train the model """ 188 | 189 | args.train_batch_size = args.per_gpu_train_batch_size * max(1, args.n_gpu) 190 | train_sampler = RandomSampler(train_dataset) if args.local_rank == -1 else DistributedSampler(train_dataset) 191 | train_dataloader = DataLoader(train_dataset, sampler=train_sampler, batch_size=args.train_batch_size) 192 | 193 | if args.max_steps > 0: 194 | t_total = args.max_steps 195 | args.num_train_epochs = args.max_steps // (len(train_dataloader) // args.gradient_accumulation_steps) + 1 196 | else: 197 | t_total = len(train_dataloader) // args.gradient_accumulation_steps * args.num_train_epochs 198 | 199 | # Prepare optimizer and schedule (linear warmup and decay) 200 | no_decay = ['bias', 'LayerNorm.weight'] 201 | optimizer_grouped_parameters = [ 202 | {'params': [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)], 'weight_decay': args.weight_decay}, 203 | {'params': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0} 204 | ] 205 | optimizer = AdamW(optimizer_grouped_parameters, lr=args.learning_rate, eps=args.adam_epsilon) 206 | scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=args.warmup_steps, num_training_steps=t_total) 207 | if args.fp16: 208 | try: 209 | from apex import amp 210 | except ImportError: 211 | raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use fp16 training.") 212 | model, optimizer = amp.initialize(model, optimizer, opt_level=args.fp16_opt_level) 213 | 214 | # multi-gpu training (should be after apex fp16 initialization) 215 | if args.n_gpu > 1: 216 | model = torch.nn.DataParallel(model) 217 | 218 | # Distributed training (should be after apex fp16 initialization) 219 | if args.local_rank != -1: 220 | model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank], 221 | output_device=args.local_rank, 222 | find_unused_parameters=True) 223 | 224 | # Train! 225 | logger.info("***** Running training *****") 226 | logger.info(" Num examples = %d", len(train_dataset)) 227 | logger.info(" Num Epochs = %d", args.num_train_epochs) 228 | logger.info(" Instantaneous batch size per GPU = %d", args.per_gpu_train_batch_size) 229 | logger.info(" Total train batch size (w. parallel, distributed & accumulation) = %d", 230 | args.train_batch_size * args.gradient_accumulation_steps * (torch.distributed.get_world_size() if args.local_rank != -1 else 1)) 231 | logger.info(" Gradient Accumulation steps = %d", args.gradient_accumulation_steps) 232 | logger.info(" Total optimization steps = %d", t_total) 233 | 234 | global_step = 0 235 | tr_loss, logging_loss = 0.0, 0.0 236 | model.zero_grad() 237 | train_iterator = trange(int(args.num_train_epochs), desc="Epoch", disable=args.local_rank not in [-1, 0]) 238 | set_seed(args) # Added here for reproductibility (even between python 2 and 3) 239 | for _ in train_iterator: 240 | epoch_iterator = tqdm(train_dataloader, desc="Iteration", disable=args.local_rank not in [-1, 0]) 241 | for step, batch in enumerate(epoch_iterator): 242 | model.train() 243 | batch = tuple(t.to(args.device) for t in batch) 244 | inputs = {'input_ids': batch[0], 245 | 'attention_mask': batch[1], 246 | 'labels': batch[3]} 247 | if args.model_type != 'distilbert': 248 | inputs['token_type_ids'] = batch[2] if args.model_type in ['bert', 'xlnet'] else None # XLM, DistilBERT and RoBERTa don't use segment_ids 249 | outputs = model(**inputs) 250 | loss = outputs[0] # model outputs are always tuple in transformers (see doc) 251 | 252 | if args.n_gpu > 1: 253 | loss = loss.mean() # mean() to average on multi-gpu parallel training 254 | if args.gradient_accumulation_steps > 1: 255 | loss = loss / args.gradient_accumulation_steps 256 | 257 | if args.fp16: 258 | with amp.scale_loss(loss, optimizer) as scaled_loss: 259 | scaled_loss.backward() 260 | else: 261 | loss.backward() 262 | 263 | tr_loss += loss.item() 264 | if (step + 1) % args.gradient_accumulation_steps == 0 and not args.tpu: 265 | if args.fp16: 266 | torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), args.max_grad_norm) 267 | else: 268 | torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm) 269 | 270 | optimizer.step() 271 | scheduler.step() # Update learning rate schedule 272 | model.zero_grad() 273 | global_step += 1 274 | 275 | if args.local_rank in [-1, 0] and args.logging_steps > 0 and global_step % args.logging_steps == 0: 276 | # Log metrics 277 | if args.local_rank == -1 and args.evaluate_during_training: # Only evaluate when single GPU otherwise metrics may not average well 278 | results = evaluate(args, model, tokenizer) 279 | logging_loss = tr_loss 280 | 281 | if args.local_rank in [-1, 0] and args.save_steps > 0 and global_step % args.save_steps == 0: 282 | # Save model checkpoint 283 | output_dir = os.path.join(args.output_dir, 'checkpoint-{}'.format(global_step)) 284 | if not os.path.exists(output_dir): 285 | os.makedirs(output_dir) 286 | model_to_save = model.module if hasattr(model, 'module') else model # Take care of distributed/parallel training 287 | model_to_save.save_pretrained(output_dir) 288 | torch.save(args, os.path.join(output_dir, 'training_args.bin')) 289 | logger.info("Saving model checkpoint to %s", output_dir) 290 | 291 | if args.tpu: 292 | args.xla_model.optimizer_step(optimizer, barrier=True) 293 | model.zero_grad() 294 | global_step += 1 295 | 296 | if args.max_steps > 0 and global_step > args.max_steps: 297 | epoch_iterator.close() 298 | break 299 | if args.max_steps > 0 and global_step > args.max_steps: 300 | train_iterator.close() 301 | break 302 | 303 | 304 | return global_step, tr_loss / global_step 305 | 306 | 307 | def evaluate(args, model, tokenizer, prefix=""): 308 | # Loop to handle MNLI double evaluation (matched, mis-matched) 309 | eval_task_names = ("mnli", "mnli-mm") if args.task_name == "mnli" else (args.task_name,) 310 | eval_outputs_dirs = (args.output_dir, args.output_dir + '-MM') if args.task_name == "mnli" else (args.output_dir,) 311 | 312 | results = {} 313 | for eval_task, eval_output_dir in zip(eval_task_names, eval_outputs_dirs): 314 | eval_dataset = load_and_cache_examples(args, eval_task, tokenizer, evaluate=True) 315 | 316 | if not os.path.exists(eval_output_dir) and args.local_rank in [-1, 0]: 317 | os.makedirs(eval_output_dir) 318 | 319 | args.eval_batch_size = args.per_gpu_eval_batch_size * max(1, args.n_gpu) 320 | # Note that DistributedSampler samples randomly 321 | eval_sampler = SequentialSampler(eval_dataset) if args.local_rank == -1 else DistributedSampler(eval_dataset) 322 | eval_dataloader = DataLoader(eval_dataset, sampler=eval_sampler, batch_size=args.eval_batch_size) 323 | 324 | # multi-gpu eval 325 | if args.n_gpu > 1: 326 | model = torch.nn.DataParallel(model) 327 | 328 | # Eval! 329 | logger.info("***** Running evaluation {} *****".format(prefix)) 330 | logger.info(" Num examples = %d", len(eval_dataset)) 331 | logger.info(" Batch size = %d", args.eval_batch_size) 332 | eval_loss = 0.0 333 | nb_eval_steps = 0 334 | preds = None 335 | out_label_ids = None 336 | for batch in tqdm(eval_dataloader, desc="Evaluating"): 337 | model.eval() 338 | batch = tuple(t.to(args.device) for t in batch) 339 | 340 | with torch.no_grad(): 341 | inputs = {'input_ids': batch[0], 342 | 'attention_mask': batch[1], 343 | 'labels': batch[3]} 344 | if args.model_type != 'distilbert': 345 | inputs['token_type_ids'] = batch[2] if args.model_type in ['bert', 'xlnet'] else None # XLM, DistilBERT and RoBERTa don't use segment_ids 346 | outputs = model(**inputs) 347 | tmp_eval_loss, logits = outputs[:2] 348 | 349 | eval_loss += tmp_eval_loss.mean().item() 350 | nb_eval_steps += 1 351 | if preds is None: 352 | preds = logits.detach().cpu().numpy() 353 | out_label_ids = inputs['labels'].detach().cpu().numpy() 354 | else: 355 | preds = np.append(preds, logits.detach().cpu().numpy(), axis=0) 356 | out_label_ids = np.append(out_label_ids, inputs['labels'].detach().cpu().numpy(), axis=0) 357 | 358 | eval_loss = eval_loss / nb_eval_steps 359 | if args.output_mode == "classification": 360 | preds = np.argmax(preds, axis=1) 361 | elif args.output_mode == "regression": 362 | preds = np.squeeze(preds) 363 | result = compute_metrics(eval_task, preds, out_label_ids) 364 | results.update(result) 365 | 366 | output_eval_file = os.path.join(eval_output_dir, prefix, "eval_results.txt") 367 | with open(output_eval_file, "w") as writer: 368 | logger.info("***** Eval results {} *****".format(prefix)) 369 | for key in sorted(result.keys()): 370 | logger.info(" %s = %s", key, str(result[key])) 371 | writer.write("%s = %s\n" % (key, str(result[key]))) 372 | 373 | return results 374 | 375 | 376 | def load_and_cache_examples(args, task, tokenizer, evaluate=False): 377 | if args.local_rank not in [-1, 0] and not evaluate: 378 | torch.distributed.barrier() # Make sure only the first process in distributed training process the dataset, and the others will use the cache 379 | 380 | processor = processors[task]() 381 | output_mode = output_modes[task] 382 | # Load data features from cache or dataset file 383 | cached_features_file = os.path.join(args.data_dir, 'cached_{}_{}_{}_{}'.format( 384 | 'dev' if evaluate else 'train', 385 | list(filter(None, args.model_name_or_path.split('/'))).pop(), 386 | str(args.max_seq_length), 387 | str(task))) 388 | if os.path.exists(cached_features_file) and not args.overwrite_cache: 389 | logger.info("Loading features from cached file %s", cached_features_file) 390 | features = torch.load(cached_features_file) 391 | else: 392 | logger.info("Creating features from dataset file at %s", args.data_dir) 393 | label_list = processor.get_labels() 394 | if task in ['mnli', 'mnli-mm'] and args.model_type in ['roberta']: 395 | # HACK(label indices are swapped in RoBERTa pretrained model) 396 | label_list[1], label_list[2] = label_list[2], label_list[1] 397 | examples = processor.get_dev_examples(args.data_dir) if evaluate else processor.get_train_examples(args.data_dir) 398 | """ 399 | Examples is a list of dicts like so: 400 | { 401 | "guid": "train-0", # train id 402 | "label": "1", 403 | "text_a": "RIP Roger Rabbit ... James Bond", 404 | "text_b": null 405 | } 406 | """ 407 | features = convert_examples_to_features(examples, 408 | tokenizer, 409 | label_list=label_list, 410 | max_length=args.max_seq_length, 411 | output_mode=output_mode, 412 | pad_on_left=bool(args.model_type in ['xlnet']), # pad on the left for xlnet 413 | pad_token=tokenizer.convert_tokens_to_ids([tokenizer.pad_token])[0], 414 | pad_token_segment_id=4 if args.model_type in ['xlnet'] else 0, 415 | ) 416 | if args.local_rank in [-1, 0]: 417 | logger.info("Saving features into cached file %s", cached_features_file) 418 | torch.save(features, cached_features_file) 419 | 420 | if args.local_rank == 0 and not evaluate: 421 | torch.distributed.barrier() # Make sure only the first process in distributed training process the dataset, and the others will use the cache 422 | 423 | # Convert to Tensors and build dataset 424 | all_input_ids = torch.tensor([f.input_ids for f in features], dtype=torch.long) 425 | all_attention_mask = torch.tensor([f.attention_mask for f in features], dtype=torch.long) 426 | all_token_type_ids = torch.tensor([f.token_type_ids for f in features], dtype=torch.long) 427 | if output_mode == "classification": 428 | all_labels = torch.tensor([f.label for f in features], dtype=torch.long) 429 | elif output_mode == "regression": 430 | all_labels = torch.tensor([f.label for f in features], dtype=torch.float) 431 | 432 | dataset = TensorDataset(all_input_ids, all_attention_mask, all_token_type_ids, all_labels) 433 | return dataset 434 | 435 | 436 | def main(): 437 | parser = argparse.ArgumentParser() 438 | 439 | ## Required parameters 440 | parser.add_argument("--data_dir", default=None, type=str, required=True, 441 | help="The input data dir. Should contain the .tsv files (or other data files) for the task.") 442 | parser.add_argument("--model_type", default=None, type=str, required=True, 443 | help="Model type selected in the list: " + ", ".join(MODEL_CLASSES.keys())) 444 | parser.add_argument("--model_name_or_path", default=None, type=str, required=True, 445 | help="Path to pre-trained model or shortcut name selected in the list: " + ", ".join(ALL_MODELS)) 446 | parser.add_argument("--task_name", default=None, type=str, required=True, 447 | help="The name of the task to train selected in the list: " + ", ".join(processors.keys())) 448 | parser.add_argument("--output_dir", default=None, type=str, required=True, 449 | help="The output directory where the model predictions and checkpoints will be written.") 450 | 451 | ## Other parameters 452 | parser.add_argument("--config_name", default="", type=str, 453 | help="Pretrained config name or path if not the same as model_name") 454 | parser.add_argument("--tokenizer_name", default="", type=str, 455 | help="Pretrained tokenizer name or path if not the same as model_name") 456 | parser.add_argument("--cache_dir", default="", type=str, 457 | help="Where do you want to store the pre-trained models downloaded from s3") 458 | parser.add_argument("--max_seq_length", default=128, type=int, 459 | help="The maximum total input sequence length after tokenization. Sequences longer " 460 | "than this will be truncated, sequences shorter will be padded.") 461 | parser.add_argument("--do_train", action='store_true', 462 | help="Whether to run training.") 463 | parser.add_argument("--do_eval", action='store_true', 464 | help="Whether to run eval on the dev set.") 465 | parser.add_argument("--evaluate_during_training", action='store_true', 466 | help="Rul evaluation during training at each logging step.") 467 | parser.add_argument("--do_lower_case", action='store_true', 468 | help="Set this flag if you are using an uncased model.") 469 | 470 | parser.add_argument("--per_gpu_train_batch_size", default=8, type=int, 471 | help="Batch size per GPU/CPU for training.") 472 | parser.add_argument("--per_gpu_eval_batch_size", default=8, type=int, 473 | help="Batch size per GPU/CPU for evaluation.") 474 | parser.add_argument('--gradient_accumulation_steps', type=int, default=1, 475 | help="Number of updates steps to accumulate before performing a backward/update pass.") 476 | parser.add_argument("--learning_rate", default=5e-5, type=float, 477 | help="The initial learning rate for Adam.") 478 | parser.add_argument("--weight_decay", default=0.0, type=float, 479 | help="Weight deay if we apply some.") 480 | parser.add_argument("--adam_epsilon", default=1e-8, type=float, 481 | help="Epsilon for Adam optimizer.") 482 | parser.add_argument("--max_grad_norm", default=1.0, type=float, 483 | help="Max gradient norm.") 484 | parser.add_argument("--num_train_epochs", default=3.0, type=float, 485 | help="Total number of training epochs to perform.") 486 | parser.add_argument("--max_steps", default=-1, type=int, 487 | help="If > 0: set total number of training steps to perform. Override num_train_epochs.") 488 | parser.add_argument("--warmup_steps", default=0, type=int, 489 | help="Linear warmup over warmup_steps.") 490 | 491 | parser.add_argument('--logging_steps', type=int, default=6000, 492 | help="Log every X updates steps.") 493 | parser.add_argument('--save_steps', type=int, default=6000, 494 | help="Save checkpoint every X updates steps.") 495 | parser.add_argument("--eval_all_checkpoints", action='store_true', 496 | help="Evaluate all checkpoints starting with the same prefix as model_name ending and ending with step number") 497 | parser.add_argument("--no_cuda", action='store_true', 498 | help="Avoid using CUDA when available") 499 | parser.add_argument('--overwrite_output_dir', action='store_true', 500 | help="Overwrite the content of the output directory") 501 | parser.add_argument('--overwrite_cache', action='store_true', 502 | help="Overwrite the cached training and evaluation sets") 503 | parser.add_argument('--seed', type=int, default=42, 504 | help="random seed for initialization") 505 | 506 | parser.add_argument('--tpu', action='store_true', 507 | help="Whether to run on the TPU defined in the environment variables") 508 | parser.add_argument('--tpu_ip_address', type=str, default='', 509 | help="TPU IP address if none are set in the environment variables") 510 | parser.add_argument('--tpu_name', type=str, default='', 511 | help="TPU name if none are set in the environment variables") 512 | parser.add_argument('--xrt_tpu_config', type=str, default='', 513 | help="XRT TPU config if none are set in the environment variables") 514 | 515 | parser.add_argument('--fp16', action='store_true', 516 | help="Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit") 517 | parser.add_argument('--fp16_opt_level', type=str, default='O1', 518 | help="For fp16: Apex AMP optimization level selected in ['O0', 'O1', 'O2', and 'O3']." 519 | "See details at https://nvidia.github.io/apex/amp.html") 520 | parser.add_argument("--local_rank", type=int, default=-1, 521 | help="For distributed training: local_rank") 522 | parser.add_argument('--server_ip', type=str, default='', help="For distant debugging.") 523 | parser.add_argument('--server_port', type=str, default='', help="For distant debugging.") 524 | args = parser.parse_args() 525 | 526 | if os.path.exists(args.output_dir) and os.listdir(args.output_dir) and args.do_train and not args.overwrite_output_dir: 527 | raise ValueError("Output directory ({}) already exists and is not empty. Use --overwrite_output_dir to overcome.".format(args.output_dir)) 528 | 529 | # Setup distant debugging if needed 530 | if args.server_ip and args.server_port: 531 | # Distant debugging - see https://code.visualstudio.com/docs/python/debugging#_attach-to-a-local-script 532 | import ptvsd 533 | print("Waiting for debugger attach") 534 | ptvsd.enable_attach(address=(args.server_ip, args.server_port), redirect_output=True) 535 | ptvsd.wait_for_attach() 536 | 537 | # Setup CUDA, GPU & distributed training 538 | if args.local_rank == -1 or args.no_cuda: 539 | device = torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu") 540 | args.n_gpu = torch.cuda.device_count() 541 | else: # Initializes the distributed backend which will take care of sychronizing nodes/GPUs 542 | torch.cuda.set_device(args.local_rank) 543 | device = torch.device("cuda", args.local_rank) 544 | torch.distributed.init_process_group(backend='nccl') 545 | args.n_gpu = 1 546 | args.device = device 547 | 548 | if args.tpu: 549 | if args.tpu_ip_address: 550 | os.environ["TPU_IP_ADDRESS"] = args.tpu_ip_address 551 | if args.tpu_name: 552 | os.environ["TPU_NAME"] = args.tpu_name 553 | if args.xrt_tpu_config: 554 | os.environ["XRT_TPU_CONFIG"] = args.xrt_tpu_config 555 | 556 | assert "TPU_IP_ADDRESS" in os.environ 557 | assert "TPU_NAME" in os.environ 558 | assert "XRT_TPU_CONFIG" in os.environ 559 | 560 | import torch_xla 561 | import torch_xla.core.xla_model as xm 562 | args.device = xm.xla_device() 563 | args.xla_model = xm 564 | 565 | # Setup logging 566 | logging.basicConfig(format = '%(asctime)s - %(levelname)s - %(name)s - %(message)s', 567 | datefmt = '%m/%d/%Y %H:%M:%S', 568 | level = logging.INFO if args.local_rank in [-1, 0] else logging.WARN) 569 | logger.warning("Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, 16-bits training: %s", 570 | args.local_rank, device, args.n_gpu, bool(args.local_rank != -1), args.fp16) 571 | 572 | # Set seed 573 | set_seed(args) 574 | 575 | # Prepare GLUE task 576 | args.task_name = args.task_name.lower() 577 | if args.task_name not in processors: 578 | raise ValueError("Task not found: %s" % (args.task_name)) 579 | processor = processors[args.task_name]() 580 | args.output_mode = output_modes[args.task_name] 581 | label_list = processor.get_labels() 582 | num_labels = len(label_list) 583 | 584 | # Load pretrained model and tokenizer 585 | if args.local_rank not in [-1, 0]: 586 | torch.distributed.barrier() # Make sure only the first process in distributed training will download model & vocab 587 | 588 | args.model_type = args.model_type.lower() 589 | config_class, model_class, tokenizer_class = MODEL_CLASSES[args.model_type] 590 | config = config_class.from_pretrained(args.config_name if args.config_name else args.model_name_or_path, 591 | num_labels=num_labels, 592 | finetuning_task=args.task_name, 593 | cache_dir=args.cache_dir if args.cache_dir else None) 594 | tokenizer = tokenizer_class.from_pretrained(args.tokenizer_name if args.tokenizer_name else args.model_name_or_path, 595 | do_lower_case=args.do_lower_case, 596 | cache_dir=args.cache_dir if args.cache_dir else None) 597 | model = model_class.from_pretrained(args.model_name_or_path, 598 | from_tf=bool('.ckpt' in args.model_name_or_path), 599 | config=config, 600 | cache_dir=args.cache_dir if args.cache_dir else None) 601 | 602 | if args.local_rank == 0: 603 | torch.distributed.barrier() # Make sure only the first process in distributed training will download model & vocab 604 | 605 | model.to(args.device) 606 | 607 | logger.info("Training/evaluation parameters %s", args) 608 | 609 | 610 | # Training 611 | if args.do_train: 612 | train_dataset = load_and_cache_examples(args, args.task_name, tokenizer, evaluate=False) 613 | global_step, tr_loss = train(args, train_dataset, model, tokenizer) 614 | logger.info(" global_step = %s, average loss = %s", global_step, tr_loss) 615 | 616 | 617 | # Saving best-practices: if you use defaults names for the model, you can reload it using from_pretrained() 618 | if args.do_train and (args.local_rank == -1 or torch.distributed.get_rank() == 0) and not args.tpu: 619 | # Create output directory if needed 620 | if not os.path.exists(args.output_dir) and args.local_rank in [-1, 0]: 621 | os.makedirs(args.output_dir) 622 | 623 | logger.info("Saving model checkpoint to %s", args.output_dir) 624 | # Save a trained model, configuration and tokenizer using `save_pretrained()`. 625 | # They can then be reloaded using `from_pretrained()` 626 | model_to_save = model.module if hasattr(model, 'module') else model # Take care of distributed/parallel training 627 | model_to_save.save_pretrained(args.output_dir) 628 | tokenizer.save_pretrained(args.output_dir) 629 | 630 | # Good practice: save your training arguments together with the trained model 631 | torch.save(args, os.path.join(args.output_dir, 'training_args.bin')) 632 | 633 | # Load a trained model and vocabulary that you have fine-tuned 634 | model = model_class.from_pretrained(args.output_dir) 635 | tokenizer = tokenizer_class.from_pretrained(args.output_dir) 636 | model.to(args.device) 637 | 638 | 639 | # Evaluation 640 | results = {} 641 | if args.do_eval and args.local_rank in [-1, 0]: 642 | tokenizer = tokenizer_class.from_pretrained(args.output_dir, do_lower_case=args.do_lower_case) 643 | checkpoints = [args.output_dir] 644 | if args.eval_all_checkpoints: 645 | checkpoints = list(os.path.dirname(c) for c in sorted(glob.glob(args.output_dir + '/**/' + WEIGHTS_NAME, recursive=True))) 646 | logging.getLogger("transformers.modeling_utils").setLevel(logging.WARN) # Reduce logging 647 | logger.info("Evaluate the following checkpoints: %s", checkpoints) 648 | for checkpoint in checkpoints: 649 | global_step = checkpoint.split('-')[-1] if len(checkpoints) > 1 else "" 650 | prefix = checkpoint.split('/')[-1] if checkpoint.find('checkpoint') != -1 else "" 651 | 652 | model = model_class.from_pretrained(checkpoint) 653 | model.to(args.device) 654 | result = evaluate(args, model, tokenizer, prefix=prefix) 655 | result = dict((k + '_{}'.format(global_step), v) for k, v in result.items()) 656 | results.update(result) 657 | 658 | return results 659 | 660 | 661 | if __name__ == "__main__": 662 | main() -------------------------------------------------------------------------------- /analysis/time_statistics.py: -------------------------------------------------------------------------------- 1 | import os 2 | from datetime import datetime 3 | import random 4 | 5 | import pandas as pd 6 | import numpy as np 7 | import matplotlib.pyplot as plt 8 | from vaderSentiment.vaderSentiment import SentimentIntensityAnalyzer 9 | import seaborn as sns 10 | from sklearn.decomposition import LatentDirichletAllocation as LDA 11 | 12 | from dataset_statistics import plot_n_most_common_words, EXT_TYPES 13 | 14 | random.seed(42) 15 | np.random.seed(42) 16 | 17 | def sentiment_scores(sid_obj: SentimentIntensityAnalyzer, sentence: str) -> int: 18 | sentiment_dict = sid_obj.polarity_scores(sentence) 19 | if sentiment_dict['compound'] >= 0.05 : 20 | # print("positive", sentence) 21 | return 1 22 | elif sentiment_dict['compound'] <= - 0.05 : 23 | # print("negative", sentence) 24 | return -1 25 | else: 26 | # print("neutral", sentence) 27 | return 0 28 | 29 | 30 | def get_sentiment(df, data_name): 31 | sid_obj = SentimentIntensityAnalyzer() 32 | sentiments = np.array([sentiment_scores(sid_obj, sentence) for sentence in df["joke"].dropna().values]) 33 | results = dict(zip(*np.unique(sentiments, return_counts=True))) 34 | pos_prop = results[-1] / len(sentiments) 35 | neutral_prop = results[0] / len(sentiments) 36 | neg_prop = results[1] / len(sentiments) 37 | 38 | # std for a proportion is sqrt(p*q / n) 39 | pos_std = np.sqrt((pos_prop * (1 - pos_prop)) / len(sentiments)) 40 | neutral_std = np.sqrt((neutral_prop * (1 - neutral_prop)) / len(sentiments)) 41 | neg_std = np.sqrt((neg_prop * (1 - neg_prop)) / len(sentiments)) 42 | 43 | results = {"positive_prop": pos_prop, "positive_std": pos_std, "neutral_prop": neutral_prop, "neutral_std": neutral_std, "negative_prop": neg_prop, 44 | "negative_std": neg_std} 45 | 46 | if "date" not in df.columns: 47 | return results 48 | else: 49 | list_of_results = [] 50 | for prop_type in ["positive", "neutral", "negative"]: 51 | name = prop_type + "_prop" 52 | name_std = prop_type + "_std" 53 | list_of_results.append({"prop": results[name], "sentiment": prop_type, "std": results[name_std], "date": df.date.iloc[0]}) 54 | return pd.DataFrame(list_of_results) 55 | 56 | 57 | def get_scores(df, data_name): 58 | return df[["date", "score"]] 59 | 60 | def get_jokes(df, data_name): 61 | return df[["date", "joke"]] 62 | 63 | def split_by_year(df: pd.DataFrame, func, data_name="all") -> pd.DataFrame: 64 | df["date"] = df["date"].map(lambda x: pd.to_datetime(datetime.utcfromtimestamp(x)).year) 65 | date_df = df.groupby("date").apply(lambda x: func(x, data_name)) 66 | return date_df 67 | 68 | def plot_lines_by_year(date_df: pd.DataFrame, data_name: str, name: str = "sentiment"): 69 | for ext_type in EXT_TYPES: 70 | ax = sns.lineplot(x="date", y="prop", hue="sentiment", data=date_df, palette=["C2", "C7", "C3"]) 71 | ax.errorbar(date_df["date"], date_df["prop"], yerr=date_df["std"], fmt='o', ecolor='black') 72 | plt.xlabel('Year') 73 | plt.ylabel('Proportion of Sentiment') 74 | plt.savefig(os.path.join(os.path.realpath('..'), "plots", data_name, "{}.{}".format(name, ext_type))) 75 | plt.close() 76 | 77 | def plot_distribution_by_year(date_df: pd.DataFrame, data_name: str, name: str = "score"): 78 | date_df = date_df[date_df["score"] != 0] # zeros skew the distribution 79 | for ext_type in EXT_TYPES: 80 | ax = sns.violinplot(x="date", y="score", data=date_df) 81 | plt.xlabel('Year') 82 | plt.ylabel('Distribution of Jokes') 83 | plt.savefig(os.path.join(os.path.realpath('..'), "plots", data_name, "{}.{}".format(name, ext_type))) 84 | plt.close() 85 | 86 | def plot_bar_by_year(df: pd.DataFrame, data_name: str, name: str): 87 | if name == "mean": 88 | ave = df.groupby('date').mean().reset_index(drop=True) 89 | ave.columns = ["average"] 90 | max = df.groupby('date').max().reset_index() 91 | max.columns = ["date", "score"] 92 | data = pd.concat([max, ave], axis=1) 93 | elif name == "count": 94 | data = df.groupby('date').count().reset_index() 95 | elif name == "max": 96 | data = df.groupby('date').max().reset_index() 97 | else: 98 | raise NotImplementedError() 99 | 100 | for ext_type in EXT_TYPES: 101 | ax = sns.barplot(x="date", y="score", data=data) 102 | for index, p in enumerate(ax.patches): 103 | height = p.get_height() 104 | ax.text(p.get_x()+p.get_width()/2., 105 | height + 3, 106 | # will use average for numbering in max case 107 | '{}'.format(int(data["score"].iloc[index]) if name in ["count", "max"] else int(data["average"].iloc[index])), 108 | ha="center") 109 | plt.xlabel('Year') 110 | plt.ylabel(name) 111 | plt.tight_layout() 112 | plt.savefig(os.path.join(os.path.realpath('..'), "plots", data_name, "{}.{}".format(name, ext_type))) 113 | plt.close() 114 | 115 | 116 | if __name__ == "__main__" : 117 | df = pd.read_csv(os.path.join(os.path.realpath('..'), "data", "preprocessed.csv"), index_col=None, encoding="UTF-8", keep_default_na=False) 118 | df["date"] = pd.to_numeric(df["date"]) 119 | df["score"] = pd.to_numeric(df["score"]) 120 | df = df[df["date"].isna() == False] 121 | 122 | plot_bar_by_year(split_by_year(df.copy(deep=True), get_scores), "all", "max") 123 | plot_bar_by_year(split_by_year(df.copy(deep=True), get_scores), "all", "count") 124 | plot_lines_by_year(split_by_year(df.copy(deep=True), get_sentiment, "all"), "all", "sentiment") 125 | plot_distribution_by_year(split_by_year(df.copy(deep=True), get_scores, "all"), "all", "score_distributions") 126 | # plot top entities per year 127 | jokes_by_year = split_by_year(df.copy(deep=True), get_jokes, "all") 128 | for year in jokes_by_year["date"].unique(): 129 | plot_n_most_common_words(jokes_by_year[jokes_by_year["date"] == year]["joke"].tolist(), data_name="yearbyyear", specific_year=year) 130 | -------------------------------------------------------------------------------- /data/dev.tsv.gz: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/orionw/rJokesData/d48bedd71bdacc7557b84ad697bf556e7aad7c21/data/dev.tsv.gz -------------------------------------------------------------------------------- /data/fullrjokes.json.gz: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/orionw/rJokesData/d48bedd71bdacc7557b84ad697bf556e7aad7c21/data/fullrjokes.json.gz -------------------------------------------------------------------------------- /data/preprocessed.csv.gz: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/orionw/rJokesData/d48bedd71bdacc7557b84ad697bf556e7aad7c21/data/preprocessed.csv.gz -------------------------------------------------------------------------------- /data/submissions.json.gz: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/orionw/rJokesData/d48bedd71bdacc7557b84ad697bf556e7aad7c21/data/submissions.json.gz -------------------------------------------------------------------------------- /data/test.tsv.gz: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/orionw/rJokesData/d48bedd71bdacc7557b84ad697bf556e7aad7c21/data/test.tsv.gz -------------------------------------------------------------------------------- /data/train.tsv.gz: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/orionw/rJokesData/d48bedd71bdacc7557b84ad697bf556e7aad7c21/data/train.tsv.gz -------------------------------------------------------------------------------- /download_nltk_packages.sh: -------------------------------------------------------------------------------- 1 | python3 -m nltk.downloader averaged_perceptron_tagger 2 | python3 -m nltk.downloader words 3 | python3 -m nltk.downloader stopwords 4 | python3 -m nltk.downloader maxent_ne_chunker 5 | -------------------------------------------------------------------------------- /plots/all/count.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/orionw/rJokesData/d48bedd71bdacc7557b84ad697bf556e7aad7c21/plots/all/count.pdf -------------------------------------------------------------------------------- /plots/all/count.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/orionw/rJokesData/d48bedd71bdacc7557b84ad697bf556e7aad7c21/plots/all/count.png -------------------------------------------------------------------------------- /plots/all/logdist.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/orionw/rJokesData/d48bedd71bdacc7557b84ad697bf556e7aad7c21/plots/all/logdist.pdf -------------------------------------------------------------------------------- /plots/all/logdist.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/orionw/rJokesData/d48bedd71bdacc7557b84ad697bf556e7aad7c21/plots/all/logdist.png -------------------------------------------------------------------------------- /plots/all/max.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/orionw/rJokesData/d48bedd71bdacc7557b84ad697bf556e7aad7c21/plots/all/max.pdf -------------------------------------------------------------------------------- /plots/all/max.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/orionw/rJokesData/d48bedd71bdacc7557b84ad697bf556e7aad7c21/plots/all/max.png -------------------------------------------------------------------------------- /plots/all/percentiles.csv: -------------------------------------------------------------------------------- 1 | ,percentile,value 2 | 0,0,0.0 3 | 1,10,0.0 4 | 2,25,1.0 5 | 3,50,5.0 6 | 4,75,20.0 7 | 5,90,103.0 8 | 6,100,1541018587.0 9 | -------------------------------------------------------------------------------- /plots/all/score_distributions.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/orionw/rJokesData/d48bedd71bdacc7557b84ad697bf556e7aad7c21/plots/all/score_distributions.pdf -------------------------------------------------------------------------------- /plots/all/score_distributions.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/orionw/rJokesData/d48bedd71bdacc7557b84ad697bf556e7aad7c21/plots/all/score_distributions.png -------------------------------------------------------------------------------- /plots/all/sentiment.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/orionw/rJokesData/d48bedd71bdacc7557b84ad697bf556e7aad7c21/plots/all/sentiment.pdf -------------------------------------------------------------------------------- /plots/all/sentiment.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/orionw/rJokesData/d48bedd71bdacc7557b84ad697bf556e7aad7c21/plots/all/sentiment.png -------------------------------------------------------------------------------- /plots/all/statistics.txt: -------------------------------------------------------------------------------- 1 | ,ave_punchline_len,ave_body_len,ave_joke_len,std_punch,std_body,std_joke,total_tokens 2 | 0,47.84236006473576,191.93603856241978,239.7849140577041,25.90976722809782,502.55412523216035,501.56866351178525,256619 3 | -------------------------------------------------------------------------------- /plots/all/top_wordcounts.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/orionw/rJokesData/d48bedd71bdacc7557b84ad697bf556e7aad7c21/plots/all/top_wordcounts.pdf -------------------------------------------------------------------------------- /plots/all/top_wordcounts.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/orionw/rJokesData/d48bedd71bdacc7557b84ad697bf556e7aad7c21/plots/all/top_wordcounts.png -------------------------------------------------------------------------------- /plots/all/total_score_distribution.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/orionw/rJokesData/d48bedd71bdacc7557b84ad697bf556e7aad7c21/plots/all/total_score_distribution.pdf -------------------------------------------------------------------------------- /plots/all/total_score_distribution.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/orionw/rJokesData/d48bedd71bdacc7557b84ad697bf556e7aad7c21/plots/all/total_score_distribution.png -------------------------------------------------------------------------------- /plots/all/wordCloud.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/orionw/rJokesData/d48bedd71bdacc7557b84ad697bf556e7aad7c21/plots/all/wordCloud.pdf -------------------------------------------------------------------------------- /plots/all/wordCloud.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/orionw/rJokesData/d48bedd71bdacc7557b84ad697bf556e7aad7c21/plots/all/wordCloud.png -------------------------------------------------------------------------------- /plots/all/wordCloudSmall.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/orionw/rJokesData/d48bedd71bdacc7557b84ad697bf556e7aad7c21/plots/all/wordCloudSmall.pdf -------------------------------------------------------------------------------- /plots/yearbyyear/2008_top_wordcounts.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/orionw/rJokesData/d48bedd71bdacc7557b84ad697bf556e7aad7c21/plots/yearbyyear/2008_top_wordcounts.pdf -------------------------------------------------------------------------------- /plots/yearbyyear/2008_top_wordcounts.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/orionw/rJokesData/d48bedd71bdacc7557b84ad697bf556e7aad7c21/plots/yearbyyear/2008_top_wordcounts.png -------------------------------------------------------------------------------- /plots/yearbyyear/2009_top_wordcounts.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/orionw/rJokesData/d48bedd71bdacc7557b84ad697bf556e7aad7c21/plots/yearbyyear/2009_top_wordcounts.pdf -------------------------------------------------------------------------------- /plots/yearbyyear/2009_top_wordcounts.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/orionw/rJokesData/d48bedd71bdacc7557b84ad697bf556e7aad7c21/plots/yearbyyear/2009_top_wordcounts.png -------------------------------------------------------------------------------- /plots/yearbyyear/2010_top_wordcounts.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/orionw/rJokesData/d48bedd71bdacc7557b84ad697bf556e7aad7c21/plots/yearbyyear/2010_top_wordcounts.pdf -------------------------------------------------------------------------------- /plots/yearbyyear/2010_top_wordcounts.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/orionw/rJokesData/d48bedd71bdacc7557b84ad697bf556e7aad7c21/plots/yearbyyear/2010_top_wordcounts.png -------------------------------------------------------------------------------- /plots/yearbyyear/2011_top_wordcounts.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/orionw/rJokesData/d48bedd71bdacc7557b84ad697bf556e7aad7c21/plots/yearbyyear/2011_top_wordcounts.pdf -------------------------------------------------------------------------------- /plots/yearbyyear/2011_top_wordcounts.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/orionw/rJokesData/d48bedd71bdacc7557b84ad697bf556e7aad7c21/plots/yearbyyear/2011_top_wordcounts.png -------------------------------------------------------------------------------- /plots/yearbyyear/2012_top_wordcounts.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/orionw/rJokesData/d48bedd71bdacc7557b84ad697bf556e7aad7c21/plots/yearbyyear/2012_top_wordcounts.pdf -------------------------------------------------------------------------------- /plots/yearbyyear/2012_top_wordcounts.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/orionw/rJokesData/d48bedd71bdacc7557b84ad697bf556e7aad7c21/plots/yearbyyear/2012_top_wordcounts.png -------------------------------------------------------------------------------- /plots/yearbyyear/2013_top_wordcounts.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/orionw/rJokesData/d48bedd71bdacc7557b84ad697bf556e7aad7c21/plots/yearbyyear/2013_top_wordcounts.pdf -------------------------------------------------------------------------------- /plots/yearbyyear/2013_top_wordcounts.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/orionw/rJokesData/d48bedd71bdacc7557b84ad697bf556e7aad7c21/plots/yearbyyear/2013_top_wordcounts.png -------------------------------------------------------------------------------- /plots/yearbyyear/2014_top_wordcounts.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/orionw/rJokesData/d48bedd71bdacc7557b84ad697bf556e7aad7c21/plots/yearbyyear/2014_top_wordcounts.pdf -------------------------------------------------------------------------------- /plots/yearbyyear/2014_top_wordcounts.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/orionw/rJokesData/d48bedd71bdacc7557b84ad697bf556e7aad7c21/plots/yearbyyear/2014_top_wordcounts.png -------------------------------------------------------------------------------- /plots/yearbyyear/2015_top_wordcounts.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/orionw/rJokesData/d48bedd71bdacc7557b84ad697bf556e7aad7c21/plots/yearbyyear/2015_top_wordcounts.pdf -------------------------------------------------------------------------------- /plots/yearbyyear/2015_top_wordcounts.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/orionw/rJokesData/d48bedd71bdacc7557b84ad697bf556e7aad7c21/plots/yearbyyear/2015_top_wordcounts.png -------------------------------------------------------------------------------- /plots/yearbyyear/2016_top_wordcounts.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/orionw/rJokesData/d48bedd71bdacc7557b84ad697bf556e7aad7c21/plots/yearbyyear/2016_top_wordcounts.pdf -------------------------------------------------------------------------------- /plots/yearbyyear/2016_top_wordcounts.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/orionw/rJokesData/d48bedd71bdacc7557b84ad697bf556e7aad7c21/plots/yearbyyear/2016_top_wordcounts.png -------------------------------------------------------------------------------- /plots/yearbyyear/2017_top_wordcounts.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/orionw/rJokesData/d48bedd71bdacc7557b84ad697bf556e7aad7c21/plots/yearbyyear/2017_top_wordcounts.pdf -------------------------------------------------------------------------------- /plots/yearbyyear/2017_top_wordcounts.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/orionw/rJokesData/d48bedd71bdacc7557b84ad697bf556e7aad7c21/plots/yearbyyear/2017_top_wordcounts.png -------------------------------------------------------------------------------- /plots/yearbyyear/2018_top_wordcounts.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/orionw/rJokesData/d48bedd71bdacc7557b84ad697bf556e7aad7c21/plots/yearbyyear/2018_top_wordcounts.pdf -------------------------------------------------------------------------------- /plots/yearbyyear/2018_top_wordcounts.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/orionw/rJokesData/d48bedd71bdacc7557b84ad697bf556e7aad7c21/plots/yearbyyear/2018_top_wordcounts.png -------------------------------------------------------------------------------- /plots/yearbyyear/2019_top_wordcounts.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/orionw/rJokesData/d48bedd71bdacc7557b84ad697bf556e7aad7c21/plots/yearbyyear/2019_top_wordcounts.pdf -------------------------------------------------------------------------------- /plots/yearbyyear/2019_top_wordcounts.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/orionw/rJokesData/d48bedd71bdacc7557b84ad697bf556e7aad7c21/plots/yearbyyear/2019_top_wordcounts.png -------------------------------------------------------------------------------- /prepare_data/gather_reddit_pushshift.py: -------------------------------------------------------------------------------- 1 | """ 2 | Code framework taken from: 3 | https://www.osrsbox.com/blog/2019/03/18/watercooler-scraping-an-entire-subreddit-2007scape/ 4 | """ 5 | 6 | import requests 7 | import json 8 | import re 9 | import time 10 | import datetime 11 | import os 12 | from dateutil.relativedelta import relativedelta 13 | 14 | PUSHSHIFT_REDDIT_URL = "http://api.pushshift.io/reddit" 15 | 16 | def fetchObjects(**kwargs): 17 | # Default paramaters for API query 18 | params = { 19 | "sort_type":"created_utc", 20 | "sort":"asc", 21 | "size":1000, 22 | "filter": ("id", "created_utc") 23 | } 24 | 25 | # Add additional paramters based on function arguments 26 | for key,value in kwargs.items(): 27 | params[key] = value 28 | 29 | # Print API query paramaters 30 | print(params) 31 | 32 | # Set the type variable based on function input 33 | # The type can be "comment" or "submission", default is "comment" 34 | type_post = "comment" 35 | if 'type' in kwargs and kwargs['type'].lower() == "submission": 36 | type_post = "submission" 37 | 38 | # Perform an API request 39 | r = requests.get(PUSHSHIFT_REDDIT_URL + "/" + type_post + "/search/", params=params, timeout=30) 40 | 41 | # Check the status code, if successful, process the data 42 | if r.status_code == 200: 43 | response = json.loads(r.text) 44 | data = response['data'] 45 | sorted_data_by_id = sorted(data, key=lambda x: int(x['id'],36)) 46 | return sorted_data_by_id 47 | 48 | else: 49 | return [] 50 | 51 | def extract_reddit_data(**kwargs): 52 | # Speficify the start timestamp, make sure we get everything 53 | max_created_utc = int((datetime.datetime.now() - relativedelta(years=100)).timestamp()) 54 | max_id = 0 55 | 56 | # Open a file for JSON output 57 | file = open(os.path.join(os.path.realpath('..'), "data", "submissions.json"), "w") 58 | 59 | # While loop for recursive function 60 | while 1: 61 | nothing_processed = True 62 | # Call the recursive function 63 | print("After {}".format(max_created_utc)) 64 | objects = fetchObjects(**kwargs,after=max_created_utc) 65 | 66 | # Loop the returned data, ordered by date 67 | for object in objects: 68 | id = int(object['id'], 36) 69 | if id > max_id: 70 | nothing_processed = False 71 | created_utc = object['created_utc'] 72 | max_id = id 73 | if created_utc > max_created_utc: max_created_utc = created_utc 74 | # Output JSON data to the opened file 75 | print(json.dumps(object, sort_keys=True, ensure_ascii=True), file=file) 76 | 77 | # Exit if nothing happened 78 | if nothing_processed: 79 | return 80 | max_created_utc -= 1 81 | 82 | # Sleep a little before the next recursive function call 83 | time.sleep(.5) 84 | 85 | # Start program by calling function with: 86 | # 1) Subreddit specified 87 | # 2) The type of data required (comment or submission) 88 | extract_reddit_data(subreddit="jokes",type="submission") -------------------------------------------------------------------------------- /prepare_data/preprocess.py: -------------------------------------------------------------------------------- 1 | import re 2 | import requests 3 | import json 4 | import praw 5 | import datetime 6 | import os 7 | from multiprocessing import Pool 8 | import time 9 | import csv 10 | 11 | import argparse 12 | import pandas as pd 13 | import numpy as np 14 | import nltk 15 | from sklearn.model_selection import train_test_split 16 | import seaborn as sns 17 | import matplotlib.pyplot as plt 18 | 19 | 20 | reddit = praw.Reddit(client_id='client_id', 21 | client_secret='secret_key', 22 | user_agent="Desciption Here") 23 | 24 | random.seed(42) 25 | np.random.seed(42) 26 | 27 | def processExistingPost(doc): 28 | """ A helper function to gather data from the Reddit API based on post ID """ 29 | # create a post object 30 | post = reddit.submission(doc['id']) 31 | # fetch the updated post from reddit 32 | post._fetch() 33 | return post 34 | 35 | 36 | def read_and_update_ids(file_name: str = "submissions.json", output_name: str = "fullrjokes.json"): 37 | """ 38 | Gathers the data in the data/{file_name} file and grabs the fields we want to keep, along with the updated 39 | upvote count and stores them in data/{output_name}. 40 | 41 | If you want to include other meta-data from the reddit API, you can add it to `KEEP_KEYS` list 42 | 43 | Args: 44 | file_name: the name of the file in the `data` folder that contains the Reddit post IDs 45 | output_name: the name of the file to output the gathered posts to, in the `data` folder 46 | """ 47 | print("Reading and update jokes for each id...") 48 | KEEP_KEYS = ["id", "selftext", "title", "downs", "ups", "score", "name", "created_utc"] 49 | with open(os.path.join(os.path.realpath('..'), "data", file_name), "r") as f: 50 | for line in f: 51 | line = json.loads(line) 52 | for retry in range(3): 53 | try: 54 | updated_doc = processExistingPost(line) 55 | except Exception as e: 56 | print("On retry number", retry) 57 | time.sleep(5) 58 | continue 59 | break 60 | consalidated_dict = {your_key: updated_doc.__dict__[your_key] for your_key in KEEP_KEYS} 61 | with open(os.path.join(os.path.realpath('..'), "data", output_name), "a") as fout: 62 | fout.write(json.dumps(consalidated_dict, default=str) + "\n") 63 | print("wrote out id={}".format(updated_doc)) 64 | 65 | 66 | def phrase_in_doc(phrase, updated_doc): 67 | """ A helper function to see if the phrase is in the document body (title) or punchline (selftext) """ 68 | return phrase in updated_doc["title"] or phrase in updated_doc["selftext"] 69 | 70 | 71 | def read_filter_and_preprocess(file_name: str = "fullrjokes.json", output_name: str = "preprocessed.csv"): 72 | """ 73 | Gathers the data in the data/{file_name} file, preprocesses it to exclude deleted/pic/video/removed/empty 74 | posts and stores them in data/{output_name}. 75 | 76 | It also creates the split for the log-upvote regression problem and stores them in `data/{train|dev|test}.tsv` 77 | 78 | Args: 79 | file_name: the name of the file in the `data` folder that contains the Reddit posts 80 | output_name: the name of the file to output the gathered preprocessed files to 81 | """ 82 | START_OF_2020 = 1577836801 # GMT time, 1/1/2020 83 | list_of_jokes = [] 84 | skipped_deleted = 0 85 | skipped_removed = 0 86 | skipped_format = 0 87 | TEXT_COLS = ["body", "punchline", "joke"] 88 | print("Reading and preprocessing all jokes") 89 | with open(os.path.join(os.path.realpath('..'), "data", file_name), "r") as f: 90 | for index, line in enumerate(f): 91 | updated_doc = json.loads(line) 92 | if updated_doc["created_utc"] >= START_OF_2020: 93 | # NOTE: if analyzing data after 2020, remove this 94 | continue 95 | if updated_doc["title"] == "" and updated_doc["selftext"] == "": 96 | skipped_format += 1 97 | continue 98 | if phrase_in_doc("[deleted]", updated_doc): 99 | skipped_deleted += 1 100 | continue 101 | if phrase_in_doc("[removed]", updated_doc): 102 | skipped_removed += 1 103 | continue 104 | if pd.isnull(updated_doc["score"]) or type(updated_doc["score"]) == str and updated_doc["score"] == "": 105 | skipped_format += 1 106 | continue 107 | if phrase_in_doc("[pic]", updated_doc) or phrase_in_doc("VIDEO", updated_doc): 108 | skipped_format += 1 109 | continue 110 | 111 | joke = updated_doc["title"] + " " + updated_doc["selftext"] 112 | list_of_jokes.append({"joke": joke, "body": updated_doc["selftext"], 113 | "punchline": updated_doc["title"], "score": updated_doc["score"], "date": updated_doc["created_utc"]}) 114 | 115 | df = pd.DataFrame(list_of_jokes) 116 | 117 | # get rid of newlines 118 | df["joke"] = df["joke"].replace(r'\n',' ',regex=True) 119 | df["joke"] = df["joke"].replace(r'\r',' ',regex=True) 120 | assert df["joke"].isnull().any() == False, "there were NaNs in the joke columns, ERROR" 121 | assert df["score"].isnull().any() == False, "there were NaNs in the score columns, ERROR" 122 | df = df.dropna() # drop all jokes that have something funny 123 | print("Writing df with shape", df.shape, "to file path", os.path.join(os.path.realpath('..'), "data", output_name)) 124 | print("SKIPPED: {} were removed, {} deleted, {} format".format(skipped_removed, skipped_deleted, skipped_format)) 125 | df.to_csv(os.path.join(os.path.realpath('..'), "data", output_name), index=None, encoding="UTF-8") 126 | 127 | # format the data for a regression task 128 | # remove data before 1/1/2016 GMT 129 | year2016 = 1451606401 130 | df = df[df["date"] > year2016] 131 | df["score"] = np.log(df["score"]) # use a log scale, make 0's 0 still 132 | df = df.replace([np.inf, -np.inf], np.nan) 133 | df["score"] = df["score"].fillna(0).astype(int) 134 | print("Values are between", df["score"].value_counts()) 135 | print("For log prediction task, {} unique jokes".format(df.shape)) 136 | df = df[["score", "joke"]] 137 | 138 | # save log-distribution plots 139 | plt.tight_layout() 140 | ax = sns.distplot(df["score"], kde=False) 141 | ax.set(xlabel='Log Upvotes', ylabel='Frequency') 142 | plt.savefig(os.path.join("..", "plots", "all", "logdist.png"), bbox_inches = "tight") 143 | plt.savefig(os.path.join("..", "plots", "all", "logdist.pdf"), bbox_inches = "tight") 144 | plt.close() 145 | 146 | train, dev = train_test_split(df, test_size=0.2, random_state=42) 147 | dev, test = train_test_split(dev, test_size=0.5, random_state=42) 148 | print("writing train, dev, and test with shapes", train.shape, dev.shape, test.shape) 149 | train.to_csv(os.path.join("..", "data", "train.tsv"), sep="\t", quoting=csv.QUOTE_NONE, escapechar='\\', index=None, header=None, encoding="UTF-8") 150 | dev.to_csv(os.path.join("..", "data", "dev.tsv"), sep="\t", quoting=csv.QUOTE_NONE, escapechar='\\', index=None, header=None, encoding="UTF-8") 151 | test.to_csv(os.path.join("..", "data", "test.tsv"), sep="\t", quoting=csv.QUOTE_NONE, escapechar='\\', index=None, header=None, encoding="UTF-8") 152 | 153 | 154 | 155 | if __name__ == "__main__": 156 | parser = argparse.ArgumentParser() 157 | parser.add_argument("--update", action="store_true", default=False, help="updates the ids for all files in the given file_path") 158 | parser.add_argument("--preprocess", action="store_true", default=False, help="preprocesses the given id for all docs in the given file_path") 159 | parser.add_argument("--file_name_update", type=str, default="submissions.json", help="updates the ids for all files in the given file_name") 160 | parser.add_argument("--output_name_update", type=str, default="fullrjokes.json", help="updates the ids for all files in the given output_name") 161 | parser.add_argument("--file_name_preprocess", type=str, default="fullrjokes.json", help="preprocesses for all files in the given file_name") 162 | parser.add_argument("--output_name_preprocess", type=str, default="preprocessed.csv", help="preprocesses for all files in the given output_name") 163 | args = parser.parse_args() 164 | 165 | if args.update: 166 | read_and_update_ids(file_name=args.file_name_update, output_name=args.output_name_update) 167 | if args.preprocess: 168 | read_filter_and_preprocess(file_name=args.file_name_preprocess, output_name=args.output_name_preprocess) 169 | if not args.update and not args.preprocess: 170 | raise Exception("Did not give a command to execute...") 171 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | pandas==0.25.3 2 | seaborn==0.9.0 3 | wordcloud==1.5.0 4 | scikit-learn==0.21.3 5 | requests==2.22.0 6 | transformers==2.2.0 7 | torch==1.2.0 8 | praw==6.4.0 9 | pymongo==3.9.0 10 | nltk==3.4.5 11 | vaderSentiment==3.2.1 12 | requests==2.22.0 13 | --------------------------------------------------------------------------------