Uses NATS to publish responses of LLM inference to the subject inference.requests

llm_worker.py

# Launch nats-server
# wget https://huggingface.co/remyxai/stablelm-zephyr-3B_localmentor/resolve/main/ggml-model-q4_0.gguf -o stablelm-localmentor_2.gguf
import nats
import asyncio
from llama_cpp import Llama


async def llm_runner(nats_url, model_path, subject):
    nc = await nats.connect(nats_url)
    llm = Llama(model_path)

    async def inference_handler(msg):
        data = msg.data.decode()
        response = llm(data, max_tokens=2048, stop=["###", "\n\n"], echo=True)
        r = response["choices"][0]["text"]
        await nc.publish(msg.reply, str(r).encode())

    await nc.subscribe(subject, cb=inference_handler)
    await asyncio.Future()


if __name__ == "__main__":
    asyncio.run(
        llm_runner(
            "nats://localhost:4222", "stablelm-localmentor.gguf", "inference.requests"
        )
    )