test_csv.py

"""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 first becomes available after all inference and application decisions.
    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'])))