blue-machines/Gemma3-270M-sentence-intent-v1

probs = softmax(intent_logits)
pred  = argmax(probs)
conf  = max(probs)
if conf < 0.55:
    pred = unclear   # safe abstain

import json
from pathlib import Path

import numpy as np
from transformers import AutoTokenizer, CanineTokenizer

MODEL_DIR = Path(".")  # or huggingface_hub.snapshot_download(...)
MAX_LEN = 64
CONF_THRESHOLD = 0.55  # optimal gate (coverage>=0.97, max selective acc)
UNCLEAR_ID = 6

# tokenizer — Gemma tokenizer (max length 64).
tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

with open(MODEL_DIR / "label_map.json", encoding="utf-8") as f:
    maps = json.load(f)
id2intent = {int(k): v for k, v in maps["head4_intent"]["id2label"].items()}

import onnxruntime as ort

session = ort.InferenceSession(
    str(MODEL_DIR / "model.onnx"),  # or model_fp32.onnx
    providers=["CPUExecutionProvider"],
)

def predict_intent(sentence: str) -> dict:
    """Pass a single bare sentence; returns label + confidence."""
    enc = tokenizer(
        sentence,
        return_tensors="np",
        truncation=True,
        max_length=MAX_LEN,
        padding="max_length",
    )
    feed = {
        "input_ids": enc["input_ids"].astype(np.int64),
        "attention_mask": enc["attention_mask"].astype(np.int64),
    }
    logits = session.run(None, feed)[0][0]

    x = logits.astype(np.float64)
    x = x - x.max()
    probs = np.exp(x)
    probs = probs / probs.sum()
    pred_id = int(probs.argmax())
    conf = float(probs.max())
    abstained = conf < CONF_THRESHOLD
    if abstained:
        pred_id = UNCLEAR_ID
    return {
        "text": sentence,
        "intent": id2intent[pred_id],
        "confidence": conf,
        "abstained": abstained,
        "threshold": CONF_THRESHOLD,
        "probs": {id2intent[i]: float(probs[i]) for i in range(7)},
    }

for s in [
    "Haan, ye sahi hai. Proceed karo.",
    "What is the minimum balance required?",
    "Nahi, maine ye payment nahi kiya.",
]:
    print(predict_intent(s))