FrontiersMind/Lumma-0.6B-Tool

import torch
import json
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_PATH = "FrontiersMind/Lumma-0.6B-Tool"


tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_PATH, trust_remote_code=True, torch_dtype=torch.bfloat16, device_map="auto"
)
model.eval()


def generate(max_new_tokens,messages):
    prompt = tokenizer.apply_chat_template(
        messages,
        tools_list=tools_list,   # pass every turn
        tokenize=False,
        add_generation_prompt=True,
    )
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    with torch.no_grad():
        out = model.generate(
            **inputs, max_new_tokens=max_new_tokens, do_sample=False,
            eos_token_id=tokenizer.eos_token_id, pad_token_id=tokenizer.pad_token_id,
        )
    text = tokenizer.decode(out[0, inputs["input_ids"].shape[1]:], skip_special_tokens=False)
    return text.split("<|endoftext|>")[0].strip()


tools = [
    {
        "name": "Quotes by Keywords",
        "description": "Returns a list of quotes containing the specified keyword.",
        "parameters": {
            "type": "dict",
            "properties": {"word": {"description": "The keyword to search for in quotes.", "type": "string"}},
            "required": ["word"],
        },
        "required": None,
    },
    {
        "name": "Get Zip Code Information",
        "description": "Retrieve information about a specific zip code in the United States.",
        "parameters": {
            "type": "dict",
            "properties": {
                "country": {"description": "The country code (default: 'us')", "type": "string"},
                "postal_code": {"description": "The zip code (default: '90210')", "type": "string"},
            },
            "required": ["country", "postal_code"],
        },
        "required": None,
    },
]

tools_list = json.dumps(tools)
messages = []

# --- Turn 1: user → tool call ---
messages.append({"role": "user", "content": 'Find quotes about "inspiration".'})
reply = generate(max_new_tokens=512,messages=messages)
messages.append({"role": "assistant", "content": reply})
print("Turn 1:", reply)

# --- Turn 2: tool result → answer ---
tool_call_response = """[{"name": "Quotes by Keywords", "results": {"quotes": [{"text": "Keep going.", "author": "Sam Levenson"}]}}]"""
messages.append({"role": "tool", "content": tool_call_response})
reply = generate(max_new_tokens=512,messages=messages)
messages.append({"role": "assistant", "content": reply})
print("Turn 2:", reply)