Branching

import asyncio

from inferlet import chat, model
from inferlet.eta import ForwardPass, Pipeline, StateGeometry, WorkingSet, intrinsics, reduce_argmax
# generate() as in "Your first inferlet"


async def branch(ws, ending):
    pipe = Pipeline()
    words = [t async for t in generate(ws, pipe, model.encode(ending), max_tokens=24)]
    pipe.close()
    return ending + model.decode(words)


async def main(input: dict) -> dict:
    # 1. Prefill and freeze the shared prompt once.
    root, setup = WorkingSet(), Pipeline()
    fwd = ForwardPass()
    fwd.embed(chat.prefix() + model.encode("Write one sentence about a cat"))
    fwd.state([root], StateGeometry.folding())
    done = fwd.epilogue(lambda: reduce_argmax(intrinsics.output(0, "logits")))
    setup.submit(fwd)
    await done

    # 2. Fork three child working sets and run them concurrently.
    endings = [" at dawn:", " at night:", " in the rain:"]
    children = [root.fork(setup) for _ in endings]
    setup.close()
    texts = await asyncio.gather(*(branch(ws, e) for ws, e in zip(children, endings)))
    return {"branches": texts}