"""Classify construction text into MasterFormat level-2 groups.

    pip install transformers torch
    python predict.py "4000 psi concrete slab on grade" "TPO roofing, 60 mil, fully adhered"

Without arguments it classifies one example line item.
"""
import sys

from transformers import pipeline

MODEL = "constructelligence/masterformat-classifier"


def main():
    texts = sys.argv[1:] or ["4000 psi concrete slab on grade"]
    clf = pipeline("text-classification", model=MODEL, top_k=3)
    for text, preds in zip(texts, clf(texts)):
        print(f"\n{text}")
        for p in preds:
            print(f"  {p['label']:64s} {p['score']:.3f}")


if __name__ == "__main__":
    main()