Views
No views yet
google/functiongemma-270m-it1from transformers import AutoModelForCausalLM, AutoTokenizer
2import json
3import re
4
5# Load model
6model = AutoModelForCausalLM.from_pretrained(
7 "gnumanth/xkcd-functiongemma",
8 device_map="auto",
9 torch_dtype="auto"
10)
11tokenizer = AutoTokenizer.from_pretrained("gnumanth/xkcd-functiongemma")
12
13# Define tools
14TOOLS = [{
15 "type": "function",
16 "function": {
17 "name": "search_xkcd",
18 "description": "Search XKCD comics by topic",
19 "parameters": {
20 "type": "object",
21 "properties": {"query": {"type": "string"}},
22 "required": ["query"]
23 }
24 }
25}]
26
27# Generate function call
28messages = [{"role": "user", "content": "Find xkcd about programming"}]
29text = tokenizer.apply_chat_template(messages, tools=TOOLS, tokenize=False, add_generation_prompt=True)
30inputs = tokenizer(text, return_tensors="pt").to(model.device)
31outputs = model.generate(**inputs, max_new_tokens=128, do_sample=False)
32response = tokenizer.decode(outputs[0], skip_special_tokens=True)
33
34print(response)
35# Output: <start_function_call>call:search_xkcd{"query": "programming"}<end_function_call>1def parse_function_call(output: str) -> dict | None:
2 """Extract function name and arguments from model output."""
3 match = re.search(r'call:(\w+)\s*\{(.+)\}', output, re.DOTALL)
4 if not match:
5 return None
6
7 func_name = match.group(1)
8 args_raw = match.group(2).strip()
9
10 # Handle double braces from training format
11 args_raw = re.sub(r'^\s*\{', '', args_raw)
12 if args_raw.endswith('}'):
13 args_raw = args_raw[:-1]
14
15 try:
16 return {"function": func_name, "arguments": json.loads('{' + args_raw + '}')}
17 except json.JSONDecodeError:
18 return None
19
20# Usage
21call = parse_function_call(response)
22# {'function': 'search_xkcd', 'arguments': {'query': 'programming'}}