import json
import re
from transformers import AutoModelForCausalLM, AutoTokenizer
model_id = "LiquidAI/LFM3-1.7B"
model = AutoModelForCausalLM.from_pretrained(model_id, dtype="bfloat16", device_map="auto", trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
def tool(name, description, properties, required, defer_loading=False):
schema = {
"type": "function",
"function": {
"name": name,
"description": description,
"parameters": {"type": "object", "properties": properties, "required": required},
},
}
if defer_loading:
schema["defer_loading"] = True
return schema
# The search tool is always loaded; the catalog tools are deferred
tools = [
tool("search_tools", "Search the tool catalog and return the tools that match a query.",
{"query": {"type": "string", "description": "What the tool should do"}}, ["query"]),
tool("get_weather", "Get the current weather for a city.",
{"city": {"type": "string", "description": "City name"}}, ["city"], defer_loading=True),
tool("convert_currency", "Convert an amount of money from one currency to another.",
{"amount": {"type": "number"}, "from_currency": {"type": "string"}, "to_currency": {"type": "string"}},
["amount", "from_currency", "to_currency"], defer_loading=True),
]
# Mock functions (replace with actual API calls in production)
functions = {
"get_weather": lambda city: {"city": city, "condition": "light rain", "temperature_c": 11},
"convert_currency": lambda amount, from_currency, to_currency: {"amount": round(float(amount) * 1.17, 2), "currency": to_currency},
}
def search_catalog(query):
words = set(query.lower().split())
return [
t["function"]["name"] for t in tools
if t.get("defer_loading") and words & set(t["function"]["description"].lower().rstrip(".").split())
]
def parse_tool_calls(text):
calls = []
for name, body in re.findall(r"<function=([^>\n]+)>\n(.*?)</function>", text, re.S):
arguments = dict(re.findall(r"<parameter=([^>\n]+)>\n(.*?)\n</parameter>", body, re.S))
calls.append({"name": name, "arguments": arguments})
return calls
messages = [
{"role": "system", "content": "You only have one tool right now: search_tools, which searches a catalog of many more tools. If none of your current tools can handle a request, call search_tools to find one, then call the tool it returns."},
{"role": "user", "content": "Do I need an umbrella in Edinburgh today?"},
]
for _ in range(5):
inputs = tokenizer.apply_chat_template(
messages, tools=tools, add_generation_prompt=True, return_dict=True, return_tensors="pt"
)
outputs = model.generate(
**inputs.to(model.device),
do_sample=True,
temperature=0.1,
top_k=50,
repetition_penalty=1.1,
max_new_tokens=1024,
)
text = tokenizer.decode(outputs[0][len(inputs["input_ids"][0]):], skip_special_tokens=False)
reasoning, _, answer = text.partition("</think>")
calls = parse_tool_calls(answer)
if not calls:
print(answer.replace("<|im_end|>", "").strip())
break
messages.append({
"role": "assistant",
"reasoning": reasoning,
"tool_calls": [{"type": "function", "function": call} for call in calls],
})
for call in calls:
if call["name"] == "search_tools":
names = search_catalog(call["arguments"]["query"])
content = [{"type": "text", "text": f"Found {len(names)} tool(s)."}]
content += [{"type": "tool_reference", "name": name} for name in names]
else:
content = json.dumps(functions[call["name"]](**call["arguments"]))
messages.append({"role": "tool", "content": content})
