examples/usage.py

"""Run the s1d1 router on a few prompts and print its three outputs (Python >= 3.10).

    pip install "transformers>=5" torch tokenizers huggingface_hub
    python examples/usage.py [repo-or-local-folder] [device]      # device: cpu (default), mps or cuda
"""
import os
import sys

sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), ".."))
from s1d1_router import Router  # noqa: E402

REPO = "scombat/s1d1-router"

if __name__ == "__main__":
    prompts = [
        "Quelle est la capitale de l'Australie ?",
        "Write a Python function that merges overlapping intervals and explain its complexity.",
        "Search the web for today's EUR/USD rate and put it in my spreadsheet.",
    ]
    router = Router(sys.argv[1] if len(sys.argv) > 1 else REPO, device=sys.argv[2] if len(sys.argv) > 2 else "cpu")
    for prompt, result in zip(prompts, router.route(prompts), strict=True):
        print(f"{prompt[:60]!r:64} -> {result}")