| """Gradio demo: bag-of-words sentiment via ONNX Runtime.""" |
|
|
| import json |
| from pathlib import Path |
|
|
| import gradio as gr |
| import numpy as np |
| import onnxruntime as ort |
|
|
| from preprocess import text_to_bow |
|
|
| VOCAB_PATH = Path(__file__).resolve().parent / "vocab.json" |
| vocab_data = json.loads(VOCAB_PATH.read_text()) |
| VOCAB: list[str] = vocab_data["vocab"] |
| LABELS: list[str] = vocab_data["labels"] |
|
|
| session = ort.InferenceSession("sentiment.onnx") |
|
|
|
|
| def predict(text: str) -> dict[str, float]: |
| if not text.strip(): |
| return {label: 0.0 for label in LABELS} |
|
|
| bow = text_to_bow(text, VOCAB)[np.newaxis, :] |
| logits = session.run(None, {"bow": bow})[0][0] |
| exp_logits = np.exp(logits - logits.max()) |
| probs = exp_logits / exp_logits.sum() |
|
|
| return {label: float(prob) for label, prob in zip(LABELS, probs)} |
|
|
|
|
| demo = gr.Interface( |
| fn=predict, |
| inputs=gr.Textbox(label="Text", placeholder="I love this product"), |
| outputs=gr.Label(num_top_classes=2, label="Sentiment"), |
| title="ONNX Sentiment Demo", |
| description=( |
| "Tiny bag-of-words classifier trained from scratch in PyTorch " |
| "(no pretrained model). Exported to ONNX for inference." |
| ), |
| examples=[ |
| ["I love this"], |
| ["This is awful"], |
| ["Amazing product"], |
| ["So disappointed"], |
| ], |
| ) |
|
|
| demo.launch() |
|
|
