Quick python script to try structured output with vLLM

example.py

from openai import OpenAI
import argparse
client = OpenAI(
    base_url="http://localhost:8000/v1",
    api_key="-",
)

parser = argparse.ArgumentParser(description="Query local vLLM")
parser.add_argument("prompt", help="Prompt to send to the model")
args = parser.parse_args()

chat_template = (
    "{% for message in messages %}"
    "{% if message['role'] == 'system' %}"
    "<|system|>{{ message['content'] }}\n"
    "{% elif message['role'] == 'user' %}"
    "<|user|>{{ message['content'] }}\n"
    "{% elif message['role'] == 'assistant' %}"
    "<|assistant|>{{ message['content'] }}\n"
    "{% endif %}"
    "{% endfor %}"
    "<|assistant|>"
)

completion = client.chat.completions.create(
    model="mistralai/Mistral-7B-v0.1",
    messages=[
	{"role": "user", "content": "Classify if this is rad or bogus as if you're some ski bro from Vermont, dude: " + args.prompt}
    ],
    extra_body={
        "chat_template": chat_template,
        "guided_choice": ["rad", "bogus"]
    },
)


print(completion.choices[0].message.content)