"""
Example: Tool-calling inference with Gemma 4 E4B Tool-Call v02

Usage:
    python example_inference.py
"""
import torch
from transformers import AutoProcessor, AutoModelForMultimodalLM

MODEL_ID = "roshangrewal/gemma4-e4b-toolcall-v02"

print("Loading model...")
model = AutoModelForMultimodalLM.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16, device_map="auto")
processor = AutoProcessor.from_pretrained(MODEL_ID)
print("Ready!\n")

tools = [
    {"type": "function", "function": {
        "name": "get_weather", "description": "Get current weather for a city",
        "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}
    }},
    {"type": "function", "function": {
        "name": "search_flights", "description": "Search available flights",
        "parameters": {"type": "object", "properties": {"origin": {"type": "string"}, "destination": {"type": "string"}, "date": {"type": "string"}}, "required": ["origin", "destination", "date"]}
    }},
    {"type": "function", "function": {
        "name": "send_email", "description": "Send an email",
        "parameters": {"type": "object", "properties": {"to": {"type": "string"}, "subject": {"type": "string"}, "body": {"type": "string"}}, "required": ["to", "subject", "body"]}
    }},
]

queries = [
    "What's the weather in Mumbai?",
    "Find flights from Delhi to London on March 20",
    "Hi, how are you today?",
    "Email john@company.com about the project deadline",
]

for query in queries:
    messages = [{"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": query}]
    text = processor.apply_chat_template(messages, tools=tools, tokenize=False, add_generation_prompt=True)
    inputs = processor(text=text, return_tensors="pt").to(model.device)
    with torch.no_grad():
        output = model.generate(**inputs, max_new_tokens=200, do_sample=False)
    response = processor.decode(output[0][inputs["input_ids"].shape[1]:], skip_special_tokens=False)
    response_clean = response.split("<eos>")[0].split("<turn|>")[0].strip()
    print(f"User: {query}")
    print(f"Model: {response_clean}\n")