Spaces:
Sleeping
Sleeping
File size: 1,722 Bytes
3d16511 25ac059 3d16511 25ac059 efefc6e 3d16511 293b192 efefc6e 3d16511 8cbc163 3d16511 efefc6e 3d16511 8cbc163 3d16511 efefc6e 3d16511 efefc6e 25ac059 3d16511 efefc6e 3d16511 efefc6e 8cbc163 efefc6e 293b192 0ec8839 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 | import os, json, gradio as gr, torch
from transformers import AutoProcessor, AutoModelForCausalLM
hf_token = os.getenv("HF_TOKEN")
model_id = "google/functiongemma-270m-it"
# 1. Load for CPU specifically
processor = AutoProcessor.from_pretrained(model_id, token=hf_token)
model = AutoModelForCausalLM.from_pretrained(
model_id,
torch_dtype=torch.float32, # CPU prefers float32
device_map={"": "cpu"}, # Forces everything onto CPU, avoiding "meta device"
token=hf_token
)
def process_request(user_prompt, developer_prompt, tools_json):
try:
tools = json.loads(tools_json) if tools_json.strip() else []
# FunctionGemma format
messages = [
{"role": "developer", "content": developer_prompt},
{"role": "user", "content": user_prompt}
]
# 2. Ensure inputs are on CPU
inputs = processor.apply_chat_template(
messages, tools=tools, add_generation_prompt=True,
return_dict=True, return_tensors="pt"
).to("cpu")
with torch.no_grad():
outputs = model.generate(**inputs, max_new_tokens=128, do_sample=False)
input_len = inputs.input_ids.shape[1]
decoded = processor.decode(outputs[0][input_len:], skip_special_tokens=True)
return decoded if decoded.strip() else "Model returned an empty string."
except Exception as e:
return f"Error: {str(e)}"
demo = gr.Interface(
fn=process_request,
inputs=[gr.Textbox(label="User Prompt"), gr.Textbox(label="Developer Prompt"), gr.Textbox(label="Tools (JSON Array)")],
outputs=gr.Code(label="Model Output"),
title="FunctionGemma CPU Fixed"
)
demo.launch(share=True) |