examples/inference.py

from vllm import SamplingParams
from prefixLLM import PrefixLLM

SYSTEM_PROMPT = "You are a helpful and harmless assistant.You should think step-by-step and put your final answer within \\boxed{{}}."

if __name__ == "__main__":
    llm = PrefixLLM(model="amandaa/AutoL2S-Plus-7b")
    max_tokens, temp = 32768, 0.7
    sampling_params= SamplingParams(max_tokens=max_tokens, temperature=temp)
    
    question = "Melissa works as a pet groomer. This week, she has 8 dogs that need to be bathed, 5 cats that need their nails clipped, 3 birds that need their wings trimmed, and 12 horses that need to be brushed. If she splits the grooming jobs evenly over the days, how many animals will she groom each day of the week?"
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": question}
    ]
    responses = llm.chat(messages=messages, sampling_params=sampling_params, use_tqdm=True)

    print(responses[0].outputs[0].text)