| """CSV inference and visual review. Run with the training script's Python dependencies.""" |
| from __future__ import annotations |
|
|
| import argparse |
| import csv |
| import html |
| import json |
| import random |
| import time |
| import unicodedata |
| from typing import cast |
| from dataclasses import asdict |
| from pathlib import Path |
| import xml.etree.ElementTree as ET |
|
|
| import torch |
| from train_tabfix import (BUSINESS_CATEGORIES as CATEGORIES, Consumer, Edit, Record, generate_repair, load_checkpoint, predict_edits, deterministic_findings, deterministic_repair, conditional_perplexity, replace_cell_xml, read_json, mapping, number) |
|
|
|
|
| def read_csv(path: Path) -> list[list[str]]: |
| with path.open(encoding="utf-8-sig", newline="") as stream: |
| rows = list(csv.reader(stream)) |
| if not rows or not rows[0] or any(len(r) != len(rows[0]) for r in rows): |
| raise ValueError("Expected a rectangular, comma-delimited UTF-8 CSV with a header") |
| return rows |
|
|
|
|
| def write_csv(path: Path, rows: list[list[str]]) -> None: |
| with path.open("w", encoding="utf-8", newline="") as stream: |
| csv.writer(stream).writerows(rows) |
|
|
|
|
| def context_xml(rows: list[list[str]], selected: list[int], schema: str) -> str: |
| chunks = [f'<table rows="{len(rows)-1}" columns="{len(rows[0])}" encoding="utf-8">', schema, '<rows>'] |
| for i in selected: |
| chunks.append(f'<row index="{i}">') |
| for header, value in zip(rows[0], rows[i+1]): |
| chunks.append(f'<cell column="{html.escape(header, quote=True)}" xml:space="preserve">'+(html.escape(value) if value else '<empty/>')+'</cell>') |
| chunks.append('</row>') |
| return '\n'.join(chunks+['</rows>', '</table>']) |
|
|
|
|
| def review(rows: list[list[str]], result: list[list[str]], reference: list[list[str]] | None, |
| audit: list[dict[str, object]], stats: dict[str, object]) -> str: |
| def show(value: str) -> str: |
| quoted = json.dumps(value, ensure_ascii=False) |
| visible = ''.join(f"\\u{ord(c):04x}" if unicodedata.category(c) == 'Cf' else c for c in quoted) |
| return html.escape(visible) |
| details: list[str] = [] |
| for i, row in enumerate(rows[1:]): |
| for j, value in enumerate(row): |
| actual = result[i+1][j] |
| expected = reference[i+1][j] if reference else None |
| relevant = actual != value or (expected is not None and expected != value) or any(a['row']==i and a['column']==j for a in audit) |
| if not relevant: |
| continue |
| status = 'Changed' if actual != value else 'Unchanged' |
| if expected is not None: |
| status = ('Repaired' if actual == expected else 'Missed / incorrect') if value != expected else ('Unwanted change' if actual != value else 'Preserved') |
| color = 'good' if status in ('Repaired','Preserved') else 'bad' |
| categories = ', '.join(str(a['category']) for a in audit if a['row']==i and a['column']==j) or 'No detection' |
| details.append(f'<tr><td>{i+1}</td><td>{html.escape(rows[0][j])}</td><td>{show(value)}</td><td>{show(actual)}</td><td>{show(expected) if expected is not None else "—"}</td><td class="{color}">{status}<small>{html.escape(categories)}</small></td></tr>') |
| full = '<tr><th>Data row</th>'+''.join('<th>'+html.escape(h)+'</th>' for h in rows[0])+'</tr>' |
| for i,row in enumerate(result[1:]): |
| full += f'<tr><td>{i+1}</td>'+''.join('<td class="'+('changed' if v!=rows[i+1][j] else '')+'">'+html.escape(v)+'</td>' for j,v in enumerate(row))+'</tr>' |
| cards = ''.join('<div class="card"><strong>'+html.escape(str(stats.get(key,'—')))+'</strong><span>'+label+'</span></div>' for key,label in [('cells','Cells scanned'),('initial_faulty_cells','Errors in input'),('exactly_repaired_cells','Errors repaired'),('valid_cells_changed','Valid cells changed')]) |
| return '''<!doctype html><meta charset="utf-8"><title>TabFix — end-to-end CSV review</title><style> |
| .cards{display:flex;gap:16px;flex-wrap:wrap}.card{background:white;padding:18px 24px;border-radius:12px;min-width:130px}.card strong{display:block;font-size:32px}.card span{color:#536175}body{font:16px system-ui;background:#f3f5f8;color:#172238;margin:36px}h1{font-size:30px}p{max-width:950px;line-height:1.6}table{border-collapse:collapse;width:100%;background:white;margin:20px 0}td,th{padding:12px;text-align:left;border-bottom:1px solid #dbe1e8;white-space:pre-wrap;overflow-wrap:anywhere}th{background:#e6ebf2}small{display:block;color:#536175;margin-top:6px}.good{color:#15704b}.bad{color:#b13a37}table.full{min-width:1050px}.changed{background:#fff0c7}.scroll{overflow:auto}pre{white-space:pre-wrap;background:white;padding:20px}details{margin-top:24px}</style> |
| <h1>Does TabFix fix this CSV?</h1><p>Actual CSV → schema and context → predicted error locations → generated replacements → corrected CSV. The clean reference is used only after inference. No gold error locations or replacement text are supplied to the model. Row numbers below count data rows, excluding the header.</p> |
| '''+'<div class="cards">'+cards+'</div><p>CSV inference review · Threshold '+str(stats['threshold'])+' · CPU inference '+str(stats['inference_seconds'])+' seconds (model loading excluded).</p><h2>Errors and proposed changes</h2><div class="scroll"><table><tr><th>Row</th><th>Column</th><th>Input</th><th>Model output</th><th>Reference</th><th>Result</th></tr>'+''.join(details)+'</table></div><p>Strings are quoted to make spaces and empty values visible. Yellow cells below were changed. All other cells are included so unwanted edits remain visible.</p><h2>Complete model output</h2><div class="scroll"><table class="full">'+full+'</table></div><details><summary>Every detection and application decision</summary><pre>'+html.escape(json.dumps(audit,ensure_ascii=False,indent=2))+'</pre></details>' |
|
|
|
|
| def run(csv_path: Path, checkpoint: Path, output: Path, schema_path: Path | None, |
| reference_path: Path | None, threshold: float, enabled: set[str]) -> None: |
| if output.resolve() == csv_path.parent.resolve(): |
| raise ValueError("Use a separate output directory to preserve inputs") |
| output.mkdir(parents=True, exist_ok=True) |
| rows = read_csv(csv_path) |
| schema = schema_path.read_text() if schema_path else '<schema>\n'+'\n'.join(f'<column index="{i}" name="{html.escape(h,quote=True)}" type="text" nullable="unknown"/>' for i,h in enumerate(rows[0]))+'\n</schema>' |
| element = ET.fromstring(schema) |
| if element.tag != 'schema' or [c.get('name') for c in element.findall('column')] != rows[0]: |
| raise ValueError("Schema columns must match the CSV header exactly") |
| torch.set_num_threads(4) |
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
| begin = time.perf_counter() |
| model,tok,_ = load_checkpoint(checkpoint,device) |
| load_seconds = time.perf_counter()-begin |
| consumer = Consumer(tok) |
| result = [r.copy() for r in rows] |
| audit: list[dict[str,object]] = [] |
| n = len(rows)-1 |
| begin = time.perf_counter() |
| if device.type == 'cuda': |
| torch.cuda.reset_peak_memory_stats() |
| for start in range(0,n,12): |
| targets = set(range(start,min(n,start+12))) |
| context = targets | set(range(min(n,2))) | set(range(max(0,start-2),min(n,start+14))) |
| other = sorted(set(range(n))-context) |
| context.update(random.Random(start+42).sample(other,min(4,len(other)))) |
| xml = context_xml(rows,sorted(context),schema) |
| record = Record('csv',xml,[],{'target_rows':[], 'supervised_columns':[], 'label_scope':'unknown'}) |
| prepared = consumer.prepare(record) |
| violations = [f for f in deterministic_findings(xml) if f.status == 'invalid' and f.category in enabled and f.row in targets] |
| proposed = {(e.row, e.column): e for e in predict_edits(model, consumer, prepared, device, enabled, threshold) if e.row in targets} |
| grouped = {(f.row, c) for f in violations if len(f.columns) > 1 for c in f.columns} |
| for finding in violations: |
| if len(finding.columns) == 1: |
| c = finding.columns[0] |
| proposed[finding.row, c] = Edit(finding.row, c, 0, len(rows[finding.row+1][c]), '', finding.category) |
| else: |
| audit.append({'row':finding.row, 'column':finding.columns[-1], 'category':finding.category, 'status':'relation contradiction: ambiguous faulty member; unchanged'}) |
| columns = element.findall('column') |
| metrics = read_json(checkpoint / 'metrics.json') |
| calibration = mapping(mapping(metrics.get('validation', {})).get('perplexity_calibration', {})) |
| max_perplexity = number(calibration.get('threshold', 0.0)) |
| for (r, c), proposal in proposed.items(): |
| value = rows[r+1][c] |
| entry: dict[str, object] = dict(asdict(proposal)) |
| if (r, c) in grouped: |
| entry['status'] = 'abstained: relation does not identify faulty member' |
| audit.append(entry) |
| continue |
| replacement = deterministic_repair(value, columns[c]) |
| route = 'deterministic' |
| if replacement is None and enabled != set(CATEGORIES): |
| entry['status'] = 'abstained: whole-cell model cannot guarantee disabled categories remain unchanged' |
| audit.append(entry) |
| continue |
| if replacement is None: |
| route = 'model' |
| full = Edit(r, c, 0, len(value), '', proposal.category) |
| replacement = generate_repair(model, consumer, prepared, full, device) |
| if replacement is not None: |
| score = conditional_perplexity(model, consumer, prepared, full, replacement, device) |
| entry['perplexity'] = score |
| if score > max_perplexity: |
| replacement = None |
| entry['status'] = 'abstained: perplexity exceeds validation threshold' |
| if replacement is not None: |
| candidate_xml = replace_cell_xml(xml, r, c, replacement) |
| remaining = [f for f in deterministic_findings(candidate_xml) if f.status == 'invalid' and f.row == r and c in f.columns and f.category in enabled] |
| if remaining: |
| entry['status'] = 'abstained: output violates a declared rule' |
| else: |
| result[r+1][c] = replacement |
| entry['status'] = 'preserved' if replacement == value else 'applied' |
| else: |
| entry.setdefault('status', 'abstained: no terminated answer') |
| entry.update(route=route, replacement=replacement) |
| audit.append(entry) |
| proposals = list(proposed.values()) |
| print(json.dumps({'processed_rows':min(n,start+12),'detections':len(proposals)}),flush=True) |
| if device.type=='cuda': |
| torch.cuda.synchronize() |
| elapsed=time.perf_counter()-begin |
| write_csv(output/'corrected.csv',result) |
| |
| reference=read_csv(reference_path) if reference_path else None |
| if reference is not None and (reference[0]!=rows[0] or len(reference)!=len(rows)): |
| raise ValueError('Reference shape/header mismatch') |
| stats: dict[str,object]={'device':str(device),'rows':n,'cells':n*len(rows[0]),'threshold':threshold,'load_seconds':round(load_seconds,2),'inference_seconds':round(elapsed,2),'cells_per_second':round(n*len(rows[0])/elapsed,2),'peak_gpu_bytes':torch.cuda.max_memory_allocated() if device.type=='cuda' else None,'detected_spans':len(audit),'changed_cells':sum(a!=b for x,y in zip(rows[1:],result[1:]) for a,b in zip(x,y))} |
| if reference: |
| triples=[(a,b,c) for x,y,z in zip(rows[1:],result[1:],reference[1:]) for a,b,c in zip(x,y,z)] |
| stats.update(initial_faulty_cells=sum(a!=c for a,_b,c in triples),exactly_repaired_cells=sum(a!=c and b==c for a,b,c in triples),remaining_faulty_cells=sum(b!=c for _a,b,c in triples),valid_cells_changed=sum(a==c and b!=c for a,b,c in triples),valid_cells=sum(a==c for a,_b,c in triples)) |
| (output/'audit.json').write_text(json.dumps({'metrics':stats,'proposals':audit},ensure_ascii=False,indent=2)) |
| (output/'review.html').write_text(review(rows,result,reference,audit,stats)) |
| print(json.dumps(stats),flush=True) |
|
|
|
|
| if __name__=='__main__': |
| parser=argparse.ArgumentParser(description=__doc__) |
| parser.add_argument('csv',type=Path) |
| parser.add_argument('--checkpoint',type=Path,required=True) |
| parser.add_argument('--output',type=Path,required=True) |
| parser.add_argument('--schema',type=Path) |
| parser.add_argument('--reference',type=Path) |
| parser.add_argument('--threshold',type=float,default=.9) |
| parser.add_argument('--categories',nargs='+',choices=CATEGORIES,default=CATEGORIES) |
| args=cast(dict[str, object], vars(parser.parse_args())) |
| run(cast(Path,args['csv']),cast(Path,args['checkpoint']),cast(Path,args['output']), |
| cast(Path | None,args['schema']),cast(Path | None,args['reference']), |
| cast(float,args['threshold']),set(cast(list[str],args['categories']))) |
|
|
