Views
No views yet
1#!/usr/bin/env python3
2import time
3from vllm import LLM, SamplingParams
4
5def main():
6 # Hard-coded model and tensor parallel configuration.
7 model_path = "miike-ai/r1-8b-fp8"
8 tensor_parallel_size = 1
9
10 # Define sampling parameters with an increased max_tokens and a stop string.
11 sampling_params = SamplingParams(
12 temperature=0.0,
13 top_p=0.95,
14 max_tokens=32000, # Increase this to allow longer responses.
15 stop=["\nUser:"], # Stop when the model outputs a new user marker.
16 )
17
18 print(f"Loading model '{model_path}' ...")
19 model = LLM(
20 model=model_path,
21 enforce_eager=True,
22 dtype="auto",
23 tensor_parallel_size=tensor_parallel_size,
24 )
25 print("Model loaded. You can now chat!")
26 print("Type 'exit' or 'quit' to end the conversation.\n")
27
28 conversation = ""
29 while True:
30 try:
31 user_input = input("User: ").strip()
32 except (KeyboardInterrupt, EOFError):
33 print("\nExiting chat.")
34 break
35
36 if user_input.lower() in {"exit", "quit"}:
37 print("Exiting chat.")
38 break
39
40 # Append the user's input to the conversation history.
41 conversation += f"User: {user_input}\nBot: "
42 print("Bot: ", end="", flush=True)
43
44 # Generate a response using the conversation history and sampling parameters.
45 response = model.generate(conversation, sampling_params=sampling_params)
46 # Extract the generated reply.
47 bot_reply = response[0].outputs[0].text.strip()
48
49 # Simulate streaming by printing one character at a time.
50 for char in bot_reply:
51 print(char, end="", flush=True)
52 time.sleep(0.02) # Adjust delay (in seconds) as desired.
53 print() # Newline after bot reply.
54
55 # Append the bot reply to conversation history.
56 conversation += bot_reply + "\n"
57
58if __name__ == "__main__":
59 main()