python/examples/quickstart.py

"""The quick start from the model card, in Python.

    python examples/quickstart.py [model_dir]

model_dir holds meta.json and model.bin. The default is the folder above python/, which is where
they sit in a downloaded TheREZOR/TinyDecide repo.
"""
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))   # run from a checkout without installing
from tinydecide import TinyDecide  # noqa: E402

model_dir = sys.argv[1] if len(sys.argv) > 1 else Path(__file__).resolve().parent.parent.parent
model = TinyDecide.load(model_dir)

r = model.answer(
    "Book a table for 4 at an Italian place near the station on Friday at 7:30",
    [
        {"type": "choice", "text": "Which app should handle this?",
         "options": ["reminders", "music", "calendar", "restaurants", "weather"]},
        {"type": "noul", "text": "The message is urgent."},
        {"type": "score", "text": "How positive is the tone?", "options": ["negative", "neutral", "positive"]},
        {"type": "span", "text": "Extract the time."},
    ],
)

app, urgent, tone, time = r["answers"]
options = ["reminders", "music", "calendar", "restaurants", "weather"]
print(f"app:    {options[app['pick']]} (p = {app['probs'][app['pick']]:.2f})")
print(f"urgent: p = {urgent['p']:.2f}")
print(f"tone:   score = {tone['score']:.2f}")
print(f"time:   {time['text']!r} (p_present = {time['p_present']:.2f})")
print(f"{r['tokens']['total']} tokens, {r['ms']:.0f} ms")