# Copyright 2021 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import json
import logging
import statistics
import torch
from os import mkdir
from os.path import exists, isdir
from os.path import join as pjoin
import nltk
import numpy as np
import pandas as pd
import plotly
import plotly.express as px
import plotly.figure_factory as ff
import plotly.graph_objects as go
import pyarrow.feather as feather
from datasets import load_from_disk
from nltk.corpus import stopwords
from sklearn.feature_extraction.text import CountVectorizer
from .dataset_utils import (
CNT,
DEDUP_TOT,
EMBEDDING_FIELD,
LENGTH_FIELD,
OUR_LABEL_FIELD,
OUR_TEXT_FIELD,
PROP,
TEXT_NAN_CNT,
TOKENIZED_FIELD,
TXT_LEN,
VOCAB,
WORD,
extract_field,
load_truncated_dataset,
)
from .embeddings import Embeddings
from .npmi import nPMI
from .zipf import Zipf
pd.options.display.float_format = "{:,.3f}".format
logs = logging.getLogger(__name__)
logs.setLevel(logging.WARNING)
logs.propagate = False
if not logs.handlers:
# Logging info to log file
file = logging.FileHandler("./log_files/dataset_statistics.log")
fileformat = logging.Formatter("%(asctime)s:%(message)s")
file.setLevel(logging.INFO)
file.setFormatter(fileformat)
# Logging debug messages to stream
stream = logging.StreamHandler()
streamformat = logging.Formatter("[data_measurements_tool] %(message)s")
stream.setLevel(logging.WARNING)
stream.setFormatter(streamformat)
logs.addHandler(file)
logs.addHandler(stream)
# TODO: Read this in depending on chosen language / expand beyond english
nltk.download("stopwords")
_CLOSED_CLASS = (
stopwords.words("english")
+ [
"t",
"n",
"ll",
"d",
"wasn",
"weren",
"won",
"aren",
"wouldn",
"shouldn",
"didn",
"don",
"hasn",
"ain",
"couldn",
"doesn",
"hadn",
"haven",
"isn",
"mightn",
"mustn",
"needn",
"shan",
"would",
"could",
"dont",
"u",
]
+ [str(i) for i in range(0, 21)]
)
_IDENTITY_TERMS = [
"man",
"woman",
"non-binary",
"gay",
"lesbian",
"queer",
"trans",
"straight",
"cis",
"she",
"her",
"hers",
"he",
"him",
"his",
"they",
"them",
"their",
"theirs",
"himself",
"herself",
]
# treating inf values as NaN as well
pd.set_option("use_inf_as_na", True)
_MIN_VOCAB_COUNT = 10
_TREE_DEPTH = 12
_TREE_MIN_NODES = 250
# as long as we're using sklearn - already pushing the resources
_MAX_CLUSTER_EXAMPLES = 5000
_NUM_VOCAB_BATCHES = 2000
_CVEC = CountVectorizer(token_pattern="(?u)\\b\\w+\\b", lowercase=True)
num_rows = 200000
class DatasetStatisticsCacheClass:
def __init__(
self,
cache_dir,
dset_name,
dset_config,
split_name,
text_field,
label_field,
label_names,
calculation=None,
):
# This is only used for standalone runs for each kind of measurement.
self.calculation = calculation
self.our_text_field = OUR_TEXT_FIELD
self.our_length_field = LENGTH_FIELD
self.our_label_field = OUR_LABEL_FIELD
self.our_tokenized_field = TOKENIZED_FIELD
self.our_embedding_field = EMBEDDING_FIELD
self.cache_dir = cache_dir
### What are we analyzing?
# name of the Hugging Face dataset
self.dset_name = dset_name
# name of the dataset config
self.dset_config = dset_config
# name of the split to analyze
self.split_name = split_name
# which text fields are we analysing?
self.text_field = text_field
# which label fields are we analysing?
self.label_field = label_field
# what are the names of the classes?
self.label_names = label_names
## Hugging Face dataset objects
self.dset = None # original dataset
# HF dataset with all of the self.text_field instances in self.dset
self.text_dset = None
# HF dataset with text embeddings in the same order as self.text_dset
self.embeddings_dset = None
# HF dataset with all of the self.label_field instances in self.dset
self.label_dset = None
## Data frames
# Tokenized text
self.tokenized_df = []
# save sentence length histogram in the class so it doesn't ge re-computed
self.fig_tok_length = None
# Data Frame version of self.label_dset
self.label_df = None
# save label pie chart in the class so it doesn't ge re-computed
self.fig_labels = None
# Vocabulary with word counts in the dataset
self.vocab_counts_df = None
# Vocabulary filtered to remove stopwords
self.vocab_counts_filtered_df = None
## General statistics and duplicates
# Number of NaN values (NOT empty strings)
self.text_nan_count = 0
# Number of text items that appear more than once in the dataset
self.dedup_total = 0
# Duplicated text items along with their number of occurences ("count")
self.text_dup_counts_df = None
self.avg_length = None
self.std_length = None
self.general_stats_dict = None
# clustering text by embeddings
# the hierarchical clustering tree is represented as a list of nodes,
# the first is the root
self.node_list = []
# save tree figure in the class so it doesn't ge re-computed
self.fig_tree = None
# keep Embeddings object around to explore clusters
self.embeddings = None
# nPMI
# Holds a nPMIStatisticsCacheClass object
self.npmi_stats = None
# TODO: Users ideally can type in whatever words they want.
self.termlist = _IDENTITY_TERMS
# termlist terms that are available more than _MIN_VOCAB_COUNT times
self.available_terms = _IDENTITY_TERMS
# TODO: Have lowercase be an option for a user to set.
self.to_lowercase = True
# The minimum amount of times a word should occur to be included in
# word-count-based calculations (currently just relevant to nPMI)
self.min_vocab_count = _MIN_VOCAB_COUNT
# zipf
self.z = None
self.zipf_fig = None
self.cvec = _CVEC
# File definitions
# path to the directory used for caching
if not isinstance(text_field, str):
text_field = "-".join(text_field)
if isinstance(label_field, str):
label_field = label_field
else:
label_field = "-".join(label_field)
self.cache_path = pjoin(
self.cache_dir,
f"{dset_name}_{dset_config}_{split_name}_{text_field}_{label_field}",
)
if not isdir(self.cache_path):
logs.warning("Creating cache directory %s." % self.cache_path)
mkdir(self.cache_path)
self.dset_fid = pjoin(self.cache_path, "base_dset")
self.text_dset_fid = pjoin(self.cache_path, "text_dset")
self.tokenized_df_fid = pjoin(self.cache_path, "tokenized_df.feather")
self.label_dset_fid = pjoin(self.cache_path, "label_dset")
self.vocab_counts_df_fid = pjoin(self.cache_path, "vocab_counts.feather")
self.general_stats_fid = pjoin(self.cache_path, "general_stats.json")
self.text_duplicate_counts_df_fid = pjoin(
self.cache_path, "text_dup_counts_df.feather"
)
self.fig_tok_length_fid = pjoin(self.cache_path, "fig_tok_length.json")
self.fig_labels_fid = pjoin(self.cache_path, "fig_labels.json")
self.node_list_fid = pjoin(self.cache_path, "node_list.th")
self.fig_tree_fid = pjoin(self.cache_path, "fig_tree.json")
self.zipf_fid = pjoin(self.cache_path, "zipf_basic_stats.json")
self.zipf_fig_fid = pjoin(self.cache_path, "zipf_fig.json")
def get_base_dataset(self):
"""Gets a pointer to the truncated base dataset object."""
if not self.dset:
self.dset = load_truncated_dataset(
self.dset_name,
self.dset_config,
self.split_name,
cache_name=self.dset_fid,
use_cache=True,
use_streaming=True,
)
def get_dataset_peek(self):
self.get_base_dataset()
return self.dset[:100]
def load_or_prepare_general_stats(self, use_cache=False):
"""Data structures used in calculating general statistics and duplicates"""
# TODO: These probably don't need to be feather files, could be csv.
# General statistics
if (
use_cache
and exists(self.general_stats_fid)
and exists(self.text_duplicate_counts_df_fid)
):
self.load_general_stats(
self.general_stats_fid, self.text_duplicate_counts_df_fid
)
else:
(
self.text_nan_count,
self.dedup_total,
self.text_dup_counts_df,
) = self.prepare_general_text_stats()
self.general_stats_dict = {
TEXT_NAN_CNT: self.text_nan_count,
DEDUP_TOT: self.dedup_total,
}
write_df(self.text_dup_counts_df, self.text_duplicate_counts_df_fid)
write_json(self.general_stats_dict, self.general_stats_fid)
def load_or_prepare_text_lengths(self, use_cache=False, save=True):
# TODO: Everything here can be read from cache; it's in a transitory
# state atm where just the fig is cached. Clean up.
if use_cache and exists(self.fig_tok_length_fid):
self.fig_tok_length = read_plotly(self.fig_tok_length_fid)
if len(self.tokenized_df) == 0:
self.tokenized_df = self.do_tokenization()
self.tokenized_df[LENGTH_FIELD] = self.tokenized_df[TOKENIZED_FIELD].apply(len)
self.avg_length = round(
sum(self.tokenized_df[self.our_length_field])
/ len(self.tokenized_df[self.our_length_field]),
1,
)
self.std_length = round(
statistics.stdev(self.tokenized_df[self.our_length_field]), 1
)
self.fig_tok_length = make_fig_lengths(self.tokenized_df, self.our_length_field)
if save:
write_plotly(self.fig_tok_length, self.fig_tok_length_fid)
def load_or_prepare_embeddings(self, use_cache=False, save=True):
if use_cache and exists(self.node_list_fid) and exists(self.fig_tree_fid):
self.node_list = torch.load(self.node_list_fid)
self.fig_tree = read_plotly(self.fig_tree_fid)
elif use_cache and exists(self.node_list_fid):
self.node_list = torch.load(self.node_list_fid)
self.fig_tree = make_tree_plot(self.node_list,
self.text_dset)
if save:
write_plotly(self.fig_tree, self.fig_tree_fid)
else:
self.embeddings = Embeddings(self, use_cache=use_cache)
self.embeddings.make_hierarchical_clustering()
self.node_list = self.embeddings.node_list
self.fig_tree = make_tree_plot(self.node_list,
self.text_dset)
if save:
torch.save(self.node_list, self.node_list_fid)
write_plotly(self.fig_tree, self.fig_tree_fid)
# get vocab with word counts
def load_or_prepare_vocab(self, use_cache=True, save=True):
"""
Calculates the vocabulary count from the tokenized text.
The resulting dataframes may be used in nPMI calculations, zipf, etc.
:param use_cache:
:return:
"""
if (
use_cache
and exists(self.vocab_counts_df_fid)
):
logs.info("Reading vocab from cache")
self.load_vocab()
self.vocab_counts_filtered_df = filter_vocab(self.vocab_counts_df)
else:
logs.info("Calculating vocab afresh")
if len(self.tokenized_df) == 0:
self.tokenized_df = self.do_tokenization()
if save:
logs.info("Writing out.")
write_df(self.tokenized_df, self.tokenized_df_fid)
word_count_df = count_vocab_frequencies(self.tokenized_df)
logs.info("Making dfs with proportion.")
self.vocab_counts_df = calc_p_word(word_count_df)
self.vocab_counts_filtered_df = filter_vocab(self.vocab_counts_df)
if save:
logs.info("Writing out.")
write_df(self.vocab_counts_df, self.vocab_counts_df_fid)
logs.info("unfiltered vocab")
logs.info(self.vocab_counts_df)
logs.info("filtered vocab")
logs.info(self.vocab_counts_filtered_df)
def load_or_prepare_npmi_terms(self, use_cache=False):
self.npmi_stats = nPMIStatisticsCacheClass(self, use_cache=use_cache)
self.npmi_stats.load_or_prepare_npmi_terms()
def load_or_prepare_zipf(self, use_cache=False, save=True):
# TODO: Current UI only uses the fig, meaning the self.z here is irrelevant
# when only reading from cache. Either the UI should use it, or it should
# be removed when reading in cache
if use_cache and exists(self.zipf_fig_fid) and exists(self.zipf_fid):
with open(self.zipf_fid, "r") as f:
zipf_dict = json.load(f)
self.z = Zipf()
self.z.load(zipf_dict)
self.zipf_fig = read_plotly(self.zipf_fig_fid)
elif use_cache and exists(self.zipf_fid):
# TODO: Read zipf data so that the vocab is there.
with open(self.zipf_fid, "r") as f:
zipf_dict = json.load(f)
self.z = Zipf()
self.z.load(zipf_dict)
self.zipf_fig = make_zipf_fig(self.vocab_counts_df, self.z)
if save:
write_plotly(self.zipf_fig, self.zipf_fig_fid)
else:
self.z = Zipf(self.vocab_counts_df)
self.zipf_fig = make_zipf_fig(self.vocab_counts_df, self.z)
if save:
write_zipf_data(self.z, self.zipf_fid)
write_plotly(self.zipf_fig, self.zipf_fig_fid)
def prepare_general_text_stats(self):
text_nan_count = int(self.tokenized_df.isnull().sum().sum())
dup_df = self.tokenized_df[self.tokenized_df.duplicated([self.our_text_field])]
dedup_df = pd.DataFrame(
dup_df.pivot_table(
columns=[self.our_text_field], aggfunc="size"
).sort_values(ascending=False),
columns=[CNT],
)
dedup_df.index = dedup_df.index.map(str)
dedup_df[OUR_TEXT_FIELD] = dedup_df.index
dedup_total = sum(dedup_df[CNT])
return text_nan_count, dedup_total, dedup_df
def load_general_stats(self, general_stats_fid, text_duplicate_counts_df_fid):
general_stats = json.load(open(general_stats_fid, encoding="utf-8"))
self.text_nan_count = general_stats[TEXT_NAN_CNT]
self.dedup_total = general_stats[DEDUP_TOT]
with open(text_duplicate_counts_df_fid, "rb") as f:
self.text_dup_counts_df = feather.read_feather(f)
def load_or_prepare_dataset(self, use_cache=True, save=True):
"""
Prepares the HF datasets and data frames containing the untokenized and
tokenized text as well as the label values.
self.tokenized_df is used further for calculating text lengths,
word counts, etc.
Args:
use_cache: Used stored data if there; otherwise calculate afresh
save: Store the calculated data to disk.
Returns:
"""
self.load_or_prepare_text_dset(save, use_cache)
self.load_or_prepare_tokenized_df(save, use_cache)
def load_or_prepare_tokenized_df(self, save, use_cache):
if (use_cache and exists(self.tokenized_df_fid)):
self.tokenized_df = feather.read_feather(self.tokenized_df_fid)
else:
# tokenize all text instances
self.tokenized_df = self.do_tokenization()
if save:
# save tokenized text
write_df(self.tokenized_df, self.tokenized_df_fid)
def load_or_prepare_text_dset(self, save, use_cache):
if (use_cache and exists(self.text_dset_fid)):
# load extracted text
self.text_dset = load_from_disk(self.text_dset_fid)
logs.warning("Loaded dataset from disk")
logs.info(self.text_dset)
# ...Or load it from the server and store it anew
else:
self.get_base_dataset()
# extract all text instances
self.text_dset = self.dset.map(
lambda examples: extract_field(
examples, self.text_field, OUR_TEXT_FIELD
),
batched=True,
remove_columns=list(self.dset.features),
)
if save:
# save extracted text instances
logs.warning("Saving dataset to disk")
self.text_dset.save_to_disk(self.text_dset_fid)
def do_tokenization(self):
"""
Tokenizes the dataset
:return:
"""
sent_tokenizer = self.cvec.build_tokenizer()
def tokenize_batch(examples):
# TODO: lowercase should be an option
res = {
TOKENIZED_FIELD: [
tuple(sent_tokenizer(text.lower()))
for text in examples[OUR_TEXT_FIELD]
]
}
res[LENGTH_FIELD] = [len(tok_text) for tok_text in res[TOKENIZED_FIELD]]
return res
tokenized_dset = self.text_dset.map(
tokenize_batch,
batched=True,
# remove_columns=[OUR_TEXT_FIELD], keep around to print
)
tokenized_df = pd.DataFrame(tokenized_dset)
return tokenized_df
def set_label_field(self, label_field="label"):
"""
Setter for label_field. Used in the CLI when a user asks for information
about labels, but does not specify the field;
'label' is assumed as a default.
"""
self.label_field = label_field
def load_or_prepare_labels(self, use_cache=False, save=True):
# TODO: This is in a transitory state for creating fig cache.
# Clean up to be caching and reading everything correctly.
"""
Extracts labels from the Dataset
:param use_cache:
:return:
"""
# extracted labels
if len(self.label_field) > 0:
if use_cache and exists(self.fig_labels_fid):
self.fig_labels = read_plotly(self.fig_labels_fid)
elif use_cache and exists(self.label_dset_fid):
# load extracted labels
self.label_dset = load_from_disk(self.label_dset_fid)
self.label_df = self.label_dset.to_pandas()
self.fig_labels = make_fig_labels(
self.label_df, self.label_names, OUR_LABEL_FIELD
)
if save:
write_plotly(self.fig_labels, self.fig_labels_fid)
else:
self.get_base_dataset()
self.label_dset = self.dset.map(
lambda examples: extract_field(
examples, self.label_field, OUR_LABEL_FIELD
),
batched=True,
remove_columns=list(self.dset.features),
)
self.label_df = self.label_dset.to_pandas()
self.fig_labels = make_fig_labels(
self.label_df, self.label_names, OUR_LABEL_FIELD
)
if save:
# save extracted label instances
self.label_dset.save_to_disk(self.label_dset_fid)
write_plotly(self.fig_labels, self.fig_labels_fid)
def load_vocab(self):
with open(self.vocab_counts_df_fid, "rb") as f:
self.vocab_counts_df = feather.read_feather(f)
# Handling for changes in how the index is saved.
self.vocab_counts_df = self._set_idx_col_names(self.vocab_counts_df)
def _set_idx_col_names(self, input_vocab_df):
if input_vocab_df.index.name != VOCAB and VOCAB in input_vocab_df.columns:
input_vocab_df = input_vocab_df.set_index([VOCAB])
input_vocab_df[VOCAB] = input_vocab_df.index
return input_vocab_df
class nPMIStatisticsCacheClass:
""" "Class to interface between the app and the nPMI class
by calling the nPMI class with the user's selections."""
def __init__(self, dataset_stats, use_cache=False):
self.dstats = dataset_stats
self.pmi_cache_path = pjoin(self.dstats.cache_path, "pmi_files")
if not isdir(self.pmi_cache_path):
logs.warning("Creating pmi cache directory %s." % self.pmi_cache_path)
# We need to preprocess everything.
mkdir(self.pmi_cache_path)
self.joint_npmi_df_dict = {}
self.termlist = self.dstats.termlist
logs.info(self.termlist)
self.use_cache = use_cache
# TODO: Let users specify
self.open_class_only = True
self.min_vocab_count = self.dstats.min_vocab_count
self.subgroup_files = {}
self.npmi_terms_fid = pjoin(self.dstats.cache_path, "npmi_terms.json")
self.available_terms = self.dstats.available_terms
logs.info(self.available_terms)
def load_or_prepare_npmi_terms(self, use_cache=False):
"""
Figures out what identity terms the user can select, based on whether
they occur more than self.min_vocab_count times
:param use_cache:
:return: Identity terms occurring at least self.min_vocab_count times.
"""
# TODO: Add the user's ability to select subgroups.
# TODO: Make min_vocab_count here value selectable by the user.
if (
use_cache
and exists(self.npmi_terms_fid)
and json.load(open(self.npmi_terms_fid))["available terms"] != []
):
available_terms = json.load(open(self.npmi_terms_fid))["available terms"]
else:
true_false = [
term in self.dstats.vocab_counts_df.index for term in self.termlist
]
word_list_tmp = [x for x, y in zip(self.termlist, true_false) if y]
true_false_counts = [
self.dstats.vocab_counts_df.loc[word, CNT] >= self.min_vocab_count
for word in word_list_tmp
]
available_terms = [
word for word, y in zip(word_list_tmp, true_false_counts) if y
]
logs.info(available_terms)
with open(self.npmi_terms_fid, "w+") as f:
json.dump({"available terms": available_terms}, f)
self.available_terms = available_terms
return available_terms
def load_or_prepare_joint_npmi(self, subgroup_pair, use_cache=True):
"""
Run on-the fly, while the app is already open,
as it depends on the subgroup terms that the user chooses
:param subgroup_pair:
:return:
"""
# Canonical ordering for subgroup_list
subgroup_pair = sorted(subgroup_pair)
subgroups_str = "-".join(subgroup_pair)
if not isdir(self.pmi_cache_path):
logs.warning("Creating cache")
# We need to preprocess everything.
# This should eventually all go into a prepare_dataset CLI
mkdir(self.pmi_cache_path)
joint_npmi_fid = pjoin(self.pmi_cache_path, subgroups_str + "_npmi.csv")
subgroup_files = define_subgroup_files(subgroup_pair, self.pmi_cache_path)
# Defines the filenames for the cache files from the selected subgroups.
# Get as much precomputed data as we can.
if use_cache and exists(joint_npmi_fid):
# When everything is already computed for the selected subgroups.
logs.info("Loading cached joint npmi")
joint_npmi_df = self.load_joint_npmi_df(joint_npmi_fid)
# When maybe some things have been computed for the selected subgroups.
else:
logs.info("Preparing new joint npmi")
joint_npmi_df, subgroup_dict = self.prepare_joint_npmi_df(
subgroup_pair, subgroup_files
)
# Cache new results
logs.info("Writing out.")
for subgroup in subgroup_pair:
write_subgroup_npmi_data(subgroup, subgroup_dict, subgroup_files)
with open(joint_npmi_fid, "w+") as f:
joint_npmi_df.to_csv(f)
logs.info("The joint npmi df is")
logs.info(joint_npmi_df)
return joint_npmi_df
def load_joint_npmi_df(self, joint_npmi_fid):
"""
Reads in a saved dataframe with all of the paired results.
:param joint_npmi_fid:
:return: paired results
"""
with open(joint_npmi_fid, "rb") as f:
joint_npmi_df = pd.read_csv(f)
joint_npmi_df = self._set_idx_cols_from_cache(joint_npmi_df)
return joint_npmi_df.dropna()
def prepare_joint_npmi_df(self, subgroup_pair, subgroup_files):
"""
Computs the npmi bias based on the given subgroups.
Handles cases where some of the selected subgroups have cached nPMI
computations, but other's don't, computing everything afresh if there
are not cached files.
:param subgroup_pair:
:return: Dataframe with nPMI for the words, nPMI bias between the words.
"""
subgroup_dict = {}
# When npmi is computed for some (but not all) of subgroup_list
for subgroup in subgroup_pair:
logs.info("Load or failing...")
# When subgroup npmi has been computed in a prior session.
cached_results = self.load_or_fail_cached_npmi_scores(
subgroup, subgroup_files[subgroup]
)
# If the function did not return False and we did find it, use.
if cached_results:
# FYI: subgroup_cooc_df, subgroup_pmi_df, subgroup_npmi_df = cached_results
# Holds the previous sessions' data for use in this session.
subgroup_dict[subgroup] = cached_results
logs.info("Calculating for subgroup list")
joint_npmi_df, subgroup_dict = self.do_npmi(subgroup_pair, subgroup_dict)
return joint_npmi_df.dropna(), subgroup_dict
# TODO: Update pairwise assumption
def do_npmi(self, subgroup_pair, subgroup_dict):
"""
Calculates nPMI for given identity terms and the nPMI bias between.
:param subgroup_pair: List of identity terms to calculate the bias for
:return: Subset of data for the UI
:return: Selected identity term's co-occurrence counts with
other words, pmi per word, and nPMI per word.
"""
logs.info("Initializing npmi class")
npmi_obj = self.set_npmi_obj()
# Canonical ordering used
subgroup_pair = tuple(sorted(subgroup_pair))
# Calculating nPMI statistics
for subgroup in subgroup_pair:
# If the subgroup data is already computed, grab it.
# TODO: Should we set idx and column names similarly to how we set them for cached files?
if subgroup not in subgroup_dict:
logs.info("Calculating statistics for %s" % subgroup)
vocab_cooc_df, pmi_df, npmi_df = npmi_obj.calc_metrics(subgroup)
# Store the nPMI information for the current subgroups
subgroup_dict[subgroup] = (vocab_cooc_df, pmi_df, npmi_df)
# Pair the subgroups together, indexed by all words that
# co-occur between them.
logs.info("Computing pairwise npmi bias")
paired_results = npmi_obj.calc_paired_metrics(subgroup_pair, subgroup_dict)
UI_results = make_npmi_fig(paired_results, subgroup_pair)
return UI_results, subgroup_dict
def set_npmi_obj(self):
"""
Initializes the nPMI class with the given words and tokenized sentences.
:return:
"""
npmi_obj = nPMI(self.dstats.vocab_counts_df, self.dstats.tokenized_df)
return npmi_obj
def load_or_fail_cached_npmi_scores(self, subgroup, subgroup_fids):
"""
Reads cached scores from the specified subgroup files
:param subgroup: string of the selected identity term
:return:
"""
# TODO: Ordering of npmi, pmi, vocab triple should be consistent
subgroup_npmi_fid, subgroup_pmi_fid, subgroup_cooc_fid = subgroup_fids
if (
exists(subgroup_npmi_fid)
and exists(subgroup_pmi_fid)
and exists(subgroup_cooc_fid)
):
logs.info("Reading in pmi data....")
with open(subgroup_cooc_fid, "rb") as f:
subgroup_cooc_df = pd.read_csv(f)
logs.info("pmi")
with open(subgroup_pmi_fid, "rb") as f:
subgroup_pmi_df = pd.read_csv(f)
logs.info("npmi")
with open(subgroup_npmi_fid, "rb") as f:
subgroup_npmi_df = pd.read_csv(f)
subgroup_cooc_df = self._set_idx_cols_from_cache(
subgroup_cooc_df, subgroup, "count"
)
subgroup_pmi_df = self._set_idx_cols_from_cache(
subgroup_pmi_df, subgroup, "pmi"
)
subgroup_npmi_df = self._set_idx_cols_from_cache(
subgroup_npmi_df, subgroup, "npmi"
)
return subgroup_cooc_df, subgroup_pmi_df, subgroup_npmi_df
return False
def _set_idx_cols_from_cache(self, csv_df, subgroup=None, calc_str=None):
"""
Helps make sure all of the read-in files can be accessed within code
via standardized indices and column names.
:param csv_df:
:param subgroup:
:param calc_str:
:return:
"""
# The csv saves with this column instead of the index, so that's weird.
if "Unnamed: 0" in csv_df.columns:
csv_df = csv_df.set_index("Unnamed: 0")
csv_df.index.name = WORD
elif WORD in csv_df.columns:
csv_df = csv_df.set_index(WORD)
csv_df.index.name = WORD
elif VOCAB in csv_df.columns:
csv_df = csv_df.set_index(VOCAB)
csv_df.index.name = WORD
if subgroup and calc_str:
csv_df.columns = [subgroup + "-" + calc_str]
elif subgroup:
csv_df.columns = [subgroup]
elif calc_str:
csv_df.columns = [calc_str]
return csv_df
def get_available_terms(self, use_cache=False):
return self.load_or_prepare_npmi_terms(use_cache=use_cache)
def dummy(doc):
return doc
def count_vocab_frequencies(tokenized_df):
"""
Based on an input pandas DataFrame with a 'text' column,
this function will count the occurrences of all words.
:return: [num_words x num_sentences] DataFrame with the rows corresponding to the
different vocabulary words and the column to the presence (0 or 1) of that word.
"""
cvec = CountVectorizer(
tokenizer=dummy,
preprocessor=dummy,
)
# We do this to calculate per-word statistics
# Fast calculation of single word counts
logs.info("Fitting dummy tokenization to make matrix using the previous tokenization")
cvec.fit(tokenized_df[TOKENIZED_FIELD])
document_matrix = cvec.transform(tokenized_df[TOKENIZED_FIELD])
batches = np.linspace(0, tokenized_df.shape[0], _NUM_VOCAB_BATCHES).astype(int)
i = 0
tf = []
while i < len(batches) - 1:
logs.info("%s of %s vocab batches" % (str(i), str(len(batches))))
batch_result = np.sum(
document_matrix[batches[i] : batches[i + 1]].toarray(), axis=0
)
tf.append(batch_result)
i += 1
word_count_df = pd.DataFrame(
[np.sum(tf, axis=0)], columns=cvec.get_feature_names()
).transpose()
# Now organize everything into the dataframes
word_count_df.columns = [CNT]
word_count_df.index.name = WORD
return word_count_df
def calc_p_word(word_count_df):
# p(word)
word_count_df[PROP] = word_count_df[CNT] / float(sum(word_count_df[CNT]))
vocab_counts_df = pd.DataFrame(word_count_df.sort_values(by=CNT, ascending=False))
vocab_counts_df[VOCAB] = vocab_counts_df.index
return vocab_counts_df
def filter_vocab(vocab_counts_df):
# TODO: Add warnings (which words are missing) to log file?
filtered_vocab_counts_df = vocab_counts_df.drop(_CLOSED_CLASS,
errors="ignore")
filtered_count = filtered_vocab_counts_df[CNT]
filtered_count_denom = float(sum(filtered_vocab_counts_df[CNT]))
filtered_vocab_counts_df[PROP] = filtered_count / filtered_count_denom
return filtered_vocab_counts_df
## Figures ##
def write_plotly(fig, fid):
write_json(plotly.io.to_json(fig), fid)
def read_plotly(fid):
fig = plotly.io.from_json(json.load(open(fid, encoding="utf-8")))
return fig
def make_fig_lengths(tokenized_df, length_field):
fig_tok_length = px.histogram(
tokenized_df, x=length_field, marginal="rug", hover_data=[length_field]
)
return fig_tok_length
def make_fig_labels(label_df, label_names, label_field):
labels = label_df[label_field].unique()
label_sums = [len(label_df[label_df[label_field] == label]) for label in labels]
fig_labels = px.pie(label_df, values=label_sums, names=label_names)
return fig_labels
def make_zipf_fig_ranked_word_list(vocab_df, unique_counts, unique_ranks):
ranked_words = {}
for count, rank in zip(unique_counts, unique_ranks):
vocab_df[vocab_df[CNT] == count]["rank"] = rank
ranked_words[rank] = ",".join(
vocab_df[vocab_df[CNT] == count].index.astype(str)
) # Use the hovertext kw argument for hover text
ranked_words_list = [wrds for rank, wrds in sorted(ranked_words.items())]
return ranked_words_list
def make_npmi_fig(paired_results, subgroup_pair):
subgroup1, subgroup2 = subgroup_pair
UI_results = pd.DataFrame()
if "npmi-bias" in paired_results:
UI_results["npmi-bias"] = paired_results["npmi-bias"].astype(float)
UI_results[subgroup1 + "-npmi"] = paired_results["npmi"][
subgroup1 + "-npmi"
].astype(float)
UI_results[subgroup1 + "-count"] = paired_results["count"][
subgroup1 + "-count"
].astype(int)
if subgroup1 != subgroup2:
UI_results[subgroup2 + "-npmi"] = paired_results["npmi"][
subgroup2 + "-npmi"
].astype(float)
UI_results[subgroup2 + "-count"] = paired_results["count"][
subgroup2 + "-count"
].astype(int)
return UI_results.sort_values(by="npmi-bias", ascending=True)
def make_zipf_fig(vocab_counts_df, z):
zipf_counts = z.calc_zipf_counts(vocab_counts_df)
unique_counts = z.uniq_counts
unique_ranks = z.uniq_ranks
ranked_words_list = make_zipf_fig_ranked_word_list(
vocab_counts_df, unique_counts, unique_ranks
)
zmin = z.get_xmin()
logs.info("zipf counts is")
logs.info(zipf_counts)
layout = go.Layout(xaxis=dict(range=[0, 100]))
fig = go.Figure(
data=[
go.Bar(
x=z.uniq_ranks,
y=z.uniq_counts,
hovertext=ranked_words_list,
name="Word Rank Frequency",
)
],
layout=layout,
)
fig.add_trace(
go.Scatter(
x=z.uniq_ranks[zmin : len(z.uniq_ranks)],
y=zipf_counts[zmin : len(z.uniq_ranks)],
hovertext=ranked_words_list[zmin : len(z.uniq_ranks)],
line=go.scatter.Line(color="crimson", width=3),
name="Zipf Predicted Frequency",
)
)
# Customize aspect
# fig.update_traces(marker_color='limegreen',
# marker_line_width=1.5, opacity=0.6)
fig.update_layout(title_text="Word Counts, Observed and Predicted by Zipf")
fig.update_layout(xaxis_title="Word Rank")
fig.update_layout(yaxis_title="Frequency")
fig.update_layout(legend=dict(yanchor="top", y=0.99, xanchor="left", x=0.10))
return fig
def make_tree_plot(node_list, text_dset):
nid_map = dict([(node["nid"], nid) for nid, node in enumerate(node_list)])
for nid, node in enumerate(node_list):
node["label"] = node.get(
"label",
f"{nid:2d} - {node['weight']:5d} items
"
+ "
".join(
[
"> " + txt[:64] + ("..." if len(txt) >= 63 else "")
for txt in list(
set(text_dset.select(node["example_ids"])[OUR_TEXT_FIELD])
)[:5]
]
),
)
# make plot nodes
# TODO: something more efficient than set to remove duplicates
labels = [node["label"] for node in node_list]
root = node_list[0]
root["X"] = 0
root["Y"] = 0
def rec_make_coordinates(node):
total_weight = 0
add_weight = len(node["example_ids"]) - sum(
[child["weight"] for child in node["children"]]
)
for child in node["children"]:
child["X"] = node["X"] + total_weight
child["Y"] = node["Y"] - 1
total_weight += child["weight"] + add_weight / len(node["children"])
rec_make_coordinates(child)
rec_make_coordinates(root)
E = [] # list of edges
Xn = []
Yn = []
Xe = []
Ye = []
for nid, node in enumerate(node_list):
Xn += [node["X"]]
Yn += [node["Y"]]
for child in node["children"]:
E += [(nid, nid_map[child["nid"]])]
Xe += [node["X"], child["X"], None]
Ye += [node["Y"], child["Y"], None]
# make figure
fig = go.Figure()
fig.add_trace(
go.Scatter(
x=Xe,
y=Ye,
mode="lines",
line=dict(color="rgb(210,210,210)", width=1),
hoverinfo="none",
)
)
fig.add_trace(
go.Scatter(
x=Xn,
y=Yn,
mode="markers",
name="nodes",
marker=dict(
symbol="circle-dot",
size=18,
color="#6175c1",
line=dict(color="rgb(50,50,50)", width=1)
# '#DB4551',
),
text=labels,
hoverinfo="text",
opacity=0.8,
)
)
return fig
## Input/Output ###
def define_subgroup_files(subgroup_list, pmi_cache_path):
"""
Sets the file ids for the input identity terms
:param subgroup_list: List of identity terms
:return:
"""
subgroup_files = {}
for subgroup in subgroup_list:
# TODO: Should the pmi, npmi, and count just be one file?
subgroup_npmi_fid = pjoin(pmi_cache_path, subgroup + "_npmi.csv")
subgroup_pmi_fid = pjoin(pmi_cache_path, subgroup + "_pmi.csv")
subgroup_cooc_fid = pjoin(pmi_cache_path, subgroup + "_vocab_cooc.csv")
subgroup_files[subgroup] = (
subgroup_npmi_fid,
subgroup_pmi_fid,
subgroup_cooc_fid,
)
return subgroup_files
## Input/Output ##
def intersect_dfs(df_dict):
started = 0
new_df = None
for key, df in df_dict.items():
if df is None:
continue
for key2, df2 in df_dict.items():
if df2 is None:
continue
if key == key2:
continue
if started:
new_df = new_df.join(df2, how="inner", lsuffix="1", rsuffix="2")
else:
new_df = df.join(df2, how="inner", lsuffix="1", rsuffix="2")
started = 1
return new_df.copy()
def write_df(df, df_fid):
feather.write_feather(df, df_fid)
def write_json(json_dict, json_fid):
with open(json_fid, "w", encoding="utf-8") as f:
json.dump(json_dict, f)
def write_subgroup_npmi_data(subgroup, subgroup_dict, subgroup_files):
"""
Saves the calculated nPMI statistics to their output files.
Includes the npmi scores for each identity term, the pmi scores, and the
co-occurrence counts of the identity term with all the other words
:param subgroup: Identity term
:return:
"""
subgroup_fids = subgroup_files[subgroup]
subgroup_npmi_fid, subgroup_pmi_fid, subgroup_cooc_fid = subgroup_fids
subgroup_dfs = subgroup_dict[subgroup]
subgroup_cooc_df, subgroup_pmi_df, subgroup_npmi_df = subgroup_dfs
with open(subgroup_npmi_fid, "w+") as f:
subgroup_npmi_df.to_csv(f)
with open(subgroup_pmi_fid, "w+") as f:
subgroup_pmi_df.to_csv(f)
with open(subgroup_cooc_fid, "w+") as f:
subgroup_cooc_df.to_csv(f)
def write_zipf_data(z, zipf_fid):
zipf_dict = {}
zipf_dict["xmin"] = int(z.xmin)
zipf_dict["xmax"] = int(z.xmax)
zipf_dict["alpha"] = float(z.alpha)
zipf_dict["ks_distance"] = float(z.distance)
zipf_dict["p-value"] = float(z.ks_test.pvalue)
zipf_dict["uniq_counts"] = [int(count) for count in z.uniq_counts]
zipf_dict["uniq_ranks"] = [int(rank) for rank in z.uniq_ranks]
with open(zipf_fid, "w+", encoding="utf-8") as f:
json.dump(zipf_dict, f)