import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
model_id = "alibaba-pai/SearchQwen3-8B"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
model_id,
torch_dtype=torch.bfloat16,
device_map="auto",
).eval()
messages = [
{"role": "system", "content": "You are a search agent. Use tools before answering."},
{"role": "user", "content": "Which city is the birthplace of the author of The Old Man and the Sea?"},
]
tools = [
{
"type": "function",
"function": {
"name": "search",
"description": "Search the web.",
"parameters": {
"type": "object",
"properties": {"query": {"type": "string"}},
"required": ["query"],
},
},
}
]
inputs = tokenizer.apply_chat_template(
messages,
tools=tools,
add_generation_prompt=True,
tokenize=True,
return_tensors="pt",
return_dict=True,
enable_thinking=False,
).to(model.device)
output = model.generate(**inputs, max_new_tokens=256, do_sample=False)
print(tokenizer.decode(output[0, inputs["input_ids"].shape[-1]:], skip_special_tokens=False))