smoke_runtime.py

"""Run only the published runtime, without any source model weights."""

import asyncio
import json
import math

from transformers import PreTrainedModel

from coreai_d3 import D3CoreAI


def forbid_weights(*args, **kwargs):
    raise AssertionError("Runtime attempted to load PyTorch model weights")


async def main():
    PreTrainedModel.from_pretrained = forbid_weights
    engine = await D3CoreAI.load(".")
    result = await engine.system_one(
        state="The customer was charged twice and has a receipt.",
        questions={
            "team": {
                "type": "choice",
                "instructions": "Which team?",
                "criteria": {"billing": "Payments", "shipping": "Delivery"},
            },
            "receipt": {
                "type": "noul",
                "instructions": "Does the customer have a receipt?",
            },
            "severity": {
                "type": "score",
                "instructions": "How severe?",
                "criteria": ["Minor", "Major", "Complete"],
            },
        },
    )
    assert result["answers"]["team"]["choice"] in {"billing", "shipping"}
    assert 0 <= result["answers"]["receipt"]["noul"] <= 1
    assert math.isfinite(result["answers"]["severity"]["score"])
    print(json.dumps(result, indent=2))
    open("runtime-smoke.json", "w").write(
        json.dumps(
            {"passed": True, "source_weights_loaded": False, "result": result}, indent=2
        )
        + "\n"
    )


if __name__ == "__main__":
    asyncio.run(main())