# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # **Identify_Incorrect_labels.ipynb**
#
# Use text extensions for pandas cleaning library functionality to rapidly and easily identify potentially incorrect labels in the english language EWT corpus

# +
import numpy as np
import pandas as pd

try:
    import text_extensions_for_pandas as tp
except ModuleNotFoundError as e:
    raise Exception("text_extensions_for_pandas package not found on the Jupyter "
                    "kernel's path. Please either run:\n"
                    "   ln -s ../../text_extensions_for_pandas .\n"
                    "from the directory containing this notebook, or use a Python "
                    "environment on which you have used `pip` to install the package.")
from text_extensions_for_pandas import cleaning
import matplotlib.pyplot as plt


# init some globals we need later on 
BASE_DIR = 'Corpus_docs/'
ewt_base_url = "https://raw.githubusercontent.com/UniversalDependencies/UD_English-EWT/master/en_ewt-ud-"
# -

#import (and possibly download) our corpus: 
fold_paths = {"test":  tp.io.conll.maybe_download_dataset_data(BASE_DIR, ewt_base_url + "test.conllu"),
              "dev":   tp.io.conll.maybe_download_dataset_data(BASE_DIR, ewt_base_url + "dev.conllu"),
              "train": tp.io.conll.maybe_download_dataset_data(BASE_DIR, ewt_base_url + "train.conllu")}
# then use text-extensions conll module import functions to bring them into our notebook
fold_docs = {}
for fold,fold_path in fold_paths.items(): 
    fold_docs[fold] = tp.io.conll.conll_u_to_dataframes(fold_path)
    print(f"converted fold: '{fold}' to list of {len(fold_docs[fold])} dataframes")
fold_docs['dev'][0].head(7)[['span','lemma','upostag','xpostag','features','head','deprel','deps','misc','sentence','line_num']]

# ## Preprocessing
# Lets suppose we wish to verify parts of speech using the xpostag label
#
# so we will need to preprocess using that column. 
#

corpus_df, postag_dtype, postag_list, postag_dict = cleaning.preprocess_documents(fold_docs, label_col='upostag',iob_format=False,carry_cols=['lemma'])

corpus_df.head()

## Train an ensemble of models:
# start by prepping features
kernel_sizes = [16, 32, 64,128,768]
seeds = [1, 2, 3,4]
train_fold = corpus_df[corpus_df.fold == 'train']
train_fold.head()

# now use cleaning util to create set of models 
model_ensemble = cleaning.train_model_ensemble(train_fold,'upostag_id',model_sizes=kernel_sizes, model_seeds=seeds,max_iters=10000)


# +
## now run infererence using those models
dev_and_test_df = corpus_df[corpus_df.fold != 'train']
predictions = {name: cleaning.infer_and_extract_raw_entites(dev_and_test_df,postag_list,model,keep_cols=['upostag','upostag_id','raw_span']) for name,model in model_ensemble.items()}

a_model = list(model_ensemble.keys())[0]
predictions[a_model].head()
# -

cleaning.create_f1_score_report(predictions[a_model],'upostag','predicted_class').transpose()

in_gold, not_in_gold = cleaning.flag_suspicious_labels(predictions,'upostag','predicted_class',label_name='postag')
in_gold.head(10)

not_in_gold.head(10)

# # Try alternate 'cleanlab like' heuristic for flagging incorrect entities: 
#
# We will attempt a few alternate error flagging heuristics: 
#
# 1. sum self confidence in each of the chosen classes to act as a 'confidence weighted count' 
# 1. sum self confidence over all models (chosen or not) 
# 1. single model constant threshold flagging 
# 1. single model flagging with dynamic thresholding (cleanlab)

# try alternate flagging method that uses each label's probability as a weight
gold_df = predictions[a_model][['fold','doc_num','raw_span_id','upostag','raw_span']].copy()
gold_df['in_gold'] = True
gold_df['models'] = 'GOLD'
gold_df['weighted_count'] =0
gold_df['corpus_tag'] = gold_df['upostag'].copy()
# other models 
feats_dfs = [gold_df]
for model_name, model_vals in predictions.items(): 
    df = model_vals[['fold','doc_num','raw_span_id','raw_span','predicted_class','upostag']].copy()
    df.rename(columns={'upostag':'corpus_tag'},inplace=True)
    df.rename(columns={'predicted_class':'upostag'},inplace=True)
    df['in_gold']=False
    df['models'] = model_name
    df['weighted_count'] = np.stack(model_vals.raw_output.to_numpy()).max(axis=1)
    feats_dfs.append(df)
all_feats = pd.concat(feats_dfs)

# +
aggby = {col: 'first' for col in ['raw_span','corpus_tag'] if col in all_feats.columns}
aggby.update({'in_gold':'any','weighted_count':'sum','models':'count'})

res = all_feats.groupby(['fold','doc_num','raw_span_id','upostag']).agg(aggby).reset_index().sort_values(['weighted_count'],ascending=False, kind="mergesort")
res[~res.in_gold].head(10)
# -

res.shape

res[res.in_gold].sort_values(['weighted_count'],kind='mergesort',ascending=True).head(10)

# ### Attempt 'cleanlab-like' heuristic version 2'
# Here we attempt to sum raw outputs regardless of whether or not they were the most likely of each model
#
# hopefully this gives us a higher resolution than the last heuristic 

# +
#get prediction matrices - use numpy for speed
arrs = [np.stack(prediction.raw_output.to_numpy()) for prediction in predictions.values()]
sum_arr = np.stack(arrs).sum(axis=0)

combined_probs = pd.DataFrame({'predicted_class':None,'pred_id':np.argmax(sum_arr,axis=1),'raw_outputs':tp.TensorArray(sum_arr),'weighted_sum':np.max(sum_arr,axis=1)})
combined_probs['predicted_class']= combined_probs['pred_id'].apply(lambda x: postag_list[x]).astype(postag_dtype)
combined_probs
# -

# reuse the gold df from earlier and join our lines 
combined_df = combined_probs.join(gold_df[['fold','doc_num','raw_span_id','raw_span','corpus_tag']].copy(),how='outer')
not_in_gold = combined_df[combined_df['corpus_tag'] != combined_df['predicted_class']].copy()
not_in_gold.sort_values('weighted_sum',ascending=False,inplace=True)
not_in_gold.head(20)[['predicted_class','weighted_sum','corpus_tag','raw_span','fold','doc_num','raw_span_id']]

in_gold = combined_df[['corpus_tag','raw_outputs','raw_span','fold','doc_num','raw_span_id']].copy()
in_gold['weighted_score'] = [in_gold.loc[i,'raw_outputs'].to_numpy()[postag_dict[in_gold.loc[i,'corpus_tag']]] for i in in_gold.index]
in_gold.sort_values('weighted_score',ascending=True,inplace=True)
in_gold.head(10)[['weighted_score','corpus_tag','raw_span','fold','doc_num','raw_span_id']]

# ### Attempt actual cleanlab heuristics (from paper)
#
# Create a confident joint using methodology from the cleanlab paper,
# then, use three different error flagging approaches to flag suspicious entities
#
# 2. **Confident Joint:** use off diagonals from the confusion matrix 
# 3. **Prune by Class:** use n*(row_sum(Q) ) and choose the n lowest self_confidence values from each labeled class
# 4. **Prune by Noise Rate:** choose n*Q<sub>i,j</sub> for each bin: (y&#771; =i, y<super> *</super>=j)

# +
#calculate confident joint 
model_predictions = predictions[list(predictions.keys())[0]].copy()
print(f"using model: {list(predictions.keys())[0]}")
# first extract our model values into numpy
model_raw_preds = np.stack(model_predictions.raw_output.to_numpy())
# just for good measure jot down our predicted measures
model_predictions['pred_confidence'] = model_raw_preds.max(axis=1)
corpus_labels = model_predictions['upostag_id'].to_numpy()
gold_conf = model_raw_preds[range(model_raw_preds.shape[0]),model_predictions.upostag_id.to_numpy()]


#first calculate thresholds
threshold_mask=np.zeros_like(model_raw_preds)
threshold_mask[range(len(corpus_labels)),corpus_labels] = 1
thresholds = np.mean(model_raw_preds*threshold_mask,axis=0)
# model
mask = (model_raw_preds - thresholds)>=0 # p(y=j) > t_j
preds_masked = np.where(mask,model_raw_preds,0) # j = argmax(p(y=l,x)|p(y=l>t_l))
model_predictions['CL_pred'] = preds_masked.argmax(axis=1)
model_predictions['cl_confidence'] =preds_masked.max(axis=1)
model_predictions['normalized_margin'] = gold_conf - model_predictions.cl_confidence
model_predictions['self_confidence'] = gold_conf
# just out of interest show the places where the cleanlab flagged models differ from 
# the raw model predictions. this isn't very common; only 11 spots over the whole test fold
model_predictions[model_predictions.CL_pred != model_predictions.predicted_id]
# -

# display results from Confident Joint off diagonals
CL_method_2 = model_predictions[model_predictions.upostag_id != model_predictions.CL_pred].sort_values('normalized_margin')[['fold','doc_num','upostag','raw_span','CL_pred','normalized_margin']]
CL_method_2.head(20)

# calculate the C matrix
counts_per = model_predictions.groupby(['upostag_id','CL_pred']).agg({'fold':'count'}).reset_index()
C_mat = counts_per.pivot(index='upostag_id',columns='CL_pred',values='fold').fillna(0)#.astype('int')
C_mat

# calculate Q matrix 
c_mat = C_mat.to_numpy()
# calculate row by row correction to match observed
c_row_sums = c_mat.sum(axis=0)
obseved_sums =model_predictions.groupby('upostag_id').agg({'fold':'count'}).fold.to_numpy()
row_correction = (c_row_sums/obseved_sums)
q_mat = c_mat*row_correction[:,np.newaxis]
#normalize so that all values sum to one 
q_mat /= np.sum(q_mat)
# number of elements in each class to prune. 
q_n_mat = q_mat*model_predictions.shape[0]
plt.imshow(np.where(np.eye(q_n_mat.shape[0]),300,q_n_mat))
pd.DataFrame(q_n_mat).astype(int)


# calculate outputs from prune by noise rate
model_predictions.sort_values('self_confidence', ascending=True,inplace = True)
flagged_list = []
# TODO: there's probably a groupby way to do this thats better than O(M^2 N)
# iterate by cleanlab bin and put the correct number of entities in each bin
# into the 'suspicious' pile 
display_n = 0
for label in range(c_mat.shape[0]):
    for cl_label in range(c_mat.shape[1]):
        if label == cl_label: continue 
        flagged=(model_predictions[(model_predictions.upostag_id==label) & (model_predictions.CL_pred==cl_label)]
                    .sort_values('self_confidence',ascending=True,)
                .head(int(q_n_mat[label,cl_label])))
        if display_n: 
            display(label, cl_label, flagged.shape)
            display(flagged[['fold','doc_num','raw_span_id','upostag','upostag_id','CL_pred','normalized_margin','self_confidence']])
            display_n -=1
        flagged_list.append(flagged)
flagged_PNBR = pd.concat(flagged_list)    
flagged_PNBR.sort_values('normalized_margin',ascending=True,inplace=True)
flagged_PNBR[['fold','doc_num','upostag','upostag_id','raw_span','CL_pred','normalized_margin']]

# +
#calculate outputs from prune by class

n_prune_by_class = np.where(np.eye(q_n_mat.shape[0]),0,q_n_mat).sum(axis=0).astype(np.int32)
model_predictions.sort_values('self_confidence', ascending=True,inplace = True)
flagged_pbc_list =[]

display_n = 0
for class_label in range(len(n_prune_by_class)):
    flagged_row = model_predictions[model_predictions.upostag_id ==class_label].sort_values('self_confidence',ascending=True).head(n_prune_by_class[class_label])
    if display_n: 
        display(class_label, flagged_row.shape,n_prune_by_class[class_label])
        display(flagged_row.head(3)[['fold','doc_num','raw_span_id','upostag','upostag_id','CL_pred','normalized_margin','self_confidence']])
        display_n -=1 
    flagged_pbc_list.append(flagged_row)
flagged_PBC = pd.concat(flagged_pbc_list)
flagged_PBC.sort_values('normalized_margin',ascending=True,inplace=True)
flagged_PBC[['fold','doc_num','upostag','upostag_id','raw_span','CL_pred','normalized_margin','self_confidence']]
# -





StarcoderdataJupyterScriptsDedupFiltered