Views
No views yet
1
2import time
3from unsloth import FastLanguageModel, PatchFastRL
4PatchFastRL("GRPO", FastLanguageModel)
5from unsloth import is_bfloat16_supported
6import torch
7
8# Define model parameters
9max_seq_length = 13000 # Increase for longer contexts if needed
10lora_rank = 32 # Larger rank = potentially better performance but slower inference
11
12# Define system prompt and response format
13SYSTEM_PROMPT = """
14Respond in the following format:
15<reasoning>
16...
17</reasoning>
18<answer>
19...
20</answer>
21"""
22
23# Load the base model in 16-bit mode (required for LoRA with PEFT)
24model, tokenizer = FastLanguageModel.from_pretrained(
25 model_name="meta-llama/Llama-3.1-8B-Instruct",
26 max_seq_length=max_seq_length,
27 load_in_4bit=True, # Use 16-bit mode for LoRA support
28 fast_inference=True, # Enable vLLM fast inference
29 max_lora_rank=lora_rank,
30 gpu_memory_utilization=0.6, # Adjust if you run into memory issues
31)
32
33# Load the LoRA weights using PEFT
34from peft import PeftModel
35model = PeftModel.from_pretrained(model, "miike-ai/Llama-3.1-8b-gsm8k-r")
36
37# Import sampling parameters from vLLM
38from vllm import SamplingParams
39
40print("Model fully loaded into memory. Ready to chat!")
41print("Type 'exit' or 'quit' to stop.\n")
42
43# Interactive chat loop
44while True:
45 try:
46 user_input = input("User: ")
47 except KeyboardInterrupt:
48 print("\nExiting...")
49 break
50
51 # Exit the chat loop if the user types 'exit' or 'quit'
52 if user_input.strip().lower() in {"exit", "quit"}:
53 print("Exiting...")
54 break
55
56 # Prepare the prompt with the system prompt and user input
57 prompt = tokenizer.apply_chat_template(
58 [
59 {"role": "system", "content": SYSTEM_PROMPT},
60 {"role": "user", "content": user_input},
61 ],
62 tokenize=False,
63 add_generation_prompt=True
64 )
65
66 # Set sampling parameters for generation
67 sampling_params = SamplingParams(
68 temperature=0.8,
69 top_p=0.95,
70 max_tokens=1024
71 )
72
73 # Stream the response
74 start_time = time.time()
75 print("\nAssistant: ", end="", flush=True)
76
77 # Generate initial response
78 outputs = model.fast_generate(prompt, sampling_params=sampling_params)
79 response_text = outputs[0].outputs[0].text
80
81 # Stream the response token by token
82 current_text = ""
83 for i in range(len(response_text)):
84 new_token = response_text[i]
85 print(new_token, end="", flush=True)
86 current_text += new_token
87 time.sleep(0.02) # Small delay for readability
88
89 inference_time = time.time() - start_time
90 print(f"\nInference time: {inference_time:.2f} seconds\n")
91