Update app.py

app.py

@@ -4,52 +4,46 @@ from transformers import AutoProcessor, AutoModelForCausalLM
 hf_token = os.getenv("HF_TOKEN")
 model_id = "google/functiongemma-270m-it"
 
-# Load model
+# 1. Load for CPU specifically
 processor = AutoProcessor.from_pretrained(model_id, token=hf_token)
 model = AutoModelForCausalLM.from_pretrained(
     model_id, 
-    torch_dtype=torch.float16, 
-    device_map="auto", 
-    token=hf_token,
-    low_cpu_mem_usage=True
+    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, tools_json):
     try:
         tools = json.loads(tools_json) if tools_json.strip() else []
+        
+        # FunctionGemma format
         messages = [
             {"role": "developer", "content": "You are a model that can do function calling with the following functions"},
             {"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(model.device)
+        ).to("cpu")
 
         with torch.no_grad():
             outputs = model.generate(**inputs, max_new_tokens=128, do_sample=False)
         
         input_len = inputs.input_ids.shape[1]
-        return processor.decode(outputs[0][input_len:], skip_special_tokens=True)
+        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)}"
 
-# Create the UI using Blocks for better API support
-with gr.Blocks() as demo:
-    gr.Markdown("# FunctionGemma API Server")
-    with gr.Row():
-        prompt_input = gr.Textbox(label="User Prompt")
-        tools_input = gr.Textbox(label="Tools (JSON Array)")
-    output_text = gr.Code(label="Model Output")
-    
-    submit_btn = gr.Button("Submit")
-    
-    # CRITICAL: This 'api_name' must match what your client expects
-    submit_btn.click(
-        fn=process_request, 
-        inputs=[prompt_input, tools_input], 
-        outputs=output_text,
-        api_name="predict" 
-    )
+demo = gr.Interface(
+    fn=process_request,
+    inputs=[gr.Textbox(label="User Prompt"), gr.Textbox(label="Tools (JSON Array)")],
+    outputs=gr.Code(label="Model Output"),
+    title="FunctionGemma CPU Fixed"
+)
 
 demo.launch()