| ## How to use this model |
|
|
| To use this model: |
|
|
| ```python |
| import torch |
| import transformers |
| import pyreft |
| |
| device = "cuda" |
| |
| # Load the base model |
| model_name_or_path = "meta-llama/Meta-Llama-3-8B" |
| model = transformers.AutoModelForCausalLM.from_pretrained( |
| model_name_or_path, torch_dtype=torch.bfloat16, device_map={"": device} |
| ) |
| |
| # Load the ReFT model |
| reft_model = pyreft.ReftModel.load( |
| "./CiscoDevNetSandboxRunningConfig", model |
| ) |
| |
| # Ensure the ReFT model components are also on the GPU |
| reft_model.set_device(device) |
| |
| # Define the prompt template |
| prompt_no_input_template = """<s>[INST] <<SYS>> |
| You are a computer networking expert specialized in Cisco IOS XE running configurations. |
| <</SYS>> |
| |
| %s [/INST] |
| """ |
| |
| # Load the tokenizer |
| tokenizer = transformers.AutoTokenizer.from_pretrained( |
| model_name_or_path, model_max_length=2048, |
| padding_side="right", use_fast=False |
| ) |
| |
| # Set pad_token as eos_token |
| tokenizer.pad_token = tokenizer.eos_token |
| |
| def generate_response(instruction): |
| # Tokenize and prepare the input |
| prompt = prompt_no_input_template % instruction |
| prompt = tokenizer(prompt, return_tensors="pt").to(device) |
| |
| base_unit_location = prompt["input_ids"].shape[-1] - 1 # Last position |
| |
| # Move all relevant tensors and operations to the GPU |
| prompt = {key: value.to(device) for key, value in prompt.items()} # Ensure the prompt is on the GPU |
| |
| # Generate the response using the reft_model |
| _, reft_response = reft_model.generate( |
| prompt, unit_locations={"sources->base": (None, [[[base_unit_location]]])}, |
| intervene_on_prompt=True, max_new_tokens=512, do_sample=True, |
| eos_token_id=tokenizer.eos_token_id, early_stopping=True |
| ) |
| |
| fine_tuned_answer = tokenizer.decode(reft_response[0], skip_special_tokens=True) |
| return fine_tuned_answer |
| |
| while True: |
| instruction = input("You: ") |
| |
| if instruction.lower() == "exit": |
| print("Goodbye!") |
| break |
| |
| # Generate and print the response |
| fine_tuned_answer = generate_response(instruction) |
| print(f"Question: {instruction}") |
| print(f"Fine-tuned Answer: {fine_tuned_answer}") |
| |
| # Log intervention details |
| print(f"Intervention applied: {reft_model.interventions}") |
| ``` |
