File size: 2,111 Bytes
c85c557 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 | import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import PeftModel
BASE_MODEL_ID = "google/gemma-4-E2B-it"
ADAPTER_DIR = "./final_adapter"
print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(ADAPTER_DIR)
# Configure 4-bit quantization
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16,
bnb_4bit_use_double_quant=True,
)
print("Loading base model in 4-bit...")
base_model = AutoModelForCausalLM.from_pretrained(
BASE_MODEL_ID,
quantization_config=bnb_config,
device_map={"": 0},
dtype=torch.bfloat16,
low_cpu_mem_usage=True,
)
print("Loading LoRA adapter...")
model = PeftModel.from_pretrained(base_model, ADAPTER_DIR)
model.eval()
print("Model and LoRA adapter successfully loaded!")
# Define available system tools
tools = [
{
"type": "function",
"function": {
"name": "execute_bash",
"description": "Execute a bash command on the local system.",
"parameters": {
"type": "object",
"properties": {
"command": {
"type": "string",
"description": "The shell command to run."
}
},
"required": ["command"]
}
}
}
]
messages = [
{
"role": "user",
"content": "Can you check the current status of the repository and list any modified files?"
}
]
print("\nFormatting prompt with tools...")
prompt = tokenizer.apply_chat_template(
messages,
tools=tools,
tokenize=False,
add_generation_prompt=True
)
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
print("Generating response...")
with torch.no_grad():
outputs = model.generate(
**inputs,
max_new_tokens=256,
do_sample=False
)
response = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=False)
print("\n--- Model Output ---")
print(response)
|