File size: 3,195 Bytes
e7eaeed |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 |
import pyrootutils
root = pyrootutils.setup_root(
search_from=__file__,
indicator=[".project-root"],
pythonpath=True,
dotenv=False,
)
import argparse
from typing import List, Optional
import pandas as pd
import plotly.figure_factory as ff
from pie_datasets import DatasetDict
pd.options.plotting.backend = "plotly"
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="Show score distribution of annotations per layer"
)
# --data-dir predictions/default/2025-02-26_14-28-17
parser.add_argument(
"--data-dir", type=str, required=True, help="Path to the dataset directory"
)
parser.add_argument("--split", type=str, default="test", help="Dataset split to use")
parser.add_argument(
"--layers",
nargs="+",
default=["labeled_spans", "binary_relations"],
help="Annotation layers to use",
)
# --layer-captions ADUs "Argumentative Relations"
parser.add_argument(
"--layer-captions", nargs="+", help="Captions for the figure traces per layer"
)
# --layer-colors "rgb(31,119,180)" "rgb(255,127,14)"
parser.add_argument("--layer-colors", nargs="+", help="Colors for the figure traces per layer")
args = parser.parse_args()
# Load the dataset
ds = DatasetDict.from_json(data_dir=args.data_dir)[args.split]
# get scores per annotation layer and label
layers = args.layers
all_scores = []
all_scores_idx = []
for doc in ds:
for layer in layers:
for ann in doc[layer].predictions:
all_scores.append(ann.score)
all_scores_idx.append((doc.id, layer, getattr(ann, "label", None)))
scores = pd.Series(
all_scores,
index=pd.MultiIndex.from_tuples(all_scores_idx, names=["doc_id", "layer", "label"]),
name="score",
)
if args.layer_captions is not None:
if len(args.layer_captions) < len(layers):
raise ValueError("Not enough captions provided for all layers")
name_mapping = dict(zip(layers, args.layer_captions))
else:
name_mapping = dict(zip(layers, layers))
colors: Optional[List[str]] = None
if args.layer_colors is not None:
if len(args.layer_colors) < len(layers):
raise ValueError("Not enough colors provided for all layers")
color_mapping = dict(zip(layers, args.layer_colors))
colors = [color_mapping[layer] for layer in layers]
else:
colors = None
score_groups = {layer: scores.xs(layer, level="layer").to_numpy() for layer in layers}
group_labels, hist_data = zip(*score_groups.items())
group_labels_renamed = [name_mapping[label] for label in group_labels]
fig = ff.create_distplot(
hist_data,
group_labels=group_labels_renamed,
show_hist=True,
colors=colors,
bin_size=0.025,
)
fig.update_layout(
height=600,
width=800,
title_text="Score Distribution per Annotation Layer",
title_x=0.5,
barmode="group",
)
fig.update_layout(legend=dict(yanchor="top", y=0.99, xanchor="left", x=0.01))
fig.show()
print("done")
|