predict.py

"""Classify construction text into MasterFormat level-2 groups (171 classes, 32 divisions).

Examples
--------
    pip install transformers torch
    python predict.py "4000 psi concrete slab on grade" "TPO roofing, 60 mil, fully adhered"
    python predict.py --top-k 5 --divisions "8\" CMU wall, grout filled at 32\" o.c."
    echo "Cat 6 data cabling and jacks" | python predict.py -
    python predict.py --file items.txt --json > out.json

ONNX (no torch, ~34 MB int8 weights):
    pip install onnxruntime transformers
    python predict.py --onnx onnx/model_quantized.onnx "Panelboards 208/120V 42 circuit"

Without arguments it classifies one built-in example.
"""
import argparse
import json
import sys

MODEL = "constructelligence/masterformat-classifier"

# Division code -> name. Kept local so --divisions needs no extra package.
DIVISIONS = {
    "01": "General Requirements", "02": "Existing Conditions", "03": "Concrete",
    "04": "Masonry", "05": "Metals", "06": "Wood, Plastics, and Composites",
    "07": "Thermal and Moisture Protection", "08": "Openings", "09": "Finishes",
    "10": "Specialties", "11": "Equipment", "12": "Furnishings",
    "13": "Special Construction", "14": "Conveying Equipment", "21": "Fire Suppression",
    "22": "Plumbing", "23": "HVAC", "25": "Integrated Automation", "26": "Electrical",
    "27": "Communications", "28": "Electronic Safety and Security", "31": "Earthwork",
    "32": "Exterior Improvements", "33": "Utilities", "34": "Transportation",
    "35": "Waterway and Marine Construction", "40": "Process Interconnections",
    "41": "Material Processing and Handling Equipment",
    "43": "Process Gas and Liquid Handling, Purification, and Storage Equipment",
    "44": "Pollution and Waste Control Equipment", "46": "Water and Wastewater Equipment",
    "48": "Electrical Power Generation",
}

EXAMPLE = "4000 psi concrete slab on grade"


def split_label(label):
    """'03 30 00 Cast-in-Place Concrete' -> ('03 30 00', 'Cast-in-Place Concrete')."""
    code, name = label[:8].strip(), label[8:].strip()
    return code, (name or label)


def division_rollup(preds):
    """Sum level-2 probabilities by the first two digits (division)."""
    totals = {}
    for p in preds:
        div = p["code"][:2]
        totals[div] = totals.get(div, 0.0) + p["score"]
    return [
        {"code": d, "name": DIVISIONS.get(d, d), "score": s}
        for d, s in sorted(totals.items(), key=lambda kv: -kv[1])
    ]


def hf_scorer(model, top_k):
    from transformers import pipeline

    clf = pipeline("text-classification", model=model, top_k=top_k)
    return lambda texts: [[d for d in out] for out in clf(texts)]


def onnx_scorer(path, top_k):
    """Run onnx/model*.onnx directly. Mirrors scripts/export_onnx.py: inputs
    input_ids / attention_mask / token_type_ids, output logits [batch, num_labels]."""
    import numpy as np
    import onnxruntime as ort
    from pathlib import Path
    from transformers import AutoTokenizer

    # The labels and tokenizer live beside onnx/ in the repo; fall back to the Hub.
    root = Path(path).resolve().parent.parent
    src = str(root) if (root / "tokenizer.json").exists() or (root / "vocab.txt").exists() else MODEL
    tok = AutoTokenizer.from_pretrained(src)
    id2label = None
    try:
        c = json.loads((root / "config.json").read_text())
        id2label = {int(k): v for k, v in c.get("id2label", {}).items()}
    except Exception:
        id2label = None
    sess = ort.InferenceSession(path, providers=["CPUExecutionProvider"])
    names = {i.name for i in sess.get_inputs()}

    def score(texts):
        enc = tok(texts, truncation=True, max_length=128, padding=True, return_tensors="np")
        feed = {n: enc[n].astype(np.int64) for n in ("input_ids", "attention_mask", "token_type_ids") if n in names}
        logits = sess.run(None, feed)[0]
        idx = np.argsort(-logits, 1)[:, :top_k]
        shifted = logits - logits.max(1, keepdims=True)
        probs = np.exp(shifted) / np.exp(shifted).sum(1, keepdims=True)
        out = []
        for row, cols in zip(probs, idx):
            out.append([{"label": id2label.get(int(c), str(int(c))), "score": float(row[c])} for c in cols])
        return out

    return score


def predict(texts, score, top_k, want_div):
    results = []
    for text, preds in zip(texts, score(texts)):
        pl = []
        for p in preds:
            code, name = split_label(p["label"])
            pl.append({"code": code, "label": p["label"], "name": name, "score": round(float(p["score"]), 4)})
        item = {"text": text, "predictions": pl}
        if want_div:
            item["divisions"] = [dict(d, score=round(d["score"], 4)) for d in division_rollup(pl)]
        results.append(item)
    return results


def main(argv=None):
    ap = argparse.ArgumentParser(description="MasterFormat level-2 classifier (171 groups, 32 divisions).")
    ap.add_argument("text", nargs="*", help="text to classify; use '-' to read lines from stdin")
    ap.add_argument("--model", default=MODEL, help="HF repo id or local checkpoint dir")
    ap.add_argument("--onnx", metavar="PATH", help="classify with an ONNX model instead of PyTorch")
    ap.add_argument("--top-k", type=int, default=3, help="number of level-2 predictions to show")
    ap.add_argument("--divisions", action="store_true", help="also show the level-1 (division) roll-up")
    ap.add_argument("--file", help="read newline-separated inputs from a file")
    ap.add_argument("--json", action="store_true", help="emit JSON instead of a table")
    a = ap.parse_args(argv)

    texts = list(a.text)
    if a.file:
        texts += [l.rstrip("\n") for l in open(a.file, encoding="utf-8") if l.strip()]
    if "-" in texts:
        texts = [t for t in texts if t != "-"] + [l.rstrip("\n") for l in sys.stdin if l.strip()]
    if not texts:
        texts = [EXAMPLE]
    texts = [t for t in texts if t.strip()]
    if not texts:
        ap.error("no input text")

    score = onnx_scorer(a.onnx, a.top_k) if a.onnx else hf_scorer(a.model, a.top_k)
    results = predict(texts, score, a.top_k, a.divisions)

    if a.json:
        json.dump(results, sys.stdout, indent=2, ensure_ascii=False)
        print()
        return
    for r in results:
        print(f"\n{r['text']}")
        for p in r["predictions"]:
            print(f"  {p['code']} {p['name']:<48.48} {p['score']:.3f}")
        if a.divisions:
            print("  -- divisions --")
            for d in r.get("divisions", []):
                print(f"  {d['code']}      {d['name']:<48.48} {d['score']:.3f}")


if __name__ == "__main__":
    main()