<jupyter_start><jupyter_text>Annotating pivotsUsing the pigeon library.<jupyter_code>%%capture --no-display
from dask.distributed import Client
client = Client(address = 'tcp://192.168.62.207:8686')
client
%%time
import dask.bag as db
from takco.util import robust_json_loads_lines
from takco.table import from_tabel
import takco

steps = takco.config.build('step', load=['resources/graphs/wikidata.toml','resources/pipelines/wikitables.toml'])
prefix_header_rules = steps[0]['prefix_header_rules']
unpivot_heuristics = {h.name: h for h in steps[0]['unpivot_heuristics']}

# fnames = 'hdfs://bricks07:9000/user/kruit/tabel/2-*'
fnames = '/export/scratch1/home/kruit/scratch/tabel/*'
bag = db.read_text(fnames).map_partitions(robust_json_loads_lines)\
    .map_partitions(lambda ts: list(takco.reshape.restructure(ts, prefix_header_rules=prefix_header_rules)))\
    .persist()
bag.count().compute()
headers = bag.distinct(key=takco.reshape.table_get_headerId).map(lambda t: t.get('tableHeaders')).persist()
headers.count().compute()
from dask.bag.random import sample

for up in bag.map_partitions(lambda ts: list(takco.reshape.build_heuristics(ts, unpivot_heuristics.values()))):
    unpivot_heuristics[up.name].merge(up)

hsample = sample(headers, 100).compute()
header_tables = {}
for h in hsample:
    h = {'tableHeaders': h}
    header_tables[takco.reshape.table_get_headerId(h)] = h

for p in takco.reshape.yield_pivots(hsample, unpivot_heuristics.values()):
    header_tables[p['headerId']]['pivot'] = p


def show_pivot(h):
    df = takco.table.from_tabel(h)
    if h.get('pivot') and not h['pivot'].get('discard'):
        print(h.get('pivot'))
        display(df.takco.highlight_pivot(**h.get('pivot')))
    else:
        display(df.takco)

from pigeon import annotate
annotations = annotate(
    header_tables.values(),
    options=['good', 'bad'],
    display_fn = show_pivot
)
import json
    
outpath = 'resources/pivot-annotation-recall.json'
if os.path.exists(outpath):
    annotations = json.load(open(outpath))
else:
    with open(outpath, 'w') as fw:
        json.dump(annotations, fw)
    
import pandas as pd
p = pd.Series([bool(h.get('pivot') and not h['pivot'].get('discard')) for h,a in annotations])
ok = pd.Series([a == 'good' for h,a in annotations])

tp = sum(p & ok)
fp = sum(p & (~ok))
fn = sum((~p) & (~ok))
print(f"TP: {tp}  FP: {fp}  FN: {fn}  Total: {len(annotations)}")
print(f"Precision: {(tp/(tp+fp)):.2f}, Recall: {(tp/(tp+fn)):.2f}")
from dask.bag.random import sample
import tqdm

hsample = tqdm.tqdm(headers.compute())

header_tables = {}
for h in hsample:
    h = {'tableHeaders': h}
    header_tables[takco.reshape.table_get_headerId(h)] = h

heuristic_pivots = {}
for p in takco.reshape.yield_pivots(hsample, unpivot_heuristics.values()):
    h = header_tables[p['headerId']]
    h['pivot'] = p
    if not h['pivot'].get('discard'):
        heuristic_pivots.setdefault(p.get('heuristic'), []).append(h)
    
{h:len(ps) for h,ps in heuristic_pivots.items()}
def show_pivot(h):
    df = takco.table.from_tabel(h)
    if h.get('pivot') and not h['pivot'].get('discard'):
        print(h.get('pivot'))
        display(df.takco.highlight_pivot(**h.get('pivot')))
    else:
        display(df.takco)

per_heuristic_tables = [p for h,ps in heuristic_pivots.items() for p in ps[:50]]
        
from pigeon import annotate
per_heuristic_annotations = annotate(
    per_heuristic_tables,
    options=['good', 'bad'],
    display_fn = show_pivot
)
import pandas as pd
import json
import os

outpath = 'resources/pivot-annotation-precision.json'
if os.path.exists(outpath):
    per_heuristic_annotations = json.load(open(outpath))
else:
    with open(outpath, 'w') as fw:
        json.dump(per_heuristic_annotations, fw)

heuristic_annotation = {}
for h,a in per_heuristic_annotations:
    hname = h['pivot']['heuristic']
    heuristic_annotation.setdefault(hname, []).append((h,a))
    
precision = {}
for hname, anns in heuristic_annotation.items():
    p = pd.Series([bool(h.get('pivot')) for h,a in anns])
    ok = pd.Series([a == 'good' for h,a in anns])

    tp = sum(p & ok)
    fp = sum(p & (~ok))
    fn = sum((~p) & (~ok))
    precision[hname] = (tp/(tp+fp))
    
s = pd.Series(precision)
s.index.name = 'Heuristic'
df = s.sort_values(ascending=False).to_frame('Precision').reset_index()

texpath = 'fig/pivot-precision.tex'
df.to_latex(texpath, index=False, float_format="%.2f")
df<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.16307277628032346, 'fraction_numerical': 0.007412398921832884, 'mean_word_length': 4.622474747474747, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 1, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '428353', 'license': 'MIT', 'n_tokens_mistral': 1739, 'n_tokens_neox': 1672, 'n_words': 316, 'path': 'docs/benchmarks/pivots.ipynb', 'repo_name': 'karmaresearch/takco', 'size': '4452'}