Views
No views yet
1import torch
2import transformers
3import pyreft
4
5device = "cuda"
6
7# Load the base model
8model_name_or_path = "meta-llama/Meta-Llama-3-8B"
9model = transformers.AutoModelForCausalLM.from_pretrained(
10 model_name_or_path, torch_dtype=torch.bfloat16, device_map={"": device}
11)
12
13# Load the ReFT model
14reft_model = pyreft.ReftModel.load(
15 "./CiscoDevNetSandboxRunningConfig", model
16)
17
18# Ensure the ReFT model components are also on the GPU
19reft_model.set_device(device)
20
21# Define the prompt template
22prompt_no_input_template = """<s>[INST] <<SYS>>
23You are a computer networking expert specialized in Cisco IOS XE running configurations.
24<</SYS>>
25
26%s [/INST]
27"""
28
29# Load the tokenizer
30tokenizer = transformers.AutoTokenizer.from_pretrained(
31 model_name_or_path, model_max_length=2048,
32 padding_side="right", use_fast=False
33)
34
35# Set pad_token as eos_token
36tokenizer.pad_token = tokenizer.eos_token
37
38def generate_response(instruction):
39 # Tokenize and prepare the input
40 prompt = prompt_no_input_template % instruction
41 prompt = tokenizer(prompt, return_tensors="pt").to(device)
42
43 base_unit_location = prompt["input_ids"].shape[-1] - 1 # Last position
44
45 # Move all relevant tensors and operations to the GPU
46 prompt = {key: value.to(device) for key, value in prompt.items()} # Ensure the prompt is on the GPU
47
48 # Generate the response using the reft_model
49 _, reft_response = reft_model.generate(
50 prompt, unit_locations={"sources->base": (None, [[[base_unit_location]]])},
51 intervene_on_prompt=True, max_new_tokens=512, do_sample=True,
52 eos_token_id=tokenizer.eos_token_id, early_stopping=True
53 )
54
55 fine_tuned_answer = tokenizer.decode(reft_response[0], skip_special_tokens=True)
56 return fine_tuned_answer
57
58while True:
59 instruction = input("You: ")
60
61 if instruction.lower() == "exit":
62 print("Goodbye!")
63 break
64
65 # Generate and print the response
66 fine_tuned_answer = generate_response(instruction)
67 print(f"Question: {instruction}")
68 print(f"Fine-tuned Answer: {fine_tuned_answer}")
69
70 # Log intervention details
71 print(f"Intervention applied: {reft_model.interventions}")