inferlet (Python)

The Python library for writing Pie inferlets — the small programs that run next to the model.

from inferlet import model
from inferlet.eta import *

async def main(input: dict) -> dict:
    prompt = model.encode(input.get("prompt", "The capital of France is"))
    n = len(prompt)
    ws = WorkingSet()
    ws.reserve(max(-(-(n + 9) // kv_page_size()), 1))

    tokens = Channel.from_(prompt, dtype.i32)
    indptr_ch = Channel.from_([0, n], dtype.u32)
    kv_len = Channel.from_([n], dtype.u32)
    tok_out = Channel([1], dtype.i32)
    # ...pages / page_indptr / w_slot / w_off / positions as in
    # examples/text-completion-py/main.py

    fwd = ForwardPass()                    # picks the model's pass kind
    fwd.embed(tokens, indptr_ch)
    fwd.bind_state(ws, KvGeometry(kv_len=kv_len, ...))

    @fwd.epilogue                          # traced once, runs on the device
    def _():
        tok_out.put(reshape(reduce_argmax(intrinsics.logits()), [1]))

    pipe = Pipeline()
    fwd.submit(pipe)
    first = await tok_out.take_scalar()
    pipe.close()
    return {"text": model.decode([first])}

What is here

Spelling differences from Rust: Channel.from_([...], dtype.u32) needs a dtype for an integer sequence (the way a Rust literal needs a suffix); scalars in arithmetic take the partner tensor's dtype; integer division is x // y (ETA div truncates — / emits the same op, but // says so); </<=/>/>= compare elementwise, while == stays Python identity — elementwise equality is eq(a, b).

Running

A Python inferlet is its source: there is no build. pie run main.py (or pie inferlet install main.py, or client.install_program(main_py, pie_toml)) hands the file to the server, which runs it under the Python language component -- one wasm that bundles CPython, this package and the standard library, built once by language/build.sh with stock componentize-py >= 0.25 against crates/inferlet/wit, and installed at ~/.pie/languages/python.wasm. The world is component-model async, so main may be async and host reads are awaited (await ch.take_host()).

From a client, a function can go straight up: decorate it with pie_client.inferlet and hand it to PieClient.run (see python/client).

Tests

PYTHONPATH=src python -m pytest tests

The stub wit_world in tests/conftest.py covers the unit tests; the end-to-end twins live in examples/*-py in the pie repo (tests/examples/test_twins.py --attach ws://127.0.0.1:8080).