predict.py

"""Classify voice-note transcripts with Voice-Note-Intent-Classifier v0.1.

    pip install sentence-transformers numpy
    python predict.py "um add milk eggs and the oat milk to the list"

Or from Python:
    from predict import IntentClassifier
    clf = IntentClassifier(".")          # folder containing head.json
    clf.predict(["dentist thursday at four thirty"])

Output per transcript follows the router contract: scores for every label, the
labels above their thresholds, and `abstain` when none clears its threshold.
"""

import json
import sys
from pathlib import Path

import numpy as np
from sentence_transformers import SentenceTransformer


class IntentClassifier:
    def __init__(self, path="."):
        cfg = json.loads((Path(path) / "head.json").read_text())
        self.cfg = cfg
        self.labels = cfg["labels"]
        self.coef = np.array(cfg["coef"])
        self.intercept = np.array(cfg["intercept"])
        self.thresholds = np.array([cfg["thresholds"][l] for l in self.labels])
        self.encoder = SentenceTransformer(cfg["encoder"], device="cpu")
        self.encoder.max_seq_length = cfg["max_seq_length"]

    def predict(self, transcripts):
        emb = self.encoder.encode(list(transcripts), normalize_embeddings=self.cfg["normalize_embeddings"])
        probs = 1 / (1 + np.exp(-(emb @ self.coef.T + self.intercept)))
        out = []
        for p in probs:
            labels = [l for l, s, t in zip(self.labels, p, self.thresholds) if s >= t]
            out.append({
                "model_version": self.cfg["model_version"],
                "taxonomy_version": self.cfg["taxonomy_version"],
                "scores": {l: round(float(s), 4) for l, s in zip(self.labels, p)},
                "labels": labels,
                "abstain": not labels,
            })
        return out


if __name__ == "__main__":
    clf = IntentClassifier(Path(__file__).parent)
    print(json.dumps(clf.predict(sys.argv[1:] or ["need to get milk eggs and dish soap"]), indent=2))