import torchimport jsonfrom transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfigfrom peft import PeftModel model_id = "openbmb/MiniCPM5-2B"adapter_id = "antontuzovAI/MiniCPM5-2B-xLAM-60k-Production-Agent" # 1. Setup 4-bit Quantizationbnb_config = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=True) # 2. Load Base Model & Tokenizertokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)tokenizer.pad_token = tokenizer.eos_token base_model = AutoModelForCausalLM.from_pretrained( model_id, quantization_config=bnb_config, device_map="auto", trust_remote_code=True, torch_dtype=torch.bfloat16) # 3. Load the Agent LoRAmodel = PeftModel.from_pretrained(base_model, adapter_id)model.eval() # 4. Define Toolstools = [ { "name": "get_weather", "description": "Get current weather for a city.", "parameters": { "city": {"type": "string", "description": "City name", "required": True}, "units": {"type": "string", "enum": ["celsius", "fahrenheit"]} } }] # 5. Build Prompttools_str = json.dumps(tools, indent=2)system_msg = ( "You are a strict tool-calling agent.\n" "Output ONLY a valid JSON array of tool calls.\n" 'Format: [{"name": "tool_name", "arguments": {"arg": "val"}}]\n' "If no tool is needed, output: []\n\n" f"Available tools:\n{tools_str}") messages = [ {"role": "system", "content": system_msg}, {"role": "user", "content": "What is the weather in Tokyo in celsius?"}] prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) # 6. Generateinputs = tokenizer(prompt, return_tensors="pt").to(model.device)with torch.no_grad(): outputs = model.generate(**inputs, max_new_tokens=128, do_sample=False, pad_token_id=tokenizer.eos_token_id) response = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)print("Agent Output:\n", response)