from vllm import LLM, SamplingParams

model_name = "BSC-LT/salamandra-2b-instruct-fp8"
llm = LLM(model=model_name)

messages = []

while True:
    user_input = input("user >> ")
    if user_input.lower() == "exit":
        print("Chat ended.")
        break

    messages.append({'role': 'user', 'content': user_input})

    outputs = llm.chat(messages,
                       sampling_params=SamplingParams(
                           temperature=0.5,
                           stop_token_ids=[5],
                           max_tokens=200)
                       )[0].outputs
    
    model_output = outputs[0].text
    print(f'assistant >> {model_output}')
    messages.append({'role': 'assistant', 'content': model_output})