functiongemma-mobile-actions

import torch
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer

model_id = "dousery/functiongemma-mobile-actions"  
device = "cuda" if torch.cuda.is_available() else "cpu"

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
    device_map="auto" if device == "cuda" else None,
    trust_remote_code=True,
).eval()

tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
if device == "cpu":
    model = model.to(device)

dataset = load_dataset("google/mobile-actions", split="train")
text = tokenizer.apply_chat_template(
    dataset[0]["messages"][:2],
    tools=dataset[0]["tools"],
    tokenize=False,
    add_generation_prompt=True,
).removeprefix("<bos>")

inputs = tokenizer(text, return_tensors="pt").to(device)
with torch.no_grad():
    _ = model.generate(
        **inputs,
        max_new_tokens=256,
        streamer=TextStreamer(tokenizer, skip_prompt=True),
        top_p=0.95,
        top_k=64,
        temperature=1.0,
    )